实现多工程数据存储支持
This commit is contained in:
@@ -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}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user