diff --git a/src/utils.py b/src/utils.py index 0f0559a..2237382 100644 --- a/src/utils.py +++ b/src/utils.py @@ -9,6 +9,9 @@ from langchain_ollama import OllamaEmbeddings # Chroma collection name COLLECTION_NAME = "faq" +# Persistence directory from config +from src.config import CHROMA_PERSIST_DIR + def load_faq_to_chroma(md_dir: str = "data") -> None: """Load all .md files from md_dir into a Chroma vector store. @@ -22,21 +25,23 @@ def load_faq_to_chroma(md_dir: str = "data") -> None: all_docs.extend(docs) splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) docs = splitter.split_documents(all_docs) - embeddings = OllamaEmbeddings(model="nomic-embed-text") - # Create or update Chroma collection + embeddings = OllamaEmbeddings("nomic-embed-text") + # Create or update Chroma collection with persistence Chroma.from_documents( docs, embeddings, collection_name=COLLECTION_NAME, + persist_directory=CHROMA_PERSIST_DIR, ) def search_course_docs(query: str, k: int = 3) -> List[str]: """Return top k document snippets from the persisted Chroma store.""" - embeddings = OllamaEmbeddings(model="nomic-embed-text") + embeddings = OllamaEmbeddings("nomic-embed-text") store = Chroma( collection_name=COLLECTION_NAME, embedding_function=embeddings, + persist_directory=CHROMA_PERSIST_DIR, ) results = store.similarity_search(query, k=k) return [doc.page_content for doc in results]