Update agent.py

agent.py

@@ -1,22 +1,18 @@
-import os
+from langgraph.graph import StateGraph, END, START
+from shared_store import url_time
 import time
-import json
-from dotenv import load_dotenv
-from typing import TypedDict, Annotated, List
-
-from langgraph.graph import StateGraph, START, END
-from langgraph.prebuilt import ToolNode
-from langgraph.graph.message import add_messages
-from langchain.chat_models import init_chat_model
 from langchain_core.rate_limiters import InMemoryRateLimiter
-from langchain_core.messages import trim_messages
-
+from langgraph.prebuilt import ToolNode
 from tools import (
-    get_rendered_html, download_file, post_request, run_code,
-    add_dependencies, ocr_image_tool, transcribe_audio, encode_image_to_base64
+    get_rendered_html, download_file, post_request,
+    run_code, add_dependencies, ocr_image_tool, transcribe_audio, encode_image_to_base64
 )
-from shared_store import url_time
-
+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
+import os
+from dotenv import load_dotenv
 load_dotenv()
 
 EMAIL = os.getenv("EMAIL")
@@ -25,161 +21,110 @@ SECRET = os.getenv("SECRET")
 RECURSION_LIMIT = 5000
 MAX_TOKENS = 180000
 
-# ==============================================================  
-# STATE  
-# ==============================================================  
 
+# -------------------------------------------------
+# STATE
+# -------------------------------------------------
 class AgentState(TypedDict):
     messages: Annotated[List, add_messages]
 
-TOOLS = [
-    run_code, get_rendered_html, download_file, post_request,
-    add_dependencies, ocr_image_tool, transcribe_audio, encode_image_to_base64
-]
-
-# ==============================================================  
-# FALLBACK LLM  
-# ==============================================================
 
-FALLBACK_MODELS = [
-    "gemini-2.5-flash",
-    "gemini-2.5-flash-lite",
-    "gemini-2.0-flash-lite",
-    "gemini-2.0-flash",
-    
+TOOLS = [
+    run_code, get_rendered_html, download_file,
+    post_request, add_dependencies, ocr_image_tool, transcribe_audio, encode_image_to_base64
 ]
 
-def init_llm_with_fallback(tools):
-    """Initialize an LLM with automatic fallback selection."""
 
-    rate_limiter = InMemoryRateLimiter(
-        requests_per_second=7/60,
-        check_every_n_seconds=1,
-        max_bucket_size=7
-    )
-
-    for model_name in FALLBACK_MODELS:
-        try:
-            print(f"[LLM] Trying model: {model_name}")
-            llm = init_chat_model(
-                model_provider="google_genai",
-                model=model_name,
-                rate_limiter=rate_limiter
-            ).bind_tools(tools)
-
-            llm.invoke("ping")  # probe
-            print(f"[LLM] Model ready: {model_name}")
-            return llm
-
-        except Exception as e:
-            print(f"[LLM] Model failed ({model_name}): {e}")
-
-    raise RuntimeError("❌ No Gemini model available!")
-
-# global LLM
-llm = init_llm_with_fallback(TOOLS)
-
-# ==============================================================  
-# SAFE INVOKE (fallback switcher)  
-# ==============================================================
-
-def safe_llm_invoke(input_message):
-    global llm
-    try:
-        return llm.invoke(input_message)
-
-    except Exception as e:
-        err = str(e).lower()
-
-        trigger_fallback = any([
-            "429" in err,
-            "quota" in err,
-            "exceeded" in err,
-            "rate" in err,
-            "unavailable" in err,
-            "deadline" in err,
-            "resourceexhausted" in err
-        ])
+# -------------------------------------------------
+# LLM INIT (NO SYSTEM PROMPT HERE)
+# -------------------------------------------------
+rate_limiter = InMemoryRateLimiter(
+    requests_per_second=7 / 60,
+    check_every_n_seconds=1,
+    max_bucket_size=7
+)
 
-        if trigger_fallback:
-            print("\n⚠️ Quota/Rate error → switching LLM model...\n")
-            llm = init_llm_with_fallback(TOOLS)
-            return llm.invoke(input_message)
+llm = init_chat_model(
+    model_provider="google_genai",
+    model="gemini-2.5-flash",
+    rate_limiter=rate_limiter
+).bind_tools(TOOLS)
 
-        raise e
 
-# ==============================================================  
-# SYSTEM PROMPT  
-# ==============================================================
 
+# -------------------------------------------------
+# SYSTEM PROMPT (WILL BE INSERTED ONLY ONCE)
+# -------------------------------------------------
 SYSTEM_PROMPT = f"""
 You are an autonomous quiz-solving agent.
 
-Your job:
-1. Load each quiz page.
-2. Extract instructions, parameters & submit endpoint.
-3. Solve tasks EXACTLY.
+Your job is to:
+1. Load each quiz page from the given URL.
+2. Extract instructions, parameters, and submit endpoint.
+3. Solve tasks exactly.
 4. Submit answers ONLY to the correct endpoint.
-5. Follow returned URLs until none remain, then output END.
+5. Follow new URLs until none remain, then output END.
 
 Rules:
-- NEVER generate your own base64—use encode_image_to_base64 tool.
-- NEVER hallucinate endpoints.
-- NEVER shorten URLs.
+- 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.
-- Use tools for HTML, code execution, OCR, downloading, etc.
-Include in every submission:
-email = {EMAIL}
-secret = {SECRET}
+- Never stop early.
+- Use tools for HTML, downloading, rendering, OCR, or running code.
+- Include:
+    email = {EMAIL}
+    secret = {SECRET}
 """
 
-# ==============================================================  
-# AGENT NODE  
-# ==============================================================
 
+# -------------------------------------------------
+# AGENT NODE
+# -------------------------------------------------
 def agent_node(state: AgentState):
-
-    # ---- TIMEOUT ----
+    # time-handling
     cur_time = time.time()
     cur_url = os.getenv("url")
-    prev_time = url_time.get(cur_url)
-    offset = float(os.getenv("offset", "0"))
-
+    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
 
-        if diff >= 180 or (offset != 0 and (cur_time - offset) > 90):
-            print("Timeout exceeded — forcing WRONG submission.", diff)
+        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)
 
-            instruction = """
-            You exceeded allowed time.
-            Immediately call post_request with a WRONG answer 
-            for the CURRENT quiz.
+            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.
             """
 
-            return {"messages": [safe_llm_invoke(instruction)]}
-
-    # ---- NORMAL FLOW ----
+            # LLM will figure out the right endpoint + JSON structure itself
+            result = llm.invoke([
+                {"role": "user", "content": fail_instruction}
+            ])
+            return {"messages": [result]}
 
-    trimmed = trim_messages(
+    trimmed_messages = trim_messages(
         messages=state["messages"],
         max_tokens=MAX_TOKENS,
         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)
 
-    result = safe_llm_invoke(trimmed)
     return {"messages": [result]}
 
-# ==============================================================  
-# ROUTING  
-# ==============================================================
-
+# -------------------------------------------------
+# 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:
@@ -191,62 +136,72 @@ def route(state):
     if isinstance(content, str) and content.strip() == "END":
         return END
 
-    if isinstance(content, list) and len(content):
+    if isinstance(content, list) and len(content) and isinstance(content[0], dict):
         if content[0].get("text", "").strip() == "END":
             return END
 
     print("Route → agent")
     return "agent"
 
-# ==============================================================  
-# GRAPH BUILD  
-# ==============================================================
 
+
+# -------------------------------------------------
+# GRAPH
+# -------------------------------------------------
 graph = StateGraph(AgentState)
+
 graph.add_node("tools", ToolNode(TOOLS))
+
 graph.add_edge(START, "agent")
 graph.add_edge("tools", "agent")
 graph.add_conditional_edges("agent", route)
-
-graph.add_node("agent", agent_node, retry={
+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()
 
-# ==============================================================  
-# RUN AGENT  
-# ==============================================================
 
-def run_agent(url: str):
 
-    initial = [
+# -------------------------------------------------
+# 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": 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 {}
         if parsed.get("url") is None:
             print("Tasks completed successfully!")
             return
 
     except Exception:
-        pass
+        pass  # fallback below
 
-    print("Tasks completed successfully!")
+    # Default fallback
+    print("Tasks completed successfully!")