98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
"""FAQ Bot with ChromaDB and a mock MCP tool.
|
|
|
|
The agent answers questions about course materials using a local Chroma vector store.
|
|
If the question is about course metadata (e.g., schedule), it calls a simple HTTP mock that returns JSON. The agent is built with LangGraph.
|
|
|
|
Run with:
|
|
python main.py
|
|
|
|
The script will load the FAQ files, build the vector store, and then enter an interactive loop.
|
|
"""
|
|
|
|
import json
|
|
import pathlib
|
|
from pathlib import Path
|
|
|
|
from langchain_ollama import ChatOllama
|
|
from langchain_chroma import Chroma
|
|
from langchain_text_splitters import MarkdownHeaderTextSplitter
|
|
from langchain_community.document_loaders import TextLoader
|
|
from langchain_community.embeddings import OllamaEmbeddings
|
|
from langgraph.prebuilt import create_react_agent
|
|
from langgraph.graph import StateGraph, MessagesState
|
|
|
|
# ---------- 1. Load FAQ files and create Chroma store ----------
|
|
|
|
def load_faq_to_chroma(data_dir: str = "data", persist_dir: str = "./chroma_faq") -> Chroma:
|
|
"""Read all .md files in data_dir, chunk them, and persist to Chroma."""
|
|
Path(persist_dir).mkdir(parents=True, exist_ok=True)
|
|
docs = []
|
|
for md_file in Path(data_dir).glob("*.md"):
|
|
loader = TextLoader(str(md_file), encoding="utf-8")
|
|
docs.extend(loader.load())
|
|
splitter = MarkdownHeaderTextSplitter(headers_to_split_on=["#", "##", "###"])
|
|
split_docs = splitter.split_documents(docs)
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
chroma = Chroma.from_documents(
|
|
documents=split_docs,
|
|
embedding=embeddings,
|
|
persist_directory=persist_dir,
|
|
)
|
|
return chroma
|
|
|
|
# ---------- 2. MCP-style tool: fetch course metadata ----------
|
|
|
|
def fetch_course_meta(query: str) -> dict:
|
|
"""Mock MCP tool that returns course metadata.
|
|
In production this would be an HTTP call to an MCP server.
|
|
Here we simulate with a local JSON file.
|
|
"""
|
|
meta_path = Path("mock_meta.json")
|
|
if not meta_path.exists():
|
|
meta = {"schedule": {"Monday": "10:00-12:00", "Wednesday": "14:00-16:00"}, "instructor": "Prof. Smith"}
|
|
meta_path.write_text(json.dumps(meta, indent=2))
|
|
else:
|
|
meta = json.loads(meta_path.read_text())
|
|
# Return the whole meta; the agent can filter as needed
|
|
return meta
|
|
|
|
# ---------- 3. Agent setup ----------
|
|
|
|
def build_agent(chroma: Chroma):
|
|
"""Create a LangGraph agent that routes to either Chroma or the MCP tool."""
|
|
def meta_tool(query: str):
|
|
return json.dumps(fetch_course_meta(query), indent=2)
|
|
tools = {"fetch_course_meta": meta_tool}
|
|
agent = create_react_agent(ChatOllama(model="nomic-embed-text"), tools)
|
|
graph = StateGraph(MessagesState)
|
|
graph.add_node("agent", agent)
|
|
graph.set_entry_point("agent")
|
|
return graph.compile()
|
|
|
|
# ---------- 4. Interactive CLI ----------
|
|
|
|
def main():
|
|
chroma = load_faq_to_chroma()
|
|
agent = build_agent(chroma)
|
|
sample_questions = [
|
|
"What is the deadline for the assignment?",
|
|
"How many chapters are in the course?",
|
|
"What is the schedule for the next lecture?",
|
|
]
|
|
for q in sample_questions:
|
|
print("\nQuestion:", q)
|
|
response = agent.invoke({"messages": [{"role": "user", "content": q}]})
|
|
print("Answer:", response["messages"][0]["content"])
|
|
while True:
|
|
try:
|
|
user_q = input("\nAsk a question (or 'exit'): ")
|
|
except EOFError:
|
|
break
|
|
if user_q.lower() in {"exit", "quit"}:
|
|
break
|
|
response = agent.invoke({"messages": [{"role": "user", "content": user_q}]})
|
|
print("Answer:", response["messages"][0]["content"])
|
|
|
|
if __name__ == "__main__":
|
|
main()
|