Update agent.py

agent.py

@@ -11,7 +11,6 @@ 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()
@@ -37,7 +36,7 @@ TOOLS = [
 
 
 # -------------------------------------------------
-# LLM INIT
+# LLM INIT (NO SYSTEM PROMPT HERE)
 # -------------------------------------------------
 rate_limiter = InMemoryRateLimiter(
     requests_per_second=7 / 60,
@@ -47,17 +46,18 @@ rate_limiter = InMemoryRateLimiter(
 
 llm = init_chat_model(
     model_provider="google_genai",
-    model="gemini-2.5-flash-lite",
+    model="gemini-2.5-flash",
     rate_limiter=rate_limiter
 ).bind_tools(TOOLS)
 
 
 
 # -------------------------------------------------
-# SYSTEM PROMPT
+# SYSTEM PROMPT (WILL BE INSERTED ONLY ONCE)
 # -------------------------------------------------
 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,46 +66,40 @@ Your job is to:
 5. Follow new URLs until none remain, then output END.
 
 Rules:
-- For base64 generation NEVER use your own code — use "encode_image_to_base64"
+- For base64 generation of an image NEVER use your own code, always use the "encode_image_to_base64" tool that's provided
 - Never hallucinate URLs or fields.
 - Never shorten endpoints.
 - Always inspect server response.
 - Never stop early.
-- Use tools for HTML, downloading, rendering, OCR, running code.
+- Use tools for HTML, downloading, rendering, OCR, or running code.
 - Include:
     email = {EMAIL}
     secret = {SECRET}
 """
 
+
 # -------------------------------------------------
 # AGENT NODE
 # -------------------------------------------------
 def agent_node(state: AgentState):
-    """Fixes: KeyError, offset None, float(None)"""
-    
+    # time-handling
     cur_time = time.time()
-    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
-
+    cur_url = os.getenv("url")
+    prev_time = url_time[cur_url]
+    offset = os.getenv("offset")
     if prev_time is not None:
         prev_time = float(prev_time)
         diff = cur_time - prev_time
 
-        # 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):
+        if diff >= 180 or (offset != "0" and (cur_time - float(offset)) > 90):
             print("Timeout exceeded — instructing LLM to purposely submit wrong answer.", diff, "Offset=", offset)
 
             fail_instruction = """
-            You exceeded the time limit (130s).
-            Immediately call `post_request` and submit a WRONG answer for the CURRENT quiz.
+            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.
             """
 
+            # LLM will figure out the right endpoint + JSON structure itself
             result = llm.invoke([
                 {"role": "user", "content": fail_instruction}
             ])
@@ -117,18 +111,20 @@ def agent_node(state: AgentState):
         strategy="last",
         include_system=True,
         start_on="human",
-        token_counter=llm,
+        token_counter=llm,  # Use the LLM to count actual tokens, not just list length
     )
-
+    
     result = llm.invoke(trimmed_messages)
-    return {"messages": [result]}
 
+    return {"messages": [result]}
 
 # -------------------------------------------------
-# ROUTE
+# ROUTE LOGIC (YOURS WITH MINOR SAFETY IMPROVES)
 # -------------------------------------------------
 def route(state):
     last = state["messages"][-1]
+    # print("=== ROUTE DEBUG: last message type ===")
+
     tool_calls = getattr(last, "tool_calls", None)
 
     if tool_calls:
@@ -137,17 +133,6 @@ def route(state):
 
     content = getattr(last, "content", None)
 
-    # 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
-
     if isinstance(content, str) and content.strip() == "END":
         return END
 
@@ -159,6 +144,7 @@ def route(state):
     return "agent"
 
 
+
 # -------------------------------------------------
 # GRAPH
 # -------------------------------------------------
@@ -166,20 +152,17 @@ 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()
 
 
@@ -188,31 +171,37 @@ 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 content
+        parsed = json.loads(content) if isinstance(content, str) else {}
         if parsed.get("url") is None:
             print("Tasks completed successfully!")
             return
 
     except Exception:
-        pass
+        pass  # fallback below
 
+    # Default fallback
     print("Tasks completed successfully!")