Publishing solution for task 6a1864fa8a94f887e50d46f0: add main.py

This commit is contained in:
2026-06-04 12:33:48 +00:00
commit f5d34215f4
+115
View File
@@ -0,0 +1,115 @@
"""
Selfcorrecting LangGraph agent demo.
Run with:
python main.py
Requires:
pip install langgraph langchain-openai
"""
import random
from typing import TypedDict, Dict
# ---------- State definition ----------
class AgentState(TypedDict):
task: str
result: str | None
attempts: int
status: str # pending | success | failed | max_attempts
error: str | None
max_attempts: int
# ---------- Unreliable tool ----------
class UnreliableTool:
def run(self, input_: str) -> str:
if random.random() < 0.3:
raise ValueError("Simulated tool failure")
# simple eval for demo purposes
try:
return str(eval(input_))
except Exception as e:
raise ValueError(f"Eval error: {e}")
# ---------- LangGraph imports ----------
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
# ---------- Nodes ----------
def execute_task(state: AgentState) -> AgentState:
tool = UnreliableTool()
try:
result = tool.run(state["task"])
state["result"] = result
state["error"] = None
except Exception as e:
state["result"] = None
state["error"] = str(e)
return state
def verify_result(state: AgentState) -> AgentState:
# Ask LLM to judge success or failed based on result and error
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
prompt = (
f"Task: {state['task']}\n"
f"Result: {state.get('result')}\n"
f"Error: {state.get('error')}\n"
"Respond with only 'success' or 'failed'."
)
resp = llm.invoke(prompt)
verdict = resp.content.strip().lower()
if verdict not in ("success", "failed"):
verdict = "failed"
state["verdict"] = verdict
return state
def handle_error(state: AgentState) -> AgentState:
state["attempts"] += 1
if state["attempts"] >= state["max_attempts"]:
state["status"] = "max_attempts"
else:
state["status"] = "pending"
return state
# ---------- Graph ----------
builder = StateGraph(AgentState)
builder.add_node("execute_task", execute_task)
builder.add_node("verify_result", verify_result)
builder.add_node("handle_error", handle_error)
builder.set_entry_point("execute_task")
builder.add_edge("execute_task", "verify_result")
# From verify_result, branch on verdict
builder.add_conditional_edges(
"verify_result",
lambda state: state.get("verdict"),
{
"success": END,
"failed": "handle_error",
},
)
builder.add_edge("handle_error", "execute_task")
graph = builder.compile()
# ---------- Demo runner ----------
if __name__ == "__main__":
initial_state: AgentState = {
"task": "2+2",
"result": None,
"attempts": 0,
"status": "pending",
"error": None,
"max_attempts": 5,
}
for attempt in range(1, initial_state["max_attempts"] + 1):
print(f"Attempt {attempt}:")
result = graph.invoke(initial_state)
if result.get("verdict") == "success":
print(f"Success: {result['result']}")
break
else:
print(f"Failed, error: {result.get('error')}")
else:
print("Reached max attempts without success.")
""