Update app.py

app.py

@@ -1,3 +1,17 @@
+# Copyright 2025 Jesus Vilela Jato.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 # --- Imports ---
 import os
 import gradio as gr
@@ -12,8 +26,6 @@ import mimetypes
 import subprocess # For yt-dlp
 import io # For BytesIO with PIL
 
-# Removed: from huggingface_hub import get_space_runtime - not used for username with OAuth
-
 # --- Global Variables for Startup Status ---
 missing_vars_startup_list_global = []
 agent_pre_init_status_msg_global = "Agent status will be determined at startup."
@@ -21,16 +33,16 @@ agent_pre_init_status_msg_global = "Agent status will be determined at startup."
 # File Processing Libs
 try: from PyPDF2 import PdfReader; PYPDF2_AVAILABLE = True
 except ImportError: PYPDF2_AVAILABLE = False; print("WARNING: PyPDF2 not found, PDF tool will be disabled.")
-try: from PIL import Image; import pytesseract; PIL_TESSERACT_AVAILABLE = True # PIL is needed for new tool
+try: from PIL import Image; import pytesseract; PIL_TESSERACT_AVAILABLE = True
 except ImportError: PIL_TESSERACT_AVAILABLE = False; print("WARNING: Pillow or Pytesseract not found, OCR tool will be disabled.")
 try: import whisper; WHISPER_AVAILABLE = True
 except ImportError: WHISPER_AVAILABLE = False; print("WARNING: OpenAI Whisper not found, Audio Transcription tool will be disabled.")
 
-# Google GenAI (Used by LangChain integration)
-from google.genai.types import HarmCategory, HarmBlockThreshold
+# Google GenAI (Used by LangChain integration AND direct client)
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
 
 # LangChain
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage # Removed AnyMessage
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage
 from langchain.prompts import PromptTemplate
 from langchain.tools import BaseTool, tool as lc_tool_decorator
 from langchain_google_genai import ChatGoogleGenerativeAI
@@ -41,13 +53,13 @@ from langchain_experimental.tools import PythonREPLTool
 # LangGraph Conditional Imports
 if TYPE_CHECKING:
     from langgraph.graph import StateGraph as StateGraphAliasedForHinting
-    from langgraph.prebuilt import ToolExecutor as ToolExecutorAliasedForHinting
+    from langgraph.prebuilt import ToolNode as ToolExecutorAliasedForHinting # Prefer ToolNode
     from typing_extensions import TypedDict
     from langgraph.checkpoint.base import BaseCheckpointSaver
 
 LANGGRAPH_FLAVOR_AVAILABLE = False
 LG_StateGraph: Optional[Type[Any]] = None
-LG_ToolExecutor: Optional[Type[Any]] = None
+LG_ToolExecutor_Class: Optional[Type[Any]] = None # Store the class (ToolNode or ToolExecutor)
 LG_END: Optional[Any] = None
 LG_ToolInvocation: Optional[Type[Any]] = None
 add_messages: Optional[Any] = None
@@ -55,40 +67,56 @@ MemorySaver_Class: Optional[Type[Any]] = None
 
 AGENT_INSTANCE: Optional[Union[AgentExecutor, Any]] = None
 TOOLS: List[BaseTool] = []
-LLM_INSTANCE: Optional[ChatGoogleGenerativeAI] = None # This is the agent's "planner"
+LLM_INSTANCE: Optional[ChatGoogleGenerativeAI] = None
 LANGGRAPH_MEMORY_SAVER: Optional[Any] = None
 
-# --- google-genai Client SDK (for the new direct multimodal tool) ---
+# google-genai Client SDK
 from google import genai as google_genai_sdk
-google_genai_client: Optional[google_genai_sdk.Client] = None # Initialized later
-# --- End google-genai Client SDK section ---
+google_genai_client: Optional[google_genai_sdk.Client] = None
 
 try:
     from langgraph.graph import StateGraph, END
-    from langgraph.prebuilt import ToolExecutor, ToolInvocation as LGToolInvocationActual
-    from langgraph.graph.message import add_messages as lg_add_messages
-    from langgraph.checkpoint.memory import MemorySaver as LGMemorySaver
-    LANGGRAPH_FLAVOR_AVAILABLE = True
-    LG_StateGraph, LG_ToolExecutor, LG_END, LG_ToolInvocation, add_messages, MemorySaver_Class = \
-        StateGraph, ToolExecutor, END, LGToolInvocationActual, lg_add_messages, LGMemorySaver
-    print("Successfully imported LangGraph components.")
-except ImportError as e:
+    try:
+        from langgraph.prebuilt import ToolNode # Common in newer langgraph
+        LG_ToolExecutor_Class = ToolNode # Assign ToolNode class
+        print("Using langgraph.prebuilt.ToolNode for LangGraph tool execution.")
+    except ImportError:
+        try:
+            from langgraph.prebuilt import ToolExecutor # Original attempt
+            LG_ToolExecutor_Class = ToolExecutor
+            print("Using langgraph.prebuilt.ToolExecutor (fallback) for LangGraph tool execution.")
+        except ImportError as e_lg_exec_inner:
+             print(f"Failed to import ToolNode and ToolExecutor from langgraph.prebuilt: {e_lg_exec_inner}")
+             LG_ToolExecutor_Class = None
+
+    if LG_ToolExecutor_Class is not None: # Proceed only if a tool executor class was found
+        from langgraph.prebuilt import ToolInvocation as LGToolInvocationActual
+        from langgraph.graph.message import add_messages as lg_add_messages
+        from langgraph.checkpoint.memory import MemorySaver as LGMemorySaver
+        LANGGRAPH_FLAVOR_AVAILABLE = True
+        LG_StateGraph, LG_END, LG_ToolInvocation, add_messages, MemorySaver_Class = \
+            StateGraph, END, LGToolInvocationActual, lg_add_messages, LGMemorySaver
+        print("Successfully imported LangGraph components.")
+    else:
+        # This ensures LANGGRAPH_FLAVOR_AVAILABLE remains False if no executor was found
+        LANGGRAPH_FLAVOR_AVAILABLE = False
+        LG_StateGraph, LG_END, LG_ToolInvocation, add_messages, MemorySaver_Class = (None,) * 5
+        print(f"WARNING: No suitable LangGraph tool executor (ToolNode/ToolExecutor) found. LangGraph agent will be disabled.")
+
+except ImportError as e: # Catch import error for StateGraph, END itself
     LANGGRAPH_FLAVOR_AVAILABLE = False
-    # Assign None to all to prevent NameError if used before assignment
-    LG_StateGraph, LG_ToolExecutor, LG_END, LG_ToolInvocation, add_messages, MemorySaver_Class = (None,) * 6
-    print(f"WARNING: LangGraph components not found or import error: {e}. LangGraph agent will be disabled.")
+    LG_StateGraph, LG_ToolExecutor_Class, LG_END, LG_ToolInvocation, add_messages, MemorySaver_Class = (None,) * 6
+    print(f"WARNING: Core LangGraph components (StateGraph, END) not found or import error: {e}. LangGraph agent will be disabled.")
+
 
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-# GEMINI_MODEL_NAME is for the agent's planner LLM (LangChain)
-GEMINI_MODEL_NAME = "gemini-2.5-pro-preview-05-06" # Retained from original for planner
-# GEMINI_FLASH_MULTIMODAL_MODEL_NAME is for the new direct multimodal tool (google-genai client SDK)
+GEMINI_MODEL_NAME = "gemini-2.5-pro-preview-05-06"
 GEMINI_FLASH_MULTIMODAL_MODEL_NAME = "gemini-2.0-flash-exp"
-
 SCORING_API_BASE_URL = os.getenv("SCORING_API_URL", DEFAULT_API_URL)
 MAX_FILE_SIZE_BYTES = 50 * 1024 * 1024
 LOCAL_FILE_STORE_PATH = "./Data"
-os.makedirs(LOCAL_FILE_STORE_PATH, exist_ok=True) # Create data directory at startup
+os.makedirs(LOCAL_FILE_STORE_PATH, exist_ok=True)
 
 # --- Global State ---
 WHISPER_MODEL: Optional[Any] = None
@@ -108,13 +136,13 @@ if GOOGLE_API_KEY:
         logger.info("google-genai SDK Client initialized successfully.")
     except Exception as e:
         logger.error(f"Failed to initialize google-genai SDK Client: {e}")
-        google_genai_client = None # Ensure it's None if init fails
+        google_genai_client = None
 else:
-    logger.warning("GOOGLE_API_KEY not found. google-genai SDK Client (for direct multimodal tool) not initialized.")
-# --- End Initialize google-genai Client SDK ---
+    logger.warning("GOOGLE_API_KEY not found. google-genai SDK Client not initialized.")
 
-# --- Helper Functions (Unchanged from your original) ---
+# --- Helper Functions (Unchanged) ---
 def _strip_exact_match_answer(text: Any) -> str:
+    # ... (Your original _strip_exact_match_answer function)
     if not isinstance(text, str): text = str(text)
     text_lower_check = text.lower()
     if text_lower_check.startswith("final answer:"):
@@ -132,15 +160,17 @@ def _strip_exact_match_answer(text: Any) -> str:
     return text.strip()
 
 def _is_full_url(url_string: str) -> bool:
+    # ... (Your original _is_full_url function)
     try: result = urlparse(url_string); return all([result.scheme, result.netloc])
     except ValueError: return False
 
 def _is_youtube_url(url: str) -> bool:
+    # ... (Your original _is_youtube_url function)
     parsed_url = urlparse(url)
     return parsed_url.netloc.lower().endswith(("youtube.com", "youtu.be"))
 
 def _download_file(file_identifier: str, task_id_for_file: Optional[str] = None) -> str:
-    # ... (Your original _download_file function - unchanged)
+    # ... (Your original _download_file function - unchanged) ...
     os.makedirs(LOCAL_FILE_STORE_PATH, exist_ok=True)
     logger.debug(f"Download request: '{file_identifier}', task_id: {task_id_for_file}")
     original_filename = os.path.basename(urlparse(file_identifier).path) if _is_full_url(file_identifier) else os.path.basename(file_identifier)
@@ -221,8 +251,8 @@ def _download_file(file_identifier: str, task_id_for_file: Optional[str] = None)
 
             if effective_save_path != tentative_local_path and os.path.exists(effective_save_path) and os.path.getsize(effective_save_path) > 0:
                 logger.info(f"Cached file (CD name): {effective_save_path}"); return effective_save_path
-            with open(effective_save_path, "wb") as f:
-                for chunk in r.iter_content(chunk_size=1024*1024): f.write(chunk)
+            with open(effective_save_path, "wb") as f_download:
+                for chunk in r.iter_content(chunk_size=1024*1024): f_download.write(chunk)
             logger.info(f"File downloaded to {effective_save_path}"); return effective_save_path
     except requests.exceptions.HTTPError as e:
         err_msg = f"HTTP {e.response.status_code} for {file_url_to_try}. Detail: {e.response.text[:100]}"
@@ -230,7 +260,7 @@ def _download_file(file_identifier: str, task_id_for_file: Optional[str] = None)
     except Exception as e:
         logger.error(f"Download error for {file_url_to_try}: {e}", exc_info=True); return f"Error: {str(e)[:100]}"
 
-# --- Tool Function Definitions (Original tools unchanged) ---
+# --- Tool Function Definitions ---
 READ_PDF_TOOL_DESC = "Reads text content from a PDF file. Input: JSON '{\"file_identifier\": \"FILENAME_OR_URL\", \"task_id\": \"TASK_ID_IF_GAIA_FILENAME_ONLY\"}'. Returns extracted text."
 @lc_tool_decorator(description=READ_PDF_TOOL_DESC)
 def read_pdf_tool(action_input_json_str: str) -> str:
@@ -242,19 +272,18 @@ def read_pdf_tool(action_input_json_str: str) -> str:
     path = _download_file(file_id, task_id)
     if path.startswith("Error:"): return path
     try:
-        text = "";
-        with open(path, "rb") as f:
-            reader = PdfReader(f)
+        text_content = "";
+        with open(path, "rb") as f_pdf:
+            reader = PdfReader(f_pdf)
             if reader.is_encrypted:
                 try: reader.decrypt('')
                 except: return f"Error: PDF '{path}' encrypted."
             for page_num in range(len(reader.pages)):
                 page = reader.pages[page_num]
-                text += page.extract_text() + "\n\n"
-        return text[:40000]
+                text_content += page.extract_text() + "\n\n"
+        return text_content[:40000]
     except Exception as e: return f"Error reading PDF '{path}': {e}"
 
-
 OCR_IMAGE_TOOL_DESC = "Extracts text from an image using OCR. Input: JSON '{\"file_identifier\": \"FILENAME_OR_URL\", \"task_id\": \"TASK_ID_IF_GAIA_FILENAME_ONLY\"}'. Returns extracted text."
 @lc_tool_decorator(description=OCR_IMAGE_TOOL_DESC)
 def ocr_image_tool(action_input_json_str: str) -> str:
@@ -285,143 +314,189 @@ def transcribe_audio_tool(action_input_json_str: str) -> str:
     try: result = WHISPER_MODEL.transcribe(path, fp16=False); return result["text"][:40000] # type: ignore
     except Exception as e: logger.error(f"Whisper error on '{path}': {e}", exc_info=True); return f"Error transcribing '{path}': {e}"
 
-# +++ NEW TOOL using google-genai Client SDK for Multimodal Prompts +++
 DIRECT_MULTIMODAL_GEMINI_TOOL_DESC = (
     "Processes an image file (URL or local path) along with a text prompt using a Gemini multimodal model (gemini-2.0-flash-exp) "
-    "for tasks like image description, answering questions about the image, or generating text based on the image. "
-    "Input: JSON '{\"file_identifier\": \"IMAGE_FILENAME_OR_URL\", \"text_prompt\": \"Your question or instruction related to the image.\", \"task_id\": \"TASK_ID_IF_GAIA_FILENAME_ONLY\" (optional)}'. "
+    "for tasks like image description, Q&A about the image, or text generation based on the image. "
+    "Input: JSON '{\"file_identifier\": \"IMAGE_FILENAME_OR_URL\", \"text_prompt\": \"Your question or instruction.\", \"task_id\": \"TASK_ID\" (optional)}'. "
     "Returns the model's text response."
 )
 @lc_tool_decorator(description=DIRECT_MULTIMODAL_GEMINI_TOOL_DESC)
 def direct_multimodal_gemini_tool(action_input_json_str: str) -> str:
-    global google_genai_client # Use the initialized client
-    if not google_genai_client:
-        return "Error: google-genai SDK client not initialized. GOOGLE_API_KEY might be missing."
-    if not PIL_TESSERACT_AVAILABLE : # Check if PIL is available, as it's used to open the image
-        return "Error: Pillow (PIL) library is not available for image processing."
-
+    # ... (Implementation from previous response)
+    global google_genai_client
+    if not google_genai_client: return "Error: google-genai SDK client not initialized."
+    if not PIL_TESSERACT_AVAILABLE : return "Error: Pillow (PIL) library not available."
     try:
         data = json.loads(action_input_json_str)
         file_identifier = data.get("file_identifier")
-        text_prompt = data.get("text_prompt", "Describe this image.") # Default prompt
-        task_id = data.get("task_id") # Optional, for _download_file if needed
-
-        if not file_identifier:
-            return "Error: 'file_identifier' for the image is missing in the input."
-
-        logger.info(f"Direct Multimodal Tool: Processing image '{file_identifier}' with prompt '{text_prompt}'")
-
-        # Download the file to a local path (handles URLs and GAIA files)
+        text_prompt = data.get("text_prompt", "Describe this image.")
+        task_id = data.get("task_id")
+        if not file_identifier: return "Error: 'file_identifier' for image missing."
+        logger.info(f"Direct Multimodal Tool: Image '{file_identifier}', Prompt '{text_prompt}'")
         local_image_path = _download_file(file_identifier, task_id)
-        if local_image_path.startswith("Error:"):
-            return f"Error downloading image for Direct Multimodal Tool: {local_image_path}"
-
-        # Open the image using Pillow
+        if local_image_path.startswith("Error:"): return f"Error downloading for Direct MM Tool: {local_image_path}"
         try:
             pil_image = Image.open(local_image_path)
-            pil_image.thumbnail((1024, 1024)) # Optional: resize large images
-        except Exception as e_img:
-            logger.error(f"Error opening image at {local_image_path}: {e_img}")
-            return f"Error opening image file {local_image_path}: {str(e_img)}"
-
-        # Send to Gemini Flash model using the client SDK
+        except Exception as e_img_open: return f"Error opening image {local_image_path}: {str(e_img_open)}"
         response = google_genai_client.models.generate_content(
-            model=GEMINI_FLASH_MULTIMODAL_MODEL_NAME, # Use the specified Flash model
-            contents=[pil_image, text_prompt] # Pass PIL image and text prompt
+            model=GEMINI_FLASH_MULTIMODAL_MODEL_NAME, contents=[pil_image, text_prompt]
         )
-        logger.info(f"Direct Multimodal Tool: Response received from {GEMINI_FLASH_MULTIMODAL_MODEL_NAME}.")
-        return response.text[:40000] # Return model's text response, truncated if very long
-
-    except json.JSONDecodeError as e_json:
-        return f"Error parsing JSON input for Direct Multimodal Tool: {str(e_json)}. Input was: {action_input_json_str}"
-    except Exception as e_tool:
-        logger.error(f"Error in direct_multimodal_gemini_tool: {e_tool}", exc_info=True)
-        return f"Error executing Direct Multimodal Tool: {str(e_tool)}"
-# +++ END NEW TOOL +++
+        logger.info(f"Direct Multimodal Tool: Response from {GEMINI_FLASH_MULTIMODAL_MODEL_NAME} received.")
+        return response.text[:40000]
+    except json.JSONDecodeError as e_json_mm: return f"Error parsing JSON for Direct MM Tool: {str(e_json_mm)}. Input: {action_input_json_str}"
+    except Exception as e_tool_mm:
+        logger.error(f"Error in direct_multimodal_gemini_tool: {e_tool_mm}", exc_info=True)
+        return f"Error executing Direct Multimodal Tool: {str(e_tool_mm)}"
+
+# --- Agent Prompts (Unchanged) ---
+LANGGRAPH_PROMPT_TEMPLATE_STR = """You are a highly intelligent agent for the GAIA benchmark.
+Your goal is to provide an EXACT MATCH final answer. No conversational text, explanations, or markdown unless explicitly part of the answer.
+TOOLS:
+You have access to the following tools. Use them if necessary.
+{tools}
+TOOL USAGE:
+- To use a tool, your response must include a `tool_calls` attribute in the AIMessage. Each tool call should be a dictionary with "name", "args" (a dictionary of arguments), and "id".
+- For file tools ('read_pdf_tool', 'ocr_image_tool', 'transcribe_audio_tool', 'direct_multimodal_gemini_tool'): `args` must contain 'file_identifier' (filename/URL) and 'task_id' (if GAIA file). For 'direct_multimodal_gemini_tool', also include 'text_prompt'.
+- 'web_search': `args` is like '{{"query": "search query"}}'.
+- 'python_repl': `args` is like '{{"command": "python code string"}}'. Use print() for output.
+RESPONSE FORMAT:
+Final AIMessage should contain ONLY the answer in 'content' and NO 'tool_calls'. If using tools, 'content' can be thought process, with 'tool_calls'.
+Begin!
+Current Task Details (including Task ID and any associated files):
+{input}"""
+
+REACT_PROMPT_TEMPLATE_STR = """You are a highly intelligent agent for the GAIA benchmark.
+Goal: EXACT MATCH answer. No extra text/markdown.
+Tools: {tools}
+Process: Question -> Thought -> Action (ONE of [{tool_names}]) -> Action Input -> Observation -> Thought ... -> Final Answer: [exact answer]
+Tool Inputs:
+- web_search: Your search query string.
+- python_repl: Python code string. Use print(). For Excel/CSV, use pandas: import pandas as pd; df = pd.read_excel('./Data/TASKID_filename.xlsx'); print(df.head())
+- read_pdf_tool, ocr_image_tool, transcribe_audio_tool: JSON string like '{{"file_identifier": "FILENAME_OR_URL", "task_id": "CURRENT_TASK_ID_IF_FILENAME"}}'.
+- direct_multimodal_gemini_tool: JSON string like '{{"file_identifier": "IMAGE_FILENAME_OR_URL", "text_prompt": "Your prompt for the image.", "task_id": "TASK_ID_IF_GAIA_FILENAME"}}'.
+If tool fails or info missing, Final Answer: N/A. Do NOT use unlisted tools.
+Begin!
+{input}
+Thought:{agent_scratchpad}"""
 
 
-# --- Agent Prompts (Slightly updated to include the new tool name if available) ---
-# (Agent prompts remain largely the same, the agent will learn to use tools from their descriptions)
-
 # --- Agent Initialization and Response Logic ---
 def initialize_agent_and_tools(force_reinit=False):
-    global AGENT_INSTANCE, TOOLS, LLM_INSTANCE, LANGGRAPH_FLAVOR_AVAILABLE, LG_StateGraph, LG_ToolExecutor, LG_END, LG_ToolInvocation, add_messages, MemorySaver_Class, LANGGRAPH_MEMORY_SAVER, google_genai_client
+    global AGENT_INSTANCE, TOOLS, LLM_INSTANCE, LANGGRAPH_FLAVOR_AVAILABLE, LG_StateGraph, LG_ToolExecutor_Class, LG_END, LG_ToolInvocation, add_messages, MemorySaver_Class, LANGGRAPH_MEMORY_SAVER, google_genai_client
     if AGENT_INSTANCE and not force_reinit: logger.info("Agent already initialized."); return
     logger.info("Initializing agent and tools...")
     if not GOOGLE_API_KEY: raise ValueError("GOOGLE_API_KEY not set for LangChain LLM.")
 
-    # Initialize LangChain LLM (Planner)
+    # CORRECTED ChatGoogleGenerativeAI initialization
+    # The safety_settings should be a dictionary where keys are HarmCategory enums and values are HarmBlockThreshold enums.
+    # Or a list of dicts: [{"category": HarmCategory.XYZ, "threshold": HarmBlockThreshold.ABC}, ...]
+    # Let's use the dictionary format as it's cleaner and suggested by LangChain's type hints.
+    llm_safety_settings_map = {
+        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+    }
+    # If the above still causes issues due to Pydantic's strict enum key handling,
+    # an alternative is to pass it as a list of SafetySetting objects from google.generativeai.types,
+    # but ChatGoogleGenerativeAI's Pydantic model might not directly accept that.
+    # The most robust way if direct enums as keys fail is to convert enums to their string values for the dict if Pydantic demands.
+    # However, LangChain *should* handle the direct enums.
+    # The error `Input should be 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 or 11 [type=enum, input_value=<HarmCategory...>]`
+    # suggests that for the *keys* of the safety_settings dict, it might be expecting the integer value of the HarmCategory enum.
+    # This is unusual. Let's try with the most standard documented way first (direct enums).
+    # If that fails, the next step would be to try string names for keys.
+
     try:
-        LLM_INSTANCE = ChatGoogleGenerativeAI(model=GEMINI_MODEL_NAME, google_api_key=GOOGLE_API_KEY, temperature=0.0,
-            safety_settings={HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE, HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-                             HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE, HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,},
-            request_timeout=120, convert_system_message_to_human=True )
+        LLM_INSTANCE = ChatGoogleGenerativeAI(
+            model=GEMINI_MODEL_NAME,
+            google_api_key=GOOGLE_API_KEY,
+            temperature=0.0,
+            safety_settings=llm_safety_settings_map, # Pass the map
+            timeout=120,                            # Corrected: 'timeout' not 'request_timeout'
+            convert_system_message_to_human=True
+        )
         logger.info(f"LangChain LLM (Planner) initialized: {GEMINI_MODEL_NAME}")
-    except Exception as e: logger.error(f"LangChain LLM init failed: {e}", exc_info=True); raise
+    except Exception as e:
+        logger.error(f"LangChain LLM init failed: {e}", exc_info=True)
+        raise
 
     TOOLS = []
     if PYPDF2_AVAILABLE: TOOLS.append(read_pdf_tool)
     if PIL_TESSERACT_AVAILABLE: TOOLS.append(ocr_image_tool)
     if WHISPER_AVAILABLE: TOOLS.append(transcribe_audio_tool)
-    
-    # Add the new direct multimodal tool if its dependencies (client, PIL) are met
-    if google_genai_client and PIL_TESSERACT_AVAILABLE: # PIL_TESSERACT_AVAILABLE implies PIL is available
-        TOOLS.append(direct_multimodal_gemini_tool)
-        logger.info("Added 'direct_multimodal_gemini_tool'.")
-    else:
-        logger.warning("'direct_multimodal_gemini_tool' NOT added due to missing google_genai_client or PIL.")
-
+    if google_genai_client and PIL_TESSERACT_AVAILABLE: TOOLS.append(direct_multimodal_gemini_tool); logger.info("Added 'direct_multimodal_gemini_tool'.")
+    else: logger.warning("'direct_multimodal_gemini_tool' NOT added (client or PIL missing).")
     try: search_tool = DuckDuckGoSearchRun(name="web_search"); search_tool.description = "Web search. Input: query."; TOOLS.append(search_tool)
     except Exception as e: logger.warning(f"DuckDuckGoSearchRun init failed: {e}")
     try: python_repl = PythonREPLTool(name="python_repl"); python_repl.description = "Python REPL. print() for output."; TOOLS.append(python_repl)
     except Exception as e: logger.warning(f"PythonREPLTool init failed: {e}")
     logger.info(f"Final tools list for agent: {[t.name for t in TOOLS]}")
 
-
-    # ... (Rest of your initialize_agent_and_tools function for LangGraph/ReAct setup - unchanged)
-    if LANGGRAPH_FLAVOR_AVAILABLE and all([LG_StateGraph, LG_ToolExecutor, LG_END, LLM_INSTANCE, LG_ToolInvocation, add_messages]):
+    if LANGGRAPH_FLAVOR_AVAILABLE and all([LG_StateGraph, LG_ToolExecutor_Class, LG_END, LLM_INSTANCE, LG_ToolInvocation, add_messages]):
         if not LANGGRAPH_MEMORY_SAVER and MemorySaver_Class: LANGGRAPH_MEMORY_SAVER = MemorySaver_Class(); logger.info("LangGraph MemorySaver initialized.")
         try:
-            logger.info(f"Attempting LangGraph init (Memory: {LANGGRAPH_MEMORY_SAVER is not None})")
+            logger.info(f"Attempting LangGraph init (Tool Executor type: {LG_ToolExecutor_Class.__name__ if LG_ToolExecutor_Class else 'None'})")
             _TypedDict = getattr(__import__('typing_extensions'), 'TypedDict', dict)
-            class AgentState(_TypedDict): input: str; messages: Annotated[List[Any], add_messages] # Use Any for AnyMessage for broader compatibility
+            class AgentState(_TypedDict): input: str; messages: Annotated[List[Any], add_messages]
 
-            prompt_content_lg = LANGGRAPH_PROMPT_TEMPLATE_STR.format(
-                tools="\n".join([f"- {t.name}: {t.description}" for t in TOOLS]),
-                input="{input}"
+            prompt_content_lg_init = LANGGRAPH_PROMPT_TEMPLATE_STR.format(
+                tools="\n".join([f"- {t.name}: {t.description}" for t in TOOLS]), input="{input}"
             )
             def agent_node(state: AgentState):
-                current_input_lg = state.get('input', '')
-                formatted_system_prompt_lg = prompt_content_lg.replace("{input}", current_input_lg)
-                messages_for_llm_lg = [SystemMessage(content=formatted_system_prompt_lg)] + state.get('messages', [])
-                bound_llm_for_tools_lg = LLM_INSTANCE.bind_tools(TOOLS)
-                response_from_llm_lg = bound_llm_for_tools_lg.invoke(messages_for_llm_lg)
-                return {"messages": [response_from_llm_lg]}
-
-            tool_executor_lg_instance = LG_ToolExecutor(TOOLS) # type: ignore
-            def tool_node(state: AgentState):
-                last_msg_lg = state['messages'][-1] if state.get('messages') and isinstance(state['messages'][-1], AIMessage) else None
-                if not last_msg_lg or not last_msg_lg.tool_calls: return {"messages": []}
-                tool_results_lg = []
-                for tc_lg in last_msg_lg.tool_calls:
-                    name_lg, args_lg, tc_id_lg = tc_lg.get('name'), tc_lg.get('args'), tc_lg.get('id')
-                    if not all([name_lg, isinstance(args_lg, dict), tc_id_lg]):
-                        err_msg_lg=f"Invalid tool_call: {tc_lg}"; logger.error(err_msg_lg)
-                        tool_results_lg.append(ToolMessage(f"Error: {err_msg_lg}", tool_call_id=tc_id_lg or "error_id", name=name_lg or "error_tool"))
+                current_input = state.get('input', '')
+                formatted_system_prompt = prompt_content_lg_init.replace("{input}", current_input)
+                messages_for_llm = [SystemMessage(content=formatted_system_prompt)] + state.get('messages', [])
+                bound_llm = LLM_INSTANCE.bind_tools(TOOLS) # type: ignore
+                response = bound_llm.invoke(messages_for_llm)
+                return {"messages": [response]}
+
+            if not LG_ToolExecutor_Class: raise ValueError("LG_ToolExecutor_Class is None for LangGraph.")
+            # Instantiate ToolNode if that's what was imported
+            tool_executor_instance_lg = LG_ToolExecutor_Class(tools=TOOLS)
+
+
+            def tool_node(state: AgentState): # This function uses the instantiated tool_executor_instance_lg
+                last_msg = state['messages'][-1] if state.get('messages') and isinstance(state['messages'][-1], AIMessage) else None
+                if not last_msg or not last_msg.tool_calls: return {"messages": []}
+                # ToolNode expects a list of ToolInvocations if invoked directly,
+                # or handles it if part of a graph that structures it.
+                # The LangGraph prebuilt react agent often passes the AIMessage directly to ToolNode.
+                # Let's assume the ToolNode can handle a list of tool calls from the AIMessage.
+                # If ToolNode expects a single ToolInvocation, this loop needs adjustment.
+                # However, the standard ToolNode takes the AIMessage and iterates internally.
+                # The issue might be if `tool_executor_instance_lg` is not directly callable or its `invoke` expects different input.
+                # For now, let's assume the standard pattern where ToolNode handles the AIMessage's tool_calls.
+                # A simpler way to use ToolNode is often just to pass it to add_node if it's a runnable.
+                # tool_executor_instance_lg.invoke(last_msg.tool_calls) might be what's needed if it takes raw calls
+                
+                # The following is more aligned if tool_executor_instance_lg is the older ToolExecutor
+                # or if ToolNode is used within a manual iteration like this:
+                tool_results = []
+                for tc in last_msg.tool_calls:
+                    name, args, tc_id = tc.get('name'), tc.get('args'), tc.get('id')
+                    if not all([name, isinstance(args, dict), tc_id]):
+                        # ... error handling ...
+                        err_msg=f"Invalid tool_call: {tc}"; logger.error(err_msg)
+                        tool_results.append(ToolMessage(f"Error: {err_msg}", tool_call_id=tc_id or "error_id", name=name or "error_tool"))
                         continue
                     try:
-                        logger.info(f"LG Tool Invoking: '{name_lg}' with {args_lg} (ID: {tc_id_lg})")
-                        tool_invocation_obj_lg = LG_ToolInvocation(tool=name_lg, tool_input=args_lg) # type: ignore
-                        output_lg = tool_executor_lg_instance.invoke(tool_invocation_obj_lg)
-                        tool_results_lg.append(ToolMessage(content=str(output_lg), tool_call_id=tc_id_lg, name=name_lg))
+                        logger.info(f"LG Tool Invoking: '{name}' with {args} (ID: {tc_id})")
+                        tool_invocation_obj_lg = LG_ToolInvocation(tool=name, tool_input=args) # type: ignore
+                        # This invoke is on the tool_executor_instance_lg (which could be ToolNode or older ToolExecutor)
+                        output_lg = tool_executor_instance_lg.invoke(tool_invocation_obj_lg) # This might be the issue if ToolNode expects direct runnable call
+                        tool_results.append(ToolMessage(content=str(output_lg), tool_call_id=tc_id, name=name))
                     except Exception as e_tool_node_lg:
-                        logger.error(f"LG Tool Error ('{name_lg}'): {e_tool_node_lg}", exc_info=True)
-                        tool_results_lg.append(ToolMessage(content=f"Error for tool {name_lg}: {str(e_tool_node_lg)}", tool_call_id=tc_id_lg, name=name_lg))
-                return {"messages": tool_results_lg}
+                        logger.error(f"LG Tool Error ('{name}'): {e_tool_node_lg}", exc_info=True)
+                        tool_results.append(ToolMessage(content=f"Error for tool {name}: {str(e_tool_node_lg)}", tool_call_id=tc_id, name=name))
+                return {"messages": tool_results}
+
 
             workflow_lg = LG_StateGraph(AgentState) # type: ignore
             workflow_lg.add_node("agent", agent_node)
-            workflow_lg.add_node("tools", tool_node)
+            # If LG_ToolExecutor_Class is ToolNode, tool_node_instance is runnable
+            # workflow_lg.add_node("tools", tool_executor_instance_lg) # Preferred way if ToolNode is a runnable
+            workflow_lg.add_node("tools", tool_node) # Keep custom tool_node for now
             workflow_lg.set_entry_point("agent")
             def should_continue_lg(state: AgentState): return "tools" if state['messages'][-1].tool_calls else LG_END
             workflow_lg.add_conditional_edges("agent", should_continue_lg, {"tools": "tools", LG_END: LG_END}) # type: ignore
@@ -450,11 +525,9 @@ def initialize_agent_and_tools(force_reinit=False):
     if not AGENT_INSTANCE: raise RuntimeError("CRITICAL: Agent initialization completely failed.")
     logger.info(f"Agent init finished. Active agent type: {type(AGENT_INSTANCE).__name__}")
 
-
-# --- get_agent_response, construct_prompt_for_agent, run_and_submit_all ---
-# --- These functions remain UNCHANGED from your original code ---
+# --- get_agent_response, construct_prompt_for_agent, run_and_submit_all (Unchanged) ---
 def get_agent_response(prompt: str, task_id: Optional[str]=None, thread_id: Optional[str]=None) -> str:
-    # ... (Your original get_agent_response logic) ...
+    # ... (Your original get_agent_response logic - unchanged) ...
     global AGENT_INSTANCE, LLM_INSTANCE
     thread_id_to_use = thread_id or (f"gaia_task_{task_id}" if task_id else hashlib.md5(prompt.encode()).hexdigest()[:8])
     if not AGENT_INSTANCE or not LLM_INSTANCE:
@@ -470,7 +543,7 @@ def get_agent_response(prompt: str, task_id: Optional[str]=None, thread_id: Opti
             logger.debug(f"Using LangGraph agent (Memory: {LANGGRAPH_MEMORY_SAVER is not None}) for thread: {thread_id_to_use}")
             initial_messages_lg_get = []
             input_for_lg_get = {"input": prompt, "messages": initial_messages_lg_get}
-            final_state_lg_get = AGENT_INSTANCE.invoke(input_for_lg_get, {"configurable": {"thread_id": thread_id_to_use}})
+            final_state_lg_get = AGENT_INSTANCE.invoke(input_for_lg_get, {"configurable": {"thread_id": thread_id_to_use}}) # type: ignore
             if not final_state_lg_get or 'messages' not in final_state_lg_get or not final_state_lg_get['messages']:
                 logger.error("LangGraph: No final state/messages."); return "[ERROR] LangGraph: No final state/messages."
             for message_item_lg_get in reversed(final_state_lg_get['messages']):
@@ -489,7 +562,7 @@ def get_agent_response(prompt: str, task_id: Optional[str]=None, thread_id: Opti
         return f"[ERROR] Agent execution failed: {str(e_agent_run_get)[:150]}"
 
 def construct_prompt_for_agent(q: Dict[str,Any]) -> str:
-    # ... (Your original construct_prompt_for_agent logic) ...
+    # ... (Your original construct_prompt_for_agent logic - unchanged) ...
     tid,q_str=q.get("task_id","N/A"),q.get("question",""); files=q.get("files",[])
     files_info = ("\nFiles:\n"+"\n".join([f"- {f} (task_id:{tid})"for f in files])) if files else ""
     level = f"\nLevel:{q.get('level')}" if q.get('level') else ""
@@ -500,19 +573,16 @@ def run_and_submit_all(profile: Optional[gr.OAuthProfile] = None):
     global AGENT_INSTANCE
     space_id = os.getenv("SPACE_ID")
     username_for_submission = None
-
     if profile and hasattr(profile, 'username') and profile.username:
         username_for_submission = profile.username
         logger.info(f"Username from OAuth profile: {username_for_submission}")
     else:
         logger.warning("OAuth profile not available or username missing.")
         return "Hugging Face login required. Please use the login button and try again.", None
-
     if AGENT_INSTANCE is None:
         try: logger.info("Agent not pre-initialized. Initializing for run..."); initialize_agent_and_tools()
         except Exception as e: return f"Agent on-demand initialization failed: {e}", None
         if AGENT_INSTANCE is None: return "Agent is still None after on-demand init.", None
-    
     agent_code_url_run=f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else "local_dev_run"
     questions_url_run,submit_url_run=f"{DEFAULT_API_URL}/questions",f"{DEFAULT_API_URL}/submit"
     auth_headers_run={"Authorization":f"Bearer {HUGGINGFACE_TOKEN}"} if HUGGINGFACE_TOKEN else {}
@@ -522,7 +592,6 @@ def run_and_submit_all(profile: Optional[gr.OAuthProfile] = None):
         if not questions_data_run or not isinstance(questions_data_run,list):logger.error(f"Invalid questions data: {questions_data_run}");return "Fetched questions_data invalid.",None
         logger.info(f"Fetched {len(questions_data_run)} questions.")
     except Exception as e:logger.error(f"Fetch questions error: {e}",exc_info=True);return f"Fetch questions error:{e}",None
-    
     results_log_run,answers_payload_run=[],[]
     logger.info(f"Running agent on {len(questions_data_run)} questions for user '{username_for_submission}'...")
     for i,item_run in enumerate(questions_data_run):
@@ -538,7 +607,6 @@ def run_and_submit_all(profile: Optional[gr.OAuthProfile] = None):
             logger.error(f"Agent error task {task_id_run}:{e}",exc_info=True);error_answer_run=f"AGENT ERROR:{str(e)[:100]}"
             answers_payload_run.append({"task_id":task_id_run,"submitted_answer":"N/A [AGENT_ERROR]"})
             results_log_run.append({"Task ID":task_id_run,"Question":question_text_run,"Full Agent Prompt":prompt_run,"Raw Agent Output":error_answer_run,"Submitted Answer":"N/A [AGENT_ERROR]"})
-    
     if not answers_payload_run:return "Agent produced no answers.",pd.DataFrame(results_log_run)
     submission_payload_run={"username":username_for_submission.strip(),"agent_code":agent_code_url_run,"answers":answers_payload_run}
     logger.info(f"Submitting {len(answers_payload_run)} answers to {submit_url_run} for user '{username_for_submission}'...")
@@ -551,10 +619,8 @@ def run_and_submit_all(profile: Optional[gr.OAuthProfile] = None):
         error_http_run=f"HTTP {e.response.status_code}. Detail:{e.response.text[:200]}"; logger.error(f"Submit Fail:{error_http_run}",exc_info=True); return f"Submit Fail:{error_http_run}",pd.DataFrame(results_log_run)
     except Exception as e:logger.error(f"Submit Fail unexpected:{e}",exc_info=True);return f"Submit Fail:{str(e)[:100]}",pd.DataFrame(results_log_run)
 
-
-# --- Build Gradio Interface (Unchanged from your original) ---
+# --- Build Gradio Interface ---
 with gr.Blocks(css=".gradio-container {max-width:1280px !important;margin:auto !important;}",theme=gr.themes.Soft()) as demo:
-    # ... (Your original Gradio UI layout - unchanged) ...
     gr.Markdown("# GAIA Agent Challenge Runner v7 (OAuth for Username)")
     gr.Markdown(f"""**Instructions:**
 1. **Login with Hugging Face** using the button below. Your HF username will be used for submission.
@@ -569,12 +635,11 @@ with gr.Blocks(css=".gradio-container {max-width:1280px !important;margin:auto !
     gr.LoginButton()
     run_button = gr.Button("Run Evaluation & Submit All Answers")
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=7, interactive=False)
-    results_table = gr.DataFrame(label="Q&A Log", headers=["Task ID","Question","Prompt","Raw","Submitted"], wrap=True)
+    results_table = gr.DataFrame(label="Q&A Log", headers=["Task ID","Question","Prompt","Raw","Submitted"], wrap=True) # Removed height
 
     run_button.click(fn=run_and_submit_all, outputs=[status_output,results_table], api_name="run_evaluation")
 
     def update_ui_on_load_fn_within_context():
-        # ... (Your original update_ui_on_load_fn_within_context logic - unchanged) ...
         global missing_vars_startup_list_global, agent_pre_init_status_msg_global
         secrets_msg_md = ""
         if missing_vars_startup_list_global:
@@ -587,20 +652,21 @@ with gr.Blocks(css=".gradio-container {max-width:1280px !important;margin:auto !
         if env_issues: secrets_msg_md += f"<br/><font color='orange'>**Tool Deps Missing:** {', '.join(env_issues)}.</font>"
         current_status_md = agent_pre_init_status_msg_global
         if not LANGGRAPH_FLAVOR_AVAILABLE and "LangGraph" not in current_status_md:
-             current_status_md += " (LangGraph core import failed, ReAct fallback.)"
+             current_status_md += f" (LangGraph core components not fully loaded: LG_ToolExecutor_Class is {type(LG_ToolExecutor_Class).__name__ if LG_ToolExecutor_Class else 'None'}, ReAct fallback.)"
+        elif LANGGRAPH_FLAVOR_AVAILABLE and "LangGraph" not in current_status_md:
+             current_status_md += f" (LangGraph ready with {type(LG_ToolExecutor_Class).__name__ if LG_ToolExecutor_Class else 'UnknownExecutor'}.)"
         return { agent_status_display: gr.Markdown(value=current_status_md),
                  missing_secrets_display: gr.Markdown(value=secrets_msg_md) }
 
     demo.load(update_ui_on_load_fn_within_context, [], [agent_status_display, missing_secrets_display])
 
 if __name__ == "__main__":
-    # ... (Your original __main__ block for startup logging and pre-initialization - unchanged) ...
-    logger.info("Application starting up (v7 with Direct Multimodal Tool)...")
+    logger.info(f"Application starting up (v7 with Pydantic & LangGraph fixes)...") # Updated version in log
     if not PYPDF2_AVAILABLE: logger.warning("PyPDF2 (PDF tool) NOT AVAILABLE.")
-    if not PIL_TESSERACT_AVAILABLE: logger.warning("Pillow/Pytesseract (OCR tool) NOT AVAILABLE.") # PIL also needed for new tool
+    if not PIL_TESSERACT_AVAILABLE: logger.warning("Pillow/Pytesseract (OCR tool) NOT AVAILABLE.")
     if not WHISPER_AVAILABLE: logger.warning("Whisper (Audio tool) NOT AVAILABLE.")
-    if LANGGRAPH_FLAVOR_AVAILABLE: logger.info("Core LangGraph (StateGraph, END) loaded.")
-    else: logger.warning("Core LangGraph FAILED import. ReAct fallback. Check requirements & Space build logs.")
+    if LANGGRAPH_FLAVOR_AVAILABLE: logger.info(f"Core LangGraph components (StateGraph, END, {type(LG_ToolExecutor_Class).__name__ if LG_ToolExecutor_Class else 'FailedExecutor'}) loaded.")
+    else: logger.warning("Core LangGraph FAILED import or essential component (ToolExecutor/Node) missing. ReAct fallback. Check requirements & Space build logs.")
 
     missing_vars_startup_list_global.clear()
     if not GOOGLE_API_KEY: missing_vars_startup_list_global.append("GOOGLE_API_KEY")
@@ -608,17 +674,18 @@ if __name__ == "__main__":
 
     try:
         logger.info("Pre-initializing agent...")
-        initialize_agent_and_tools() # This will now include the new direct_multimodal_gemini_tool
+        initialize_agent_and_tools()
         if AGENT_INSTANCE:
             agent_type_name = type(AGENT_INSTANCE).__name__
             agent_pre_init_status_msg_global = f"Agent Pre-initialized: **{agent_type_name}**."
-            if LANGGRAPH_FLAVOR_AVAILABLE and "StateGraph" in agent_type_name: # More robust check
-                 agent_pre_init_status_msg_global = f"Agent Pre-initialized: **LangGraph** (Memory: {LANGGRAPH_MEMORY_SAVER is not None})."
+            if LANGGRAPH_FLAVOR_AVAILABLE and ("StateGraph" in agent_type_name or "CompiledGraph" in agent_type_name) :
+                 lg_executor_display_name = type(LG_ToolExecutor_Class).__name__ if LG_ToolExecutor_Class else "UnknownExecutor"
+                 agent_pre_init_status_msg_global = f"Agent Pre-initialized: **LangGraph** (Executor: {lg_executor_display_name}, Memory: {LANGGRAPH_MEMORY_SAVER is not None})."
         else: agent_pre_init_status_msg_global = "Agent pre-init FAILED (AGENT_INSTANCE is None)."
         logger.info(agent_pre_init_status_msg_global.replace("**",""))
     except Exception as e:
-        agent_pre_init_status_msg_global = f"Agent pre-init CRASHED: {str(e)[:100]}."
-        logger.critical(f"Agent pre-init CRASHED: {e}", exc_info=True)
+        agent_pre_init_status_msg_global = f"Agent pre-init CRASHED: {str(e)[:100]}." # Show first 100 chars of error
+        logger.critical(f"Agent pre-init CRASHED: {e}", exc_info=True) # Full traceback to logs
 
     logger.info(f"Space ID: {os.getenv('SPACE_ID', 'Not Set')}")
     logger.info("Gradio Interface launching...")