diff --git a/vector_store.py b/vector_store.py new file mode 100644 index 0000000..e1a466c --- /dev/null +++ b/vector_store.py @@ -0,0 +1,63 @@ +"""Vector store utilities for the FAQ bot. + +This module handles loading markdown files into a persistent Chroma vector store +and provides a tool that performs semantic search over the stored documents. +""" + +import os +from pathlib import Path +from typing import List + +from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma +from langchain_core.documents import Document + +from config import DATA_DIR, CHROMA_PERSIST_DIR, CHROMA_TOP_K +from chunker import chunk_text + +# Initialize embeddings once +embeddings = OllamaEmbeddings(model="nomic-embed-text") + +# Create or load the Chroma collection +vector_store = Chroma( + collection_name="faq_collection", + embedding_function=embeddings, + persist_directory=CHROMA_PERSIST_DIR, +) + +def load_faq_to_chroma() -> None: + """Load all markdown files from DATA_DIR into the Chroma vector store. + + The function reads each .md file, splits it into chunks, creates Document + objects, and adds them to the vector store. The store is persisted to disk. + """ + docs: List[Document] = [] + for md_file in Path(DATA_DIR).glob("*.md"): + text = md_file.read_text(encoding="utf-8") + chunks = chunk_text(text) + for i, chunk in enumerate(chunks): + doc = Document(page_content=chunk, metadata={"source": md_file.name, "chunk_id": i}) + docs.append(doc) + if docs: + vector_store.add_documents(docs) + vector_store.persist() + +def search_course_docs(query: str, k: int = CHROMA_TOP_K) -> str: + """Perform a semantic search over the FAQ collection. + + Parameters + ---------- + query: str + The user question. + k: int, optional + Number of top results to return. + + Returns + ------- + str + Concatenated top-k documents' page content. + """ + results = vector_store.similarity_search_with_score(query, k=k) + # results is list of (Document, score) + snippets = [doc.page_content for doc, _ in results] + return "\n\n---\n\n".join(snippets) \ No newline at end of file