fix(storage): fix wrong config in VectorStoreConnector

eosphoros-ai · May 20, 2024 · d39a785 · d39a785
1 parent f56f20e
commit d39a785
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 26 deletions.
diff --git a/dbgpt/storage/vector_store/connector.py b/dbgpt/storage/vector_store/connector.py
@@ -96,36 +96,25 @@ def __init__(
             raise Exception(f"Vector store {vector_store_type} not supported")
 
         logger.info(f"VectorStore:{self.connector_class}")
+
+        if vector_store_config is None:
+            vector_store_config = self.config_class()
+
         self._vector_store_type = vector_store_type
         self._embeddings = (
             vector_store_config.embedding_fn if vector_store_config else None
         )
-
         try:
-            if vector_store_config is not None:
-                config: IndexStoreConfig = self.config_class()
-                config.name = getattr(vector_store_config, "name", "default_name")
-                config.embedding_fn = getattr(vector_store_config, "embedding_fn", None)
-                config.max_chunks_once_load = getattr(
-                    vector_store_config, "max_chunks_once_load", 5
-                )
-                config.max_threads = getattr(vector_store_config, "max_threads", 4)
-                config.user = getattr(vector_store_config, "user", None)
-                config.password = getattr(vector_store_config, "password", None)
-
-                # extra
-                config_dict = vector_store_config.dict()
-                config.llm_client = config_dict.get("llm_client", None)
-                config.model_name = config_dict.get("model_name", None)
-                if (
-                    vector_store_type in pools
-                    and config.name in pools[vector_store_type]
-                ):
-                    self.client = pools[vector_store_type][config.name]
-                else:
-                    client = self.connector_class(config)
-                    pools[vector_store_type][config.name] = self.client = client
-                    self.client = client
+            if (
+                vector_store_type in pools
+                and vector_store_config.name in pools[vector_store_type]
+            ):
+                self.client = pools[vector_store_type][vector_store_config.name]
+            else:
+                client = self.connector_class(vector_store_config)
+                pools[vector_store_type][
+                    vector_store_config.name
+                ] = self.client = client
         except Exception as e:
             logger.error("connect vector store failed: %s", e)
             raise e

diff --git a/examples/rag/graph_rag_example.py b/examples/rag/graph_rag_example.py
@@ -6,6 +6,7 @@
 from dbgpt.rag import ChunkParameters
 from dbgpt.rag.assembler import EmbeddingAssembler
 from dbgpt.rag.knowledge import KnowledgeFactory
+from dbgpt.storage.knowledge_graph.knowledge_graph import BuiltinKnowledgeGraphConfig
 from dbgpt.storage.vector_store.base import VectorStoreConfig
 from dbgpt.storage.vector_store.connector import VectorStoreConnector
 
@@ -32,7 +33,7 @@ def _create_kg_connector():
     """Create knowledge graph connector."""
     return VectorStoreConnector(
         vector_store_type="KnowledgeGraph",
-        vector_store_config=VectorStoreConfig(
+        vector_store_config=BuiltinKnowledgeGraphConfig(
             name="graph_rag_test",
             embedding_fn=None,
             llm_client=OpenAILLMClient(),