Update agent.py

agent.py

@@ -1,18 +1,22 @@
-from langgraph.graph import StateGraph, END, START
-from shared_store import url_time
+import os
 import time
-from langchain_core.rate_limiters import InMemoryRateLimiter
+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 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 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
+from shared_store import url_time
+
 load_dotenv()
 
 EMAIL = os.getenv("EMAIL")
@@ -21,110 +25,161 @@ 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
+    run_code, get_rendered_html, download_file, post_request,
+    add_dependencies, ocr_image_tool, transcribe_audio, encode_image_to_base64
 ]
 
+# ==============================================================  
+# FALLBACK LLM  
+# ==============================================================
 
-# -------------------------------------------------
-# LLM INIT (NO SYSTEM PROMPT HERE)
-# -------------------------------------------------
-rate_limiter = InMemoryRateLimiter(
-    requests_per_second=7 / 60,
-    check_every_n_seconds=1,
-    max_bucket_size=7
-)
+FALLBACK_MODELS = [
+    "gemini-2.5-flash",
+    "gemini-2.5-flash-lite",
+    "gemini-2.0-flash-lite",
+    "gemini-2.0-flash",
+    
+]
+
+def init_llm_with_fallback(tools):
+    """Initialize an LLM with automatic fallback selection."""
 
-llm = init_chat_model(
-    model_provider="google_genai",
-    model="gemini-2.5-flash",
-    rate_limiter=rate_limiter
-).bind_tools(TOOLS)
+    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
+        ])
+
+        if trigger_fallback:
+            print("\n⚠️ Quota/Rate error → switching LLM model...\n")
+            llm = init_llm_with_fallback(TOOLS)
+            return llm.invoke(input_message)
+
+        raise e
+
+# ==============================================================  
+# 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.
-3. Solve tasks exactly.
+Your job:
+1. Load each quiz page.
+2. Extract instructions, parameters & submit endpoint.
+3. Solve tasks EXACTLY.
 4. Submit answers ONLY to the correct endpoint.
-5. Follow new URLs until none remain, then output END.
+5. Follow returned 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
-- Never hallucinate URLs or fields.
-- Never shorten endpoints.
+- NEVER generate your own base64—use encode_image_to_base64 tool.
+- NEVER hallucinate endpoints.
+- NEVER shorten URLs.
 - Always inspect server response.
-- Never stop early.
-- Use tools for HTML, downloading, rendering, OCR, or running code.
-- Include:
-    email = {EMAIL}
-    secret = {SECRET}
+- Use tools for HTML, code execution, OCR, downloading, etc.
+Include in every submission:
+email = {EMAIL}
+secret = {SECRET}
 """
 
+# ==============================================================  
+# AGENT NODE  
+# ==============================================================
 
-# -------------------------------------------------
-# AGENT NODE
-# -------------------------------------------------
 def agent_node(state: AgentState):
-    # time-handling
+
+    # ---- TIMEOUT ----
     cur_time = time.time()
     cur_url = os.getenv("url")
-    prev_time = url_time[cur_url]
-    offset = os.getenv("offset")
+    prev_time = url_time.get(cur_url)
+    offset = float(os.getenv("offset", "0"))
+
     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):
-            print("Timeout exceeded — instructing LLM to purposely submit wrong answer.", diff, "Offset=", offset)
+        if diff >= 180 or (offset != 0 and (cur_time - offset) > 90):
+            print("Timeout exceeded — forcing WRONG submission.", diff)
 
-            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.
+            instruction = """
+            You exceeded allowed time.
+            Immediately call post_request with 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}
-            ])
-            return {"messages": [result]}
+            return {"messages": [safe_llm_invoke(instruction)]}
+
+    # ---- NORMAL FLOW ----
 
-    trimmed_messages = trim_messages(
+    trimmed = trim_messages(
         messages=state["messages"],
         max_tokens=MAX_TOKENS,
         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 = safe_llm_invoke(trimmed)
     return {"messages": [result]}
 
-# -------------------------------------------------
-# ROUTE LOGIC (YOURS WITH MINOR SAFETY IMPROVES)
-# -------------------------------------------------
+# ==============================================================  
+# ROUTING  
+# ==============================================================
+
 def route(state):
     last = state["messages"][-1]
-    # print("=== ROUTE DEBUG: last message type ===")
-
     tool_calls = getattr(last, "tool_calls", None)
 
     if tool_calls:
@@ -136,73 +191,62 @@ def route(state):
     if isinstance(content, str) and content.strip() == "END":
         return END
 
-    if isinstance(content, list) and len(content) and isinstance(content[0], dict):
+    if isinstance(content, list) and len(content):
         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)
-robust_retry = {
+
+graph.add_node("agent", agent_node, 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  
+# ==============================================================
 
-
-# -------------------------------------------------
-# RUNNER
-# -------------------------------------------------
 def run_agent(url: str):
-    # system message is seeded ONCE here
-    initial_messages = [
+
+    initial = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": url}
     ]
 
-    # run agent and CAPTURE the output
     result = app.invoke(
-        {"messages": initial_messages},
+        {"messages": initial},
         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  # fallback below
+        pass
 
-    # Default fallback
     print("Tasks completed successfully!")
-