diff --git a/vector_store.py b/vector_store.py deleted file mode 100644 index adbb43f..0000000 --- a/vector_store.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Vector store implementation using ChromaDB. - -Provides methods to add documents and perform similarity search. -""" -import os -from pathlib import Path -from typing import List, Dict - -from chromadb import Client as ChromadbClient -from chromadb.config import Settings -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_ollama import OllamaEmbeddings - -# Ensure persistent directory exists -CHROMA_DIR = Path("./chroma_db") -CHROMA_DIR.mkdir(parents=True, exist_ok=True) - -# 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) - -# Chunker from chunker.py -from chunker import chunker - -class VectorStore: - def add_documents(self, docs: List[str], metadatas: List[Dict]): - """Add documents to the 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)) - - # 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]: - """Return top-k similar chunks with metadata.""" - results = col.query( - query_texts=[query], - n_results=k, - include=['documents', 'metadatas'], - ) - # 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)]