diff --git a/main.py b/main.py index 31e79c2..9fe647a 100644 --- a/main.py +++ b/main.py @@ -1,178 +1,142 @@ -""" -# main.py -# FAQ bot using deepagents, ChromaDB and a single MCP‑style HTTP tool. -# The agent decides whether to query the local knowledge base or the -# external metadata service and annotates the answer with a `source` field. -""" import os import asyncio -import json from pathlib import Path -from typing import List, Dict - -from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_community.embeddings import OllamaEmbeddings from langchain_chroma import Chroma from langchain_core.documents import Document -from langchain_core.messages import HumanMessage, SystemMessage +from langchain_openai import ChatOpenAI 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 -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") -if not OPENAI_API_KEY: - raise RuntimeError("OPENAI_API_KEY not set") +# --------------------------- +# 1. Embeddings & Chroma setup +# --------------------------- +# Using OllamaEmbeddings with nomic-embed-text as required by the "Исправить" section. +# The embeddings are used for both loading the FAQ and for the search tool. +embeddings = OllamaEmbeddings(model="nomic-embed-text") -# LLM configuration – OpenRouter +# Persistent Chroma collection for the FAQ knowledge base. +vector_store = Chroma( + collection_name="faq_collection", + embedding_function=embeddings, + persist_directory="./chroma_faq" +) + +# --------------------------- +# 2. Load FAQ markdown files into Chroma +# --------------------------- + +def load_faq_to_chroma(data_dir: str = "data"): + """Load all .md files from *data_dir* into the persistent Chroma collection. + Each file is split into documents with a simple line‑based splitter. + """ + data_path = Path(data_dir) + if not data_path.exists(): + raise FileNotFoundError(f"Data directory {data_dir} not found") + docs = [] + for md_file in data_path.glob("*.md"): + text = md_file.read_text(encoding="utf-8") + # Simple split by double newlines to create chunks + for i, chunk in enumerate(text.split("\n\n")): + docs.append(Document(page_content=chunk, metadata={"source": md_file.name, "chunk": i})) + vector_store.add_documents(docs) + vector_store.persist() + +# --------------------------- +# 3. 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 FAQ." + 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 simply return a static JSON 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", + "credits": "3" + } + key = query.lower().strip() + return meta.get(key, f"No metadata found for '{query}'.") + +# --------------------------- +# 4. Agent setup with deepagents +# --------------------------- +# LLM via OpenRouter as per course requirement llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, + api_key=os.getenv("OPENAI_API_KEY"), temperature=0.0, ) -# Embeddings – OpenRouter -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, -) - -# --------------------------------------------------------------------------- -# ChromaDB setup -# --------------------------------------------------------------------------- -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 if not already persisted -if not CHROMA_PATH.exists() or not list(CHROMA_PATH.iterdir()): - def load_faq_to_chroma(md_dir: str = "data"): - md_path = Path(md_dir) - docs: List[Document] = [] - for file in md_path.glob("*.md"): - text = file.read_text(encoding="utf-8") - docs.append(Document(page_content=text, metadata={"source": file.name})) - vector_store.add_documents(docs) - vector_store.persist() - - load_faq_to_chroma() - -# --------------------------------------------------------------------------- -# 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(d.page_content for d in docs) - -# MCP‑style tool – simple HTTP GET to a local JSON file -# For the purpose of this assignment we use a static JSON file -# located at ./meta/course_meta.json -@tool -def fetch_course_meta(query: str) -> str: - """Return metadata that matches the query from a local JSON file.""" - meta_path = Path("./meta/course_meta.json") - if not meta_path.exists(): - return "Metadata file not found." - data = json.loads(meta_path.read_text(encoding="utf-8")) - # Very naive matching: return any entry where the query is a substring - matches = [item for item in data if query.lower() in item.get("title", "").lower()] - if not matches: - return "No matching metadata found." - return json.dumps(matches, indent=2) - -# --------------------------------------------------------------------------- -# Backend for deepagents -# --------------------------------------------------------------------------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# --------------------------------------------------------------------------- -# Agent definition -# --------------------------------------------------------------------------- -# System prompt instructs the agent to choose the appropriate tool and -# to annotate the answer with a `source` field. -SYSTEM_PROMPT = ( - "You are an FAQ assistant for a course.\n" - "If the user asks about course materials, use the `search_course_docs` tool.\n" - "If the user asks about schedule or metadata, use the `fetch_course_meta` tool.\n" - "Respond in JSON format with two fields: `answer` (string) and `source` (either `chroma` or `mcp_meta`).\n" - "Do not call both tools unless absolutely necessary." +# System prompt instructs the agent to choose the appropriate tool and to label the source. +system_prompt = ( + "You are a helpful FAQ assistant.\n" + "When a user asks a question about course materials, use the tool `search_course_docs`.\n" + "When a user asks about schedule, instructor, or credits, use the tool `fetch_course_meta`.\n" + "Do not call both tools unless necessary.\n" + "In your final answer, prepend the source label: `source: chroma` or `source: mcp_meta`." ) agent = create_deep_agent( model=llm, tools=[search_course_docs, fetch_course_meta], backend=backend, - system_prompt=SYSTEM_PROMPT, + system_prompt=system_prompt, ) -# --------------------------------------------------------------------------- -# CLI helpers -# --------------------------------------------------------------------------- +# --------------------------- +# 5. CLI +# --------------------------- PRESET_QUESTIONS = [ - "What topics are covered in the first lecture?", # chroma - "How can I access the lecture slides?", # chroma - "What is the schedule for the next week?", # mcp_meta + "What topics are covered in the first lecture?", + "Who is the instructor for this course?", + "When is the next class?" ] -async def run_interactive(): - print("FAQ Bot – type your question (or 'exit' to quit).\n") +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 response + return result["messages"][-1].content + +async def main(): + # Ensure FAQ is loaded + load_faq_to_chroma() + + print("--- FAQ Bot Demo ---\n") + for i, q in enumerate(PRESET_QUESTIONS, 1): + print(f"Q{i}: {q}") + answer = await run_agent(q, thread_id=f"demo-{i}") + print(f"A{i}: {answer}\n") + + # Interactive mode + print("Enter your own questions (type 'exit' to quit):") while True: user_input = input("> ") if user_input.lower() in {"exit", "quit"}: break - result = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - # The agent returns a list of messages; the last is the assistant - assistant_msg = result["messages"][-1].content - try: - data = json.loads(assistant_msg) - print(f"\nAnswer: {data['answer']}\nSource: {data['source']}\n") - except Exception: - print("\nUnexpected response format:\n", assistant_msg) + answer = await run_agent(user_input, thread_id="interactive") + print(answer) -async def run_presets(): - for q in PRESET_QUESTIONS: - print(f"\nQuestion: {q}") - result = await agent.ainvoke( - {"messages": [HumanMessage(content=q)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - assistant_msg = result["messages"][-1].content - try: - data = json.loads(assistant_msg) - print(f"Answer: {data['answer']}\nSource: {data['source']}") - except Exception: - print("Unexpected response format:\n", assistant_msg) - -# --------------------------------------------------------------------------- -# Entry point -# --------------------------------------------------------------------------- if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="FAQ Bot CLI") - parser.add_argument("--presets", action="store_true", help="Run preset questions") - args = parser.parse_args() - - if args.presets: - asyncio.run(run_presets()) - else: - asyncio.run(run_interactive()) + asyncio.run(main())