diff --git a/main.py b/main.py index b2e3232..bdbd063 100644 --- a/main.py +++ b/main.py @@ -1,146 +1,134 @@ import os import asyncio -import argparse from pathlib import Path - 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 -# -------------------- 1. LLM and Embeddings -------------------- -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, -) +# ===================== CONFIG ===================== +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + raise RuntimeError("OPENAI_API_KEY is not set in environment") +# ===================== EMBEDDINGS & VECTOR STORE ===================== embeddings = OpenAIEmbeddings( model="text-embedding-3-small", base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), + api_key=OPENAI_API_KEY, ) -# -------------------- 2. Chroma DB -------------------- -CHROMA_PATH = Path("./chroma_faq") -CHROMA_COLLECTION = "course_faq" - vector_store = Chroma( - collection_name=CHROMA_COLLECTION, + collection_name="faq_knowledge", embedding_function=embeddings, - persist_directory=str(CHROMA_PATH), + persist_directory="./chroma_faq", ) -# Persist changes -vector_store.persist() - -# -------------------- 3. Tools -------------------- +# ===================== TOOL: SEARCH IN CHROMA ===================== @tool -def search_knowledge(query: str) -> str: - """Search the knowledge base for relevant information.""" - docs = vector_store.similarity_search(query, k=3) - return "\n".join(d.page_content for d in docs) if docs else "No results found in course materials." +def search_course_docs(query: str, k: int = 3) -> str: + """Search the local FAQ knowledge base for relevant passages.""" + docs: list[Document] = 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('title', 'Untitled')}**\n{d.page_content}" for d in docs) + +# ===================== TOOL: FETCH METADATA (MCP‑STYLE) ===================== +# For the purpose of this assignment we use a static JSON file as the mock MCP server response. +METADATA_JSON = { + "schedule": { + "Monday": "Lecture 1: Introduction", + "Wednesday": "Lecture 2: Advanced Topics", + "Friday": "Lecture 3: Practical Applications" + }, + "instructors": { + "Dr. Smith": "smith@example.com", + "Prof. Doe": "doe@example.com" + } +} @tool def fetch_course_meta(query: str) -> str: - """Fetch course metadata (schedule, syllabus, etc.) from a static JSON file.""" - meta_path = Path("meta.json") - if not meta_path.exists(): - return "Metadata file not found." - import json - data = json.loads(meta_path.read_text()) - query_lower = query.lower() - if "schedule" in query_lower: - return f"Course schedule: {data.get('schedule', 'Not available')}" - if "syllabus" in query_lower: - return f"Syllabus URL: {data.get('syllabus_url', 'Not available')}" - if "instructor" in query_lower: - return f"Instructor: {data.get('instructor', 'Not available')}" - # Default: return all - return json.dumps(data, indent=2) + """Mock MCP tool that returns course metadata based on the query. + In production this would be an HTTP GET to an MCP server. + """ + query = query.lower() + if "schedule" in query: + return "\n".join(f"{day}: {info}" for day, info in METADATA_JSON["schedule"].items()) + if "instructor" in query or "email" in query: + return "\n".join(f"{name}: {email}" for name, email in METADATA_JSON["instructors"].items()) + return "No metadata matches your query." + +# ===================== LOAD FAQ TO CHROMA ===================== +MD_DIR = Path("data") +if not MD_DIR.exists(): + MD_DIR.mkdir(parents=True, exist_ok=True) + # Create example markdown files if none exist + (MD_DIR / "faq1.md").write_text("# FAQ 1\nWhat is the course about?\nThe course covers advanced AI techniques.") + (MD_DIR / "faq2.md").write_text("# FAQ 2\nHow to install dependencies?\nRun `pip install -r requirements.txt`.") + (MD_DIR / "faq3.md").write_text("# FAQ 3\nWhere to find the schedule?\nCheck the course website.") + +# Chunking and persisting +from langchain_text_splitters import RecursiveCharacterTextSplitter + +text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + +for md_file in MD_DIR.glob("*.md"): + content = md_file.read_text(encoding="utf-8") + docs = [Document(page_content=chunk, metadata={"title": md_file.stem}) for chunk in text_splitter.split_text(content)] + vector_store.add_documents(docs) + +vector_store.persist() + +# ===================== 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, +) -# -------------------- 4. Backend -------------------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# -------------------- 5. Agent -------------------- -SYSTEM_PROMPT = ( - "You are a helpful FAQ bot for the course.\n" - "Use the 'search_knowledge' tool to answer questions about course materials.\n" - "Use the 'fetch_course_meta' tool for questions about schedule, syllabus, instructor, etc.\n" - "Do not call both tools unnecessarily.\n" - "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate which tool provided the information." -) - agent = create_deep_agent( model=llm, - tools=[search_knowledge, fetch_course_meta], + tools=[search_course_docs, fetch_course_meta], backend=backend, - system_prompt=SYSTEM_PROMPT, + system_prompt="You are a helpful FAQ assistant. Use only the provided tools. In your final answer, prefix the source with `source: chroma` or `source: mcp_meta` accordingly.", ) -# -------------------- 6. Data Loader -------------------- +# ===================== CLI ===================== +PRESET_QUESTIONS = [ + "What is the course about?", # chroma + "How to install dependencies?", # chroma + "What is the lecture schedule?", # mcp_meta +] -def load_faq_to_chroma(): - data_dir = Path("data") - if not data_dir.exists(): - data_dir.mkdir() - # Create sample markdown files if missing - (data_dir / "file1.md").write_text( - "## Course Overview\nThis course covers advanced topics in AI. Topics include machine learning, deep learning, and natural language processing." - ) - (data_dir / "file2.md").write_text( - "## FAQ\nQ: What is the schedule?\nA: Sessions are held on Mondays and Wednesdays.\nQ: Where can I find the syllabus?\nA: Syllabus is available on the course website." - ) - docs = [] - for md_file in data_dir.glob("*.md"): - text = md_file.read_text() - docs.append(Document(page_content=text, metadata={"source": md_file.name})) - vector_store.add_documents(docs) - vector_store.persist() - print(f"Loaded {len(docs)} documents into Chroma collection '{CHROMA_COLLECTION}'.") +async def run_cli(): + print("=== FAQ Assistant ===") + print("Type your question or 'exit' to quit.") + for i, q in enumerate(PRESET_QUESTIONS, 1): + print(f"\nPreset {i}: {q}") + await handle_question(q) + while True: + user_input = input("\nYour question: ") + if user_input.lower() in {"exit", "quit"}: + break + await handle_question(user_input) -# -------------------- 7. CLI -------------------- -PRESET_QUESTIONS = { - "1": "What topics are covered in the course?", - "2": "Where can I find the syllabus?", - "3": "What is the course schedule?", -} - -async def run_question(question: str, thread_id: str = "session-1"): +async def handle_question(question: str): result = await agent.ainvoke( - {"messages": ["HumanMessage(content=\"{}\")".format(question)]}, - {"configurable": {"thread_id": thread_id}}, + {"messages": [HumanMessage(content=question)]}, + {"configurable": {"thread_id": "session-1"}}, ) - # The agent returns a dict with 'messages'; extract last message - content = result["messages"][-1].content - print(content) - -async def main(): - # Ensure data is loaded - load_faq_to_chroma() - - parser = argparse.ArgumentParser(description="FAQ Bot CLI") - parser.add_argument("--preset", choices=["1", "2", "3"], help="Run a preset question") - args = parser.parse_args() - - if args.preset: - question = PRESET_QUESTIONS[args.preset] - print(f"Preset question {args.preset}: {question}") - await run_question(question) - else: - print("Enter your question (type 'exit' to quit):") - while True: - q = input("> ") - if q.lower() in {"exit", "quit"}: - break - if q.strip(): - await run_question(q) + answer = result["messages"][-1].content + print("\nAnswer:\n", answer) if __name__ == "__main__": - asyncio.run(main()) + asyncio.run(run_cli())