diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..e6e7129 --- /dev/null +++ b/agent.py @@ -0,0 +1,195 @@ +"""Self‑correcting LangGraph agent. + +The agent receives a *task* string. It executes the task via an +`unreliable_tool` that sometimes raises a ValueError. After execution +it asks the LLM (OpenAI or Ollama) to judge whether the *result* is +correct. If the judge says ``failed`` the agent retries until +``max_attempts`` is reached. + +The implementation uses LangGraph's low‑level API: a StateGraph +with three nodes – execute_task, verify_result, handle_error – and a +simple loop. + +Run the script with: + + python agent.py + +It will ask for a task, then show the attempts and final status. +""" + +from __future__ import annotations + +import random +import sys +from typing import TypedDict + +from langchain_openai import ChatOpenAI +from langgraph.graph import StateGraph, END + +# --------------------------------------------------------------------------- +# 1. 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 + +# --------------------------------------------------------------------------- +# 2. Unreliable tool – 30 % chance of raising ValueError +# --------------------------------------------------------------------------- + +def unreliable_tool(task: str) -> str: + """Simulate a tool that sometimes fails. + + The function simply returns ``task`` reversed (as a dummy result) but + raises a ValueError 30 % of the time. + """ + if random.random() < 0.3: + raise ValueError("Simulated tool failure") + return task[::-1] # dummy "computation" + +# --------------------------------------------------------------------------- +# 3. LLM judge – asks for "success" or "failed" +# --------------------------------------------------------------------------- + +llm = ChatOpenAI(temperature=0, model="gpt-4o-mini") # or use Ollama + +async def verify_result(state: AgentState) -> AgentState: + """Ask the LLM whether the result is correct. + + The prompt forces the model to answer only "success" or "failed". + """ + if state["result"] is None: + # Should not happen – guard + state["status"] = "failed" + return state + + prompt = ( + f"Task: {state['task']}\n" + f"Result: {state['result']}\n" + "Is this result correct? Respond with only 'success' or 'failed'." + ) + response = llm.invoke(prompt) + verdict = response.content.strip().lower() + if verdict == "success": + state["status"] = "success" + else: + state["status"] = "failed" + return state + +# --------------------------------------------------------------------------- +# 4. Execute task node +# --------------------------------------------------------------------------- + +async def execute_task(state: AgentState) -> AgentState: + """Run the unreliable tool and capture errors.""" + try: + result = unreliable_tool(state["task"]) + state["result"] = result + state["error"] = None + except Exception as exc: # catch ValueError + state["result"] = None + state["error"] = str(exc) + return state + +# --------------------------------------------------------------------------- +# 5. Handle error / retry node +# --------------------------------------------------------------------------- + +async def handle_error(state: AgentState) -> AgentState: + """Increment attempt counter and decide whether to retry.""" + state["attempts"] += 1 + if state["attempts"] >= state["max_attempts"]: + state["status"] = "max_attempts" + else: + # Reset result and error for next try + state["result"] = None + state["error"] = None + return state + +# --------------------------------------------------------------------------- +# 6. Build the graph +# --------------------------------------------------------------------------- + +graph = StateGraph(AgentState) + +# Add nodes +graph.add_node("execute_task", execute_task) +graph.add_node("verify_result", verify_result) +graph.add_node("handle_error", handle_error) + +# Define edges +# Start → execute_task +graph.set_entry_point("execute_task") + +# After execution, go to verification +graph.add_edge("execute_task", "verify_result") + +# Verification outcomes +# success → END +# failed → check attempts +# max_attempts → END + +# We use a conditional edge on the status field + +def verify_cond(state: AgentState): + return state["status"] + +# Map status to next node +graph.add_conditional_edges( + "verify_result", + verify_cond, + { + "success": END, + "failed": "handle_error", + "max_attempts": END, + }, +) + +# From handle_error back to execute_task +graph.add_edge("handle_error", "execute_task") + +# Compile the graph into a runnable chain +agent = graph.compile() + +# --------------------------------------------------------------------------- +# 7. CLI entry point +# --------------------------------------------------------------------------- + +def main() -> None: + print("Self‑correcting LangGraph agent demo") + task = input("Enter a task: ") + if not task: + print("No task provided. Exiting.") + sys.exit(0) + + # Initial state + state: AgentState = { + "task": task, + "result": None, + "attempts": 1, + "status": "pending", + "error": None, + "max_attempts": 5, + } + + # Run the chain + final_state = agent.invoke(state) + + print("\n--- Result ---") + print(f"Task: {final_state['task']}") + print(f"Attempts: {final_state['attempts']}") + print(f"Status: {final_state['status']}") + if final_state['result']: + print(f"Result: {final_state['result']}") + if final_state['error']: + print(f"Last error: {final_state['error']}") + + +if __name__ == "__main__": + main() +" \ No newline at end of file