From 30b31a9de2074afe5dfafa53b29002f83504003e Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Mon, 15 Jun 2026 12:16:06 +0000 Subject: [PATCH] add: main.py --- main.py | 120 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..b84e6e6 --- /dev/null +++ b/main.py @@ -0,0 +1,120 @@ +import os +import asyncio +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 + +# ----------------- Configuration ----------------- +# Load API key from .env or environment variable +os.environ.setdefault("OPENAI_API_KEY", os.getenv("OPENAI_API_KEY", "")) + +# LLM via OpenRouter +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, +) + +# Embeddings for Chroma +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), +) + +# ----------------- Chroma DB ----------------- +CHROMA_PATH = Path("./chroma_faq") +CHROMA_COLLECTION = "faq_collection" +vector_store = Chroma( + collection_name=CHROMA_COLLECTION, + embedding_function=embeddings, + persist_directory=str(CHROMA_PATH), +) + +# Load markdown files into Chroma (idempotent) +DATA_DIR = Path("./data") +if not CHROMA_PATH.exists() or not any(CHROMA_PATH.iterdir()): + 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) -> str: + """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." + return "\n---\n".join(f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in results) + +@tool +def fetch_course_meta(query: str) -> str: + """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. + """ + meta = { + "schedule": "Mon 10-12, Wed 14-16, Fri 9-11", + "instructor": "Dr. Ivanov", + "credits": 3, + } + # Simple keyword matching + if "schedule" in query.lower(): + return f"Course schedule: {meta['schedule']}" + if "instructor" in query.lower(): + return f"Instructor: {meta['instructor']}" + if "credits" in query.lower(): + return f"Credits: {meta['credits']}" + return "No metadata matches your query." + +# ----------------- Backend ----------------- +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) + +# ----------------- 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. Use the search_course_docs tool for questions about lecture materials and fetch_course_meta for questions about schedule, instructor, or credits. Do not use both tools unless necessary. In your final answer, prefix the source with 'source: chroma' or 'source: mcp_meta'.", +) + +# ----------------- CLI ----------------- +PRESET_QUESTIONS = [ + "What topics are covered in lecture 3?", + "When is the next class?", + "Who is the instructor for this course?", +] + +async def run_agent(question: str, thread_id: str = "session-1"): + result = await agent.ainvoke( + {"messages": ["HumanMessage(content=\"{}\")".format(question)]}, + {"configurable": {"thread_id": thread_id}}, + ) + # The agent returns a dict with 'messages'; extract last content + content = result["messages"][-1].content + print(f"\nQ: {question}\nA: {content}\n") + +async def main(): + print("--- FAQ Bot Demo ---") + for i, q in enumerate(PRESET_QUESTIONS, 1): + await run_agent(q, thread_id=f"demo-{i}") + print("Enter your own question (or 'exit' to quit):") + while True: + q = input("> ") + if q.lower() in {"exit", "quit"}: + break + await run_agent(q, thread_id="interactive") + +if __name__ == "__main__": + asyncio.run(main())