Create agents.py

agents.py

@@ -0,0 +1,420 @@
+import os
+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)
+
+# --- Google Gemini ---
+from google import genai
+from google.genai import types
+# It's good practice to also import potential exceptions
+
+
+# --- Mistral AI ---
+from mistralai.async_client import MistralAsyncClient # (Keep if needed)
+
+# --- 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
+
+# --- Helper Function & Base Class (Assuming they are defined above) ---
+def normalize_name(name: str) -> str:
+    """Lowercase and remove non-alphanumeric characters."""
+    return "".join(filter(str.isalnum, name)).lower()
+
+STANDARD_TOOL_SCHEMA = {
+    "choose_move": {
+        "name": "choose_move",
+        "description": "Selects and executes an available attacking or status move.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "move_name": {
+                    "type": "string",
+                    "description": "The exact name or ID (e.g., 'thunderbolt', 'swordsdance') of the move to use. Must be one of the available moves.",
+                },
+            },
+            "required": ["move_name"],
+        },
+    },
+    "choose_switch": {
+        "name": "choose_switch",
+        "description": "Selects an available Pokémon from the bench to switch into.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "pokemon_name": {
+                    "type": "string",
+                    "description": "The exact name of the Pokémon species to switch to (e.g., 'Pikachu', 'Charizard'). Must be one of the available switches.",
+                },
+            },
+            "required": ["pokemon_name"],
+        },
+    },
+}
+
+
+class LLMAgentBase(Player): # Make sure this base class exists
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.standard_tools = STANDARD_TOOL_SCHEMA
+        self.battle_history = [] # Example attribute
+
+    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))}) " \
+                           f"HP: {active_pkmn.current_hp_fraction * 100:.1f}% " \
+                           f"Status: {active_pkmn.status.name if active_pkmn.status else 'None'} " \
+                           f"Boosts: {active_pkmn.boosts}"
+
+        opponent_pkmn = battle.opponent_active_pokemon
+        opp_info_str = "Unknown"
+        if opponent_pkmn:
+            opp_info_str = f"{opponent_pkmn.species} " \
+                           f"(Type: {'/'.join(map(str, opponent_pkmn.types))}) " \
+                           f"HP: {opponent_pkmn.current_hp_fraction * 100:.1f}% " \
+                           f"Status: {opponent_pkmn.status.name if opponent_pkmn.status else 'None'} " \
+                           f"Boosts: {opponent_pkmn.boosts}"
+        opponent_pkmn_info = f"Opponent's active Pokemon: {opp_info_str}"
+
+        available_moves_info = "Available moves:\n"
+        if battle.available_moves:
+            available_moves_info += "\n".join(
+                [f"- {move.id} (Type: {move.type}, BP: {move.base_power}, Acc: {move.accuracy}, PP: {move.current_pp}/{move.max_pp}, Cat: {move.category.name})"
+                 for move in battle.available_moves]
+            )
+        else:
+             available_moves_info += "- None (Must switch or Struggle)"
+
+        available_switches_info = "Available switches:\n"
+        if battle.available_switches:
+              available_switches_info += "\n".join(
+                  [f"- {pkmn.species} (HP: {pkmn.current_hp_fraction * 100:.1f}%, Status: {pkmn.status.name if pkmn.status else 'None'})"
+                   for pkmn in battle.available_switches]
+              )
+        else:
+            available_switches_info += "- None"
+
+        state_str = f"{active_pkmn_info}\n" \
+                    f"{opponent_pkmn_info}\n\n" \
+                    f"{available_moves_info}\n\n" \
+                    f"{available_switches_info}\n\n" \
+                    f"Weather: {battle.weather}\n" \
+                    f"Terrains: {battle.fields}\n" \
+                    f"Your Side Conditions: {battle.side_conditions}\n" \
+                    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)
+        normalized_name = normalize_name(move_name)
+        # Prioritize exact ID match
+        for move in battle.available_moves:
+            if move.id == normalized_name:
+                return move
+        # Fallback: Check display name (less reliable)
+        for move in battle.available_moves:
+            if move.name.lower() == move_name.lower():
+                print(f"Warning: Matched move by display name '{move.name}' instead of ID '{move.id}'. Input was '{move_name}'.")
+                return move
+        return None
+
+    def _find_pokemon_by_name(self, battle: Battle, pokemon_name: str) -> Pokemon | None:
+        # (Implementation as provided in the question)
+        normalized_name = normalize_name(pokemon_name)
+        for pkmn in battle.available_switches:
+            # Normalize the species name for comparison
+            if normalize_name(pkmn.species) == normalized_name:
+                return pkmn
+        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")
+        error_message = decision_result.get("error")
+        action_taken = False
+        fallback_reason = ""
+
+        if decision:
+            function_name = decision.get("name")
+            args = decision.get("arguments", {})
+            if function_name == "choose_move":
+                move_name = args.get("move_name")
+                if move_name:
+                    chosen_move = self._find_move_by_name(battle, move_name)
+                    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
+                        return self.create_order(chosen_move)
+                    else:
+                        fallback_reason = f"LLM chose unavailable/invalid move '{move_name}'."
+                else:
+                     fallback_reason = "LLM 'choose_move' called without 'move_name'."
+            elif function_name == "choose_switch":
+                pokemon_name = args.get("pokemon_name")
+                if pokemon_name:
+                    chosen_switch = self._find_pokemon_by_name(battle, pokemon_name)
+                    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
+                        return self.create_order(chosen_switch)
+                    else:
+                        fallback_reason = f"LLM chose unavailable/invalid switch '{pokemon_name}'."
+                else:
+                    fallback_reason = "LLM 'choose_switch' called without 'pokemon_name'."
+            else:
+                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 error_message:
+                     fallback_reason = f"API Error: {error_message}"
+                 elif decision is None: # Model didn't call a function or response was bad
+                      fallback_reason = "LLM did not provide a valid function call."
+                 else: # Should not happen if logic above is correct
+                      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
+
+    async def _get_llm_decision(self, battle_state: str) -> dict:
+        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
+        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')
+        )
+        # --- 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:
+        """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. "
+            "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 functions to indicate your choice.\n\n"
+            f"Current Battle State:\n{battle_state}\n\n"
+            "Choose the best action by calling the appropriate function ('choose_move' or 'choose_switch')."
+        )
+
+        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
+            )
+            # --- 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}"}
+
+            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
+                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 ---
+        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
+            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):
+        super().__init__(*args, **kwargs)
+        self.model = model
+        used_api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        if not used_api_key:
+            raise ValueError("OpenAI API key not provided or found in OPENAI_API_KEY env var.")
+        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()]
+
+    async def _get_llm_decision(self, battle_state: str) -> dict:
+        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 functions 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.openai_client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt},
+                ],
+                functions=STANDARD_TOOL_SCHEMA,
+                function_call="auto",
+                temperature=0.5,
+            )
+            message = response.choices[0].message
+            print("MESSAGE BACK : ", message)
+            if message.function_call:
+                function_name = message.function_call.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}}
+                    else:
+                         return {"error": f"Model called unknown function '{function_name}'."}
+                except json.JSONDecodeError:
+                    return {"error": f"Error decoding function call arguments: {message.function_call.arguments}"}
+            else:
+                # Model decided not to call a function (or generated text instead)
+                return {"error": f"OpenAI did not return a function call. Response: {message.content}"}
+
+        except APIError as e:
+            print(f"Error during OpenAI API call: {e}")
+            return {"error": f"OpenAI API Error: {e.status_code} - {e.message}"}
+        except Exception as e:
+            print(f"Unexpected error during OpenAI API call: {e}")
+            return {"error": f"Unexpected error: {e}"}
+
+
+# --- 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):
+        super().__init__(*args, **kwargs)
+        self.model = model
+        used_api_key = api_key or os.environ.get("MISTRAL_API_KEY")
+        if not used_api_key:
+            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
+
+    async def _get_llm_decision(self, battle_state: str) -> dict:
+        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."
+        )
+        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(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": f"{system}"},
+                    {"role": "user", "content": f"{user_prompt}"}
+                ],
+                tools=self.mistral_tools,
+                tool_choice="auto", # Let the model choose
+                temperature=0.5,
+            )
+            
+            message = response.choices[0].message
+            # Mistral returns tool_calls as a list
+            if message.tool_calls:
+                tool_call = response.choices[0].message.tool_calls[0]
+                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}}
+
+            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}"}
+
+        # 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
+            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}"}