186 lines
7.6 KiB
Python
186 lines
7.6 KiB
Python
"""
|
||
Main entry point for the FAQ‑bot assignment.
|
||
|
||
The bot answers questions about course materials stored in a local ChromaDB
|
||
vector store and, when necessary, calls an external MCP‑style tool that
|
||
returns metadata such as the course schedule. The implementation uses
|
||
LangChain v1+ (``langchain>=1.0.0``) together with ``langchain-chroma``
|
||
and ``langchain-ollama`` for embeddings.
|
||
|
||
The file contains:
|
||
* Utility functions to load markdown files into ChromaDB.
|
||
* Two tools – one that searches the vector store and another that mocks an
|
||
MCP call.
|
||
* A simple LangChain agent that decides which tool to use based on the
|
||
content of the user query.
|
||
* Three example queries executed when the script is run as ``__main__``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import json
|
||
import pathlib
|
||
from typing import List, Dict
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Dependencies – all are declared in requirements.txt
|
||
# ---------------------------------------------------------------------------
|
||
from langchain_ollama import ChatOllama
|
||
from langchain_chroma import Chroma
|
||
from langchain_core.messages import HumanMessage
|
||
from langchain.agents import create_agent
|
||
from langchain.tools import tool
|
||
from langchain.embeddings.ollama import OllamaEmbeddings
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration constants
|
||
# ---------------------------------------------------------------------------
|
||
DATA_DIR = pathlib.Path("data") # directory with .md files
|
||
CHROMA_PATH = pathlib.Path("./chroma_faq") # persistent store location
|
||
MCP_ENDPOINT = "http://localhost:8000/course_meta" # mock endpoint
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper – load markdown files into ChromaDB
|
||
# ---------------------------------------------------------------------------
|
||
@tool(name="load_faq_to_chroma", description="Load all .md files from data/ into a persistent Chroma vector store.")
|
||
def load_faq_to_chroma() -> str:
|
||
"""Create or update the Chroma collection with local markdown files.
|
||
|
||
The function is idempotent – it will overwrite existing documents
|
||
with the same IDs. It returns a short JSON string describing how many
|
||
documents were added.
|
||
"""
|
||
# Ensure data directory exists
|
||
if not DATA_DIR.exists():
|
||
return json.dumps({"error": f"{DATA_DIR} does not exist."})
|
||
|
||
# Gather all markdown files
|
||
md_files = list(DATA_DIR.glob("*.md"))
|
||
if not md_files:
|
||
return json.dumps({"error": "No .md files found in data/"})
|
||
|
||
# Prepare documents for Chroma
|
||
docs: List[Dict[str, str]] = []
|
||
for file_path in md_files:
|
||
content = file_path.read_text(encoding="utf-8")
|
||
docs.append({"id": file_path.stem, "content": content})
|
||
|
||
# Create embeddings and store
|
||
embedding = OllamaEmbeddings(model="nomic-embed-text")
|
||
chroma = Chroma.from_documents(
|
||
documents=[doc["content"] for doc in docs],
|
||
ids=[doc["id"] for doc in docs],
|
||
embedding=embedding,
|
||
persist_directory=str(CHROMA_PATH),
|
||
)
|
||
|
||
return json.dumps({"status": "loaded", "documents": len(docs)})
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tool – search the vector store
|
||
# ---------------------------------------------------------------------------
|
||
@tool(name="search_course_docs", description="Search local course documents for a query. Returns top k results.")
|
||
def search_course_docs(query: str, k: int = 3) -> str:
|
||
"""Return the most relevant chunks from the Chroma store.
|
||
|
||
The function loads the persistent collection and performs a similarity
|
||
search. Results are returned as a numbered list of snippets.
|
||
"""
|
||
if not CHROMA_PATH.exists():
|
||
return json.dumps({"error": f"Chroma store {CHROMA_PATH} not found."})
|
||
|
||
embedding = OllamaEmbeddings(model="nomic-embed-text")
|
||
chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embedding)
|
||
docs = chroma.similarity_search(query, k=k)
|
||
if not docs:
|
||
return json.dumps({"result": "No relevant documents found."})
|
||
|
||
snippets = [f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)]
|
||
return json.dumps({"source": "chroma", "snippets": snippets})
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tool – mock MCP call (could be replaced with a real HTTP request)
|
||
# ---------------------------------------------------------------------------
|
||
@tool(name="fetch_course_meta", description="Retrieve course metadata such as schedule. Returns JSON.")
|
||
def fetch_course_meta(query: str) -> str:
|
||
"""Simulate an external MCP tool by returning static data.
|
||
|
||
In a real deployment this would perform an HTTP GET to the MCP server.
|
||
For the purposes of the assignment we simply return a hard‑coded JSON
|
||
payload that mimics what an MCP endpoint might provide.
|
||
"""
|
||
# Static mock data – in practice replace with httpx.get(MCP_ENDPOINT)
|
||
mock_data = {
|
||
"schedule": [
|
||
{"day": "Monday", "time": "10:00-12:00", "topic": "Introduction"},
|
||
{"day": "Wednesday", "time": "14:00-16:00", "topic": "Advanced Topics"},
|
||
],
|
||
"instructor": "Prof. Alexei Petrov",
|
||
}
|
||
return json.dumps({"source": "mcp_meta", "data": mock_data})
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent – decides which tool to use based on the query content
|
||
# ---------------------------------------------------------------------------
|
||
# System prompt guides the agent not to call both tools unnecessarily.
|
||
SYSTEM_PROMPT = (
|
||
"You are an assistant that answers questions about course materials."
|
||
" If the question is about lecture notes or FAQs, use the search_course_docs tool."
|
||
" If the question asks for schedule, instructor info, or other metadata, use fetch_course_meta."
|
||
)
|
||
|
||
# Create the agent once – it will be reused in examples.
|
||
llm = ChatOllama(model="llama3", temperature=0.0)
|
||
agent = create_agent(
|
||
llm=llm,
|
||
tools=[search_course_docs, fetch_course_meta],
|
||
system_prompt=SYSTEM_PROMPT,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Example usage – three predefined queries and an interactive prompt
|
||
# ---------------------------------------------------------------------------
|
||
EXAMPLES: List[str] = [
|
||
"What topics are covered in the first lecture?",
|
||
"When is the next class scheduled?",
|
||
"Can you give me a summary of the course material?",
|
||
]
|
||
|
||
def run_example(query: str) -> None:
|
||
print(f"\nUser: {query}")
|
||
result = agent.invoke(
|
||
{"messages": [HumanMessage(content=query)]},
|
||
{"configurable": {"thread_id": "example-session"}},
|
||
)
|
||
# The agent returns a list of messages – the last one is the answer.
|
||
answer = result["messages"][-1].content
|
||
print(f"Assistant: {answer}")
|
||
|
||
if __name__ == "__main__":
|
||
# Ensure data directory exists for demo purposes
|
||
if not DATA_DIR.exists():
|
||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||
# Create a tiny FAQ file as an example.
|
||
(DATA_DIR / "faq.md").write_text(
|
||
"# Course FAQ\n" \
|
||
"## What is the course about?\n" \
|
||
"The course covers advanced topics in AI and machine learning."
|
||
)
|
||
# Load data into Chroma (idempotent)
|
||
print(load_faq_to_chroma())
|
||
|
||
for q in EXAMPLES:
|
||
run_example(q)
|
||
|
||
# Interactive mode – optional
|
||
while True:
|
||
try:
|
||
user_input = input("\nAsk a question (or 'quit'): ")
|
||
except EOFError:
|
||
break
|
||
if user_input.lower() in {"quit", "exit"}:
|
||
break
|
||
run_example(user_input)
|
||
"""
|