diff --git a/main.py b/main.py index 0dbeef9..47eeb16 100644 --- a/main.py +++ b/main.py @@ -1,19 +1,20 @@ import os import asyncio import json -import httpx from pathlib import Path +from typing import List + from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_chroma import Chroma from langchain_core.documents import Document +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_core.messages import HumanMessage -# --------------------------- +# --------------------------------------------------------------------------- # Configuration -# --------------------------- +# --------------------------------------------------------------------------- OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") if not OPENAI_API_KEY: raise RuntimeError("OPENAI_API_KEY not set in environment") @@ -26,93 +27,111 @@ llm = ChatOpenAI( temperature=0.0, ) -# Embeddings for Chroma +# Embeddings via OpenRouter embeddings = OpenAIEmbeddings( model="text-embedding-3-small", base_url="https://openrouter.ai/api/v1", api_key=OPENAI_API_KEY, ) -# --------------------------- -# Chroma DB utilities -# --------------------------- +# --------------------------------------------------------------------------- +# Chroma persistence +# --------------------------------------------------------------------------- CHROMA_DIR = Path("./chroma_faq") CHROMA_DIR.mkdir(parents=True, exist_ok=True) vector_store = Chroma( collection_name="faq_collection", - persist_directory=str(CHROMA_DIR), embedding_function=embeddings, + persist_directory=str(CHROMA_DIR), ) - -def load_faq_to_chroma(md_dir: str = "data"): - """Load all .md files from md_dir into ChromaDB. - Each file is split into chunks and added to the vector store. +# --------------------------------------------------------------------------- +# Utility: load markdown files into Chroma +# --------------------------------------------------------------------------- +@tool +def load_faq_to_chroma() -> str: + """Load all .md files from data/ into the Chroma vector store. + This tool is idempotent – it will overwrite existing collection. """ - md_path = Path(md_dir) - if not md_path.exists(): - raise FileNotFoundError(f"Markdown directory {md_dir} not found") - for md_file in md_path.glob("*.md"): + data_dir = Path("data") + if not data_dir.exists(): + return "Data directory not found." + docs: List[Document] = [] + for md_file in data_dir.glob("*.md"): text = md_file.read_text(encoding="utf-8") - # Simple chunking: split by double newlines - chunks = [c.strip() for c in text.split("\n\n") if c.strip()] - docs = [Document(page_content=c, metadata={"source": md_file.name}) for c in chunks] + docs.append(Document(page_content=text, metadata={"source": md_file.name})) + if docs: + vector_store.delete_collection() vector_store.add_documents(docs) - vector_store.persist() + vector_store.persist() + return f"Loaded {len(docs)} documents into Chroma." + return "No markdown files found." -# --------------------------- -# Tools -# --------------------------- +# --------------------------------------------------------------------------- +# Tool: search knowledge base +# --------------------------------------------------------------------------- @tool def search_course_docs(query: str, k: int = 3) -> str: - """Search the local FAQ ChromaDB for relevant passages.""" + """Search the FAQ collection for relevant passages. + Returns a string with the top k passages and a source tag. + """ 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"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs) + return "No relevant information found in the FAQ." + results = "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs)) + return f"source: chroma\n{results}" + +# --------------------------------------------------------------------------- +# MCP-style tool: fetch course metadata +# --------------------------------------------------------------------------- +# For simplicity we use a static JSON file in the repo. In production this +# would be an HTTP call to an MCP server. +METADATA_FILE = Path("course_meta.json") @tool def fetch_course_meta(query: str) -> str: - """Mock MCP-style tool that fetches course metadata from a local JSON file. - In production this would be an HTTP call to an MCP server. + """Return metadata that matches the query. + The function performs a simple keyword search in the static JSON. """ - # For simplicity, we use a local JSON file. In a real scenario, replace with httpx.get. - meta_path = Path("course_meta.json") - if not meta_path.exists(): - return "Course metadata not available." - data = json.loads(meta_path.read_text(encoding="utf-8")) - # Very naive search: return any entry where query is a substring of title or description - results = [f"{item['title']}: {item['description']}" for item in data if query.lower() in item.get('title', '').lower() or query.lower() in item.get('description', '').lower()] - return "\n".join(results) if results else "No matching metadata found." + if not METADATA_FILE.exists(): + return "Metadata file not found." + data = json.loads(METADATA_FILE.read_text(encoding="utf-8")) + matches = [item for item in data if query.lower() in item.get("title", "").lower()] + if not matches: + return "No metadata matches the query." + return f"source: mcp_meta\n" + json.dumps(matches, indent=2) -# --------------------------- -# Agent setup -# --------------------------- +# --------------------------------------------------------------------------- +# Backend setup +# --------------------------------------------------------------------------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) +# --------------------------------------------------------------------------- +# Agent definition +# --------------------------------------------------------------------------- agent = create_deep_agent( model=llm, - tools=[search_course_docs, fetch_course_meta], + tools=[load_faq_to_chroma, search_course_docs, fetch_course_meta], backend=backend, system_prompt=( - "You are a helpful FAQ bot for the course.\n" - "When a user asks about course content, use the search_course_docs tool.\n" - "When a user asks about schedule, metadata, or other non-content info, use fetch_course_meta.\n" - "Do not call both tools unless absolutely necessary.\n" - "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate the origin." + "You are a helpful FAQ bot for a course. " + "Use the search_course_docs tool for questions about lecture materials. " + "Use fetch_course_meta for questions about schedule or metadata. " + "Do not call both tools unless necessary. " + "Always prefix your answer with the source tag (chroma or mcp_meta)." ), ) -# --------------------------- -# CLI -# --------------------------- +# --------------------------------------------------------------------------- +# CLI helpers +# --------------------------------------------------------------------------- PRESET_QUESTIONS = [ - "What is the deadline for the final project?", # likely in metadata - "Explain the concept of tokenization in NLP.", # content - "How many lectures are there in the first module?", # content + "What is the deadline for the final project?", # should hit metadata + "Explain the concept of polymorphism in OOP.", # should hit FAQ + "How many lectures are there in the first module?", # metadata ] async def run_agent(question: str, thread_id: str = "session-1"): @@ -120,24 +139,25 @@ async def run_agent(question: str, thread_id: str = "session-1"): {"messages": [HumanMessage(content=question)]}, {"configurable": {"thread_id": thread_id}}, ) - # The last message is the agent's reply - reply = result["messages"][-1].content - print(f"\nQ: {question}\nA: {reply}\n") + return result["messages"][-1].content async def main(): - # Load data into Chroma if not already persisted - if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()): - print("Loading FAQ data into ChromaDB...") - load_faq_to_chroma() - print("\n--- Predefined questions ---") + # Ensure FAQ is loaded once + await agent.ainvoke( + {"messages": [HumanMessage(content="load_faq_to_chroma()")]}, + {"configurable": {"thread_id": "init"}}, + ) + print("\n--- Preset questions ---") for q in PRESET_QUESTIONS: - await run_agent(q) + ans = await run_agent(q) + print(f"Q: {q}\nA: {ans}\n") print("\n--- Interactive mode (type 'exit' to quit) ---") while True: - user_input = input("You: ") + user_input = input("> ") if user_input.lower() in {"exit", "quit"}: break - await run_agent(user_input) + ans = await run_agent(user_input) + print(ans) if __name__ == "__main__": asyncio.run(main())