Fix custom vector store notebook

Author: natoverse

docs/examples_notebooks/custom_vector_store.ipynb

@@ -61,6 +61,7 @@
     "import numpy as np\n",
     "import yaml\n",
     "\n",
+    "from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig\n",
     "from graphrag.data_model.types import TextEmbedder\n",
     "\n",
     "# GraphRAG vector store components\n",
@@ -147,14 +148,12 @@
     "        self.vectors: dict[str, np.ndarray] = {}\n",
     "        self.connected = False\n",
     "\n",
-    "        print(\n",
-    "            f\"🚀 SimpleInMemoryVectorStore initialized for collection: {self.collection_name}\"\n",
-    "        )\n",
+    "        print(f\"🚀 SimpleInMemoryVectorStore initialized for index: {self.index_name}\")\n",
     "\n",
     "    def connect(self, **kwargs: Any) -> None:\n",
     "        \"\"\"Connect to the vector storage (no-op for in-memory store).\"\"\"\n",
     "        self.connected = True\n",
-    "        print(f\"✅ Connected to in-memory vector store: {self.collection_name}\")\n",
+    "        print(f\"✅ Connected to in-memory vector store: {self.index_name}\")\n",
     "\n",
     "    def load_documents(\n",
     "        self, documents: list[VectorStoreDocument], overwrite: bool = True\n",
@@ -250,7 +249,7 @@
     "    def get_stats(self) -> dict[str, Any]:\n",
     "        \"\"\"Get statistics about the vector store (custom method).\"\"\"\n",
     "        return {\n",
-    "            \"collection_name\": self.collection_name,\n",
+    "            \"index_name\": self.index_name,\n",
     "            \"document_count\": len(self.documents),\n",
     "            \"vector_count\": len(self.vectors),\n",
     "            \"connected\": self.connected,\n",
@@ -353,11 +352,11 @@
    "outputs": [],
    "source": [
     "# Test creating vector store using the factory\n",
-    "vector_store_config = {\"collection_name\": \"test_collection\"}\n",
+    "schema = VectorStoreSchemaConfig(index_name=\"test_collection\")\n",
     "\n",
     "# Create vector store instance using factory\n",
     "vector_store = VectorStoreFactory.create_vector_store(\n",
-    "    CUSTOM_VECTOR_STORE_TYPE, vector_store_config\n",
+    "    CUSTOM_VECTOR_STORE_TYPE, vector_store_schema_config=schema\n",
     ")\n",
     "\n",
     "print(f\"✅ Created vector store instance: {type(vector_store).__name__}\")\n",
@@ -486,9 +485,13 @@
     "    print(\"🚀 Simulating GraphRAG pipeline with custom vector store...\\n\")\n",
     "\n",
     "    # 1. GraphRAG creates vector store using factory\n",
-    "    config = {\"collection_name\": \"graphrag_entities\", \"similarity_threshold\": 0.3}\n",
+    "    schema = VectorStoreSchemaConfig(index_name=\"graphrag_entities\")\n",
     "\n",
-    "    store = VectorStoreFactory.create_vector_store(CUSTOM_VECTOR_STORE_TYPE, config)\n",
+    "    store = VectorStoreFactory.create_vector_store(\n",
+    "        CUSTOM_VECTOR_STORE_TYPE,\n",
+    "        vector_store_schema_config=schema,\n",
+    "        similarity_threshold=0.3,\n",
+    "    )\n",
     "    store.connect()\n",
     "\n",
     "    print(\"✅ Step 1: Vector store created and connected\")\n",
@@ -549,7 +552,8 @@
     "    # Test 1: Basic functionality\n",
     "    print(\"Test 1: Basic functionality\")\n",
     "    store = VectorStoreFactory.create_vector_store(\n",
-    "        CUSTOM_VECTOR_STORE_TYPE, {\"collection_name\": \"test\"}\n",
+    "        CUSTOM_VECTOR_STORE_TYPE,\n",
+    "        vector_store_schema_config=VectorStoreSchemaConfig(index_name=\"test\"),\n",
     "    )\n",
     "    store.connect()\n",
     "\n",
@@ -597,7 +601,8 @@
     "    # Test 5: Error handling\n",
     "    print(\"\\nTest 5: Error handling\")\n",
     "    disconnected_store = VectorStoreFactory.create_vector_store(\n",
-    "        CUSTOM_VECTOR_STORE_TYPE, {\"collection_name\": \"test2\"}\n",
+    "        CUSTOM_VECTOR_STORE_TYPE,\n",
+    "        vector_store_schema_config=VectorStoreSchemaConfig(index_name=\"test2\"),\n",
     "    )\n",
     "\n",
     "    try:\n",
@@ -653,7 +658,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "graphrag-venv (3.10.18)",
+   "display_name": "graphrag",
    "language": "python",
    "name": "python3"
   },
@@ -667,7 +672,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.18"
+   "version": "3.12.10"
   }
  },
  "nbformat": 4,

graphrag/index/operations/embed_text/embed_text.py

@@ -210,7 +210,7 @@ def _create_vector_store(
     vector_store = VectorStoreFactory().create_vector_store(
         vector_store_schema_config=single_embedding_config,
         vector_store_type=vector_store_type,
-        kwargs=vector_store_config,
+        **vector_store_config,
     )
 
     vector_store.connect(**vector_store_config)

graphrag/utils/api.py

@@ -130,7 +130,7 @@ def get_embedding_store(
         embedding_store = VectorStoreFactory().create_vector_store(
             vector_store_type=vector_store_type,
             vector_store_schema_config=single_embedding_config,
-            kwargs={**store},
+            **store,
         )
         embedding_store.connect(**store)
         # If there is only a single index, return the embedding store directly

graphrag/vector_stores/factory.py

@@ -53,7 +53,7 @@ def create_vector_store(
         cls,
         vector_store_type: str,
         vector_store_schema_config: VectorStoreSchemaConfig,
-        kwargs: dict,
+        **kwargs: dict,
     ) -> BaseVectorStore:
         """Create a vector store object from the provided type.
 

tests/integration/vector_stores/test_factory.py

@@ -81,9 +81,7 @@ def test_register_and_create_custom_vector_store():
     )
 
     vector_store = VectorStoreFactory.create_vector_store(
-        vector_store_type="custom",
-        vector_store_schema_config=VectorStoreSchemaConfig(),
-        kwargs={},
+        vector_store_type="custom", vector_store_schema_config=VectorStoreSchemaConfig()
     )
 
     assert custom_vector_store_class.called
@@ -109,7 +107,6 @@ def test_create_unknown_vector_store():
         VectorStoreFactory.create_vector_store(
             vector_store_type="unknown",
             vector_store_schema_config=VectorStoreSchemaConfig(),
-            kwargs={},
         )
 
 
@@ -162,7 +159,6 @@ def search_by_id(self, id):
     vector_store = VectorStoreFactory.create_vector_store(
         vector_store_type="custom_class",
         vector_store_schema_config=VectorStoreSchemaConfig(),
-        kwargs={"collection_name": "test"},
     )
 
     assert isinstance(vector_store, CustomVectorStore)