diff --git a/main.py b/main.py index 05960f2..86c9f44 100644 --- a/main.py +++ b/main.py @@ -1,7 +1,5 @@ import os import asyncio -import json -import httpx from pathlib import Path from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_chroma import Chroma @@ -10,46 +8,49 @@ from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -# --------------------- Configuration --------------------- -BASE_DIR = Path(__file__).parent -DATA_DIR = BASE_DIR / "data" -CHROMA_DIR = BASE_DIR / "chroma_faq" -META_JSON = BASE_DIR / "course_meta.json" +# ----------------- Configuration ----------------- +# Load OpenRouter API key from .env or environment variable +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + raise RuntimeError("OPENAI_API_KEY not set in environment") -# --------------------- LLM & Embeddings --------------------- +# ----------------- 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"), + api_key=OPENAI_API_KEY, temperature=0.0, ) 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, ) -# --------------------- Chroma DB --------------------- +# ----------------- ChromaDB setup ----------------- +CHROMA_PATH = Path("./chroma_faq") +CHROMA_COLLECTION = "faq_collection" vector_store = Chroma( - collection_name="faq_collection", + collection_name=CHROMA_COLLECTION, embedding_function=embeddings, - persist_directory=str(CHROMA_DIR), + persist_directory=str(CHROMA_PATH), ) # Load markdown files into Chroma if not already loaded -if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()): +if not CHROMA_PATH.exists() or not any(CHROMA_PATH.iterdir()): + data_dir = Path("data") docs = [] - for md_file in DATA_DIR.glob("*.md"): + 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 --------------------- +# ----------------- Tools ----------------- @tool def search_course_docs(query: str) -> str: - """Search the FAQ knowledge base for relevant information.""" + """Search the local FAQ collection for relevant passages.""" results = vector_store.similarity_search(query, k=3) if not results: return "No relevant information found in the course materials." @@ -57,67 +58,63 @@ def search_course_docs(query: str) -> str: @tool def fetch_course_meta(query: str) -> str: - """Fetch course metadata (e.g., schedule) from a local JSON mock.""" - if not META_JSON.exists(): - return "Metadata file not found." - data = json.loads(META_JSON.read_text(encoding="utf-8")) - # Simple keyword search in the metadata - matches = [f"{k}: {v}" for k, v in data.items() if query.lower() in k.lower() or query.lower() in str(v).lower()] - return "\n".join(matches) if matches else "No metadata matches the query." + """Mock MCP-style tool that returns course metadata. + In production this would be an HTTP call to an MCP server. + Here we return a static JSON-like string based on the query. + """ + # Simple static mapping for demo purposes + meta = { + "schedule": "Monday 10:00-12:00, Wednesday 14:00-16:00", + "instructor": "Dr. Ivanov", + "location": "Room 101", + } + key = query.lower().strip() + return meta.get(key, f"No metadata found for '{query}'.") -# --------------------- Backend --------------------- +# ----------------- Backend ----------------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# --------------------- Agent --------------------- +# ----------------- Agent ----------------- agent = create_deep_agent( model=llm, tools=[search_course_docs, fetch_course_meta], backend=backend, system_prompt=( - "You are a helpful FAQ assistant 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" + "You are a helpful FAQ bot for the course.\n" + "When a user asks about course materials, use the search_course_docs tool.\n" + "When a user asks about schedule, instructor, or location, use the fetch_course_meta tool.\n" "Do not use both tools unless absolutely necessary.\n" - "In your final answer, prefix the source with 'source: chroma' or 'source: mcp_meta'." + "In your final answer, prefix the response with 'source: chroma' or 'source: mcp_meta' to indicate where the information came from." ), ) -# --------------------- CLI --------------------- -SAMPLE_QUESTIONS = [ - "What topics are covered in the first lecture?", # should hit chroma - "Explain the concept of tokenization in NLP.", # chroma - "When is the next class scheduled?", # meta +# ----------------- CLI ----------------- +PRESET_QUESTIONS = [ + "What topics are covered in the first lecture?", + "When is the next class?", + "Who is the instructor?", ] -async def run_interactive(): - print("Welcome to the Course FAQ Bot! Type 'exit' to quit.") +async def run_cli(): + print("Welcome to the Course FAQ Bot!\n") + for i, q in enumerate(PRESET_QUESTIONS, 1): + print(f"{i}. {q}") + print("\nEnter your own question or type 'exit' to quit.") while True: - user_input = input("\nYou: ") + user_input = input("\n> ") if user_input.lower() in {"exit", "quit"}: + print("Goodbye!") break result = await agent.ainvoke( - {"messages": [{"role": "user", "content": user_input}]}, - {"configurable": {"thread_id": "interactive-session"}}, + {"messages": ["HumanMessage(content=\"{}\")".format(user_input)]}, + {"configurable": {"thread_id": "session-1"}}, ) - print("\nAssistant:", result["messages"][-1]["content"]) - -async def run_samples(): - for q in SAMPLE_QUESTIONS: - print("\nQuestion:", q) - result = await agent.ainvoke( - {"messages": [{"role": "user", "content": q}]}, - {"configurable": {"thread_id": "sample-session"}}, - ) - print("Answer:", result["messages"][-1]["content"]) - -async def main(): - # Run sample questions first - await run_samples() - # Then enter interactive mode - await run_interactive() + # The agent returns a dict with 'messages'; take the last one + content = result["messages"][-1].content + print(content) if __name__ == "__main__": - asyncio.run(main()) + asyncio.run(run_cli())