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

This commit is contained in:
+63 -91
View File
@@ -9,144 +9,116 @@ from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langchain_core.messages import HumanMessage from langchain_core.messages import HumanMessage
# --------------------------- # --------------------- Configuration ---------------------
# Configuration BASE_DIR = Path(__file__).parent
# --------------------------- DATA_DIR = BASE_DIR / "data"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") CHROMA_DIR = BASE_DIR / "chroma_faq"
if not OPENAI_API_KEY: MOCK_META_FILE = BASE_DIR / "course_meta.json"
raise RuntimeError("OPENAI_API_KEY not set in environment")
# LLM via OpenRouter # --------------------- LLM and Embeddings ---------------------
llm = ChatOpenAI( llm = ChatOpenAI(
model="openai/gpt-oss-20b:free", model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1", base_url="https://openrouter.ai/api/v1",
api_key=OPENAI_API_KEY, api_key=os.getenv("OPENAI_API_KEY"),
temperature=0.0, temperature=0.0,
) )
# Embeddings via OpenRouter
embeddings = OpenAIEmbeddings( embeddings = OpenAIEmbeddings(
model="text-embedding-3-small", model="text-embedding-3-small",
base_url="https://openrouter.ai/api/v1", base_url="https://openrouter.ai/api/v1",
api_key=OPENAI_API_KEY, api_key=os.getenv("OPENAI_API_KEY"),
) )
# Chroma vector store (persisted) # --------------------- Chroma Vector Store ---------------------
CHROMA_PATH = Path("./chroma_faq")
vector_store = Chroma( vector_store = Chroma(
collection_name="faq_collection", collection_name="faq_collection",
embedding_function=embeddings, embedding_function=embeddings,
persist_directory=str(CHROMA_PATH), persist_directory=str(CHROMA_DIR),
) )
# --------------------------- # --------------------- Tools ---------------------
# Data loading
# ---------------------------
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:
print("No markdown files found to load.")
# ---------------------------
# Tools
# ---------------------------
@tool @tool
def search_course_docs(query: str, k: int = 3) -> str: def search_course_docs(query: str, k: int = 3) -> str:
"""Search the FAQ knowledge base for relevant information.""" """Search the local FAQ collection for relevant passages."""
docs = vector_store.similarity_search(query, k=k) docs = vector_store.similarity_search(query, k=k)
if not docs: if not docs:
return "No relevant information found in the course materials." 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) return "\n\n---\n\n".join(f"**{doc.metadata.get('title', 'Document')}**\n{doc.page_content}" for doc in docs)
@tool @tool
def fetch_course_meta(query: str) -> str: def fetch_course_meta(query: str) -> str:
"""Mock MCP-style tool that returns course metadata. """Mock MCP-style tool that returns course metadata.
In production this would perform an HTTP GET to an MCP server. In production this would be an HTTP call to an MCP server.
Here we return a static JSON-like string for simplicity. Here we read a local JSON file for simplicity.
""" """
# Static mock data import json
meta = { if not MOCK_META_FILE.exists():
"schedule": { return "Metadata source not available."
"Monday": "Lecture 1: Introduction", with open(MOCK_META_FILE, "r", encoding="utf-8") as f:
"Wednesday": "Lecture 2: Advanced Topics", data = json.load(f)
"Friday": "Lab Session" # Simple keyword search in the metadata
}, results = [f"{k}: {v}" for k, v in data.items() if query.lower() in k.lower() or query.lower() in str(v).lower()]
"instructor": "Dr. Jane Doe", return "\n".join(results) if results else "No metadata matches your query."
"credits": 3
}
return f"Course metadata: {meta}"
# --------------------------- # --------------------- Backend ---------------------
# Backend setup
# ---------------------------
backend = CompositeBackend([ backend = CompositeBackend([
LocalShellBackend(workspace_dir="./workspace"), LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(), FilesystemBackend(),
]) ])
# --------------------------- # --------------------- Agent ---------------------
# Agent creation
# ---------------------------
agent = create_deep_agent( agent = create_deep_agent(
model=llm, model=llm,
tools=[search_course_docs, fetch_course_meta], tools=[search_course_docs, fetch_course_meta],
backend=backend, 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.", system_prompt=(
"You are a helpful FAQ bot for the course.\n"
"Use the search_course_docs tool for questions about lecture materials.\n"
"Use the fetch_course_meta tool for questions about schedule or metadata.\n"
"Do not call both tools unless absolutely necessary.\n"
"In your final answer, prefix the source with 'source: chroma' or 'source: mcp_meta'."
),
) )
# --------------------------- # --------------------- Data Loading ---------------------
# CLI async def load_faq_to_chroma():
# --------------------------- """Load all .md files from data/ into the Chroma collection."""
if not DATA_DIR.exists():
print("Data directory not found.")
return
docs = []
for md_file in DATA_DIR.glob("*.md"):
text = md_file.read_text(encoding="utf-8")
docs.append(Document(page_content=text, metadata={"title": md_file.stem}))
if docs:
vector_store.add_documents(docs)
vector_store.persist()
print(f"Loaded {len(docs)} documents into Chroma.")
else:
print("No markdown files found in data/.")
# --------------------- CLI ---------------------
PRESET_QUESTIONS = [ PRESET_QUESTIONS = [
"What topics are covered in Lecture 1?", # should hit chroma "What is the deadline for the final project?", # chroma
"Explain the concept of tokenization in NLP.", # chroma "Explain the concept of tokenization in NLP.", # chroma
"When is the next lab session?", # should hit mcp_meta "When is the next lecture scheduled?", # mcp_meta
] ]
async def run_agent(question: str, thread_id: str = "session-1"): async def run_cli():
result = await agent.ainvoke( await load_faq_to_chroma()
{"messages": [HumanMessage(content=question)]}, print("\n--- FAQ Bot CLI ---\n")
{"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): for i, q in enumerate(PRESET_QUESTIONS, 1):
await run_agent(q, thread_id=f"preset-{i}") print(f"{i}. {q}")
print("\nEnter your own question (or 'exit' to quit):")
# Interactive mode
print("Enter your own questions (type 'exit' to quit):")
while True: while True:
user_input = input("> ") user_input = input("> ")
if user_input.lower() in {"exit", "quit"}: if user_input.lower() in {"exit", "quit"}:
break break
await run_agent(user_input, thread_id="interactive") result = await agent.ainvoke(
{"messages": [HumanMessage(content=user_input)]},
{"configurable": {"thread_id": "session-1"}},
)
print(result["messages"][-1].content)
if __name__ == "__main__": if __name__ == "__main__":
asyncio.run(main()) asyncio.run(run_cli())