Files
brojs-task-6a1d75c5fd30e81c…/src/utils.py
T

48 lines
2.0 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import os
from pathlib import Path
import json
import httpx
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
CHROMA_PATH = Path("./chroma_faq")
# 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))
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)
# MCPstyle 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