Update agents.py

agents.py

@@ -56,6 +56,44 @@ STANDARD_TOOL_SCHEMA = {
     },
 }
 
+# --- OpenAI Tools Schema (with 'type' field) ---
+OPENAI_TOOL_SCHEMA = {
+    "choose_move": {
+        "type": "function",
+        "function": {
+            "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": {
+        "type": "function",
+        "function": {
+            "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):
@@ -202,26 +240,15 @@ class GeminiAgent(LLMAgentBase):
         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)
+        # Initialize Gemini client using the correct API
+        self.genai_client = genai.Client(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
-        )
+        # Configure the tools for function calling
+        self.function_declarations = list(self.standard_tools.values())
 
     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). "
@@ -231,49 +258,40 @@ class GeminiAgent(LLMAgentBase):
         )
 
         try:
-            # Use the async API for Gemini
-            response = await self.model.generate_content_async(
-                prompt,
-                generation_config={"temperature": 0.5}
+            # Configure tools using the Gemini API format
+            tools = genai.types.Tool(function_declarations=self.function_declarations)
+            config = genai.types.GenerateContentConfig(tools=[tools])
+            
+            # Send request to the model
+            response = self.genai_client.models.generate_content(
+                model=self.model_name,
+                contents=prompt,
+                config=config
             )
-            print("GEMINI RESPONSE : ",response)
-            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}"}
+            print("GEMINI RESPONSE : ", response)
+            
+            # Check for function calls in the response
+            if (hasattr(response, 'candidates') and 
+                response.candidates and 
+                hasattr(response.candidates[0], 'content') and
+                hasattr(response.candidates[0].content, 'parts') and
+                response.candidates[0].content.parts and
+                hasattr(response.candidates[0].content.parts[0], 'function_call')):
+                
+                function_call = response.candidates[0].content.parts[0].function_call
+                function_name = function_call.name
+                # Get arguments
+                arguments = {}
+                if hasattr(function_call, 'args'):
+                    arguments = function_call.args
+                
+                if function_name in self.standard_tools:
+                    return {"decision": {"name": function_name, "arguments": arguments}}
+                else:
+                    return {"error": f"Model called unknown function '{function_name}'."}
             
-            # 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]}..."}
+            # No function call found
+            return {"error": "Gemini did not return a function call."}
             
         except Exception as e:
             print(f"Unexpected error during Gemini processing: {e}")
@@ -293,8 +311,8 @@ class OpenAIAgent(LLMAgentBase):
             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())
+        # Use the OpenAI-specific schema with type field
+        self.openai_tools = list(OPENAI_TOOL_SCHEMA.values())
 
     async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
         system_prompt = (
@@ -354,8 +372,17 @@ class MistralAgent(LLMAgentBase):
             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())
+        # Convert standard schema to Mistral's tool format with "function" wrapper
+        self.mistral_tools = []
+        for tool_name, tool_schema in self.standard_tools.items():
+            self.mistral_tools.append({
+                "type": "function",
+                "function": {
+                    "name": tool_schema["name"],
+                    "description": tool_schema["description"],
+                    "parameters": tool_schema["parameters"]
+                }
+            })
 
     async def _get_llm_decision(self, battle_state: str) -> Dict[str, Any]:
         system_prompt = (
@@ -368,23 +395,29 @@ class MistralAgent(LLMAgentBase):
         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(
+            # Create the messages array
+            messages = [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_prompt}
+            ]
+            
+            # Call the Mistral API with tool_choice set to "any" to force tool usage
+            response = self.mistral_client.chat.complete(
                 model=self.model,
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt}
-                ],
+                messages=messages,
                 tools=self.mistral_tools,
-                tool_choice="auto",  # Let the model choose
-                temperature=0.5,
+                tool_choice="any",  # Force the model to use a tool
+                temperature=0.3,
             )
-            print("Mistral RESPONSE : ",response)
-            message = response.choices[0].message
+            print("Mistral RESPONSE : ", response)
+            
             # Check for tool calls in the response
-            if message.tool_calls:
+            message = response.choices[0].message
+            if hasattr(message, 'tool_calls') and message.tool_calls:
                 tool_call = message.tool_calls[0]  # Get the first tool call
                 function_name = tool_call.function.name
                 try:
+                    # Parse the function arguments from JSON string
                     arguments = json.loads(tool_call.function.arguments or '{}')
                     if function_name in self.standard_tools:
                         return {"decision": {"name": function_name, "arguments": arguments}}
@@ -393,9 +426,11 @@ class MistralAgent(LLMAgentBase):
                 except json.JSONDecodeError:
                     return {"error": f"Error decoding function arguments: {tool_call.function.arguments}"}
             else:
-                # Model decided not to call a tool
+                # Model did not return a tool call
                 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}")
+            import traceback
+            traceback.print_exc()
             return {"error": f"Unexpected error: {str(e)}"}