Update agent.py

agent.py

@@ -11,6 +11,7 @@ from typing import TypedDict, Annotated, List
 from langchain_core.messages import trim_messages
 from langchain.chat_models import init_chat_model
 from langgraph.graph.message import add_messages
+from langgraph.pregel.retry import RetryPolicy
 import os
 from dotenv import load_dotenv
 load_dotenv()
@@ -36,7 +37,7 @@ TOOLS = [
 
 
 # -------------------------------------------------
-# LLM INIT (NO SYSTEM PROMPT HERE)
+# LLM INIT
 # -------------------------------------------------
 rate_limiter = InMemoryRateLimiter(
     requests_per_second=7 / 60,
@@ -53,11 +54,10 @@ llm = init_chat_model(
 
 
 # -------------------------------------------------
-# SYSTEM PROMPT (WILL BE INSERTED ONLY ONCE)
+# SYSTEM PROMPT
 # -------------------------------------------------
 SYSTEM_PROMPT = f"""
 You are an autonomous quiz-solving agent.
-
 Your job is to:
 1. Load each quiz page from the given URL.
 2. Extract instructions, parameters, and submit endpoint.
@@ -66,40 +66,46 @@ Your job is to:
 5. Follow new URLs until none remain, then output END.
 
 Rules:
-- For base64 generation of an image NEVER use your own code, always use the "encode_image_to_base64" tool that's provided
+- For base64 generation NEVER use your own code — use "encode_image_to_base64"
 - Never hallucinate URLs or fields.
 - Never shorten endpoints.
 - Always inspect server response.
 - Never stop early.
-- Use tools for HTML, downloading, rendering, OCR, or running code.
+- Use tools for HTML, downloading, rendering, OCR, running code.
 - Include:
     email = {EMAIL}
     secret = {SECRET}
 """
 
-
 # -------------------------------------------------
 # AGENT NODE
 # -------------------------------------------------
 def agent_node(state: AgentState):
-    # time-handling
+    """Fixes: KeyError, offset None, float(None)"""
+    
     cur_time = time.time()
-    cur_url = os.getenv("url")
-    prev_time = url_time[cur_url]
-    offset = os.getenv("offset")
+    cur_url = os.getenv("url") or ""    # FIX 1: safe load
+    prev_time = url_time.get(cur_url)   # FIX 2: no KeyError
+    offset = os.getenv("offset") or "0" # FIX 3: no None issues
+
     if prev_time is not None:
         prev_time = float(prev_time)
         diff = cur_time - prev_time
 
-        if diff >= 180 or (offset != "0" and (cur_time - float(offset)) > 90):
+        # timeout logic unchanged, only made safe
+        try:
+            offset_f = float(offset)
+        except:
+            offset_f = 0.0
+
+        if diff >= 180 or (offset_f != 0 and (cur_time - offset_f) > 90):
             print("Timeout exceeded — instructing LLM to purposely submit wrong answer.", diff, "Offset=", offset)
 
             fail_instruction = """
-            You have exceeded the time limit for this task (over 130 seconds).
-            Immediately call the `post_request` tool and submit a WRONG answer for the CURRENT quiz.
+            You exceeded the time limit (130s).
+            Immediately call `post_request` and submit a WRONG answer for the CURRENT quiz.
             """
 
-            # LLM will figure out the right endpoint + JSON structure itself
             result = llm.invoke([
                 {"role": "user", "content": fail_instruction}
             ])
@@ -111,15 +117,15 @@ def agent_node(state: AgentState):
         strategy="last",
         include_system=True,
         start_on="human",
-        token_counter=llm,  # Use the LLM to count actual tokens, not just list length
+        token_counter=llm,
     )
-    
-    result = llm.invoke(trimmed_messages)
 
+    result = llm.invoke(trimmed_messages)
     return {"messages": [result]}
 
+
 # -------------------------------------------------
-# ROUTE LOGIC (YOURS WITH MINOR SAFETY IMPROVES)
+# ROUTE
 # -------------------------------------------------
 def route(state):
     last = state["messages"][-1]
@@ -131,7 +137,13 @@ def route(state):
 
     content = getattr(last, "content", None)
 
-    # NEW FIX
+    # allow message dicts (post_request returns dicts)
+    if isinstance(content, dict):
+        if content.get("url") == "" or content.get("correct") is False:
+            # not final, continue agent
+            print("Route → agent (dict content)")
+            return "agent"
+
     if content is None:
         print("Content is None → END")
         return END
@@ -147,8 +159,6 @@ def route(state):
     return "agent"
 
 
-
-
 # -------------------------------------------------
 # GRAPH
 # -------------------------------------------------
@@ -156,17 +166,20 @@ graph = StateGraph(AgentState)
 
 graph.add_node("tools", ToolNode(TOOLS))
 
+# FIX 4 — LangGraph retry policy MUST be RetryPolicy(...)
+retry_policy = RetryPolicy(
+    max_attempts=10,
+    initial_interval=1,
+    backoff_factor=2,
+    max_interval=60
+)
+
+graph.add_node("agent", agent_node, retry=retry_policy)
+
 graph.add_edge(START, "agent")
 graph.add_edge("tools", "agent")
 graph.add_conditional_edges("agent", route)
-robust_retry = {
-    "initial_interval": 1,
-    "backoff_factor": 2,
-    "max_interval": 60,
-    "max_attempts": 10
-}
-
-graph.add_node("agent", agent_node, retry=robust_retry)
+
 app = graph.compile()
 
 
@@ -175,37 +188,31 @@ app = graph.compile()
 # RUNNER
 # -------------------------------------------------
 def run_agent(url: str):
-    # system message is seeded ONCE here
     initial_messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": url}
     ]
 
-    # run agent and CAPTURE the output
     result = app.invoke(
         {"messages": initial_messages},
         config={"recursion_limit": RECURSION_LIMIT}
     )
 
-    # Try to detect final server response if present
     try:
         last = result["messages"][-1]
         content = getattr(last, "content", "")
 
-        # If LLM already output END – good
         if isinstance(content, str) and content.strip() == "END":
             print("Tasks completed successfully!")
             return
 
-        # If the last content is JSON from server submission
         import json
-        parsed = json.loads(content) if isinstance(content, str) else {}
+        parsed = json.loads(content) if isinstance(content, str) else content
         if parsed.get("url") is None:
             print("Tasks completed successfully!")
             return
 
     except Exception:
-        pass  # fallback below
+        pass
 
-    # Default fallback
     print("Tasks completed successfully!")