tooluniverse 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

tooluniverse/agentic_tool.py

@@ -1,27 +1,57 @@
 from __future__ import annotations
 
 import os
-import re
-import time
+import json
 from datetime import datetime
 from typing import Any, Dict, List, Optional
 
-import openai
-from google import genai
-
 from .base_tool import BaseTool
 from .tool_registry import register_tool
 from .logging_config import get_logger
+from .llm_clients import AzureOpenAIClient, GeminiClient, OpenRouterClient
+
+
+# Global default fallback configuration
+DEFAULT_FALLBACK_CHAIN = [
+    {"api_type": "CHATGPT", "model_id": "gpt-4o-1120"},
+    {"api_type": "OPENROUTER", "model_id": "openai/gpt-4o"},
+    {"api_type": "GEMINI", "model_id": "gemini-2.0-flash"},
+]
+
+# API key environment variable mapping
+API_KEY_ENV_VARS = {
+    "CHATGPT": ["AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT"],
+    "OPENROUTER": ["OPENROUTER_API_KEY"],
+    "GEMINI": ["GEMINI_API_KEY"],
+}
 
 
 @register_tool("AgenticTool")
 class AgenticTool(BaseTool):
     """Generic wrapper around LLM prompting supporting JSON-defined configs with prompts and input arguments."""
 
+    @staticmethod
+    def has_any_api_keys() -> bool:
+        """
+        Check if any API keys are available across all supported API types.
+
+        Returns:
+            bool: True if at least one API type has all required keys, False otherwise
+        """
+        for _api_type, required_vars in API_KEY_ENV_VARS.items():
+            all_keys_present = True
+            for var in required_vars:
+                if not os.getenv(var):
+                    all_keys_present = False
+                    break
+            if all_keys_present:
+                return True
+        return False
+
     def __init__(self, tool_config: Dict[str, Any]):
         super().__init__(tool_config)
-        self.logger = get_logger("AgenticTool")  # Initialize logger
-        self.name: str = tool_config.get("name", "")  # Add name attribute
+        self.logger = get_logger("AgenticTool")
+        self.name: str = tool_config.get("name", "")
         self._prompt_template: str = tool_config.get("prompt", "")
         self._input_arguments: List[str] = tool_config.get("input_arguments", [])
 
@@ -35,8 +65,6 @@ class AgenticTool(BaseTool):
         for arg in self._input_arguments:
             if arg not in self._required_arguments:
                 prop_info = properties.get(arg, {})
-
-                # First check if there's an explicit "default" field
                 if "default" in prop_info:
                     self._argument_defaults[arg] = prop_info["default"]
 
@@ -50,12 +78,30 @@ class AgenticTool(BaseTool):
         # LLM configuration
         self._api_type: str = get_config("api_type", "CHATGPT")
         self._model_id: str = get_config("model_id", "o1-mini")
-        self._temperature: float = get_config("temperature", 0.1)
-        self._max_new_tokens: int = get_config("max_new_tokens", 2048)
+        self._temperature: Optional[float] = get_config("temperature", 0.1)
+        # Ignore configured max_new_tokens; client will resolve per model/env
+        self._max_new_tokens: Optional[int] = None
         self._return_json: bool = get_config("return_json", False)
         self._max_retries: int = get_config("max_retries", 5)
         self._retry_delay: int = get_config("retry_delay", 5)
         self.return_metadata: bool = get_config("return_metadata", True)
+        self._validate_api_key: bool = get_config("validate_api_key", True)
+
+        # API fallback configuration
+        self._fallback_api_type: Optional[str] = get_config("fallback_api_type", None)
+        self._fallback_model_id: Optional[str] = get_config("fallback_model_id", None)
+
+        # Global fallback configuration
+        self._use_global_fallback: bool = get_config("use_global_fallback", True)
+        self._global_fallback_chain: List[Dict[str, str]] = (
+            self._get_global_fallback_chain()
+        )
+
+        # Gemini model configuration (optional; env override)
+        self._gemini_model_id: str = get_config(
+            "gemini_model_id",
+            __import__("os").getenv("GEMINI_MODEL_ID", "gemini-2.0-flash"),
+        )
 
         # Validation
         if not self._prompt_template:
@@ -65,8 +111,11 @@ class AgenticTool(BaseTool):
                 "AgenticTool requires 'input_arguments' in the configuration."
             )
 
-        # Validate temperature range
-        if not 0 <= self._temperature <= 2:
+        # Validate temperature range (skip if None)
+        if (
+            isinstance(self._temperature, (int, float))
+            and not 0 <= self._temperature <= 2
+        ):
             self.logger.warning(
                 f"Temperature {self._temperature} is outside recommended range [0, 2]"
             )
@@ -74,226 +123,163 @@ class AgenticTool(BaseTool):
         # Validate model compatibility
         self._validate_model_config()
 
-        # Initialize the LLM model
-        try:
-            self._model, self._tokenizer = self._init_llm(
-                api_type=self._api_type, model_id=self._model_id
-            )
-            self.logger.debug(
-                f"Successfully initialized {self._api_type} model: {self._model_id}"
-            )
-        except Exception as e:
-            self.logger.error(f"Failed to initialize LLM model: {str(e)}")
-            raise
+        # Initialize the provider client
+        self._llm_client = None
+        self._initialization_error = None
+        self._is_available = False
+        self._current_api_type = None
+        self._current_model_id = None
+
+        # Try primary API first, then fallback if configured
+        self._try_initialize_api()
+
+    def _get_global_fallback_chain(self) -> List[Dict[str, str]]:
+        """Get the global fallback chain from environment or use default."""
+        # Check environment variable for custom fallback chain
+        env_chain = os.getenv("AGENTIC_TOOL_FALLBACK_CHAIN")
+        if env_chain:
+            try:
+                chain = json.loads(env_chain)
+                if isinstance(chain, list) and all(
+                    isinstance(item, dict) and "api_type" in item and "model_id" in item
+                    for item in chain
+                ):
+                    return chain
+                else:
+                    self.logger.warning(
+                        "Invalid fallback chain format in environment variable"
+                    )
+            except json.JSONDecodeError:
+                self.logger.warning(
+                    "Invalid JSON in AGENTIC_TOOL_FALLBACK_CHAIN environment variable"
+                )
 
-    # ------------------------------------------------------------------ LLM utilities -----------
-    def _validate_model_config(self):
-        """Validate model configuration parameters."""
-        supported_api_types = ["CHATGPT", "GEMINI"]
-        if self._api_type not in supported_api_types:
-            raise ValueError(
-                f"Unsupported API type: {self._api_type}. Supported types: {supported_api_types}"
-            )
+        return DEFAULT_FALLBACK_CHAIN.copy()
 
-        # Validate model-specific configurations
-        # if self._api_type == "CHATGPT":
-        #     supported_models = ["gpt-4o", "o1-mini", "o3-mini"]
-        #     if self._model_id not in supported_models:
-        #         self.logger.warning(f"Model {self._model_id} may not be supported. Supported models: {supported_models}")
+    def _try_initialize_api(self):
+        """Try to initialize the primary API, fallback to secondary if configured."""
+        # Try primary API first
+        if self._try_api(self._api_type, self._model_id):
+            return
 
-        # Validate token limits
-        if self._max_new_tokens <= 0:
-            raise ValueError("max_new_tokens must be positive")
+        # Try explicit fallback API if configured
+        if self._fallback_api_type and self._fallback_model_id:
+            self.logger.info(
+                f"Primary API {self._api_type} failed, trying explicit fallback {self._fallback_api_type}"
+            )
+            if self._try_api(self._fallback_api_type, self._fallback_model_id):
+                return
 
-        if self._max_new_tokens > 8192:  # Conservative limit
-            self.logger.warning(
-                f"max_new_tokens {self._max_new_tokens} is very high and may cause API issues"
+        # Try global fallback chain if enabled
+        if self._use_global_fallback:
+            self.logger.info(
+                f"Primary API {self._api_type} failed, trying global fallback chain"
             )
+            for fallback_config in self._global_fallback_chain:
+                fallback_api = fallback_config["api_type"]
+                fallback_model = fallback_config["model_id"]
 
-    def _init_llm(self, api_type: str, model_id: str):
-        """Initialize the LLM model and tokenizer based on API type."""
-        if api_type == "CHATGPT":
-            if "gpt-4o" in model_id or model_id is None:
-                api_key = os.getenv("AZURE_OPENAI_API_KEY")
-                api_version = "2024-12-01-preview"
-            elif (
-                "o1-mini" in model_id or "o3-mini" in model_id or "o4-mini" in model_id
-            ):
-                api_key = os.getenv("AZURE_OPENAI_API_KEY")
-                api_version = "2025-03-01-preview"
-            else:
-                self.logger.error(
-                    f"Invalid model_id. Please use 'gpt-4o', 'o1-mini', or 'o3-mini'. Got: {model_id}"
-                )
-                raise ValueError(f"Unsupported model_id: {model_id}")
+                # Skip if it's the same as primary or explicit fallback
+                if (
+                    fallback_api == self._api_type and fallback_model == self._model_id
+                ) or (
+                    fallback_api == self._fallback_api_type
+                    and fallback_model == self._fallback_model_id
+                ):
+                    continue
 
-            if not api_key:
-                raise ValueError(
-                    "API key not found in environment. Please set the appropriate environment variable."
+                self.logger.info(
+                    f"Trying global fallback: {fallback_api} ({fallback_model})"
                 )
+                if self._try_api(fallback_api, fallback_model):
+                    return
 
-            azure_endpoint = os.getenv(
-                "AZURE_OPENAI_ENDPOINT", "https://azure-ai.hms.edu"
-            )
+        # If we get here, all APIs failed
+        self.logger.warning(
+            f"Tool '{self.name}' failed to initialize with all available APIs"
+        )
 
-            from openai import AzureOpenAI
+    def _try_api(self, api_type: str, model_id: str) -> bool:
+        """Try to initialize a specific API and model."""
+        try:
+            if api_type == "CHATGPT":
+                self._llm_client = AzureOpenAIClient(model_id, None, self.logger)
+            elif api_type == "OPENROUTER":
+                self._llm_client = OpenRouterClient(model_id, self.logger)
+            elif api_type == "GEMINI":
+                self._llm_client = GeminiClient(model_id, self.logger)
+            else:
+                raise ValueError(f"Unsupported API type: {api_type}")
 
-            self.logger.debug(
-                "Initializing AzureOpenAI client with endpoint:", azure_endpoint
-            )
-            self.logger.debug("Using API version:", api_version)
-            model_client = AzureOpenAI(
-                azure_endpoint=azure_endpoint,
-                api_key=api_key,
-                api_version=api_version,
-            )
-            model = {
-                "model": model_client,
-                "model_name": model_id,
-                "api_version": api_version,
-            }
-            tokenizer = None
-        elif api_type == "GEMINI":
-            api_key = os.getenv("GEMINI_API_KEY")
-            if not api_key:
-                raise ValueError("GEMINI_API_KEY not found in environment variables")
-
-            model = genai.Client(api_key=api_key)
-            tokenizer = None
-        else:
-            raise ValueError(f"Unsupported API type: {api_type}")
-
-        return model, tokenizer
-
-    def _chatgpt_infer(
-        self,
-        model: Dict[str, Any],
-        messages: List[Dict[str, str]],
-        temperature: float = 0.1,
-        max_new_tokens: int = 2048,
-        return_json: bool = False,
-        max_retries: int = 5,
-        retry_delay: int = 5,
-        custom_format=None,
-    ) -> Optional[str]:
-        """Inference function for ChatGPT models including o1-mini and o3-mini."""
-        model_client = model["model"]
-        model_name = model["model_name"]
-
-        retries = 0
-        import traceback
-
-        if custom_format is not None:
-            response_format = custom_format
-            call_function = model_client.chat.completions.parse
-        elif return_json:
-            response_format = {"type": "json_object"}
-            call_function = model_client.chat.completions.create
-        else:
-            response_format = None
-            call_function = model_client.chat.completions.create
-        while retries < max_retries:
-            try:
-                if "gpt-4o" in model_name:
-                    responses = call_function(
-                        model=model_name,
-                        messages=messages,
-                        temperature=temperature,
-                        max_tokens=max_new_tokens,
-                        response_format=response_format,
-                    )
-                elif (
-                    "o1-mini" in model_name
-                    or "o3-mini" in model_name
-                    or "o4-mini" in model_name
-                ):
-                    responses = call_function(
-                        model=model_name,
-                        messages=messages,
-                        max_completion_tokens=max_new_tokens,
-                        response_format=response_format,
-                    )
-                if custom_format is not None:
-                    response = responses.choices[0].message.parsed.model_dump()
-                else:
-                    response = responses.choices[0].message.content
-                # print("\033[92m" + response + "\033[0m")
-                # usage = responses.usage
-                # print("\033[95m" + str(usage) + "\033[0m")
-                return response
-            except openai.RateLimitError:
-                self.logger.warning(
-                    f"Rate limit exceeded. Retrying in {retry_delay} seconds..."
-                )
-                retries += 1
-                time.sleep(retry_delay * retries)
-            except Exception as e:
-                self.logger.error(f"An error occurred: {e}")
-                traceback.print_exc()
-                break
-        self.logger.error("Max retries exceeded. Unable to complete the request.")
-        return None
-
-    def _gemini_infer(
-        self,
-        model: Any,
-        messages: List[Dict[str, str]],
-        temperature: float = 0.1,
-        max_new_tokens: int = 2048,
-        return_json: bool = False,
-        max_retries: int = 5,
-        retry_delay: int = 5,
-        model_name: str = "gemini-2.0-flash",
-    ) -> Optional[str]:
-        """Inference function for Gemini models."""
-        retries = 0
-        contents = ""
-        for message in messages:
-            if message["role"] == "user" or message["role"] == "system":
-                contents += f"{message['content']}\n"
-            elif message["role"] == "assistant":
-                raise ValueError(
-                    "Gemini model does not support assistant role in messages for now in the code."
+            # Test API key validity after initialization (if enabled)
+            if self._validate_api_key:
+                self._llm_client.test_api()
+                self.logger.debug(
+                    f"Successfully initialized {api_type} model: {model_id}"
                 )
             else:
-                raise ValueError(
-                    "Invalid role in messages. Only 'user' and 'system' roles are supported."
+                self.logger.info("API key validation skipped (validate_api_key=False)")
+
+            self._is_available = True
+            self._current_api_type = api_type
+            self._current_model_id = model_id
+            self._initialization_error = None
+
+            if api_type != self._api_type or model_id != self._model_id:
+                self.logger.info(
+                    f"Using fallback API: {api_type} with model {model_id} "
+                    f"(originally configured: {self._api_type} with {self._model_id})"
                 )
 
-        if return_json:
+            return True
+
+        except Exception as e:
+            error_msg = f"Failed to initialize {api_type} model {model_id}: {str(e)}"
+            self.logger.warning(error_msg)
+            self._initialization_error = error_msg
+            return False
+
+    # ------------------------------------------------------------------ LLM utilities -----------
+    def _validate_model_config(self):
+        supported_api_types = ["CHATGPT", "OPENROUTER", "GEMINI"]
+        if self._api_type not in supported_api_types:
             raise ValueError(
-                "Gemini model does not support JSON format for now in the code."
+                f"Unsupported API type: {self._api_type}. Supported types: {supported_api_types}"
             )
-
-        while retries < max_retries:
-            try:
-                response = model.models.generate_content(
-                    model=model_name,
-                    contents=contents,
-                    config=genai.types.GenerateContentConfig(
-                        max_output_tokens=max_new_tokens,
-                        temperature=temperature,
-                    ),
-                )
-                return response.text
-            except openai.RateLimitError:
-                self.logger.warning(
-                    f"Rate limit exceeded. Retrying in {retry_delay} seconds..."
-                )
-                retries += 1
-                time.sleep(retry_delay * retries)
-            except Exception as e:
-                self.logger.error(f"An error occurred: {e}")
-                break
-        return None
+        if self._max_new_tokens is not None and self._max_new_tokens <= 0:
+            raise ValueError("max_new_tokens must be positive or None")
 
     # ------------------------------------------------------------------ public API --------------
     def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
-        """Execute the tool by formatting the prompt with input arguments and querying the LLM."""
         start_time = datetime.now()
 
+        # Check if tool is available before attempting to run
+        if not self._is_available:
+            error_msg = f"Tool '{self.name}' is not available due to initialization error: {self._initialization_error}"
+            self.logger.error(error_msg)
+            if self.return_metadata:
+                return {
+                    "success": False,
+                    "error": error_msg,
+                    "error_type": "ToolUnavailable",
+                    "metadata": {
+                        "prompt_used": "Tool unavailable",
+                        "input_arguments": {
+                            arg: arguments.get(arg) for arg in self._input_arguments
+                        },
+                        "model_info": {
+                            "api_type": self._api_type,
+                            "model_id": self._model_id,
+                        },
+                        "execution_time_seconds": 0,
+                        "timestamp": start_time.isoformat(),
+                    },
+                }
+            else:
+                return f"error: {error_msg} error_type: ToolUnavailable"
+
         try:
-            # Validate that all required input arguments are provided
+            # Validate required args
             missing_required_args = [
                 arg for arg in self._required_arguments if arg not in arguments
             ]
@@ -302,25 +288,27 @@ class AgenticTool(BaseTool):
                     f"Missing required input arguments: {missing_required_args}"
                 )
 
-            # Add default values for optional arguments that are missing
+            # Fill defaults for optional args
             for arg in self._input_arguments:
                 if arg not in arguments:
-                    if arg in self._argument_defaults:
-                        arguments[arg] = self._argument_defaults[arg]
-                    else:
-                        arguments[arg] = ""  # Default to empty string for optional args
+                    arguments[arg] = self._argument_defaults.get(arg, "")
 
-            # Validate argument types and content
             self._validate_arguments(arguments)
-
-            # Format the prompt template with the provided arguments
             formatted_prompt = self._format_prompt(arguments)
 
-            # Prepare messages for the LLM
             messages = [{"role": "user", "content": formatted_prompt}]
             custom_format = arguments.get("response_format", None)
-            # Call the appropriate LLM function based on API type
-            response = self._call_llm(messages, custom_format=custom_format)
+
+            # Delegate to client; client handles provider-specific logic
+            response = self._llm_client.infer(
+                messages=messages,
+                temperature=self._temperature,
+                max_tokens=None,  # client resolves per-model defaults/env
+                return_json=self._return_json,
+                custom_format=custom_format,
+                max_retries=self._max_retries,
+                retry_delay=self._retry_delay,
+            )
 
             end_time = datetime.now()
             execution_time = (end_time - start_time).total_seconds()
@@ -350,13 +338,10 @@ class AgenticTool(BaseTool):
                 }
             else:
                 return response
-
         except Exception as e:
             end_time = datetime.now()
             execution_time = (end_time - start_time).total_seconds()
-
             self.logger.error(f"Error executing {self.name}: {str(e)}")
-
             if self.return_metadata:
                 return {
                     "success": False,
@@ -384,94 +369,28 @@ class AgenticTool(BaseTool):
 
     # ------------------------------------------------------------------ helpers -----------------
     def _validate_arguments(self, arguments: Dict[str, Any]):
-        """Validate input arguments for common issues."""
         for arg_name, value in arguments.items():
             if arg_name in self._input_arguments:
-                # Check for empty strings only for required arguments
                 if isinstance(value, str) and not value.strip():
                     if arg_name in self._required_arguments:
                         raise ValueError(
                             f"Required argument '{arg_name}' cannot be empty"
                         )
-                    # Optional arguments can be empty, so we skip the check
-
-                # Check for extremely long inputs that might cause issues - silent validation
-                if (
-                    isinstance(value, str) and len(value) > 100000
-                ):  # 100k character limit
-                    pass  # Could potentially cause API issues but no need to spam output
+                if isinstance(value, str) and len(value) > 100000:
+                    pass
 
     def _format_prompt(self, arguments: Dict[str, Any]) -> str:
-        """Format the prompt template with the provided arguments."""
         prompt = self._prompt_template
-
-        # Track which placeholders we actually replace
-        replaced_placeholders = set()
-
-        # Replace placeholders in the format {argument_name} with actual values
         for arg_name in self._input_arguments:
             placeholder = f"{{{arg_name}}}"
             value = arguments.get(arg_name, "")
-
             if placeholder in prompt:
-                replaced_placeholders.add(arg_name)
-                # Handle special characters and formatting
-                if isinstance(value, str):
-                    # Simple replacement without complex escaping that was causing issues
-                    prompt = prompt.replace(placeholder, str(value))
-                else:
-                    prompt = prompt.replace(placeholder, str(value))
-
-        # Check for unreplaced expected placeholders (only check our input arguments)
-        # _unreplaced_expected = [
-        #     arg for arg in self._input_arguments if arg not in replaced_placeholders
-        # ]
-
-        # Silent handling - no debug output needed for template patterns in JSON content
-
+                prompt = prompt.replace(placeholder, str(value))
         return prompt
 
-    def _call_llm(self, messages: List[Dict[str, str]], custom_format=None) -> str:
-        """Make the actual LLM API call using the appropriate function."""
-        if self._api_type == "CHATGPT":
-            response = self._chatgpt_infer(
-                model=self._model,
-                messages=messages,
-                temperature=self._temperature,
-                max_new_tokens=self._max_new_tokens,
-                return_json=self._return_json,
-                max_retries=self._max_retries,
-                retry_delay=self._retry_delay,
-                custom_format=custom_format,
-            )
-            if response is None:
-                raise Exception("LLM API call failed after maximum retries")
-            return response
-
-        elif self._api_type == "GEMINI":
-            response = self._gemini_infer(
-                model=self._model,
-                messages=messages,
-                temperature=self._temperature,
-                max_new_tokens=self._max_new_tokens,
-                return_json=self._return_json,
-                max_retries=self._max_retries,
-                retry_delay=self._retry_delay,
-            )
-            if response is None:
-                raise Exception("Gemini API call failed after maximum retries")
-            return response
-
-        else:
-            raise ValueError(f"Unsupported API type: {self._api_type}")
-
     def get_prompt_preview(self, arguments: Dict[str, Any]) -> str:
-        """Preview how the prompt will look with the given arguments (useful for debugging)."""
         try:
-            # Create a copy to avoid modifying the original arguments
             args_copy = arguments.copy()
-
-            # Validate that all required input arguments are provided
             missing_required_args = [
                 arg for arg in self._required_arguments if arg not in args_copy
             ]
@@ -479,21 +398,14 @@ class AgenticTool(BaseTool):
                 raise ValueError(
                     f"Missing required input arguments: {missing_required_args}"
                 )
-
-            # Add default values for optional arguments that are missing
             for arg in self._input_arguments:
                 if arg not in args_copy:
-                    if arg in self._argument_defaults:
-                        args_copy[arg] = self._argument_defaults[arg]
-                    else:
-                        args_copy[arg] = ""  # Default to empty string for optional args
-
+                    args_copy[arg] = self._argument_defaults.get(arg, "")
             return self._format_prompt(args_copy)
         except Exception as e:
             return f"Error formatting prompt: {str(e)}"
 
     def get_model_info(self) -> Dict[str, Any]:
-        """Get comprehensive information about the configured model."""
         return {
             "api_type": self._api_type,
             "model_id": self._model_id,
@@ -502,59 +414,85 @@ class AgenticTool(BaseTool):
             "return_json": self._return_json,
             "max_retries": self._max_retries,
             "retry_delay": self._retry_delay,
+            "validate_api_key": self._validate_api_key,
+            "gemini_model_id": getattr(self, "_gemini_model_id", None),
+            "is_available": self._is_available,
+            "initialization_error": self._initialization_error,
+            "current_api_type": self._current_api_type,
+            "current_model_id": self._current_model_id,
+            "fallback_api_type": self._fallback_api_type,
+            "fallback_model_id": self._fallback_model_id,
+            "use_global_fallback": self._use_global_fallback,
+            "global_fallback_chain": self._global_fallback_chain,
         }
 
+    def is_available(self) -> bool:
+        """Check if the tool is available for use."""
+        return self._is_available
+
+    def get_availability_status(self) -> Dict[str, Any]:
+        """Get detailed availability status of the tool."""
+        return {
+            "is_available": self._is_available,
+            "initialization_error": self._initialization_error,
+            "api_type": self._api_type,
+            "model_id": self._model_id,
+        }
+
+    def retry_initialization(self) -> bool:
+        """Attempt to reinitialize the tool (useful if API keys were updated)."""
+        try:
+            if self._api_type == "CHATGPT":
+                self._llm_client = AzureOpenAIClient(self._model_id, None, self.logger)
+            elif self._api_type == "OPENROUTER":
+                self._llm_client = OpenRouterClient(self._model_id, self.logger)
+            elif self._api_type == "GEMINI":
+                self._llm_client = GeminiClient(self._gemini_model_id, self.logger)
+            else:
+                raise ValueError(f"Unsupported API type: {self._api_type}")
+
+            if self._validate_api_key:
+                self._llm_client.test_api()
+                self.logger.info(
+                    f"Successfully reinitialized {self._api_type} model: {self._model_id}"
+                )
+
+            self._is_available = True
+            self._initialization_error = None
+            return True
+
+        except Exception as e:
+            self._initialization_error = str(e)
+            self.logger.warning(
+                f"Retry initialization failed for {self._api_type} model {self._model_id}: {str(e)}"
+            )
+            return False
+
     def get_prompt_template(self) -> str:
-        """Get the raw prompt template."""
         return self._prompt_template
 
     def get_input_arguments(self) -> List[str]:
-        """Get the list of required input arguments."""
         return self._input_arguments.copy()
 
     def validate_configuration(self) -> Dict[str, Any]:
-        """Validate the tool configuration and return validation results."""
         validation_results = {"valid": True, "warnings": [], "errors": []}
-
         try:
             self._validate_model_config()
         except ValueError as e:
             validation_results["valid"] = False
             validation_results["errors"].append(str(e))
-
-        # Check prompt template
         if not self._prompt_template:
             validation_results["valid"] = False
             validation_results["errors"].append("Missing prompt template")
-
-        # Check for placeholder consistency
-        placeholders_in_prompt = set(re.findall(r"\{([^}]+)\}", self._prompt_template))
-        required_args = set(self._input_arguments)
-
-        missing_in_prompt = required_args - placeholders_in_prompt
-        extra_in_prompt = placeholders_in_prompt - required_args
-
-        if missing_in_prompt:
-            validation_results["warnings"].append(
-                f"Arguments not used in prompt: {missing_in_prompt}"
-            )
-
-        if extra_in_prompt:
-            validation_results["warnings"].append(
-                f"Placeholders in prompt without corresponding arguments: {extra_in_prompt}"
-            )
-
         return validation_results
 
     def estimate_token_usage(self, arguments: Dict[str, Any]) -> Dict[str, int]:
-        """Estimate token usage for the given arguments (rough approximation)."""
         prompt = self._format_prompt(arguments)
-
-        # Rough token estimation (4 characters ≈ 1 token for English text)
         estimated_input_tokens = len(prompt) // 4
-        estimated_max_output_tokens = self._max_new_tokens
+        estimated_max_output_tokens = (
+            self._max_new_tokens if self._max_new_tokens is not None else 2048
+        )
         estimated_total_tokens = estimated_input_tokens + estimated_max_output_tokens
-
         return {
             "estimated_input_tokens": estimated_input_tokens,
             "max_output_tokens": estimated_max_output_tokens,