Publishing solution for task 6a1864fa8a94f887e50d46f0: add main.py
This commit is contained in:
@@ -0,0 +1,115 @@
|
|||||||
|
"""
|
||||||
|
Self‑correcting 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.")
|
||||||
|
""
|
||||||
Reference in New Issue
Block a user