| """ |
| rag/vectorstore.py |
| |
| Handles document ingestion, chunking, embedding, and retrieval |
| using ChromaDB as the vector store. |
| """ |
|
|
| import logging |
| from pathlib import Path |
| from typing import List, Optional |
|
|
| from langchain_core.documents import Document |
| from langchain_text_splitters import RecursiveCharacterTextSplitter |
| from langchain_community.document_loaders import ( |
| DirectoryLoader, |
| PyPDFLoader, |
| TextLoader, |
| WebBaseLoader, |
| ) |
| from langchain_chroma import Chroma |
| from langchain_huggingface import HuggingFaceEmbeddings |
|
|
| import sys |
| sys.path.append(str(Path(__file__).parent.parent)) |
| from config import cfg |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") |
| log = logging.getLogger(__name__) |
|
|
|
|
| def get_embedding_model() -> HuggingFaceEmbeddings: |
| """Load the sentence-transformer embedding model.""" |
| log.info(f"Loading embedding model: {cfg.model.embedding_model_id}") |
| return HuggingFaceEmbeddings( |
| model_name=cfg.model.embedding_model_id, |
| model_kwargs={"device": "cpu"}, |
| encode_kwargs={"normalize_embeddings": True}, |
| ) |
|
|
|
|
| def load_documents(source: str) -> List[Document]: |
| """ |
| Load documents from various sources. |
| |
| Args: |
| source: Path to directory, PDF file, text file, or URL |
| |
| Returns: |
| List of LangChain Document objects |
| """ |
| source_path = Path(source) |
|
|
| if source_path.is_dir(): |
| log.info(f"Loading documents from directory: {source}") |
| loader = DirectoryLoader( |
| str(source_path), |
| glob="**/*.{txt,pdf,md}", |
| loader_cls=TextLoader, |
| show_progress=True, |
| ) |
| elif source_path.suffix == ".pdf": |
| log.info(f"Loading PDF: {source}") |
| loader = PyPDFLoader(str(source_path)) |
| elif source_path.suffix in [".txt", ".md"]: |
| log.info(f"Loading text file: {source}") |
| loader = TextLoader(str(source_path)) |
| elif source.startswith("http"): |
| log.info(f"Loading URL: {source}") |
| loader = WebBaseLoader(source) |
| else: |
| raise ValueError(f"Unsupported source: {source}") |
|
|
| docs = loader.load() |
| log.info(f"Loaded {len(docs)} documents") |
| return docs |
|
|
|
|
| def chunk_documents(docs: List[Document]) -> List[Document]: |
| """Split documents into overlapping chunks for retrieval.""" |
| splitter = RecursiveCharacterTextSplitter( |
| chunk_size=cfg.rag.chunk_size, |
| chunk_overlap=cfg.rag.chunk_overlap, |
| separators=["\n\n", "\n", ". ", " ", ""], |
| length_function=len, |
| ) |
| chunks = splitter.split_documents(docs) |
| log.info(f"Split into {len(chunks)} chunks (size={cfg.rag.chunk_size}, overlap={cfg.rag.chunk_overlap})") |
| return chunks |
|
|
|
|
| def build_vectorstore( |
| sources: Optional[List[str]] = None, |
| docs: Optional[List[Document]] = None, |
| ) -> Chroma: |
| """ |
| Build and persist a ChromaDB vector store from documents. |
| |
| Args: |
| sources: List of file paths or URLs to index |
| docs: Pre-loaded Document objects (alternative to sources) |
| |
| Returns: |
| Initialized Chroma vector store |
| """ |
| cfg.ensure_dirs() |
|
|
| if docs is None: |
| if sources is None: |
| raise ValueError("Provide either 'sources' or 'docs'") |
| all_docs = [] |
| for src in sources: |
| all_docs.extend(load_documents(src)) |
| docs = all_docs |
|
|
| chunks = chunk_documents(docs) |
| embeddings = get_embedding_model() |
|
|
| log.info(f"Embedding {len(chunks)} chunks into ChromaDB...") |
| vectorstore = Chroma.from_documents( |
| documents=chunks, |
| embedding=embeddings, |
| persist_directory=cfg.rag.chroma_persist_dir, |
| collection_name=cfg.rag.collection_name, |
| ) |
| vectorstore.persist() |
| log.info(f"✅ Vector store saved to: {cfg.rag.chroma_persist_dir}") |
| return vectorstore |
|
|
|
|
| def load_vectorstore() -> Chroma: |
| """Load an existing ChromaDB vector store from disk.""" |
| persist_dir = cfg.rag.chroma_persist_dir |
| if not Path(persist_dir).exists(): |
| raise FileNotFoundError( |
| f"No vector store found at {persist_dir}. " |
| "Run build_vectorstore() first." |
| ) |
|
|
| embeddings = get_embedding_model() |
| vectorstore = Chroma( |
| persist_directory=persist_dir, |
| embedding_function=embeddings, |
| collection_name=cfg.rag.collection_name, |
| ) |
| log.info(f"Loaded vector store with {vectorstore._collection.count():,} chunks") |
| return vectorstore |
|
|
|
|
| def retrieve(query: str, vectorstore: Chroma, top_k: Optional[int] = None) -> List[Document]: |
| """ |
| Retrieve the most relevant document chunks for a query. |
| |
| Args: |
| query: User's question |
| vectorstore: Initialized Chroma vector store |
| top_k: Number of chunks to retrieve (defaults to cfg.rag.top_k) |
| |
| Returns: |
| List of relevant Document chunks with similarity scores |
| """ |
| k = top_k or cfg.rag.top_k |
|
|
| results = vectorstore.similarity_search_with_relevance_scores( |
| query=query, |
| k=k, |
| ) |
|
|
| |
| filtered = [ |
| doc for doc, score in results |
| if score >= cfg.rag.similarity_threshold |
| ] |
|
|
| log.debug(f"Retrieved {len(filtered)}/{k} chunks above threshold={cfg.rag.similarity_threshold}") |
| return filtered |
|
|
