diff --git a/src/utils.py b/src/utils.py new file mode 100644 index 0000000..9e3afb9 --- /dev/null +++ b/src/utils.py @@ -0,0 +1,42 @@ +import os +from pathlib import Path +from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma +from langchain_text_splitters import RecursiveCharacterTextSplitter + +CHROMA_PATH = Path("./chroma_faq") + + +def load_faq_to_chroma(md_dir: str = "data"): + """Load all .md files from md_dir into a persistent Chroma store. + The function will create or update the store at CHROMA_PATH. + """ + # Ensure directory exists + Path(md_dir).mkdir(parents=True, exist_ok=True) + # Gather all markdown files + md_files = list(Path(md_dir).glob("*.md")) + if not md_files: + raise FileNotFoundError(f"No .md files found in {md_dir}") + + # Read and split documents + texts = [] + for md_file in md_files: + text = md_file.read_text(encoding="utf-8") + splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + texts.extend(splitter.split_text(text)) + + # Create embeddings + embeddings = OllamaEmbeddings(model="nomic-embed-text") + + # Persist to Chroma + chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embeddings) + chroma.add_texts(texts) + chroma.persist() + return chroma + + +def search_course_docs(query: str, k: int = 3): + """Search the persistent Chroma store for the top k documents matching query.""" + chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=OllamaEmbeddings(model="nomic-embed-text")) + retriever = chroma.as_retriever(search_kwargs={"k": k}) + return retriever.get_relevant_documents(query)