From abb2b709c27860812b0f3f63c6f146a48b2adb14 Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Tue, 16 Jun 2026 08:10:55 +0000 Subject: [PATCH] =?UTF-8?q?fix(needs=5Ffixes):=201=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 | 135 +++++++++++++++++++++++++++++++++----------------------- 1 file changed, 79 insertions(+), 56 deletions(-) diff --git a/main.py b/main.py index b84e6e6..0dbeef9 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,7 @@ import os import asyncio +import json +import httpx from pathlib import Path from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_chroma import Chroma @@ -7,16 +9,20 @@ from langchain_core.documents import Document 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 ----------------- -# Load API key from .env or environment variable -os.environ.setdefault("OPENAI_API_KEY", os.getenv("OPENAI_API_KEY", "")) +# --------------------------- +# Configuration +# --------------------------- +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + raise RuntimeError("OPENAI_API_KEY not set in environment") # 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"), + api_key=OPENAI_API_KEY, temperature=0.0, ) @@ -24,97 +30,114 @@ llm = ChatOpenAI( 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 ----------------- -CHROMA_PATH = Path("./chroma_faq") -CHROMA_COLLECTION = "faq_collection" +# --------------------------- +# Chroma DB utilities +# --------------------------- +CHROMA_DIR = Path("./chroma_faq") +CHROMA_DIR.mkdir(parents=True, exist_ok=True) vector_store = Chroma( - collection_name=CHROMA_COLLECTION, + collection_name="faq_collection", + persist_directory=str(CHROMA_DIR), 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"): + +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. + """ + 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"): text = md_file.read_text(encoding="utf-8") - docs.append(Document(page_content=text, metadata={"source": md_file.name})) - vector_store.add_documents(docs) + # 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] + vector_store.add_documents(docs) vector_store.persist() -# ----------------- Tools ----------------- +# --------------------------- +# 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: +def search_course_docs(query: str, k: int = 3) -> str: + """Search the local FAQ ChromaDB 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".join(f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in results) + return "\n\n---\n\n".join(f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs) @tool def fetch_course_meta(query: str) -> str: - """Mock MCP-style tool that returns course metadata. + """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. - 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." + # 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." -# ----------------- Backend ----------------- +# --------------------------- +# Agent setup +# --------------------------- 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'.", + 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." + ), ) -# ----------------- CLI ----------------- +# --------------------------- +# CLI +# --------------------------- PRESET_QUESTIONS = [ - "What topics are covered in lecture 3?", - "When is the next class?", - "Who is the instructor for this course?", + "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 ] async def run_agent(question: str, thread_id: str = "session-1"): result = await agent.ainvoke( - {"messages": ["HumanMessage(content=\"{}\")".format(question)]}, + {"messages": [HumanMessage(content=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") + # The last message is the agent's reply + reply = result["messages"][-1].content + print(f"\nQ: {question}\nA: {reply}\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):") + # 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 ---") + for q in PRESET_QUESTIONS: + await run_agent(q) + print("\n--- Interactive mode (type 'exit' to quit) ---") while True: - q = input("> ") - if q.lower() in {"exit", "quit"}: + user_input = input("You: ") + if user_input.lower() in {"exit", "quit"}: break - await run_agent(q, thread_id="interactive") + await run_agent(user_input) if __name__ == "__main__": asyncio.run(main())