diff --git a/main.py b/main.py new file mode 100644 index 0000000..c2eef31 --- /dev/null +++ b/main.py @@ -0,0 +1,87 @@ +import random +from typing import TypedDict, Dict +from langgraph.graph import StateGraph +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage, AIMessage + +# Define state +class AgentState(TypedDict): + task: str + result: str + attempts: int + status: str # pending | success | failed | max_attempts + error: str | None + max_attempts: int + +# Unreliable tool +async def unreliable_tool(task: str) -> str: + if random.random() < 0.3: + raise ValueError("Tool failure") + # simple eval for arithmetic + try: + return str(eval(task)) + except Exception as e: + raise ValueError(f"Eval error: {e}") + +# Execute task node +async def execute_task(state: AgentState) -> Dict[str, str]: + try: + result = await unreliable_tool(state["task"]) + return {"result": result, "error": None} + except Exception as e: + return {"result": "", "error": str(e)} + +# Verify result node using LLM-as-judge +llm = ChatOpenAI(temperature=0) +async def verify_result(state: AgentState) -> Dict[str, str]: + prompt = f"\nTask: {state['task']}\nResult: {state['result']}\nIs this correct? Answer with 'success' or 'failed'." + msg = await llm.ainvoke([HumanMessage(content=prompt)]) + verdict = msg.content.strip().lower() + if "success" in verdict: + return {"status": "success"} + else: + return {"status": "failed"} + +# Handle error / retry node +async def handle_error(state: AgentState) -> Dict[str, str]: + attempts = state["attempts"] + 1 + if attempts >= state["max_attempts"]: + return {"status": "max_attempts", "attempts": attempts} + else: + return {"attempts": attempts, "status": "pending"} + +# Build 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_conditional_edges( + "execute_task", + lambda x: "verify_result" if x["error"] is None else "handle_error", +) +builder.add_conditional_edges( + "verify_result", + lambda x: "end_success" if x.get("status") == "success" else ( + "end_max_attempts" if x.get("status") == "max_attempts" else "handle_error" + ), +) +builder.add_edge("handle_error", "execute_task") + +builder.set_finish_nodes(["end_success", "end_max_attempts"]) + +graph = builder.compile() + +# Run demo +if __name__ == "__main__": + initial_state: AgentState = { + "task": "2+2", + "result": "", + "attempts": 0, + "status": "pending", + "error": None, + "max_attempts": 5, + } + result = graph.invoke(initial_state) + print("Final state:", result)