diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..a3e3b98 --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,61 @@ +"""ChromaDB vector store with Ollama embeddings and document chunking."" +import os +from typing import List + +from langchain_chroma import Chroma +from langchain_ollama import OllamaEmbeddings +from langchain_core.documents import Document +from langchain.text_splitter import RecursiveCharacterTextSplitter + +COLLECTION_NAME = "knowledge_base" +PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db") + +embeddings = OllamaEmbeddings(model="nomic-embed-text") + +_store: Chroma | None = None + + +def get_vector_store() -> Chroma: + """Return singleton ChromaDB instance (persisted to disk).""" + global _store + if _store is None: + _store = Chroma( + collection_name=COLLECTION_NAME, + embedding_function=embeddings, + persist_directory=PERSIST_DIR, + ) + return _store + + +def add_documents(texts: List[str], chunk_size: int = 500, chunk_overlap: int = 100) -> int: + """Split texts into chunks and index them in ChromaDB. + + Args: + texts: list of raw text strings + chunk_size: maximum chunk length in characters + chunk_overlap: overlap between consecutive chunks + + Returns: + number of chunks added to the store + """ + splitter = RecursiveCharacterTextSplitter( + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + docs: List[Document] = splitter.create_documents(texts) + if docs: + get_vector_store().add_documents(docs) + return len(docs) + + +def similarity_search(query: str, k: int = 5) -> List[Document]: + """Semantic search in ChromaDB. + + Args: + query: natural language search query + k: number of top results to return + + Returns: + list of the most relevant Document objects + """ + return get_vector_store().similarity_search(query, k=k)