Update agents.py

agents.py

@@ -3,31 +3,26 @@ import json
 import asyncio
 import random
 
-
-import os
-import json
-import asyncio
-import random
-
 # --- OpenAI ---
-from openai import AsyncOpenAI, APIError # (Keep if needed for other parts)
+from openai import AsyncOpenAI, APIError
 
 # --- Google Gemini ---
 from google import genai
 from google.genai import types
-# It's good practice to also import potential exceptions
-
+from google.api_core import exceptions as google_exceptions
 
 # --- Mistral AI ---
-from mistralai.async_client import MistralAsyncClient # (Keep if needed)
+from mistralai.async_client import MistralAsyncClient
+from mistralai.exceptions import MistralAPIError
 
 # --- Poke-Env ---
 from poke_env.player import Player
 from poke_env.environment.battle import Battle
 from poke_env.environment.move import Move
 from poke_env.environment.pokemon import Pokemon
+from typing import Optional, Dict, Any, Union
 
-# --- Helper Function & Base Class (Assuming they are defined above) ---
+# --- Helper Function & Base Class ---
 def normalize_name(name: str) -> str:
     """Lowercase and remove non-alphanumeric characters."""
     return "".join(filter(str.isalnum, name)).lower()
@@ -64,14 +59,13 @@ STANDARD_TOOL_SCHEMA = {
 }
 
 
-class LLMAgentBase(Player): # Make sure this base class exists
+class LLMAgentBase(Player):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.standard_tools = STANDARD_TOOL_SCHEMA
-        self.battle_history = [] # Example attribute
+        self.battle_history = []
 
     def _format_battle_state(self, battle: Battle) -> str:
-        # (Implementation as provided in the question)
         active_pkmn = battle.active_pokemon
         active_pkmn_info = f"Your active Pokemon: {active_pkmn.species} " \
                            f"(Type: {'/'.join(map(str, active_pkmn.types))}) " \
@@ -117,9 +111,7 @@ class LLMAgentBase(Player): # Make sure this base class exists
                     f"Opponent Side Conditions: {battle.opponent_side_conditions}"
         return state_str.strip()
 
-
-    def _find_move_by_name(self, battle: Battle, move_name: str) -> Move | None:
-        # (Implementation as provided in the question)
+    def _find_move_by_name(self, battle: Battle, move_name: str) -> Optional[Move]:
         normalized_name = normalize_name(move_name)
         # Prioritize exact ID match
         for move in battle.available_moves:
@@ -132,8 +124,7 @@ class LLMAgentBase(Player): # Make sure this base class exists
                 return move
         return None
 
-    def _find_pokemon_by_name(self, battle: Battle, pokemon_name: str) -> Pokemon | None:
-        # (Implementation as provided in the question)
+    def _find_pokemon_by_name(self, battle: Battle, pokemon_name: str) -> Optional[Pokemon]:
         normalized_name = normalize_name(pokemon_name)
         for pkmn in battle.available_switches:
             # Normalize the species name for comparison
@@ -142,7 +133,6 @@ class LLMAgentBase(Player): # Make sure this base class exists
         return None
 
     async def choose_move(self, battle: Battle) -> str:
-        # (Implementation as provided in the question - relies on _get_llm_decision)
         battle_state_str = self._format_battle_state(battle)
         decision_result = await self._get_llm_decision(battle_state_str)
         decision = decision_result.get("decision")
@@ -160,8 +150,7 @@ class LLMAgentBase(Player): # Make sure this base class exists
                     if chosen_move and chosen_move in battle.available_moves:
                         action_taken = True
                         chat_msg = f"AI Decision: Using move '{chosen_move.id}'."
-                        print(chat_msg) # Print to console for debugging
-                        # await self.send_message(chat_msg, battle=battle) # Uncomment if send_message exists
+                        print(chat_msg)
                         return self.create_order(chosen_move)
                     else:
                         fallback_reason = f"LLM chose unavailable/invalid move '{move_name}'."
@@ -174,8 +163,7 @@ class LLMAgentBase(Player): # Make sure this base class exists
                     if chosen_switch and chosen_switch in battle.available_switches:
                         action_taken = True
                         chat_msg = f"AI Decision: Switching to '{chosen_switch.species}'."
-                        print(chat_msg) # Print to console for debugging
-                        # await self.send_message(chat_msg, battle=battle) # Uncomment if send_message exists
+                        print(chat_msg)
                         return self.create_order(chosen_switch)
                     else:
                         fallback_reason = f"LLM chose unavailable/invalid switch '{pokemon_name}'."
@@ -185,55 +173,53 @@ class LLMAgentBase(Player): # Make sure this base class exists
                 fallback_reason = f"LLM called unknown function '{function_name}'."
 
         if not action_taken:
-            if not fallback_reason: # If no specific reason yet, check for API errors
+            if not fallback_reason:
                  if error_message:
                      fallback_reason = f"API Error: {error_message}"
-                 elif decision is None: # Model didn't call a function or response was bad
+                 elif decision is None:
                       fallback_reason = "LLM did not provide a valid function call."
-                 else: # Should not happen if logic above is correct
+                 else:
                       fallback_reason = "Unknown error processing LLM decision."
 
             print(f"Warning: {fallback_reason} Choosing random action.")
-            # await self.send_message(f"AI Fallback: {fallback_reason} Choosing random action.", battle=battle) # Uncomment
 
-            # Use poke-env's built-in random choice
             if battle.available_moves or battle.available_switches:
                  return self.choose_random_move(battle)
             else:
                  print("AI Fallback: No moves or switches available. Using Struggle/Default.")
-                 # await self.send_message("AI Fallback: No moves or switches available. Using Struggle.", battle=battle) # Uncomment
-                 return self.choose_default_move(battle) # Handles struggle
+                 return self.choose_default_move(battle)
 
-    async def _get_llm_decision(self, battle_state: str) -> dict:
+    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
         raise NotImplementedError("Subclasses must implement _get_llm_decision")
 
+
 # --- Google Gemini Agent ---
 class GeminiAgent(LLMAgentBase):
     """Uses Google Gemini API for decisions."""
-    def __init__(self, api_key: str | None = None, model: str = "gemini-1.5-flash", *args, **kwargs): # Default to flash for speed/cost
+    def __init__(self, api_key: str = None, model: str = "gemini-1.5-flash", *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.model_name = model
         used_api_key = api_key or os.environ.get("GOOGLE_API_KEY")
-        self.model_name=model
         if not used_api_key:
             raise ValueError("Google API key not provided or found in GOOGLE_API_KEY env var.")
-        self.client = genai.Client(
-            api_key='GEMINI_API_KEY',
-            http_options=types.HttpOptions(api_version='v1alpha')
+        
+        # Initialize Gemini client
+        genai.configure(api_key=used_api_key)
+        
+        # Configure the model with tools
+        self.gemini_tool_config = [
+            {
+                "function_declarations": list(self.standard_tools.values())
+            }
+        ]
+        
+        # Initialize the model
+        self.model = genai.GenerativeModel(
+            model_name=self.model_name,
+            tools=self.gemini_tool_config
         )
-        # --- Correct Tool Definition ---
-        # Create a list of function declaration dictionaries from the values in STANDARD_TOOL_SCHEMA
-        function_declarations = list(self.standard_tools.values())
-        # Create the Tool object expected by the API
-        self.gemini_tool_config = types.Tool(function_declarations=function_declarations)
-        # --- End Tool Definition ---
-
-        # --- Correct Model Initialization ---
-        # Pass the Tool object directly to the model's 'tools' parameter
 
-        # --- End Model Initialization ---
-
-    async def _get_llm_decision(self, battle_state: str) -> dict:
+    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
         """Sends state to the Gemini API and gets back the function call decision."""
         prompt = (
             "You are a skilled Pokemon battle AI. Your goal is to win the battle. "
@@ -246,63 +232,64 @@ class GeminiAgent(LLMAgentBase):
         )
 
         try:
-            # --- Correct API Call ---
-            # Call generate_content_async directly on the model object.
-            # Tools are already configured in the model, no need to pass config here.
-            response = await client.aio.models.generate_content(
-                model=self.model_name, 
-                contents=prompt
+            # Use the async API for Gemini
+            response = await self.model.generate_content_async(
+                prompt,
+                generation_config={"temperature": 0.5}
             )
-            # --- End API Call ---
-
-            # --- Response Parsing (Your logic was already good here) ---
-            # Check candidates and parts safely
+            
             if not response.candidates:
-                 finish_reason_str = "No candidates found"
-                 try: finish_reason_str = response.prompt_feedback.block_reason.name
-                 except AttributeError: pass
-                 return {"error": f"Gemini response issue. Reason: {finish_reason_str}"}
+                finish_reason_str = "No candidates found"
+                try: 
+                    finish_reason_str = response.prompt_feedback.block_reason.name
+                except AttributeError: 
+                    pass
+                return {"error": f"Gemini response issue. Reason: {finish_reason_str}"}
 
             candidate = response.candidates[0]
             if not candidate.content or not candidate.content.parts:
                 finish_reason_str = "Unknown"
-                try: finish_reason_str = candidate.finish_reason.name
-                except AttributeError: pass
+                try: 
+                    finish_reason_str = candidate.finish_reason.name
+                except AttributeError: 
+                    pass
                 return {"error": f"Gemini response issue. Finish Reason: {finish_reason_str}"}
 
-            part = candidate.content.parts[0]
-
-            # Check for function_call attribute
-            if hasattr(part, 'function_call') and part.function_call:
-                fc = part.function_call
-                function_name = fc.name
-                # fc.args is a proto_plus.MapComposite, convert to dict
-                arguments = dict(fc.args) if fc.args else {}
-
-                if function_name in self.standard_tools:
-                    # Valid function call found
-                    return {"decision": {"name": function_name, "arguments": arguments}}
-                else:
-                    # Model hallucinated a function name
-                    return {"error": f"Model called unknown function '{function_name}'. Args: {arguments}"}
-            elif hasattr(part, 'text'):
-                 # Handle case where the model returns text instead of a function call
-                 text_response = part.text
-                 return {"error": f"Gemini did not return a function call. Response: {text_response}"}
-            else:
-                  # Unexpected part type
-                  return {"error": f"Gemini response part type unknown. Part: {part}"}
-            # --- End Response Parsing ---
+            for part in candidate.content.parts:
+                if hasattr(part, 'function_call') and part.function_call:
+                    fc = part.function_call
+                    function_name = fc.name
+                    # Convert arguments to dict
+                    arguments = {}
+                    if fc.args:
+                        arguments = {k: v for k, v in fc.args.items()}
+                    
+                    if function_name in self.standard_tools:
+                        return {"decision": {"name": function_name, "arguments": arguments}}
+                    else:
+                        return {"error": f"Model called unknown function '{function_name}'. Args: {arguments}"}
+            
+            # If we got here, no function call was found in any part
+            text_content = " ".join([
+                part.text if hasattr(part, 'text') else str(part) 
+                for part in candidate.content.parts
+            ])
+            return {"error": f"Gemini did not return a function call. Response: {text_content[:100]}..."}
+            
+        except google_exceptions.GoogleAPIError as e:
+            print(f"Google API error: {e}")
+            return {"error": f"Google API error: {str(e)}"}
         except Exception as e:
-            # Catch any other unexpected errors during the API call or processing
             print(f"Unexpected error during Gemini processing: {e}")
             import traceback
-            traceback.print_exc() # Print stack trace for debugging
+            traceback.print_exc()
             return {"error": f"Unexpected error: {str(e)}"}
+
+
 # --- OpenAI Agent ---
 class OpenAIAgent(LLMAgentBase):
     """Uses OpenAI API for decisions."""
-    def __init__(self, api_key: str | None = None, model: str = "gpt-4o", *args, **kwargs):
+    def __init__(self, api_key: str = None, model: str = "gpt-4o", *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.model = model
         used_api_key = api_key or os.environ.get("OPENAI_API_KEY")
@@ -311,9 +298,9 @@ class OpenAIAgent(LLMAgentBase):
         self.openai_client = AsyncOpenAI(api_key=used_api_key)
 
         # Convert standard schema to OpenAI's format
-        self.openai_functions = [v for k, v in self.standard_tools.items()]
+        self.openai_tools = list(self.standard_tools.values())
 
-    async def _get_llm_decision(self, battle_state: str) -> dict:
+    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
         system_prompt = (
             "You are a skilled Pokemon battle AI. Your goal is to win the battle. "
             "Based on the current battle state, decide the best action: either use an available move or switch to an available Pokémon. "
@@ -330,25 +317,26 @@ class OpenAIAgent(LLMAgentBase):
                     {"role": "system", "content": system_prompt},
                     {"role": "user", "content": user_prompt},
                 ],
-                functions=STANDARD_TOOL_SCHEMA,
-                function_call="auto",
+                tools=self.openai_tools,
+                tool_choice="auto",  # Let the model choose
                 temperature=0.5,
             )
             message = response.choices[0].message
-            print("MESSAGE BACK : ", message)
-            if message.function_call:
-                function_name = message.function_call.name
+            
+            # Check for tool calls in the response
+            if message.tool_calls:
+                tool_call = message.tool_calls[0]  # Get the first tool call
+                function_name = tool_call.function.name
                 try:
-                    # Ensure arguments is always a dict, even if empty/null
-                    arguments = json.loads(message.function_call.arguments or '{}')
-                    if function_name in self.standard_tools: # Validate function name
-                         return {"decision": {"name": function_name, "arguments": arguments}}
+                    arguments = json.loads(tool_call.function.arguments or '{}')
+                    if function_name in self.standard_tools:
+                        return {"decision": {"name": function_name, "arguments": arguments}}
                     else:
-                         return {"error": f"Model called unknown function '{function_name}'."}
+                        return {"error": f"Model called unknown function '{function_name}'."}
                 except json.JSONDecodeError:
-                    return {"error": f"Error decoding function call arguments: {message.function_call.arguments}"}
+                    return {"error": f"Error decoding function arguments: {tool_call.function.arguments}"}
             else:
-                # Model decided not to call a function (or generated text instead)
+                # Model decided not to call a function
                 return {"error": f"OpenAI did not return a function call. Response: {message.content}"}
 
         except APIError as e:
@@ -362,7 +350,7 @@ class OpenAIAgent(LLMAgentBase):
 # --- Mistral Agent ---
 class MistralAgent(LLMAgentBase):
     """Uses Mistral AI API for decisions."""
-    def __init__(self, api_key: str | None = None, model: str = "mistral-large-latest", *args, **kwargs):
+    def __init__(self, api_key: str = None, model: str = "mistral-large-latest", *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.model = model
         used_api_key = api_key or os.environ.get("MISTRAL_API_KEY")
@@ -370,51 +358,51 @@ class MistralAgent(LLMAgentBase):
             raise ValueError("Mistral API key not provided or found in MISTRAL_API_KEY env var.")
         self.mistral_client = MistralAsyncClient(api_key=used_api_key)
 
-        # Convert standard schema to Mistral's tool format (very similar to OpenAI's)
-        self.mistral_tools = STANDARD_TOOL_SCHEMA
+        # Convert standard schema to Mistral's tool format
+        self.mistral_tools = list(self.standard_tools.values())
 
-    async def _get_llm_decision(self, battle_state: str) -> dict:
+    async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
         system_prompt = (
             "You are a skilled Pokemon battle AI. Your goal is to win the battle. "
             "Based on the current battle state, decide the best action: either use an available move or switch to an available Pokémon. "
             "Consider type matchups, HP, status conditions, field effects, entry hazards, and potential opponent actions. "
             "Only choose actions listed as available using their exact ID (for moves) or species name (for switches). "
-            "Use the provided tools/functions to indicate your choice."
+            "Use the provided tools to indicate your choice."
         )
         user_prompt = f"Current Battle State:\n{battle_state}\n\nChoose the best action by calling the appropriate function ('choose_move' or 'choose_switch')."
 
         try:
-            response = await self.mistral_client.chat.complete(
+            response = await self.mistral_client.chat(
                 model=self.model,
                 messages=[
-                    {"role": "system", "content": f"{system}"},
-                    {"role": "user", "content": f"{user_prompt}"}
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
                 ],
                 tools=self.mistral_tools,
-                tool_choice="auto", # Let the model choose
+                tool_choice="auto",  # Let the model choose
                 temperature=0.5,
             )
             
             message = response.choices[0].message
-            # Mistral returns tool_calls as a list
+            # Check for tool calls in the response
             if message.tool_calls:
-                tool_call = response.choices[0].message.tool_calls[0]
+                tool_call = message.tool_calls[0]  # Get the first tool call
                 function_name = tool_call.function.name
-                function_params = json.loads(tool_call.function.arguments)
-                print("\nfunction_name: ", function_name, "\nfunction_params: ", function_params)
-                if function_name and function_params: # Validate function name
-                    return {"decision": {"name": function_name, "arguments": arguments}}
-
+                try:
+                    arguments = json.loads(tool_call.function.arguments or '{}')
+                    if function_name in self.standard_tools:
+                        return {"decision": {"name": function_name, "arguments": arguments}}
+                    else:
+                        return {"error": f"Model called unknown function '{function_name}'."}
+                except json.JSONDecodeError:
+                    return {"error": f"Error decoding function arguments: {tool_call.function.arguments}"}
             else:
-                # Model decided not to call a tool (or generated text instead)
-                 return {"error": f"Mistral did not return a tool call. Response: {message.content}"}
+                # Model decided not to call a tool
+                return {"error": f"Mistral did not return a tool call. Response: {message.content}"}
 
-        # Mistral client might raise specific exceptions, add them here if known
-        # from mistralai.exceptions import MistralAPIException # Example
-        except Exception as e: # Catch general exceptions for now
+        except MistralAPIError as e:
+            print(f"Error during Mistral API call: {e}")
+            return {"error": f"Mistral API Error: {str(e)}"}
+        except Exception as e:
             print(f"Error during Mistral API call: {e}")
-            # Try to get specific details if it's a known exception type
-            error_details = str(e)
-            # if isinstance(e, MistralAPIException): # Example
-            #     error_details = f"{e.status_code} - {e.message}"
-            return {"error": f"Mistral API Error: {error_details}"}
+            return {"error": f"Unexpected error: {str(e)}"}