实现多工程数据存储支持

This commit is contained in:
wanyaokun
2024-08-13 13:10:52 +08:00
parent 3a81a83033
commit 7e58a1a223
9 changed files with 97 additions and 51 deletions
+12 -8
View File
@@ -5,12 +5,14 @@ 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
#collection_name = os.getenv("VECTOR_STORE_COLLECTION", "default")
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 +34,11 @@ 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")
def get_chroma_vector_store(docType:str):
#collection_name = os.getenv("VECTOR_STORE_COLLECTION", "default")
collection_name = docType
vector_store_path = os.getenv("VECTOR_STORE_PATH")
vector_store_path =os.path.join(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 +59,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}")