合并代码

backend/app/engine/vectordb.py

@@ -5,12 +5,13 @@ from qdrant_client import qdrant_client
 
 qclient = None
 
-def get_qdrant_vector_store():
-    collection_name = os.getenv("VECTOR_STORE_COLLECTION", "default")
+def get_qdrant_vector_store(docType:str):
+    collection_name = docType
     vector_store_path = os.getenv("VECTOR_STORE_PATH")
     host=os.getenv("VECTOR_STORE_HOST", "127.0.0.1"),
     port=int(os.getenv("VECTOR_STORE_PORT", "6333")),
 
+    vector_store_path =os.path.join(vector_store_path,docType)
     if not vector_store_path or not host:
         raise ValueError(
             "Please provide either VECTOR_STORE_PATH or VECTOR_STORE_HOST and VECTOR_STORE_PORT"
@@ -32,9 +33,9 @@ def get_qdrant_vector_store():
     vector_store = QdrantVectorStore(client=qclient, collection_name=collection_name)
     return vector_store
 
-def get_chroma_vector_store():
-    collection_name = os.getenv("VECTOR_STORE_COLLECTION", "default")
-    vector_store_path = os.getenv("VECTOR_STORE_PATH")
+def get_chroma_vector_store(docType:str):
+    collection_name = docType
+    vector_store_path =os.path.join(os.getenv("VECTOR_STORE_PATH"),docType)
     # if VECTOR_STORE_PATH is set, use a local ChromaVectorStore from the path
     # otherwise, use a remote ChromaVectorStore (ChromaDB Cloud is not supported yet)
     if vector_store_path:
@@ -55,16 +56,16 @@ def get_chroma_vector_store():
         )
     return store
 
-def get_vector_store():
+def get_vector_store(docType:str):
     store_type=os.getenv("VECTOR_STORE_TYPE")
 
     store = None
 
     match store_type:
         case "chroma":
-            store = get_chroma_vector_store()
+            store = get_chroma_vector_store(docType)
         case "qdrant":
-            store = get_qdrant_vector_store()
+            store = get_qdrant_vector_store(docType)
         case _:
             raise ValueError(f"Invalid vector store type: {store_type}")