129 lines
4.7 KiB
Python
129 lines
4.7 KiB
Python
import os
|
||
import asyncio
|
||
import json
|
||
import httpx
|
||
from pathlib import Path
|
||
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
||
from langchain_chroma import Chroma
|
||
from langchain_core.documents import Document
|
||
from langchain.tools import tool
|
||
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||
from langchain.agents import Tool
|
||
from langchain_community.utilities import RetrievalQA
|
||
from langchain_community.vectorstores import Chroma as ChromaStore
|
||
|
||
# ---------- Configuration ----------
|
||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||
if not OPENAI_API_KEY:
|
||
raise RuntimeError("OPENAI_API_KEY not set in environment")
|
||
|
||
# ---------- Embeddings & Vector Store ----------
|
||
embeddings = OpenAIEmbeddings(
|
||
model="text-embedding-3-small",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=OPENAI_API_KEY,
|
||
)
|
||
|
||
vector_store = ChromaStore(
|
||
collection_name="faq_collection",
|
||
embedding_function=embeddings,
|
||
persist_directory="./chroma_faq",
|
||
)
|
||
|
||
# ---------- Load FAQ into Chroma ----------
|
||
|
||
def load_faq_to_chroma(data_dir: str = "data"):
|
||
"""Load all .md files from data_dir into the Chroma vector store.
|
||
The function clears the existing collection before loading.
|
||
"""
|
||
vector_store.delete_collection()
|
||
docs = []
|
||
for md_file in Path(data_dir).glob("*.md"):
|
||
text = md_file.read_text(encoding="utf-8")
|
||
docs.append(Document(page_content=text, metadata={"source": md_file.name}))
|
||
vector_store.add_documents(docs)
|
||
vector_store.persist()
|
||
|
||
# ---------- Tools ----------
|
||
@tool
|
||
def search_course_docs(query: str, k: int = 3) -> str:
|
||
"""Search the local FAQ collection for relevant passages."""
|
||
docs = vector_store.similarity_search(query, k=k)
|
||
if not docs:
|
||
return "No relevant information found in the course materials."
|
||
return "\n\n---\n\n".join([f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in docs])
|
||
|
||
@tool
|
||
def fetch_course_meta(query: str) -> str:
|
||
"""MCP‑style tool that queries a local mock server for course metadata.
|
||
The mock server should serve a JSON file at http://localhost:8000/meta.json.
|
||
"""
|
||
url = "http://localhost:8000/meta.json"
|
||
try:
|
||
response = httpx.get(url, timeout=5.0)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
except Exception as e:
|
||
return f"Error fetching metadata: {e}"
|
||
# Simple lookup: return value if query matches a key (case‑insensitive)
|
||
key = query.strip().lower()
|
||
value = data.get(key)
|
||
if value is None:
|
||
return f"No metadata entry found for '{query}'."
|
||
return f"{key}: {value}"
|
||
|
||
# ---------- Agent Setup ----------
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=OPENAI_API_KEY,
|
||
temperature=0.0,
|
||
)
|
||
|
||
# Define the system prompt with routing rule
|
||
system_prompt = (
|
||
"You are a helpful assistant that answers questions about the course. "
|
||
"If the question is about course content, use the search_course_docs tool. "
|
||
"If the question is about schedule, metadata, or other non‑content info, "
|
||
"use the fetch_course_meta tool. Do not call both tools unless the question "
|
||
"explicitly requires both. In your final answer, prepend 'source: chroma' "
|
||
"or 'source: mcp_meta' to indicate which tool provided the information."
|
||
)
|
||
|
||
# Create Tool objects
|
||
search_tool = Tool(name="search_course_docs", func=search_course_docs, description="Search local course documents.")
|
||
meta_tool = Tool(name="fetch_course_meta", func=fetch_course_meta, description="Fetch course metadata from MCP mock server.")
|
||
|
||
# Build the agent executor
|
||
agent = create_openai_tools_agent(llm=llm, tools=[search_tool, meta_tool], system_message=system_prompt)
|
||
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=[search_tool, meta_tool], verbose=True)
|
||
|
||
# ---------- CLI ----------
|
||
async def run_cli():
|
||
# Preload data if not already present
|
||
if not Path("./chroma_faq").exists():
|
||
load_faq_to_chroma()
|
||
|
||
# Predefined questions
|
||
predefined = [
|
||
"What is the main topic of the first lecture?",
|
||
"Explain the concept of polymorphism in the course.",
|
||
"What is the schedule for the next week?",
|
||
]
|
||
print("--- Predefined questions ---")
|
||
for i, q in enumerate(predefined, 1):
|
||
print(f"{i}. {q}")
|
||
print("\nEnter a number to ask a predefined question or type your own query.")
|
||
user_input = input("> ")
|
||
if user_input.isdigit() and 1 <= int(user_input) <= len(predefined):
|
||
query = predefined[int(user_input)-1]
|
||
else:
|
||
query = user_input
|
||
|
||
result = await agent_executor.ainvoke({"input": query})
|
||
print("\n--- Answer ---")
|
||
print(result["output"])
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(run_cli())
|