Spaces:
Running
Running
| import os | |
| import json | |
| import asyncio | |
| import random | |
| # --- OpenAI --- | |
| from openai import AsyncOpenAI, APIError | |
| # --- Google Gemini --- | |
| from google import genai | |
| from google.genai import types | |
| # --- Mistral AI --- | |
| from mistralai import Mistral | |
| # --- 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 --- | |
| 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): | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.standard_tools = STANDARD_TOOL_SCHEMA | |
| self.battle_history = [] | |
| def _format_battle_state(self, battle: Battle) -> str: | |
| 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) -> Optional[Move]: | |
| 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) -> Optional[Pokemon]: | |
| 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: | |
| 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) | |
| 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) | |
| 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 error_message: | |
| fallback_reason = f"API Error: {error_message}" | |
| elif decision is None: | |
| fallback_reason = "LLM did not provide a valid function call." | |
| else: | |
| fallback_reason = "Unknown error processing LLM decision." | |
| print(f"Warning: {fallback_reason} Choosing random action.") | |
| 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.") | |
| return self.choose_default_move(battle) | |
| 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, 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") | |
| if not used_api_key: | |
| raise ValueError("Google API key not provided or found in GOOGLE_API_KEY env var.") | |
| # 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 | |
| ) | |
| 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. " | |
| "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: | |
| # Use the async API for Gemini | |
| response = await self.model.generate_content_async( | |
| prompt, | |
| generation_config={"temperature": 0.5} | |
| ) | |
| 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}"} | |
| 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 Exception as e: | |
| print(f"Unexpected error during Gemini processing: {e}") | |
| import traceback | |
| 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, 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_tools = list(self.standard_tools.values()) | |
| 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 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}, | |
| ], | |
| tools=self.openai_tools, | |
| tool_choice="auto", # Let the model choose | |
| temperature=0.5, | |
| ) | |
| message = response.choices[0].message | |
| # 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: | |
| 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 function | |
| 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, 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 = Mistral(api_key=used_api_key) | |
| # 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[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 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": system_prompt}, | |
| {"role": "user", "content": user_prompt} | |
| ], | |
| tools=self.mistral_tools, | |
| tool_choice="auto", # Let the model choose | |
| temperature=0.5, | |
| ) | |
| message = response.choices[0].message | |
| # 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: | |
| 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 | |
| return {"error": f"Mistral did not return a tool call. Response: {message.content}"} | |
| except Exception as e: | |
| print(f"Error during Mistral API call: {e}") | |
| return {"error": f"Unexpected error: {str(e)}"} |