diff --git a/vector_store.py b/vector_store.py index 912c70a..adbb43f 100644 --- a/vector_store.py +++ b/vector_store.py @@ -1,54 +1,64 @@ """ -Vector store abstraction using ChromaDB. +Vector store implementation using ChromaDB. -Provides methods to add documents with embeddings and perform similarity search. +Provides methods to add documents and perform similarity search. """ import os from pathlib import Path -from typing import List, Dict, Any +from typing import List, Dict -from chromadb import Client as ChromaClient +from chromadb import Client as ChromadbClient from chromadb.config import Settings +from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_ollama import OllamaEmbeddings -# Initialize global client (in-memory for simplicity) -client = ChromaClient(Settings(chroma_db_impl="duckdb+parquet", persist_directory=None)) -collection_name = "rag_collection" +# Ensure persistent directory exists +CHROMA_DIR = Path("./chroma_db") +CHROMA_DIR.mkdir(parents=True, exist_ok=True) -# Ensure collection exists +# Embedding model via Ollama +embeddings = OllamaEmbeddings(model="nomic-embed-text") + +# Chroma client with persistence +client = ChromadbClient(Settings(persist_directory=str(CHROMA_DIR))) +collection_name = "knowledge" if collection_name not in client.list_collections(): client.create_collection(name=collection_name) col = client.get_or_create_collection(name=collection_name) -embeddings = OllamaEmbeddings(model="nomic-embed-text") +# Chunker from chunker.py +from chunker import chunker -class ChromaVectorStore: - """Wrapper around a Chroma collection.""" +class VectorStore: + def add_documents(self, docs: List[str], metadatas: List[Dict]): + """Add documents to the collection. - def __init__(self, collection): - self.collection = collection + Parameters + ---------- + docs: list of raw text strings. + metadatas: list of metadata dicts corresponding to each doc. + """ + # Split into chunks and embed + all_chunks = [] + all_metadatas = [] + for doc, meta in zip(docs, metadatas): + chunks = chunker.split_text(doc) + all_chunks.extend(chunks) + all_metadatas.extend([meta] * len(chunks)) - def add_documents(self, documents: List[Dict[str, Any]]): - ids = [] - texts = [] - metadatas = [] - for doc in documents: - ids.append(doc.get("id", os.urandom(8).hex())) - texts.append(doc["content"]) - metadatas.append(doc.get("metadata", {})) - embeddings_list = embeddings.embed_documents(texts) - self.collection.add(ids=ids, documents=texts, embeddings=embeddings_list, metadatas=metadatas) + # Generate embeddings via Ollama + embeds = embeddings.embed_documents(all_chunks) + ids = [f"chunk_{i}" for i in range(len(all_chunks))] + col.add(ids=ids, documents=all_chunks, metadatas=all_metadatas, embeddings=embeds) - def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]: - results = self.collection.query( + def similarity_search(self, query: str, k: int = 5) -> List[Dict]: + """Return top-k similar chunks with metadata.""" + results = col.query( query_texts=[query], n_results=k, - include=['documents', 'distances', 'metadatas'], + include=['documents', 'metadatas'], ) - docs = [] - for doc, dist, meta in zip(results["documents"][0], results["distances"][0], results["metadatas"][0]): - docs.append({"content": doc, "distance": dist, "metadata": meta}) - return docs - -# Singleton instance -vector_store = ChromaVectorStore(col) + # results is dict with keys documents, metadatas + docs = results["documents"][0] + metas = results["metadatas"][0] + return [{"content": d, "metadata": m} for d, m in zip(docs, metas)]