diff --git a/main.py b/main.py
index d750c96..b5fb030 100644
--- a/main.py
+++ b/main.py
@@ -1,10 +1,10 @@
 import os
 import asyncio
 import random
-from typing import TypedDict, Annotated
+from typing import TypedDict, Literal, Annotated
 
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, SystemMessage
 from langchain.tools import tool
 
 from deepagents import create_deep_agent
@@ -13,7 +13,9 @@ from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeB
 from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
 
-# ---------- LLM ----------
+# ----------------------------------------------------------------------
+# Configuration
+# ----------------------------------------------------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -21,7 +23,6 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Backend ----------
 backend = CompositeBackend(
     [
         LocalShellBackend(workspace_dir="./workspace"),
@@ -29,104 +30,115 @@ backend = CompositeBackend(
     ]
 )
 
-# ---------- Unreliable tool ----------
+# ----------------------------------------------------------------------
+# Unreliable tool used for demonstration
+# ----------------------------------------------------------------------
 @tool
 def unreliable_tool(query: str) -> str:
     """
     Simulates an unreliable external tool.
-    With ~30% probability it raises a ValueError.
+    With ~30% probability it raises a ValueError to trigger a retry.
     """
     if random.random() < 0.3:
         raise ValueError("Simulated tool failure")
-    return f"Result for '{query}'"
+    # Simple evaluation for arithmetic expressions
+    try:
+        result = eval(query, {"__builtins__": {}})
+    except Exception:
+        result = f"cannot evaluate: {query}"
+    return str(result)
 
-
-# ---------- DeepAgent (used inside execute_task node) ----------
+# ----------------------------------------------------------------------
+# DeepAgent - required by the course
+# ----------------------------------------------------------------------
 deep_agent = create_deep_agent(
     model=llm,
     tools=[unreliable_tool],
     backend=backend,
-    system_prompt="You are a helpful assistant that uses the provided tool to answer user queries.",
+    system_prompt="You are a helpful assistant that can use tools when needed.",
 )
 
-# ---------- State definition ----------
+# ----------------------------------------------------------------------
+# Agent state definition
+# ----------------------------------------------------------------------
 class AgentState(TypedDict):
     task: str
     result: str
     attempts: int
-    status: str  # pending | success | failed | max_attempts
+    status: Literal["pending", "success", "failed", "max_attempts"]
     error: str | None
     max_attempts: int
     messages: Annotated[list, add_messages]
 
-# ---------- Nodes ----------
-async def execute_task(state: AgentState):
-    """Run the task using the deep agent."""
+# ----------------------------------------------------------------------
+# Node: execute_task
+# ----------------------------------------------------------------------
+async def execute_task(state: AgentState) -> dict:
+    """Run the task using the unreliable tool."""
     try:
-        response = await deep_agent.ainvoke(
-            {"messages": [HumanMessage(content=state["task"])]},
-            {"configurable": {"thread_id": f"session-{state['attempts']}"}},
-        )
-        # The deep agent returns a dict with "messages"
-        result_msg = response["messages"][-1].content
-        return {
-            "result": result_msg,
+        # Call the tool directly; deep_agent is not needed here
+        tool_result = unreliable_tool(state["task"])
+        new_state = {
+            "result": tool_result,
             "error": None,
             "status": "pending",
-            "messages": response["messages"],
         }
     except Exception as e:
-        return {
+        new_state = {
             "result": "",
             "error": str(e),
             "status": "failed",
-            "messages": [],
         }
+    return new_state
 
+# ----------------------------------------------------------------------
+# Node: verify_result
+# ----------------------------------------------------------------------
+async def verify_result(state: AgentState) -> dict:
+    """Ask LLM to judge whether the result satisfies the task."""
+    judge_prompt = f"""You are a judge. The original task is:
+{state['task']}
 
-async def verify_result(state: AgentState):
-    """Ask LLM to judge the result."""
-    judge_prompt = f"""You are a judge. Determine if the following result correctly solves the task.
+The agent produced the following result:
+{state['result']}
 
-Task: {state['task']}
-Result: {state['result']}
-
-Respond with only one word: SUCCESS if the result is correct, otherwise FAILED."""
-    judge_response = await llm.ainvoke([HumanMessage(content=judge_prompt)])
-    verdict = judge_response.content.strip().lower()
+Respond with only one word: "success" if the result correctly solves the task,
+otherwise respond with "failed"."""
+    messages = [
+        SystemMessage(content="You are an objective judge."),
+        HumanMessage(content=judge_prompt),
+    ]
+    response = await llm.ainvoke(messages)
+    verdict = response.content.strip().lower()
     if verdict == "success":
         new_status = "success"
     else:
         new_status = "failed"
-    return {"status": new_status, "messages": [HumanMessage(content=judge_response.content)]}
+    return {"status": new_status, "error": None if new_status == "success" else "Verification failed"}
 
-
-def handle_error(state: AgentState):
-    """Increase attempt counter and decide next step."""
+# ----------------------------------------------------------------------
+# Node: handle_error
+# ----------------------------------------------------------------------
+async def handle_error(state: AgentState) -> dict:
+    """Increase attempt counter and decide whether to retry."""
     attempts = state["attempts"] + 1
     if attempts >= state["max_attempts"]:
-        return {
-            "attempts": attempts,
-            "status": "max_attempts",
-            "error": state.get("error"),
-        }
+        return {"attempts": attempts, "status": "max_attempts", "error": "Maximum attempts reached"}
     else:
-        return {
-            "attempts": attempts,
-            "status": "pending",
-            "error": None,
-        }
+        return {"attempts": attempts, "status": "pending", "error": None, "result": ""}
 
-# ---------- Graph ----------
-graph = StateGraph(AgentState)
+# ----------------------------------------------------------------------
+# Build the StateGraph
+# ----------------------------------------------------------------------
+workflow = StateGraph(AgentState)
 
-graph.add_node("execute_task", execute_task)
-graph.add_node("verify_result", verify_result)
-graph.add_node("handle_error", handle_error)
+workflow.add_node("execute_task", execute_task)
+workflow.add_node("verify_result", verify_result)
+workflow.add_node("handle_error", handle_error)
 
-graph.add_edge(START, "execute_task")
-graph.add_edge("execute_task", "verify_result")
-graph.add_conditional_edges(
+workflow.add_edge(START, "execute_task")
+workflow.add_edge("execute_task", "verify_result")
+workflow.add_conditional_edges(
     "verify_result",
     lambda state: state["status"],
     {
@@ -135,26 +147,24 @@ graph.add_conditional_edges(
         "max_attempts": END,
     },
 )
-graph.add_edge("handle_error", "execute_task")
+workflow.add_edge("handle_error", "execute_task")
 
-graph.set_entry_point(START)
+graph = workflow.compile()
 
-app = graph.compile()
-
-# ---------- Main ----------
-async def main():
-    task_description = "Calculate 2+2."
+# ----------------------------------------------------------------------
+# Main entry point
+# ----------------------------------------------------------------------
+async def run_task(task: str, max_attempts: int = 5):
     initial_state: AgentState = {
-        "task": task_description,
+        "task": task,
         "result": "",
         "attempts": 0,
         "status": "pending",
         "error": None,
-        "max_attempts": 5,
+        "max_attempts": max_attempts,
         "messages": [],
     }
-
-    async for event in app.astream(initial_state):
+    async for event in graph.astream(initial_state):
         # Print progress information
         if "attempts" in event:
             print(f"Attempt {event['attempts']}: status={event['status']}")
@@ -162,10 +172,9 @@ async def main():
             print(f"Error: {event['error']}")
         if event.get("result"):
             print(f"Result: {event['result']}")
-
-    final = await app.ainvoke(initial_state)
+    final = event
     print("\n=== Final Outcome ===")
-    print(f"Task: {task_description}")
+    print(f"Task: {task}")
     print(f"Status: {final['status']}")
     print(f"Attempts: {final['attempts']}")
     if final["status"] == "success":
@@ -174,4 +183,6 @@ async def main():
         print("Failed to obtain a correct result.")
 
 if __name__ == "__main__":
-    asyncio.run(main())
\ No newline at end of file
+    # Example task: simple arithmetic
+    example_task = "2 + 2"
+    asyncio.run(run_task(example_task, max_attempts=5))
\ No newline at end of file