初始化提交

2024-08-13 09:37:23 +08:00
parent 4923337038
commit e112fa4e44
50 changed files with 1649 additions and 259 deletions
@@ -0,0 +1,71 @@
+import os
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from llama_index.vector_stores.qdrant import QdrantVectorStore
+from qdrant_client import qdrant_client
+
+qclient = None
+
+def get_qdrant_vector_store():
+    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")),
+
+    if not vector_store_path or not host:
+        raise ValueError(
+            "Please provide either VECTOR_STORE_PATH or VECTOR_STORE_HOST and VECTOR_STORE_PORT"
+        )
+    # if VECTOR_STORE_PATH is set, use a local QdrantVectorStore from the path
+    # otherwise, use a remote QdrantVectorStore
+    global qclient
+    if qclient == None:
+        if vector_store_path:
+            qclient = qdrant_client.QdrantClient(
+                path=vector_store_path,
+            )
+        else:
+            qclient = qdrant_client.QdrantClient(
+                host=host,
+                port=port,
+            )
+
+    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")
+    # 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:
+        store = ChromaVectorStore.from_params(
+            persist_dir=vector_store_path, collection_name=collection_name,
+            collection_kwargs={"metadata":{"hnsw:space":"cosine"}},
+        )
+    else:
+        if not os.getenv("VECTOR_STORE_HOST") or not os.getenv("VECTOR_STORE_PORT"):
+            raise ValueError(
+                "Please provide either VECTOR_STORE_PATH or VECTOR_STORE_HOST and VECTOR_STORE_PORT"
+            )
+        store = ChromaVectorStore.from_params(
+            host=os.getenv("VECTOR_STORE_HOST"),
+            port=int(os.getenv("VECTOR_STORE_PORT")),
+            collection_name=collection_name,
+            collection_kwargs={"metadata":{"hnsw:space":"cosine"}},
+        )
+    return store
+
+def get_vector_store():
+    store_type=os.getenv("VECTOR_STORE_TYPE")
+
+    store = None
+
+    match store_type:
+        case "chroma":
+            store = get_chroma_vector_store()
+        case "qdrant":
+            store = get_qdrant_vector_store()
+        case _:
+            raise ValueError(f"Invalid vector store type: {store_type}")
+
+    return store