diff --git a/main.py b/main.py new file mode 100644 index 0000000..206382c --- /dev/null +++ b/main.py @@ -0,0 +1,121 @@ +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 deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain_core.messages import HumanMessage + +# ---------- Configuration ---------- +DATA_DIR = Path("data") +CHROMA_DIR = Path("./chroma_faq") +CHROMA_COLLECTION = "faq_collection" +MOCK_META_URL = "http://localhost:8000/meta.json" # change if you use a different mock + +# ---------- Embeddings and LLM ---------- +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), +) + +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, +) + +# ---------- Chroma setup ---------- +vector_store = Chroma( + collection_name=CHROMA_COLLECTION, + embedding_function=embeddings, + persist_directory=str(CHROMA_DIR), +) + +# ---------- Load FAQ into Chroma ---------- + +def load_faq_to_chroma(): + if not CHROMA_DIR.exists(): + CHROMA_DIR.mkdir(parents=True, exist_ok=True) + # If collection already exists, skip loading + if vector_store.get_collection().count() > 0: + return + docs = [] + for md_file in 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 local course documents in Chroma and return top k snippets.""" + results = vector_store.similarity_search(query, k=k) + snippets = [f"{res.metadata.get('source', 'unknown')}\n{res.page_content[:200]}..." for res in results] + return "\n\n".join(snippets) + +@tool +def fetch_course_meta(query: str) -> str: + """Mock MCP tool that fetches course metadata from a local JSON endpoint.""" + try: + response = httpx.get(MOCK_META_URL, timeout=5) + response.raise_for_status() + data = response.json() + # Simple filtering: return items that contain the query string (case-insensitive) + matches = [item for item in data if query.lower() in json.dumps(item).lower()] + return json.dumps(matches, indent=2) if matches else "No metadata found for the query." + except Exception as e: + return f"Error fetching metadata: {e}" + +# ---------- Backend ---------- +backend = CompositeBackend( + default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True), + routes={}, +) + +# ---------- Agent ---------- +agent = create_deep_agent( + model=llm, + tools=[search_course_docs, fetch_course_meta], + backend=backend, + system_prompt="You are a helpful course assistant. Use the search_course_docs tool for questions about lecture materials and fetch_course_meta for schedule or metadata questions. In your answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate the used tool.", +) + +# ---------- CLI ---------- +PRESET_QUESTIONS = [ + "What topics are covered in Lecture 3?", # should use chroma + "Explain the concept of tokenization in NLP.", # chroma + "When is the next midterm exam scheduled?", # mcp_meta +] + +async def run_agent(question: str, thread_id: str = "session-1"): + result = await agent.ainvoke( + {"messages": [HumanMessage(content=question)]}, + {"configurable": {"thread_id": thread_id}}, + ) + return result["messages"][-1].content + +async def main(): + load_faq_to_chroma() + print("--- FAQ Bot Demo ---") + for i, q in enumerate(PRESET_QUESTIONS, 1): + print(f"\nQuestion {i}: {q}") + answer = await run_agent(q, thread_id=f"demo-{i}") + print("Answer:\n", answer) + print("\nEnter your own question (or type 'exit' to quit):") + while True: + user_q = input("> ") + if user_q.lower() in {"exit", "quit"}: + break + answer = await run_agent(user_q, thread_id="interactive") + print("Answer:\n", answer) + +if __name__ == "__main__": + asyncio.run(main())