From d747028e9ee850b9783f355571d2624f8f20565c Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Tue, 16 Jun 2026 16:38:13 +0000 Subject: [PATCH] =?UTF-8?q?fix(needs=5Ffixes):=205=20=D0=B8=D1=81=D0=BF?= =?UTF-8?q?=D1=80=D0=B0=D0=B2=D0=BB=D0=B5=D0=BD=D0=B8=D0=B9,=200=20=D0=BE?= =?UTF-8?q?=D1=82=D1=81=D1=82=D0=BE=D1=8F=D0=BD=D0=BE=20=E2=80=94=20main.p?= =?UTF-8?q?y?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 226 +++++++++++++++++++++++++++++++------------------------- 1 file changed, 125 insertions(+), 101 deletions(-) diff --git a/main.py b/main.py index 9647c64..2c08a31 100644 --- a/main.py +++ b/main.py @@ -1,78 +1,22 @@ 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 langchain.agents import AgentExecutor, create_openai_tools_agent -from langchain.agents import Tool -from langchain_community.utilities import RetrievalQA -from langchain_community.vectorstores import Chroma as ChromaStore +from deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain_core.messages import HumanMessage -# ---------- Configuration ---------- +# --------------------------- +# Configuration +# --------------------------- OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") if not OPENAI_API_KEY: raise RuntimeError("OPENAI_API_KEY not set in environment") -# ---------- Embeddings & Vector Store ---------- -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, -) - -vector_store = ChromaStore( - collection_name="faq_collection", - embedding_function=embeddings, - persist_directory="./chroma_faq", -) - -# ---------- Load FAQ into Chroma ---------- - -def load_faq_to_chroma(data_dir: str = "data"): - """Load all .md files from data_dir into the Chroma vector store. - The function clears the existing collection before loading. - """ - vector_store.delete_collection() - docs = [] - for md_file in Path(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 the local FAQ collection for relevant passages.""" - 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"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in docs]) - -@tool -def fetch_course_meta(query: str) -> str: - """MCP‑style tool that queries a local mock server for course metadata. - The mock server should serve a JSON file at http://localhost:8000/meta.json. - """ - url = "http://localhost:8000/meta.json" - try: - response = httpx.get(url, timeout=5.0) - response.raise_for_status() - data = response.json() - except Exception as e: - return f"Error fetching metadata: {e}" - # Simple lookup: return value if query matches a key (case‑insensitive) - key = query.strip().lower() - value = data.get(key) - if value is None: - return f"No metadata entry found for '{query}'." - return f"{key}: {value}" - -# ---------- Agent Setup ---------- +# LLM via OpenRouter llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -80,49 +24,129 @@ llm = ChatOpenAI( temperature=0.0, ) -# Define the system prompt with routing rule -system_prompt = ( - "You are a helpful assistant that answers questions about the course. " - "If the question is about course content, use the search_course_docs tool. " - "If the question is about schedule, metadata, or other non‑content info, " - "use the fetch_course_meta tool. Do not call both tools unless the question " - "explicitly requires both. In your final answer, prepend 'source: chroma' " - "or 'source: mcp_meta' to indicate which tool provided the information." +# Embeddings via OpenRouter +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=OPENAI_API_KEY, ) -# Create Tool objects -search_tool = Tool(name="search_course_docs", func=search_course_docs, description="Search local course documents.") -meta_tool = Tool(name="fetch_course_meta", func=fetch_course_meta, description="Fetch course metadata from MCP mock server.") +# Chroma vector store (persisted) +CHROMA_PATH = Path("./chroma_faq") +vector_store = Chroma( + collection_name="faq_collection", + embedding_function=embeddings, + persist_directory=str(CHROMA_PATH), +) -# Build the agent executor -agent = create_openai_tools_agent(llm=llm, tools=[search_tool, meta_tool], system_message=system_prompt) -agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=[search_tool, meta_tool], verbose=True) +# --------------------------- +# Data loading +# --------------------------- -# ---------- CLI ---------- -async def run_cli(): - # Preload data if not already present - if not Path("./chroma_faq").exists(): - load_faq_to_chroma() - - # Predefined questions - predefined = [ - "What is the main topic of the first lecture?", - "Explain the concept of polymorphism in the course.", - "What is the schedule for the next week?", - ] - print("--- Predefined questions ---") - for i, q in enumerate(predefined, 1): - print(f"{i}. {q}") - print("\nEnter a number to ask a predefined question or type your own query.") - user_input = input("> ") - if user_input.isdigit() and 1 <= int(user_input) <= len(predefined): - query = predefined[int(user_input)-1] +def load_faq_to_chroma(md_folder: str = "data"): + """Load all .md files from md_folder into ChromaDB. + Each file is split into chunks and added to the vector store. + """ + md_path = Path(md_folder) + if not md_path.exists(): + raise FileNotFoundError(f"Markdown folder {md_folder} not found") + docs = [] + for md_file in md_path.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()] + for i, chunk in enumerate(chunks): + docs.append(Document(page_content=chunk, metadata={"source": md_file.name, "chunk": i})) + if docs: + vector_store.add_documents(docs) + vector_store.persist() else: - query = user_input + print("No markdown files found to load.") - result = await agent_executor.ainvoke({"input": query}) - print("\n--- Answer ---") - print(result["output"]) +# --------------------------- +# Tools +# --------------------------- +@tool +def search_course_docs(query: str, k: int = 3) -> str: + """Search the FAQ knowledge base for relevant information.""" + 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"**{doc.metadata.get('source')}** (chunk {doc.metadata.get('chunk')}):\n{doc.page_content}" for doc in docs) + +@tool +def fetch_course_meta(query: str) -> str: + """Mock MCP-style tool that returns course metadata. + In production this would perform an HTTP GET to an MCP server. + Here we return a static JSON-like string for simplicity. + """ + # Static mock data + meta = { + "schedule": { + "Monday": "Lecture 1: Introduction", + "Wednesday": "Lecture 2: Advanced Topics", + "Friday": "Lab Session" + }, + "instructor": "Dr. Jane Doe", + "credits": 3 + } + return f"Course metadata: {meta}" + +# --------------------------- +# Backend setup +# --------------------------- +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) + +# --------------------------- +# Agent creation +# --------------------------- +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 or metadata. In your answer, clearly indicate the source: either 'chroma' or 'mcp_meta'. Do not use both tools unless necessary.", +) + +# --------------------------- +# CLI +# --------------------------- +PRESET_QUESTIONS = [ + "What topics are covered in Lecture 1?", # should hit chroma + "Explain the concept of tokenization in NLP.", # chroma + "When is the next lab session?", # should hit 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}}, + ) + # The last message is the agent's reply + reply = result["messages"][-1].content + print(f"\nQ: {question}\nA: {reply}\n") + +async def main(): + # Load data if not already loaded + if not CHROMA_PATH.exists() or not any(CHROMA_PATH.iterdir()): + print("Loading FAQ data into Chroma...") + load_faq_to_chroma() + else: + print("Chroma database already loaded.") + + # Run preset questions + for i, q in enumerate(PRESET_QUESTIONS, 1): + await run_agent(q, thread_id=f"preset-{i}") + + # Interactive mode + print("Enter your own questions (type 'exit' to quit):") + while True: + user_input = input("> ") + if user_input.lower() in {"exit", "quit"}: + break + await run_agent(user_input, thread_id="interactive") if __name__ == "__main__": - asyncio.run(run_cli()) + asyncio.run(main())