fix(needs_fixes): 1 исправлений, 0 отстояно — main.py

This commit is contained in:
+161 -56
View File
@@ -1,73 +1,178 @@
"""
# main.py
# FAQ bot using deepagents, ChromaDB and a single MCPstyle 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 os
import asyncio
import json import json
import httpx
from pathlib import Path from pathlib import Path
from langchain_ollama import ChatOllama from typing import List, Dict
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma from langchain_chroma import Chroma
from langchain.embeddings import OllamaEmbeddings from langchain_core.documents import Document
from langchain.schema import Document from langchain_core.messages import HumanMessage, SystemMessage
from langchain.prompts import ChatPromptTemplate from langchain.tools import tool
from langchain.chains import LLMChain from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
# 1. Load FAQ into Chroma # ---------------------------------------------------------------------------
# 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")
def load_faq_to_chroma(md_path: str, persist_dir: str = "./chroma_faq"): # LLM configuration OpenRouter
from langchain.document_loaders import TextLoader llm = ChatOpenAI(
from langchain.text_splitter import RecursiveCharacterTextSplitter model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1",
api_key=OPENAI_API_KEY,
temperature=0.0,
)
loader = TextLoader(md_path) # Embeddings OpenRouter
docs = loader.load() embeddings = OpenAIEmbeddings(
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) model="text-embedding-3-small",
texts = splitter.split_documents(docs) base_url="https://openrouter.ai/api/v1",
embeddings = OllamaEmbeddings(model="nomic-embed-text") api_key=OPENAI_API_KEY,
chroma = Chroma.from_documents(texts, embeddings, persist_directory=persist_dir) )
chroma.persist()
return chroma
# 2. Search function # ---------------------------------------------------------------------------
# ChromaDB setup
# ---------------------------------------------------------------------------
CHROMA_PATH = Path("./chroma_faq")
CHROMA_COLLECTION = "faq_collection"
def search_course_docs(query: str, k: int = 3): vector_store = Chroma(
chroma = Chroma(persist_directory="./chroma_faq", embedding_function=OllamaEmbeddings(model="nomic-embed-text")) collection_name=CHROMA_COLLECTION,
results = chroma.similarity_search(query, k=k) embedding_function=embeddings,
return [doc.page_content for doc in results] persist_directory=str(CHROMA_PATH),
)
# 3. MCP-style tool # 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()
def fetch_course_meta(query: str): load_faq_to_chroma()
# For demo, use static JSON file
meta_path = Path("meta.json") # ---------------------------------------------------------------------------
# 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)
# MCPstyle 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(): if not meta_path.exists():
return {"error": "Meta not found"} return "Metadata file not found."
meta = json.loads(meta_path.read_text()) data = json.loads(meta_path.read_text(encoding="utf-8"))
# simple search by key # Very naive matching: return any entry where the query is a substring
return meta.get(query, {}) 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)
# 4. Agent logic # ---------------------------------------------------------------------------
# Backend for deepagents
# ---------------------------------------------------------------------------
backend = CompositeBackend([
LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(),
])
def answer_question(question: str): # ---------------------------------------------------------------------------
# Simple heuristic: if question contains "schedule" or "метаданные" use meta # Agent definition
if any(word in question.lower() for word in ["schedule", "расписание", "метаданные"]): # ---------------------------------------------------------------------------
source = "mcp_meta" # System prompt instructs the agent to choose the appropriate tool and
answer = fetch_course_meta(question) # to annotate the answer with a `source` field.
else: SYSTEM_PROMPT = (
source = "chroma" "You are an FAQ assistant for a course.\n"
answer = search_course_docs(question, k=1)[0] "If the user asks about course materials, use the `search_course_docs` tool.\n"
return {"answer": answer, "source": source} "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."
)
if __name__ == "__main__": agent = create_deep_agent(
# CLI with preset questions model=llm,
preset = [ tools=[search_course_docs, fetch_course_meta],
"Что такое ChromaDB?", backend=backend,
"Как подключить Ollama embeddings?", system_prompt=SYSTEM_PROMPT,
"Когда будет расписание следующего занятия?" )
]
for q in preset: # ---------------------------------------------------------------------------
res = answer_question(q) # CLI helpers
print(f"Q: {q}\nA: {res['answer']}\nSource: {res['source']}\n") # ---------------------------------------------------------------------------
# interactive 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
]
async def run_interactive():
print("FAQ Bot type your question (or 'exit' to quit).\n")
while True: while True:
q = input("Ask a question (or 'exit'): ") user_input = input("> ")
if q.lower() == "exit": if user_input.lower() in {"exit", "quit"}:
break break
res = answer_question(q) result = await agent.ainvoke(
print(f"A: {res['answer']} (source: {res['source']})") {"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)
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())