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 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..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,12 +13,15 @@ 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 = 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