85 lines
2.7 KiB
Python
85 lines
2.7 KiB
Python
import json
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
|
|
import httpx
|
|
from langchain_community.document_loaders import TextLoader
|
|
from langchain_community.embeddings import OllamaEmbeddings
|
|
from langchain_community.vectorstores import Chroma
|
|
from langchain_core.documents import Document
|
|
from langchain_core.tools import tool
|
|
|
|
# Path to the data directory
|
|
DATA_DIR = Path(__file__).parent.parent / "data"
|
|
CHROMA_DIR = Path(__file__).parent.parent / "chroma_faq"
|
|
|
|
def load_faq_to_chroma() -> Chroma:
|
|
"""
|
|
Load all .md files from the data directory, chunk them, embed with Ollama,
|
|
and persist into a Chroma vector store.
|
|
"""
|
|
# Check if the Chroma collection already exists
|
|
if CHROMA_DIR.exists():
|
|
# Load existing collection
|
|
return Chroma(persist_directory=str(CHROMA_DIR), embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
|
|
|
|
# Gather all markdown files
|
|
md_files = list(DATA_DIR.glob("*.md"))
|
|
documents: List[Document] = []
|
|
|
|
for md_file in md_files:
|
|
loader = TextLoader(str(md_file), encoding="utf-8")
|
|
docs = loader.load()
|
|
documents.extend(docs)
|
|
|
|
# Create embeddings
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
|
|
# Create Chroma vector store
|
|
chroma = Chroma.from_documents(
|
|
documents=documents,
|
|
embedding=embeddings,
|
|
persist_directory=str(CHROMA_DIR),
|
|
)
|
|
return chroma
|
|
|
|
@tool
|
|
def search_course_docs(query: str, k: int = 3) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search the local FAQ Chroma vector store for relevant documents.
|
|
|
|
Returns a list of dictionaries containing the content and metadata.
|
|
"""
|
|
chroma = load_faq_to_chroma()
|
|
results = chroma.similarity_search(query, k=k)
|
|
output = []
|
|
for doc in results:
|
|
output.append(
|
|
{
|
|
"content": doc.page_content,
|
|
"metadata": doc.metadata,
|
|
}
|
|
)
|
|
return output
|
|
|
|
@tool
|
|
def fetch_course_meta(query: str) -> Dict[str, Any]:
|
|
"""
|
|
Simulate an MCP-style HTTP tool that returns course metadata
|
|
matching the query. The metadata is read from a local JSON file.
|
|
"""
|
|
meta_path = DATA_DIR / "course_meta.json"
|
|
with open(meta_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
# Simple keyword matching in schedule and instructor fields
|
|
results = {}
|
|
if "schedule" in query.lower():
|
|
results["schedule"] = data.get("schedule", [])
|
|
if "instructor" in query.lower() or "professor" in query.lower():
|
|
results["instructor"] = data.get("instructor", {})
|
|
if not results:
|
|
# Default to returning the whole metadata if no keyword matched
|
|
results = data
|
|
return results |