Run llamaindex tests

.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

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. "

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