Files
task-6a1d75c5fd30e81cf3126ae7/main.py
T
2026-06-04 15:58:14 +00:00

115 lines
3.8 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 asyncio, os
from pathlib import Path
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
import httpx
import json
# ---------- LLM ----------
llm = ChatOpenAI(
model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1",
api_key=os.getenv("OPENAI_API_KEY"),
temperature=0.0,
)
# ---------- Backend ----------
backend = CompositeBackend([
LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(),
])
# ---------- Chroma DB ----------
CHROMA_PATH = Path("./chroma_faq")
CHROMA_PATH.mkdir(exist_ok=True)
embeddings = OllamaEmbeddings(model="nomic-embed-text")
# Load or create vector store
if CHROMA_PATH.exists() and any(CHROMA_PATH.iterdir()):
chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embeddings)
else:
# Load markdown files
docs = []
for md_file in Path("data").glob("*.md"):
text = md_file.read_text(encoding="utf-8")
docs.append(text)
chroma = Chroma.from_texts(docs, embedding=embeddings, persist_directory=str(CHROMA_PATH))
chroma.persist()
@tool
def search_course_docs(query: str, k: int = 3) -> str:
"""Search local course documents in Chroma."""
results = chroma.similarity_search(query, k=k)
return "\n---\n".join(doc.page_content for doc in results) if results else "No relevant docs found."
# ---------- MCPstyle tool ----------
# For demo we use a static JSON file. In production this would be an HTTP call.
META_JSON = Path("meta.json")
if not META_JSON.exists():
# Create a simple mock meta file
META_JSON.write_text(json.dumps({
"schedule": "Mon 10-12, Wed 14-16, Fri 9-11",
"instructor": "Dr. Smith",
"location": "Room 101"
}))
@tool
def fetch_course_meta(query: str) -> str:
"""Return course metadata matching the query keyword."""
data = json.loads(META_JSON.read_text())
# Simple keyword search in values
for key, value in data.items():
if query.lower() in key.lower() or query.lower() in str(value).lower():
return f"{key}: {value}"
return "No metadata found for the query."
# ---------- Agent ----------
agent = create_deep_agent(
model=llm,
tools=[search_course_docs, fetch_course_meta],
backend=backend,
system_prompt=(
"You are a helpful FAQ bot for the course.\n"
"If the question is about course content, use search_course_docs.\n"
"If the question is about schedule, instructor, or location, use fetch_course_meta.\n"
"Do not call both tools unless necessary.\n"
"In your answer, prefix the source with 'source: chroma' or 'source: mcp_meta'."
),
)
# ---------- CLI ----------
PRESET_QUESTIONS = [
"What is the main topic of the first lecture?",
"How can I access the lecture slides?",
"When is the next class?"
]
async def run_cli():
print("--- FAQ Bot CLI ---")
for i, q in enumerate(PRESET_QUESTIONS, 1):
print(f"\nPreset {i}: {q}")
result = await agent.ainvoke(
{"messages": [HumanMessage(content=q)]},
{"configurable": {"thread_id": f"preset-{i}"}},
)
print(result["messages"][-1].content)
print("\nEnter your own question (or 'exit'): ")
while True:
user_q = input("> ")
if user_q.lower() in {"exit", "quit"}:
break
result = await agent.ainvoke(
{"messages": [HumanMessage(content=user_q)]},
{"configurable": {"thread_id": "interactive"}},
)
print(result["messages"][-1].content)
if __name__ == "__main__":
asyncio.run(run_cli())