Store connector (#89)

1.knowledge_init.py config update
2.DB_GPT_wechat2

Author: yhjun1026

assets/DB_GPT_wechat.png

@@ -38,14 +38,6 @@ def __init__(self) -> None:
         self.use_mac_os_tts = False
         self.use_mac_os_tts = os.getenv("USE_MAC_OS_TTS")
 
-        # milvus or zilliz cloud configuration
-        self.milvus_addr = os.getenv("MILVUS_ADDR", "localhost:19530")
-        self.milvus_username = os.getenv("MILVUS_USERNAME")
-        self.milvus_password = os.getenv("MILVUS_PASSWORD")
-        self.milvus_collection = os.getenv("MILVUS_COLLECTION", "dbgpt")
-        self.milvus_secure = os.getenv("MILVUS_SECURE") == "True"
-
-
         self.authorise_key = os.getenv("AUTHORISE_COMMAND_KEY", "y")
         self.exit_key = os.getenv("EXIT_KEY", "n")
         self.image_provider = os.getenv("IMAGE_PROVIDER", True)

pilot/configs/config.py

@@ -139,29 +139,21 @@ def init_schema_and_load(self, vector_name, documents):
         fields.append(
             FieldSchema(text_field, DataType.VARCHAR, max_length=max_length + 1)
         )
-        # create the primary key field
+        # primary key field
         fields.append(
             FieldSchema(primary_field, DataType.INT64, is_primary=True, auto_id=True)
         )
-        # create the vector field
+        # vector field
         fields.append(FieldSchema(vector_field, DataType.FLOAT_VECTOR, dim=dim))
-        # Create the schema for the collection
+        # milvus the schema for the collection
         schema = CollectionSchema(fields)
         # Create the collection
         collection = Collection(collection_name, schema)
         self.col = collection
-        # Index parameters for the collection
+        # index parameters for the collection
         index = self.index_params
-        # Create the index
+        # milvus index
         collection.create_index(vector_field, index)
-        # Create the VectorStore
-        # milvus = cls(
-        #     embedding,
-        #     kwargs.get("connection_args", {"port": 19530}),
-        #     collection_name,
-        #     text_field,
-        # )
-        # Add the texts.
         schema = collection.schema
         for x in schema.fields:
             self.fields.append(x.name)

pilot/vector_store/milvus_store.py

@@ -69,6 +69,7 @@ colorama
 playsound
 distro
 pypdf
+milvus-cli==0.3.2
 
 # Testing dependencies
 pytest

requirements.txt

@@ -1,12 +1,17 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 import argparse
+import os
+import sys
 
-from pilot.configs.model_config import DATASETS_DIR, LLM_MODEL_CONFIG, VECTOR_SEARCH_TOP_K, VECTOR_STORE_CONFIG, \
-    VECTOR_STORE_TYPE
-from pilot.source_embedding.knowledge_embedding import KnowledgeEmbedding
 
+sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
+
+from pilot.configs.config import Config
+from pilot.configs.model_config import DATASETS_DIR, LLM_MODEL_CONFIG, VECTOR_SEARCH_TOP_K
+from pilot.source_embedding.knowledge_embedding import KnowledgeEmbedding
 
+CFG = Config()
 class LocalKnowledgeInit:
     embeddings: object = None
     model_name = LLM_MODEL_CONFIG["text2vec"]
@@ -32,6 +37,7 @@ def query(self, q):
             dc, s = doc
             yield s, dc
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--vector_name", type=str, default="default")
@@ -40,8 +46,8 @@ def query(self, q):
     args = parser.parse_args()
     vector_name = args.vector_name
     append_mode = args.append
-    store_type = VECTOR_STORE_TYPE
-    vector_store_config = {"url": VECTOR_STORE_CONFIG["url"], "port": VECTOR_STORE_CONFIG["port"], "vector_store_name":vector_name}
+    store_type = CFG.VECTOR_STORE_TYPE
+    vector_store_config = {"vector_store_name": vector_name}
     print(vector_store_config)
     kv  = LocalKnowledgeInit(vector_store_config=vector_store_config)
     vector_store = kv.knowledge_persist(file_path=DATASETS_DIR, append_mode=append_mode)