diff --git a/src/utils.py b/src/utils.py index 9e3afb9..c1f0cb2 100644 --- a/src/utils.py +++ b/src/utils.py @@ -6,37 +6,23 @@ 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)