Update agent.py

agent.py

@@ -121,7 +121,60 @@ def reverse_text(text: str) -> str:
     return text[::-1]
 
 
-TOOLS = [wiki_search, web_search, fetch_page, run_python, reverse_text]
+@tool
+def youtube_transcript(url: str) -> str:
+    """
+    Fetch the transcript of a YouTube video.
+    Use this for ANY question that references a YouTube URL or asks about
+    what was said, shown, or happened in a video.
+    Provide the full YouTube URL (e.g. https://www.youtube.com/watch?v=XXXX).
+    Returns the transcript text which you can read to answer the question.
+    """
+    try:
+        from youtube_transcript_api import YouTubeTranscriptApi
+        from youtube_transcript_api._errors import TranscriptsDisabled, NoTranscriptFound
+
+        # Extract video ID from various YouTube URL formats
+        match = re.search(
+            r'(?:v=|youtu\.be/|embed/|shorts/)([A-Za-z0-9_-]{11})',
+            url
+        )
+        if not match:
+            return "Could not extract video ID from URL. Check the URL format."
+
+        video_id = match.group(1)
+
+        # Try English first, then any available language
+        try:
+            transcript_list = YouTubeTranscriptApi.get_transcript(
+                video_id, languages=["en"]
+            )
+        except NoTranscriptFound:
+            # Fall back to auto-generated or any other available transcript
+            transcripts = YouTubeTranscriptApi.list_transcripts(video_id)
+            transcript_list = transcripts.find_generated_transcript(
+                ["en", "en-US", "en-GB"]
+            ).fetch()
+
+        if not transcript_list:
+            return "Transcript is empty."
+
+        # Join all text segments into a readable block
+        full_text = " ".join(
+            entry["text"].strip()
+            for entry in transcript_list
+            if entry.get("text", "").strip()
+        )
+
+        return full_text[:8000] if full_text else "Transcript appears empty."
+
+    except TranscriptsDisabled:
+        return "Transcripts are disabled for this video. Try web_search for the video title to find a summary."
+    except Exception as e:
+        return f"Transcript unavailable: {type(e).__name__}: {e}. Try web_search for the video content instead."
+
+
+TOOLS = [wiki_search, web_search, fetch_page, run_python, reverse_text, youtube_transcript]
 
 # ==========================================================
 # MODELS — primary + ordered fallback chain
@@ -153,6 +206,9 @@ SYSTEM_PROMPT = """You are a precise benchmark task solver.
 Produce the exact correct answer — nothing more, nothing less.
 
 ## Tool use
+- Use youtube_transcript for ANY question that includes a YouTube URL or asks
+  about video content (what was said, what happened, species shown, quotes, etc.).
+  Always try youtube_transcript FIRST before web_search for YouTube questions.
 - Use wiki_search for historical facts, biographies, science, geography.
 - Use web_search for recent events, specific articles, prices, or anything time-sensitive.
 - Use fetch_page when a URL is provided or a search result points to a relevant page.
@@ -184,10 +240,10 @@ Produce the exact correct answer — nothing more, nothing less.
    - Vegetables (botanical): true vegetables are leaves (lettuce, spinach), stems (celery), roots
      (carrot, sweet potato), bulbs (onion), or flowers (broccoli, cauliflower).
    - Do NOT confuse culinary and botanical definitions. A tomato is a fruit botanically.
-10. Never abbreviate. Always write full words:
-    - City/place names in full: "Ho Chi Minh City" not "HCMC", "Saint Petersburg" not "St. Petersburg"
-    - Country names in full: "United States" not "US" or "USA", "United Kingdom" not "UK"
-    - Exception: only use an abbreviation if the question itself uses it or explicitly asks for it.
+10. Never abbreviate UNLESS the question explicitly asks for an abbreviation or code:
+    - If asked for an IOC country code: return the 3-letter code (e.g. CUB, GBR, EGY).
+    - If asked for an ISO code, airport code, or similar: return the code as-is.
+    - Otherwise write full words: "United States" not "US", "Saint Petersburg" not "St. Petersburg".
 
 ## Required final line
 Always end your response with exactly:
@@ -254,7 +310,19 @@ def _answer_looks_weak(result) -> bool:
 # INVOKE — exception fallback + content-quality fallback
 # ==========================================================
 
-def invoke(messages: list) -> object:
+def invoke(messages: list, is_final: bool = False) -> object:
+    """
+    Call models in fallback order.
+
+    is_final=False  → first turn (no tool results yet). Use Qwen only.
+                      If Qwen wants tools, return immediately.
+                      If Qwen gives a weak answer on the first turn,
+                      still return it — the graph hasn't searched yet,
+                      so a weak answer just means it will call a tool next.
+    is_final=True   → the model has already used tools and is giving its
+                      terminal answer. If that answer is weak, try the
+                      next model in the chain.
+    """
     global LAST_MODEL_USED, LAST_MODEL_FALLBACK, LAST_MODEL_ERROR
 
     LAST_MODEL_FALLBACK = "No"
@@ -279,26 +347,31 @@ def invoke(messages: list) -> object:
             result = model.bind_tools(TOOLS).invoke(messages)
             last_result = result
 
-            # If the model wants tool calls, return immediately —
-            # the graph will handle the tool execution and loop back
+            # Model wants to call a tool — return immediately regardless of turn
             tool_calls = getattr(result, "tool_calls", None)
             if tool_calls:
                 return result
 
-            # No tool calls — check if the answer is actually useful
+            # Not a final turn — don't burn fallback models on intermediate steps.
+            # Return whatever Qwen gave; the graph will loop or terminate naturally.
+            if not is_final:
+                return result
+
+            # Final turn — check answer quality and try next model if weak
             if not _answer_looks_weak(result):
                 return result
 
-            # Answer is weak (N/A or missing FINAL ANSWER) — try next model
             LAST_MODEL_FALLBACK = "Yes"
             LAST_MODEL_ERROR = f"weak answer from {key}"
+            time.sleep(0.5)   # brief pause to avoid 413/429 on rapid retries
             continue
 
         except Exception as e:
             LAST_MODEL_ERROR = str(e)
+            time.sleep(0.5)
             continue
 
-    # All models tried — return whatever the last one gave us
+    # All models tried — return the last result rather than raising
     if last_result is not None:
         return last_result
 
@@ -316,8 +389,16 @@ def assistant(state: MessagesState) -> dict:
     if direct is not None:
         return {"messages": [AIMessage(content=f"FINAL ANSWER: {direct}")]}
 
+    # Detect whether this is a final turn (tool results already in history).
+    # If so, the model should be giving its terminal answer — activate fallback
+    # chain if that answer is weak.
+    is_final = any(
+        getattr(m, "type", "") == "tool"
+        for m in state["messages"]
+    )
+
     messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
-    result = invoke(messages)
+    result = invoke(messages, is_final=is_final)
     return {"messages": [result]}
 
 
@@ -399,13 +480,21 @@ def _expand_abbreviation(answer: str) -> str:
     """
     If the entire answer is a known abbreviation, replace it with the full form.
     Also expands 'St. <Name>' → 'Saint <Name>' for city/place names.
-    Does NOT modify answers that are longer than a single token/short phrase,
-    to avoid corrupting sentences that happen to contain an abbreviation.
+
+    Does NOT expand:
+    - IOC country codes (3 uppercase letters like CUB, EGY, GBR) — these are
+      intentional when the question asks for the IOC code.
+    - Answers longer than 5 words — too risky to mutate longer text.
     """
     # Only act on short answers (≤ 5 words) to stay safe
     if len(answer.split()) > 5:
         return answer
 
+    # Never expand 3-letter ALL-CAPS strings — these are almost certainly
+    # IOC codes, ISO codes, or other intentional abbreviations the question asked for
+    if re.fullmatch(r'[A-Z]{3}', answer.strip()):
+        return answer
+
     # Whole-answer lookup (case-insensitive)
     lookup = answer.lower().strip(".")
     if lookup in _ABBREV_MAP: