test fixes

Author: Gaudy Blanco

.vscode/launch.json

@@ -45,7 +45,7 @@
 			"module": "pytest",
 			"args": [
 				"./tests/integration/vector_stores",
-				"-k", "test_lancedb"
+				"-k", "test_azure_ai_search"
 			],
 			"console": "integratedTerminal",
     		"justMyCode": false 

graphrag/vector_stores/azure_ai_search.py

@@ -159,7 +159,7 @@ def filter_by_id(self, include_ids: list[str] | list[int]) -> Any:
         # More info about odata filtering here: https://learn.microsoft.com/en-us/azure/search/search-query-odata-search-in-function
         # search.in is faster that joined and/or conditions
         id_filter = ",".join([f"{id!s}" for id in include_ids])
-        self.query_filter = f"search.in(id, '{id_filter}', ',')"
+        self.query_filter = f"search.in({self.id_field}, '{id_filter}', ',')"
 
         # Returning to keep consistency with other methods, but not needed
         # TODO: Refactor on a future PR

graphrag/vector_stores/lancedb.py

@@ -99,10 +99,10 @@ def filter_by_id(self, include_ids: list[str] | list[int]) -> Any:
         else:
             if isinstance(include_ids[0], str):
                 id_filter = ", ".join([f"'{id}'" for id in include_ids])
-                self.query_filter = f"id in ({id_filter})"
+                self.query_filter = f"{self.id_field} in ({id_filter})"
             else:
                 self.query_filter = (
-                    f"id in ({', '.join([str(id) for id in include_ids])})"
+                    f"{self.id_field} in ({', '.join([str(id) for id in include_ids])})"
                 )
         return self.query_filter
 
@@ -155,7 +155,7 @@ def search_by_id(self, id: str) -> VectorStoreDocument:
         """Search for a document by id."""
         doc = (
             self.document_collection.search()
-            .where(f"id == '{id}'", prefilter=True)
+            .where(f"{self.id_field} == '{id}'", prefilter=True)
             .to_list()
         )
         if doc:

tests/integration/vector_stores/test_azure_ai_search.py

@@ -41,8 +41,33 @@ def mock_index_client(self):
     def vector_store(self, mock_search_client, mock_index_client):
         """Create an Azure AI Search vector store instance."""
         vector_store = AzureAISearchVectorStore(
-            collection_name="test_vectors",
-            vector_store_schema_config=VectorStoreSchemaConfig(),
+            vector_store_schema_config=VectorStoreSchemaConfig(
+                index_name="test_vectors", vector_size=5
+            ),
+        )
+
+        # Create the necessary mocks first
+        vector_store.db_connection = mock_search_client
+        vector_store.index_client = mock_index_client
+
+        vector_store.connect(
+            url=TEST_AZURE_AI_SEARCH_URL,
+            api_key=TEST_AZURE_AI_SEARCH_KEY,
+        )
+        return vector_store
+
+    @pytest.fixture
+    def vector_store_custom(self, mock_search_client, mock_index_client):
+        """Create an Azure AI Search vector store instance."""
+        vector_store = AzureAISearchVectorStore(
+            vector_store_schema_config=VectorStoreSchemaConfig(
+                index_name="test_vectors",
+                id_field="id_custom",
+                text_field="text_custom",
+                attributes_field="attributes_custom",
+                vector_field="vector_custom",
+                vector_size=5,
+            ),
         )
 
         # Create the necessary mocks first
@@ -52,7 +77,6 @@ def vector_store(self, mock_search_client, mock_index_client):
         vector_store.connect(
             url=TEST_AZURE_AI_SEARCH_URL,
             api_key=TEST_AZURE_AI_SEARCH_KEY,
-            vector_size=5,
         )
         return vector_store
 
@@ -148,3 +172,72 @@ def none_embedder(text: str) -> None:
         )
         assert not mock_search_client.search.called
         assert len(results) == 0
+
+    async def test_vector_store_customization(
+        self,
+        vector_store_custom,
+        sample_documents,
+        mock_search_client,
+        mock_index_client,
+    ):
+        """Test vector store customization with Azure AI Search."""
+        # Setup mock responses
+        mock_index_client.list_index_names.return_value = []
+        mock_index_client.create_or_update_index = MagicMock()
+        mock_search_client.upload_documents = MagicMock()
+
+        search_results = [
+            {
+                vector_store_custom.id_field: "doc1",
+                vector_store_custom.text_field: "This is document 1",
+                vector_store_custom.vector_field: [0.1, 0.2, 0.3, 0.4, 0.5],
+                vector_store_custom.attributes_field: '{"title": "Doc 1", "category": "test"}',
+                "@search.score": 0.9,
+            },
+            {
+                vector_store_custom.id_field: "doc2",
+                vector_store_custom.text_field: "This is document 2",
+                vector_store_custom.vector_field: [0.2, 0.3, 0.4, 0.5, 0.6],
+                vector_store_custom.attributes_field: '{"title": "Doc 2", "category": "test"}',
+                "@search.score": 0.8,
+            },
+        ]
+        mock_search_client.search.return_value = search_results
+
+        mock_search_client.get_document.return_value = {
+            vector_store_custom.id_field: "doc1",
+            vector_store_custom.text_field: "This is document 1",
+            vector_store_custom.vector_field: [0.1, 0.2, 0.3, 0.4, 0.5],
+            vector_store_custom.attributes_field: '{"title": "Doc 1", "category": "test"}',
+        }
+
+        vector_store_custom.load_documents(sample_documents)
+        assert mock_index_client.create_or_update_index.called
+        assert mock_search_client.upload_documents.called
+
+        filter_query = vector_store_custom.filter_by_id(["doc1", "doc2"])
+        assert (
+            filter_query
+            == f"search.in({vector_store_custom.id_field}, 'doc1,doc2', ',')"
+        )
+
+        vector_results = vector_store_custom.similarity_search_by_vector(
+            [0.1, 0.2, 0.3, 0.4, 0.5], k=2
+        )
+        assert len(vector_results) == 2
+        assert vector_results[0].document.id == "doc1"
+        assert vector_results[0].score == 0.9
+
+        # Define a simple text embedder function for testing
+        def mock_embedder(text: str) -> list[float]:
+            return [0.1, 0.2, 0.3, 0.4, 0.5]
+
+        text_results = vector_store_custom.similarity_search_by_text(
+            "test query", mock_embedder, k=2
+        )
+        assert len(text_results) == 2
+
+        doc = vector_store_custom.search_by_id("doc1")
+        assert doc.id == "doc1"
+        assert doc.text == "This is document 1"
+        assert doc.attributes["title"] == "Doc 1"

tests/integration/vector_stores/test_cosmosdb.py

@@ -103,3 +103,64 @@ def test_clear():
         assert vector_store._database_exists() is False  # noqa: SLF001
     finally:
         pass
+
+
+def test_vector_store_customization():
+    """Test vector store customization with CosmosDB."""
+    vector_store = CosmosDBVectorStore(
+        vector_store_schema_config=VectorStoreSchemaConfig(
+            index_name="text-embeddings",
+            id_field="id_custom",
+            text_field="text_custom",
+            vector_field="vector_custom",
+            attributes_field="attributes_custom",
+            vector_size=5,
+        ),
+    )
+
+    try:
+        vector_store.connect(
+            connection_string=WELL_KNOWN_COSMOS_CONNECTION_STRING,
+            database_name="test_db",
+        )
+
+        docs = [
+            VectorStoreDocument(
+                id="doc1",
+                text="This is document 1",
+                vector=[0.1, 0.2, 0.3, 0.4, 0.5],
+                attributes={"title": "Doc 1", "category": "test"},
+            ),
+            VectorStoreDocument(
+                id="doc2",
+                text="This is document 2",
+                vector=[0.2, 0.3, 0.4, 0.5, 0.6],
+                attributes={"title": "Doc 2", "category": "test"},
+            ),
+        ]
+        vector_store.load_documents(docs)
+
+        vector_store.filter_by_id(["doc1"])
+
+        doc = vector_store.search_by_id("doc1")
+        assert doc.id == "doc1"
+        assert doc.text == "This is document 1"
+        assert doc.vector is not None
+        assert np.allclose(doc.vector, [0.1, 0.2, 0.3, 0.4, 0.5])
+        assert doc.attributes["title"] == "Doc 1"
+
+        # Define a simple text embedder function for testing
+        def mock_embedder(text: str) -> list[float]:
+            return [0.1, 0.2, 0.3, 0.4, 0.5]  # Return fixed embedding
+
+        vector_results = vector_store.similarity_search_by_vector(
+            [0.1, 0.2, 0.3, 0.4, 0.5], k=2
+        )
+        assert len(vector_results) > 0
+
+        text_results = vector_store.similarity_search_by_text(
+            "test query", mock_embedder, k=2
+        )
+        assert len(text_results) > 0
+    finally:
+        vector_store.clear()