From 7b3a7488414310422aba26194aceec8587460dcd Mon Sep 17 00:00:00 2001 From: Ann Zhang Date: Wed, 9 Apr 2025 17:18:49 -0700 Subject: [PATCH 1/3] add llamaindex test Signed-off-by: Ann Zhang --- .github/workflows/main.yml | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 90c88c4d..521c55b0 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -91,3 +91,24 @@ jobs: - name: Run tests run: | pytest integrations/openai/tests/unit_tests + + llamaindex_test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] + timeout-minutes: 20 + steps: + - name: Checkout code + uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + pip install . + pip install integrations/llamaindex[dev] + - name: Run tests + run: | + pytest integrations/llamaindex/tests/unit_tests From 09fe945e94a05027180a0c74e84013f6b98bd9f5 Mon Sep 17 00:00:00 2001 From: Ann Zhang Date: Wed, 9 Apr 2025 17:39:18 -0700 Subject: [PATCH 2/3] update tests Signed-off-by: Ann Zhang --- .../src/databricks_llamaindex/vector_search_retriever_tool.py | 3 +++ .../tests/unit_tests/test_vector_search_retriever_tool.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py b/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py index 322f44d8..32a0a6e8 100644 --- a/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py +++ b/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py @@ -19,6 +19,9 @@ class VectorSearchRetrieverTool(FunctionTool, VectorSearchRetrieverToolMixin): """Vector search retriever tool implementation.""" + class Config: + extra = "allow" # allow FunctionTool to set unknown attributes + text_column: Optional[str] = Field( None, description="The name of the text column to use for the embeddings. " diff --git a/integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py index e6e3ce42..6378d607 100644 --- a/integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py +++ b/integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py @@ -112,7 +112,7 @@ def test_vector_search_retriever_tool_description_generation(index_name: str) -> assert vector_search_tool.metadata.name != "" assert vector_search_tool.metadata.description != "" assert vector_search_tool.metadata.fn_schema == VectorSearchRetrieverToolInput - assert vector_search_tool.metadata.name == index_name + assert vector_search_tool.metadata.name == index_name.replace(".", "__") assert ( "A vector search-based retrieval tool for querying indexed embeddings." in vector_search_tool.metadata.description From ec1f7f85ad9e08520b910f7716e7501359a5c99f Mon Sep 17 00:00:00 2001 From: Ann Zhang Date: Wed, 9 Apr 2025 17:42:59 -0700 Subject: [PATCH 3/3] lint Signed-off-by: Ann Zhang --- .../src/databricks_llamaindex/vector_search_retriever_tool.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py b/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py index 32a0a6e8..33afad7c 100644 --- a/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py +++ b/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py @@ -13,14 +13,14 @@ from llama_index.core.embeddings import BaseEmbedding from llama_index.core.tools import FunctionTool from llama_index.core.tools.types import ToolMetadata -from pydantic import Field, PrivateAttr +from pydantic import Extra, Field, PrivateAttr class VectorSearchRetrieverTool(FunctionTool, VectorSearchRetrieverToolMixin): """Vector search retriever tool implementation.""" class Config: - extra = "allow" # allow FunctionTool to set unknown attributes + extra = Extra.allow # allow FunctionTool to set unknown attributes text_column: Optional[str] = Field( None,