diff --git a/src/utils.py b/src/utils.py index 262504e..560c085 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,47 +1,35 @@ import os -from pathlib import Path -import json -import httpx -from langchain_ollama import OllamaEmbeddings -from langchain_chroma import Chroma +import pathlib +from typing import List + +from langchain_community.document_loaders import TextLoader from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_chroma import Chroma +from langchain_ollama import OllamaEmbeddings -CHROMA_PATH = Path("./chroma_faq") +CHROMA_PATH = pathlib.Path("./chroma_faq") +CHROMA_PATH.mkdir(parents=True, exist_ok=True) -# Load FAQ markdown files into ChromaDB -def load_faq_to_chroma(md_dir: str = "data"): - Path(md_dir).mkdir(parents=True, exist_ok=True) - md_files = list(Path(md_dir).glob("*.md")) - if not md_files: - raise FileNotFoundError(f"No .md files found in {md_dir}") - 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)) + +def load_faq_to_chroma(md_dir: str = "data") -> None: + """Load all .md files from md_dir into a Chroma vector store. + The store is persisted at CHROMA_PATH. + """ + loader = TextLoader + all_docs = [] + for md_file in pathlib.Path(md_dir).glob("*.md"): + loader_obj = loader(str(md_file)) + docs = loader_obj.load() + all_docs.extend(docs) + splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + docs = splitter.split_documents(all_docs) + embeddings = OllamaEmbeddings(model="nomic-embed-text") + Chroma.from_documents(docs, embeddings, persist_directory=str(CHROMA_PATH)) + + +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") chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embeddings) - chroma.add_texts(texts) - chroma.persist() - return chroma - -# Search the ChromaDB for relevant documents -def search_course_docs(query: str, k: int = 3): - 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) - -# MCP‑style tool: fetch metadata from a mock HTTP endpoint -# In production this would be a real MCP server. Here we use a local JSON file served by http.server. -def fetch_course_meta(query: str): - url = f"http://localhost:8000/course_meta.json" - try: - response = httpx.get(url, timeout=5.0) - response.raise_for_status() - data = response.json() - except Exception: - # Fallback to local static file if server not running - data = json.loads(Path("data/course_meta.json").read_text(encoding="utf-8")) - # Simple filtering: return items where query is in title or description - results = [item for item in data if query.lower() in item.get("title", "").lower() or query.lower() in item.get("description", "").lower()] - return results + results = chroma.similarity_search(query, k=k) + return [doc.page_content for doc in results]