Обновлен vectorstore.py: update vectorstore.py

vectorstore.py

@@ -1,30 +1,25 @@
 from pathlib import Path
 from uuid import uuid4
 
-from langchain_qdrant import Qdrant
-from langchain_core.documents import Document
+import chromadb
+from chromadb.config import Settings
 from langchain_ollama import OllamaEmbeddings
+from langchain_core.documents import Document
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
-QDRANT_DIR = "./qdrant_db"
+CHROMA_DIR = "./chroma_db"
 COLLECTION_NAME = "local_kb"
 EMBED_MODEL = "nomic-embed-text"
 OLLAMA_BASE_URL = "http://127.0.0.1:11434"
 
 
-def create_vectorstore(persist_directory: str = QDRANT_DIR) -> Qdrant:
-    embeddings = OllamaEmbeddings(
-        model=EMBED_MODEL,
-        base_url=OLLAMA_BASE_URL,
-    )
-    return Qdrant(
-        collection_name=COLLECTION_NAME,
-        embedding_function=embeddings,
-        persist_directory=persist_directory,
-    )
+def create_vectorstore(persist_directory: str = CHROMA_DIR):
+    client = chromadb.Client(Settings(persist_directory=persist_directory))
+    collection = client.get_or_create_collection(name=COLLECTION_NAME)
+    return collection
 
 
-def load_documents(directory: str, vectorstore: Qdrant) -> int:
+def load_documents(directory: str, vectorstore: chromadb.Collection) -> int:
     splitter = RecursiveCharacterTextSplitter(
         chunk_size=1000,
         chunk_overlap=200,
@@ -51,5 +46,7 @@ def load_documents(directory: str, vectorstore: Qdrant) -> int:
 
     chunks = splitter.split_documents(raw_docs)
     ids = [str(uuid4()) for _ in chunks]
-    vectorstore.add_documents(documents=chunks, ids=ids)
+    embeddings = OllamaEmbeddings(model=EMBED_MODEL, base_url=OLLAMA_BASE_URL)
+    vectors = embeddings.embed_documents([doc.page_content for doc in chunks])
+    vectorstore.add(ids=ids, documents=[doc.page_content for doc in chunks], embeddings=vectors, metadatas=[doc.metadata for doc in chunks])
     return len(chunks)