Add agent.py

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()
+"