diff --git a/llama-index-integrations/llms/llama-index-llms-databricks/README.md b/llama-index-integrations/llms/llama-index-llms-databricks/README.md index 4bf0b5a03881a0d1bc08ce17f5f6d2e7f00a241c..cd5e60ba6709611efa51e8ee2f0b9faae9218729 100644 --- a/llama-index-integrations/llms/llama-index-llms-databricks/README.md +++ b/llama-index-integrations/llms/llama-index-llms-databricks/README.md @@ -2,7 +2,7 @@ ## Overview -Integrate with DataBricks LLMs APIs. +Integrate with Databricks LLMs APIs. ## Installation @@ -15,15 +15,15 @@ pip install llama-index-llms-databricks With environmental variables. ```.env -DATABRICKS_API_KEY=your_api_key -DATABRICKS_API_BASE=https://[your-work-space].cloud.databricks.com/serving-endpoints/[your-serving-endpoint] +DATABRICKS_TOKEN=your_api_key +DATABRICKS_SERVING_ENDPOINT=https://[your-work-space].cloud.databricks.com/serving-endpoints ``` ```python -from llama_index.llms.databricks import DataBricks +from llama_index.llms.databricks import Databricks -# Initialize DataBricks LLM without explicitly passing the API key and base -llm = DataBricks(model="databricks-dbrx-instruct") +# Initialize Databricks LLM without explicitly passing the API key and base +llm = Databricks(model="databricks-dbrx-instruct") # Make a query to the LLM response = llm.complete("Explain the importance of open source LLMs") @@ -34,13 +34,13 @@ print(response) Without environmental variables ```python -from llama_index.llms.databricks import DataBricks +from llama_index.llms.databricks import Databricks -# Set up the DataBricks class with the required model, API key and serving endpoint -llm = DataBricks( +# Set up the Databricks class with the required model, API key and serving endpoint +llm = Databricks( model="databricks-dbrx-instruct", api_key="your_api_key", - api_base="https://[your-work-space].cloud.databricks.com/serving-endpoints/[your-serving-endpoint]", + api_base="https://[your-work-space].cloud.databricks.com/serving-endpoints", ) # Call the complete method with a query diff --git a/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/__init__.py b/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/__init__.py index a46ccbff700657dc83b78d923da53791134ccaa9..5872eab54b2307a4638f89db031e66995ad7e8fc 100644 --- a/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/__init__.py +++ b/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/__init__.py @@ -1,3 +1,3 @@ -from llama_index.llms.databricks.base import DataBricks +from llama_index.llms.databricks.base import Databricks -__all__ = ["DataBricks"] +__all__ = ["Databricks"] diff --git a/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/base.py b/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/base.py index dc99f7feb13a776674d949a66370de7dd93884b0..2d0e35a27e470db67876a589e6e37b81367ccdfa 100644 --- a/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/base.py +++ b/llama-index-integrations/llms/llama-index-llms-databricks/llama_index/llms/databricks/base.py @@ -4,17 +4,17 @@ from typing import Any, Optional from llama_index.llms.openai_like import OpenAILike -class DataBricks(OpenAILike): - """DataBricks LLM. +class Databricks(OpenAILike): + """Databricks LLM. Examples: `pip install llama-index-llms-databricks` ```python - from llama_index.llms.databricks import DataBricks + from llama_index.llms.databricks import Databricks - # Set up the DataBricks class with the required model, API key and serving endpoint - llm = DataBricks(model="databricks-dbrx-instruct", api_key="your_api_key", api_base="https://[your-work-space].cloud.databricks.com/serving-endpoints/[your-serving-endpoint]") + # Set up the Databricks class with the required model, API key and serving endpoint + llm = Databricks(model="databricks-dbrx-instruct", api_key="your_api_key", api_base="https://[your-work-space].cloud.databricks.com/serving-endpoints") # Call the complete method with a query response = llm.complete("Explain the importance of open source LLMs") @@ -31,8 +31,8 @@ class DataBricks(OpenAILike): is_chat_model: bool = True, **kwargs: Any, ) -> None: - api_key = api_key or os.environ.get("DATABRICKS_API_KEY", None) - api_base = api_base or os.environ.get("DATABRICKS_API_BASE", None) + api_key = api_key or os.environ.get("DATABRICKS_TOKEN", None) + api_base = api_base or os.environ.get("DATABRICKS_SERVING_ENDPOINT", None) super().__init__( model=model, api_key=api_key, @@ -44,4 +44,4 @@ class DataBricks(OpenAILike): @classmethod def class_name(cls) -> str: """Get class name.""" - return "DataBricks" + return "Databricks" diff --git a/llama-index-integrations/llms/llama-index-llms-databricks/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-databricks/pyproject.toml index 631af5be36257b45afa8fbbf4b6490bbd3cd0e6c..c5d980fcc1d8105ebc4ee9ff6cc831ae4989d9bc 100644 --- a/llama-index-integrations/llms/llama-index-llms-databricks/pyproject.toml +++ b/llama-index-integrations/llms/llama-index-llms-databricks/pyproject.toml @@ -30,7 +30,7 @@ license = "MIT" name = "llama-index-llms-databricks" packages = [{include = "llama_index/"}] readme = "README.md" -version = "0.1.0" +version = "0.1.1" [tool.poetry.dependencies] python = ">=3.8.1,<4.0" diff --git a/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_integration_databricks.py b/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_integration_databricks.py index fbf6af8af5b5adb237fdd4e835906f9e06f73d89..8f0a1ca161a2b4c544f64a22b7a0a0d154921558 100644 --- a/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_integration_databricks.py +++ b/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_integration_databricks.py @@ -2,15 +2,16 @@ import os import pytest -from llama_index.llms.databricks import DataBricks +from llama_index.llms.databricks import Databricks @pytest.mark.skipif( - "DATABRICKS_API_KEY" not in os.environ or "DATABRICKS_API_BASE" not in os.environ, - reason="DATABRICKS_API_KEY or DATABRICKS_API_BASE not set in environment", + "DATABRICKS_TOKEN" not in os.environ + or "DATABRICKS_SERVING_ENDPOINT" not in os.environ, + reason="DATABRICKS_TOKEN or DATABRICKS_SERVING_ENDPOINT not set in environment", ) def test_completion(): - databricks = DataBricks( + databricks = Databricks( model="databricks-dbrx-instruct", temperature=0, max_tokens=2 ) resp = databricks.complete("hello") @@ -18,11 +19,12 @@ def test_completion(): @pytest.mark.skipif( - "DATABRICKS_API_KEY" not in os.environ or "DATABRICKS_API_BASE" not in os.environ, - reason="DATABRICKS_API_KEY or DATABRICKS_API_BASE not set in environment", + "DATABRICKS_TOKEN" not in os.environ + or "DATABRICKS_SERVING_ENDPOINT" not in os.environ, + reason="DATABRICKS_TOKEN or DATABRICKS_SERVING_ENDPOINT not set in environment", ) def test_stream_completion(): - databricks = DataBricks( + databricks = Databricks( model="databricks-dbrx-instruct", temperature=0, max_tokens=2 ) stream = databricks.stream_complete("hello") diff --git a/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_llms_databricks.py b/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_llms_databricks.py index 1bfe9b0022f4b57129d3e6bb155c81bfb629189f..e724776328dc9bddec0b1d682faf385f437547de 100644 --- a/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_llms_databricks.py +++ b/llama-index-integrations/llms/llama-index-llms-databricks/tests/test_llms_databricks.py @@ -1,7 +1,7 @@ from llama_index.core.base.llms.base import BaseLLM -from llama_index.llms.databricks import DataBricks +from llama_index.llms.databricks import Databricks def test_embedding_class(): - names_of_base_classes = [b.__name__ for b in DataBricks.__mro__] + names_of_base_classes = [b.__name__ for b in Databricks.__mro__] assert BaseLLM.__name__ in names_of_base_classes