diff --git a/vectorstore.py b/vectorstore.py index ff41888..2d919e0 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -2,7 +2,6 @@ from pathlib import Path from uuid import uuid4 from langchain_chroma import Chroma, Settings - from langchain_ollama import OllamaEmbeddings from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter @@ -14,12 +13,17 @@ OLLAMA_BASE_URL = "http://127.0.0.1:11434" def create_vectorstore(persist_directory: str = CHROMA_DIR): + """Return a Chroma collection configured for the local knowledge base. + + The collection is persisted in ``persist_directory`` and named + ``COLLECTION_NAME``. + """ client = Chroma(Settings(persist_directory=persist_directory)) collection = client.get_or_create_collection(name=COLLECTION_NAME) return collection -def load_documents(directory: str, vectorstore: Chroma) -> int: +def load_documents(directory: str, collection) -> int: splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200, @@ -48,5 +52,10 @@ def load_documents(directory: str, vectorstore: Chroma) -> int: ids = [str(uuid4()) for _ in chunks] 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]) + collection.add( + ids=ids, + documents=[doc.page_content for doc in chunks], + embeddings=vectors, + metadatas=[doc.metadata for doc in chunks], + ) return len(chunks)