177 lines
5.0 KiB
Python
177 lines
5.0 KiB
Python
import os
|
|
import asyncio
|
|
import random
|
|
from typing import TypedDict, Annotated
|
|
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.messages import HumanMessage
|
|
from langchain.tools import tool
|
|
|
|
from deepagents import create_deep_agent
|
|
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
|
|
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.graph.message import add_messages
|
|
|
|
# ---------- LLM ----------
|
|
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,
|
|
)
|
|
|
|
# ---------- Backend ----------
|
|
backend = CompositeBackend(
|
|
[
|
|
LocalShellBackend(workspace_dir="./workspace"),
|
|
FilesystemBackend(),
|
|
]
|
|
)
|
|
|
|
# ---------- Unreliable tool ----------
|
|
@tool
|
|
def unreliable_tool(query: str) -> str:
|
|
"""
|
|
Simulates an unreliable external tool.
|
|
With ~30% probability it raises a ValueError.
|
|
"""
|
|
if random.random() < 0.3:
|
|
raise ValueError("Simulated tool failure")
|
|
return f"Result for '{query}'"
|
|
|
|
|
|
# ---------- DeepAgent (used inside execute_task node) ----------
|
|
deep_agent = create_deep_agent(
|
|
model=llm,
|
|
tools=[unreliable_tool],
|
|
backend=backend,
|
|
system_prompt="You are a helpful assistant that uses the provided tool to answer user queries.",
|
|
)
|
|
|
|
# ---------- State definition ----------
|
|
class AgentState(TypedDict):
|
|
task: str
|
|
result: str
|
|
attempts: int
|
|
status: str # pending | success | failed | max_attempts
|
|
error: str | None
|
|
max_attempts: int
|
|
messages: Annotated[list, add_messages]
|
|
|
|
# ---------- Nodes ----------
|
|
async def execute_task(state: AgentState):
|
|
"""Run the task using the deep agent."""
|
|
try:
|
|
response = await deep_agent.ainvoke(
|
|
{"messages": [HumanMessage(content=state["task"])]},
|
|
{"configurable": {"thread_id": f"session-{state['attempts']}"}},
|
|
)
|
|
# The deep agent returns a dict with "messages"
|
|
result_msg = response["messages"][-1].content
|
|
return {
|
|
"result": result_msg,
|
|
"error": None,
|
|
"status": "pending",
|
|
"messages": response["messages"],
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"result": "",
|
|
"error": str(e),
|
|
"status": "failed",
|
|
"messages": [],
|
|
}
|
|
|
|
|
|
async def verify_result(state: AgentState):
|
|
"""Ask LLM to judge the result."""
|
|
judge_prompt = f"""You are a judge. Determine if the following result correctly solves the task.
|
|
|
|
Task: {state['task']}
|
|
Result: {state['result']}
|
|
|
|
Respond with only one word: SUCCESS if the result is correct, otherwise FAILED."""
|
|
judge_response = await llm.ainvoke([HumanMessage(content=judge_prompt)])
|
|
verdict = judge_response.content.strip().lower()
|
|
if verdict == "success":
|
|
new_status = "success"
|
|
else:
|
|
new_status = "failed"
|
|
return {"status": new_status, "messages": [HumanMessage(content=judge_response.content)]}
|
|
|
|
|
|
def handle_error(state: AgentState):
|
|
"""Increase attempt counter and decide next step."""
|
|
attempts = state["attempts"] + 1
|
|
if attempts >= state["max_attempts"]:
|
|
return {
|
|
"attempts": attempts,
|
|
"status": "max_attempts",
|
|
"error": state.get("error"),
|
|
}
|
|
else:
|
|
return {
|
|
"attempts": attempts,
|
|
"status": "pending",
|
|
"error": None,
|
|
}
|
|
|
|
# ---------- Graph ----------
|
|
graph = StateGraph(AgentState)
|
|
|
|
graph.add_node("execute_task", execute_task)
|
|
graph.add_node("verify_result", verify_result)
|
|
graph.add_node("handle_error", handle_error)
|
|
|
|
graph.add_edge(START, "execute_task")
|
|
graph.add_edge("execute_task", "verify_result")
|
|
graph.add_conditional_edges(
|
|
"verify_result",
|
|
lambda state: state["status"],
|
|
{
|
|
"success": END,
|
|
"failed": "handle_error",
|
|
"max_attempts": END,
|
|
},
|
|
)
|
|
graph.add_edge("handle_error", "execute_task")
|
|
|
|
graph.set_entry_point(START)
|
|
|
|
app = graph.compile()
|
|
|
|
# ---------- Main ----------
|
|
async def main():
|
|
task_description = "Calculate 2+2."
|
|
initial_state: AgentState = {
|
|
"task": task_description,
|
|
"result": "",
|
|
"attempts": 0,
|
|
"status": "pending",
|
|
"error": None,
|
|
"max_attempts": 5,
|
|
"messages": [],
|
|
}
|
|
|
|
async for event in app.astream(initial_state):
|
|
# Print progress information
|
|
if "attempts" in event:
|
|
print(f"Attempt {event['attempts']}: status={event['status']}")
|
|
if event.get("error"):
|
|
print(f"Error: {event['error']}")
|
|
if event.get("result"):
|
|
print(f"Result: {event['result']}")
|
|
|
|
final = await app.ainvoke(initial_state)
|
|
print("\n=== Final Outcome ===")
|
|
print(f"Task: {task_description}")
|
|
print(f"Status: {final['status']}")
|
|
print(f"Attempts: {final['attempts']}")
|
|
if final["status"] == "success":
|
|
print(f"Successful result: {final['result']}")
|
|
else:
|
|
print("Failed to obtain a correct result.")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |