129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
import os
|
|
import asyncio
|
|
import json
|
|
from typing import TypedDict
|
|
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.graph import StateGraph, START, END
|
|
|
|
from deepagents import create_deep_agent, tool
|
|
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
|
|
|
|
# LLM configuration - OpenRouter
|
|
llm = ChatOpenAI(
|
|
model="openai/gpt-oss-20b:free",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
temperature=0.0,
|
|
)
|
|
|
|
# State definition
|
|
class ReflectState(TypedDict):
|
|
question: str
|
|
draft: str
|
|
critique: str
|
|
verdict: str # ok | needs_revision
|
|
round: int
|
|
max_rounds: int
|
|
|
|
# Node: draft_answer
|
|
def draft_answer(state: ReflectState) -> ReflectState:
|
|
prompt = f"Write a concise answer (5-10 sentences) to the following question:\n\n{state['question']}"
|
|
response = llm.invoke([HumanMessage(content=prompt)])
|
|
state["draft"] = response.content.strip()
|
|
return state
|
|
|
|
# Node: reflect
|
|
def reflect(state: ReflectState) -> ReflectState:
|
|
prompt = f"""You are a critic evaluating the following draft answer. Provide a verdict ('ok' or 'needs_revision') and 2-3 specific points of improvement. Do not provide the revised answer. Use JSON format:
|
|
{{
|
|
"verdict": "ok" | "needs_revision",
|
|
"critique": "list of points"
|
|
}}
|
|
Draft:
|
|
{state['draft']}"""
|
|
response = llm.invoke([HumanMessage(content=prompt)])
|
|
try:
|
|
data = json.loads(response.content)
|
|
except Exception:
|
|
data = {"verdict": "needs_revision", "critique": "Could not parse critique"}
|
|
state["critique"] = data.get("critique", "")
|
|
state["verdict"] = data.get("verdict", "needs_revision")
|
|
return state
|
|
|
|
# Node: rewrite
|
|
def rewrite(state: ReflectState) -> ReflectState:
|
|
prompt = f"""You are revising the draft answer based on the following critique. Produce a revised answer (5-10 sentences). Do not include the critique. Use the critique points to improve clarity, specificity, and remove filler. Draft:\n{state['draft']}\nCritique:\n{state['critique']}"""
|
|
response = llm.invoke([HumanMessage(content=prompt)])
|
|
state["draft"] = response.content.strip()
|
|
state["round"] = state.get("round", 0) + 1
|
|
return state
|
|
|
|
# Build the graph
|
|
def build_graph() -> StateGraph:
|
|
graph = StateGraph(ReflectState)
|
|
graph.add_node("draft_answer", draft_answer)
|
|
graph.add_node("reflect", reflect)
|
|
graph.add_node("rewrite", rewrite)
|
|
graph.set_entry_point("draft_answer")
|
|
graph.add_edge("draft_answer", "reflect")
|
|
graph.add_conditional_edges(
|
|
"reflect",
|
|
lambda s: "ok" if s["verdict"] == "ok" else ("rewrite" if s["round"] < s["max_rounds"] else "END"),
|
|
{
|
|
"ok": END,
|
|
"rewrite": "rewrite",
|
|
"END": END,
|
|
},
|
|
)
|
|
graph.add_edge("rewrite", "reflect")
|
|
return graph
|
|
|
|
# Tool that runs the graph
|
|
def answer_question_tool(question: str, max_rounds: int = 2) -> str:
|
|
graph = build_graph()
|
|
initial_state: ReflectState = {
|
|
"question": question,
|
|
"draft": "",
|
|
"critique": "",
|
|
"verdict": "",
|
|
"round": 0,
|
|
"max_rounds": max_rounds,
|
|
}
|
|
final_state = graph.invoke(initial_state)
|
|
return final_state["draft"]
|
|
|
|
# DeepAgent tool
|
|
@tool
|
|
def answer_question(query: str) -> str:
|
|
"""Answer a question using a self-reflective process."""
|
|
return answer_question_tool(query)
|
|
|
|
# Backend for DeepAgent
|
|
backend = CompositeBackend(
|
|
[
|
|
LocalShellBackend(workspace_dir="./workspace"),
|
|
FilesystemBackend(),
|
|
]
|
|
)
|
|
|
|
# Create the DeepAgent
|
|
agent = create_deep_agent(
|
|
model=llm,
|
|
tools=[answer_question],
|
|
backend=backend,
|
|
system_prompt="You are an assistant that answers questions using a self-reflective process. Use the tool 'answer_question' to answer the question.",
|
|
)
|
|
|
|
# CLI demo
|
|
async def main():
|
|
question = "Объясни студенту разницу между tool и resource в MCP"
|
|
result = await agent.ainvoke(
|
|
{"messages": [HumanMessage(content=question)]},
|
|
{"configurable": {"thread_id": "session-1"}},
|
|
)
|
|
print(result["messages"][-1].content)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |