import os
import base64
import json

import backoff
import boto3
import numpy as np
from anthropic import Anthropic
from botocore.config import Config as BotoConfig
from botocore.exceptions import ClientError
from google import genai
from google.genai import errors as genai_errors
from google.genai import types
from mm_agents.vlaa_gui.utils.prompt_caching import (
    apply_bedrock_cache_control_to_recent_messages,
    build_bedrock_cache_control,
    resolve_bedrock_prompt_cache_ttl,
)
from openai import (
    APIConnectionError,
    APIError,
    AzureOpenAI,
    OpenAI,
    RateLimitError,
)
from volcenginesdkarkruntime import Ark

# from litellm import completion
DEFAULT_MAX_TOKENS = 16384 * 2
DEFAULT_MAX_TIME = 90
DEFAULT_THINKING_BUDGET = 2048
GEMINI_BACKOFF_EXCEPTIONS = (genai_errors.APIError, APIConnectionError, APIError)


def _resolve_sampling_value(configured_value, requested_value):
    return configured_value if configured_value is not None else requested_value


class BedrockRateLimitError(Exception):
    """Raised when Bedrock returns a rate-limit / throttling error.

    This intentionally does NOT inherit from ``ClientError`` so that
    ``@backoff.on_exception`` decorators targeting ``ClientError`` will
    **not** retry it.  The caller should treat this as a signal to
    abandon the current task immediately.
    """

    pass


class GeminiRateLimitError(Exception):
    """Raised when all Gemini API keys are exhausted due to rate limiting.

    This intentionally does NOT inherit from ``RateLimitError`` so that
    ``@backoff.on_exception`` decorators will **not** retry it.  The
    caller should treat this as a signal to abandon the current task
    immediately.
    """

    pass


class LLMRetryExhaustedError(Exception):
    """Raised when all LLM retry attempts are exhausted.

    The caller should treat this as a signal to abandon the current
    task immediately — the LLM API is unreachable and continuing
    would only burn steps without making progress.
    """

    pass


class LMMEngine:
    pass


class OpenAIEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-3-small",
        base_url=None,
        api_key=None,
    ):
        """Init an OpenAI Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-3-small".
            api_key (_type_, optional): Auth key from OpenAI. Defaults to None.
            base_url (_type_, optional): Base URL for the API. Defaults to None.
        """
        self.model = embedding_model
        self.api_key = api_key
        self.base_url = base_url

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> np.ndarray:
        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
            )
        client = OpenAI(api_key=api_key, base_url=self.base_url)
        response = client.embeddings.create(model=self.model, input=text)
        return np.array([data.embedding for data in response.data])


class GeminiEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-004",
        base_url=None,
        api_key=None,
    ):
        """Init an Gemini Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-004".
            api_key (_type_, optional): Auth key from Gemini. Defaults to None.
        """
        self.model = embedding_model
        self.api_key = api_key

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> np.ndarray:
        api_key = self.api_key or os.getenv("GEMINI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named GEMINI_API_KEY"
            )
        client = genai.Client(api_key=api_key)

        result = client.models.embed_content(
            model=self.model,
            contents=text,
            config=types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY"),
        )

        return np.array([i.values for i in result.embeddings])


class AzureOpenAIEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-3-small",
        api_key=None,
        api_version=None,
        base_url=None,
    ):
        """Init an Azure OpenAI Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-3-small".
            api_key (_type_, optional): Auth key from Azure OpenAI. Defaults to None.
            api_version (_type_, optional): API version. Defaults to None.
            endpoint_url (_type_, optional): Endpoint URL. Defaults to None.
        """
        self.model = embedding_model
        self.api_key = api_key
        self.api_version = api_version
        self.endpoint_url = base_url

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> np.ndarray:
        api_key = self.api_key or os.getenv("AZURE_OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named AZURE_OPENAI_API_KEY"
            )
        api_version = self.api_version or os.getenv("OPENAI_API_VERSION")
        if api_version is None:
            raise ValueError(
                "An API Version needs to be provided in either the api_version parameter or as an environment variable named OPENAI_API_VERSION"
            )
        endpoint_url = self.endpoint_url or os.getenv("AZURE_OPENAI_ENDPOINT")
        if endpoint_url is None:
            raise ValueError(
                "An Endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named AZURE_OPENAI_ENDPOINT"
            )
        client = AzureOpenAI(
            api_key=api_key,
            api_version=api_version,
            azure_endpoint=endpoint_url,
        )
        response = client.embeddings.create(input=text, model=self.model)
        return np.array([data.embedding for data in response.data])


class ArkEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "doubao-embedding-text-240715",
        api_key=None,
        base_url=None,
    ):
        """Init an Ark Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "doubao-embedding-text-240715".
            api_key (_type_, optional): Auth key from Ark. Defaults to None.
            base_url (_type_, optional): Base URL. Defaults to None.
        """
        self.model = embedding_model
        self.api_key = api_key
        self.base_url = base_url

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> np.ndarray:
        api_key = self.api_key or os.getenv("ARK_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ARK_API_KEY"
            )
        base_url = self.base_url or "https://ark.cn-beijing.volces.com/api/v3"
        if base_url is None:
            raise ValueError(
                "An Endpoint URL needs to be provided in either the base_url parameter or as an environment variable named ARK_ENDPOINT_URL"
            )
        client = Ark(base_url=base_url, api_key=api_key)
        response = client.embeddings.create(
            input=text, model=self.model, encoding_format="float"
        )
        return np.array([data.embedding for data in response.data])


class QwenEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-v4",
        api_key=None,
        base_url=None,
    ):
        """Init an Qwen Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-v4".
            api_key (_type_, optional): Auth key from Qwen. Defaults to None.
        """
        self.model = embedding_model
        self.api_key = api_key
        self.base_url = base_url

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> np.ndarray:
        api_key = self.api_key or os.getenv("DASHSCOPE_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named DASHSCOPE_API_KEY"
            )
        base_url = self.base_url or "https://dashscope.aliyuncs.com/compatible-mode/v1"
        if base_url is None:
            raise ValueError(
                "An Endpoint URL needs to be provided in either the base_url parameter or as an environment variable named DASHSCOPE_ENDPOINT_URL"
            )
        client = OpenAI(api_key=api_key, base_url=base_url)
        response = client.embeddings.create(model=self.model, input=text)
        return np.array([data.embedding for data in response.data])


class LMMEngineOpenAI(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")
        self.last_usage = None
        self.organization = kwargs.get("organization")
        self.thinking = kwargs.get("thinking", False)
        self.reasoning_effort = kwargs.get("reasoning_effort") or kwargs.get(
            "thinking_level"
        )
        self._stored_response_id = None
        self._stored_instructions = None
        self._stored_messages = []
        self._last_response_text = None

    def _get_usage_from_response(self, response) -> dict:
        """Extract token usage from API response."""
        try:
            usage = response.usage
            prompt_tokens = getattr(usage, "input_tokens", None)
            completion_tokens = getattr(usage, "output_tokens", None)
            total_tokens = getattr(usage, "total_tokens", None)
            if prompt_tokens is not None and completion_tokens is not None:
                return {
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "total_tokens": total_tokens,
                }
            return {
                "prompt_tokens": usage.prompt_tokens,
                "completion_tokens": usage.completion_tokens,
                "total_tokens": usage.total_tokens,
            }
        except AttributeError:
            return None

    def _get_client(self, api_key):
        if not self.llm_client:
            client_kwargs = {"api_key": api_key}
            if self.base_url:
                client_kwargs["base_url"] = self.base_url
            if self.organization:
                client_kwargs["organization"] = self.organization
            self.llm_client = OpenAI(**client_kwargs)
        return self.llm_client

    def _extract_text_segments(self, content):
        if content is None:
            return []
        if isinstance(content, str):
            return [content]
        if isinstance(content, dict):
            content = [content]
        if not isinstance(content, list):
            return []

        text_segments = []
        for item in content:
            if isinstance(item, str):
                text_segments.append(item)
                continue
            if not isinstance(item, dict):
                continue
            if item.get("type") == "text":
                text_segments.append(item.get("text", ""))
        return text_segments

    def _responses_content_from_message(self, content, role="user"):
        text_type = "output_text" if role == "assistant" else "input_text"
        if content is None:
            return []
        if isinstance(content, str):
            return [{"type": text_type, "text": content}]
        if isinstance(content, dict):
            content = [content]
        if not isinstance(content, list):
            return []

        response_content = []
        for item in content:
            if isinstance(item, str):
                response_content.append({"type": text_type, "text": item})
                continue
            if not isinstance(item, dict):
                continue

            item_type = item.get("type")
            if item_type == "text":
                response_content.append(
                    {"type": text_type, "text": item.get("text", "")}
                )
            elif item_type == "image_url":
                image_url = item.get("image_url", {})
                if isinstance(image_url, str):
                    url = image_url
                    detail = item.get("detail") or "original"
                else:
                    url = image_url.get("url", "")
                    detail = image_url.get("detail") or item.get("detail") or "original"
                if url:
                    response_content.append(
                        {
                            "type": "input_image",
                            "image_url": url,
                            "detail": detail,
                        }
                    )
            elif item_type == "image":
                source = item.get("source", {})
                data = source.get("data")
                if data:
                    mime_type = source.get("media_type", "image/png")
                    response_content.append(
                        {
                            "type": "input_image",
                            "image_url": f"data:{mime_type};base64,{data}",
                            "detail": item.get("detail")
                            or source.get("detail")
                            or "original",
                        }
                    )
        return response_content

    def _prepare_response_conversation(self, messages):
        instructions_parts = []
        input_items = []

        for message in messages:
            if isinstance(message, str):
                input_items.append(
                    {
                        "role": "user",
                        "content": [{"type": "input_text", "text": message}],
                    }
                )
                continue

            if not isinstance(message, dict):
                continue

            role = message.get("role") or "user"
            content = message.get("content")
            if role == "system":
                instructions_parts.extend(self._extract_text_segments(content))
                continue

            response_content = self._responses_content_from_message(content, role=role)
            if not response_content:
                text_type = "output_text" if role == "assistant" else "input_text"
                response_content = [{"type": text_type, "text": ""}]
            input_items.append({"role": role, "content": response_content})

        instructions = "\n".join(part for part in instructions_parts if part) or None
        return input_items, instructions

    def _is_matching_assistant_echo(self, item):
        if item.get("role") != "assistant" or not self._last_response_text:
            return False

        text_parts = []
        for content_part in item.get("content", []):
            if content_part.get("type") in ("input_text", "output_text"):
                text_parts.append(content_part.get("text", ""))

        return "\n".join(text_parts).strip() == self._last_response_text.strip()

    def _resolve_response_state(self, messages):
        input_items, instructions = self._prepare_response_conversation(messages)
        request_input = input_items
        previous_response_id = None

        if (
            self._stored_response_id
            and instructions == self._stored_instructions
            and len(input_items) >= len(self._stored_messages)
            and input_items[: len(self._stored_messages)] == self._stored_messages
        ):
            start_idx = len(self._stored_messages)
            if start_idx < len(input_items) and self._is_matching_assistant_echo(
                input_items[start_idx]
            ):
                start_idx += 1

            delta_items = input_items[start_idx:]
            if delta_items:
                request_input = delta_items
                previous_response_id = self._stored_response_id

        return input_items, instructions, request_input, previous_response_id

    def _build_response_request(
        self,
        messages,
        temperature=0.0,
        max_new_tokens=None,
        reasoning=None,
        **kwargs,
    ):
        (
            input_items,
            instructions,
            request_input,
            previous_response_id,
        ) = self._resolve_response_state(messages)

        resolved_temperature = (
            self.temperature if self.temperature is not None else temperature
        )
        resolved_top_p = self.top_p if self.top_p is not None else kwargs.get("top_p")
        response_kwargs = {
            "model": self.model,
            "input": request_input,
            "max_output_tokens": max_new_tokens
            if max_new_tokens
            else DEFAULT_MAX_TOKENS,
            "store": True,
            "truncation": "auto",
            **kwargs,
        }
        if instructions:
            response_kwargs["instructions"] = instructions
        if previous_response_id:
            response_kwargs["previous_response_id"] = previous_response_id
        if reasoning:
            response_kwargs["reasoning"] = reasoning
        if resolved_temperature is not None:
            response_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            response_kwargs["top_p"] = resolved_top_p

        return response_kwargs, input_items, instructions

    def _extract_response_text(self, response):
        output_text = getattr(response, "output_text", None)
        if output_text:
            return output_text

        text_parts = []
        for item in getattr(response, "output", []) or []:
            if getattr(item, "type", None) != "message":
                continue
            for content_part in getattr(item, "content", []) or []:
                text = getattr(content_part, "text", None)
                if text:
                    text_parts.append(text)
        return "\n".join(text_parts)

    def _extract_reasoning_summary(self, response):
        summaries = []
        for item in getattr(response, "output", []) or []:
            if getattr(item, "type", None) != "reasoning":
                continue
            for summary_part in getattr(item, "summary", []) or []:
                text = getattr(summary_part, "text", None)
                if text:
                    summaries.append(text)
        return "\n".join(summaries) if summaries else None

    def _store_response_state(self, input_items, instructions, response):
        self._stored_messages = input_items
        self._stored_instructions = instructions
        self._stored_response_id = getattr(response, "id", None)
        self._last_response_text = self._extract_response_text(response)

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
            )
        client = self._get_client(api_key)
        reasoning = None
        if self.thinking:
            reasoning = {"summary": "concise"}
            if self.reasoning_effort:
                reasoning["effort"] = self.reasoning_effort
        response_kwargs, input_items, instructions = self._build_response_request(
            messages=messages,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            reasoning=reasoning,
            **kwargs,
        )
        response = client.responses.create(**response_kwargs)
        self.last_usage = self._get_usage_from_response(response)
        self._store_response_state(input_items, instructions, response)
        return self._extract_response_text(response)

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate_with_thinking(
        self, messages, temperature=0.0, max_new_tokens=None, **kwargs
    ):
        if not self.thinking:
            return self.generate(
                messages,
                temperature=temperature,
                max_new_tokens=max_new_tokens,
                **kwargs,
            ), None
        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
            )
        client = self._get_client(api_key)
        reasoning = {"summary": "concise"}
        if self.reasoning_effort:
            reasoning["effort"] = self.reasoning_effort
        response_kwargs, input_items, instructions = self._build_response_request(
            messages=messages,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            reasoning=reasoning,
            **kwargs,
        )
        response = client.responses.create(**response_kwargs)
        self.last_usage = self._get_usage_from_response(response)
        self._store_response_state(input_items, instructions, response)
        return self._extract_response_text(response), self._extract_reasoning_summary(
            response
        )


class LMMEngineAnthropic(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, thinking=False, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.thinking = thinking
        self.api_key = api_key
        self.llm_client = None
        self.last_usage = None
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")

    def _get_usage_from_response(self, response) -> dict:
        """Extract token usage from API response."""
        try:
            usage = response.usage
            return {
                "prompt_tokens": usage.input_tokens,
                "completion_tokens": usage.output_tokens,
                "total_tokens": usage.input_tokens + usage.output_tokens,
            }
        except AttributeError:
            return None

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", None))
        api_key = self.api_key or os.getenv("ANTHROPIC_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ANTHROPIC_API_KEY"
            )
        if not self.llm_client:
            self.llm_client = Anthropic(api_key=api_key)
        if self.thinking:
            full_response = self.llm_client.messages.create(
                system=messages[0]["content"][0]["text"],
                model=self.model,
                messages=messages[1:],
                max_tokens=DEFAULT_MAX_TOKENS + 4096,
                thinking={"type": "enabled", "budget_tokens": DEFAULT_MAX_TOKENS},
                **kwargs,
            )
            self.last_usage = self._get_usage_from_response(full_response)
            thoughts = full_response.content[0].thinking
            print("CLAUDE 3.7 THOUGHTS:", thoughts)
            return full_response.content[1].text
        request_kwargs = {
            "system": messages[0]["content"][0]["text"],
            "model": self.model,
            "messages": messages[1:],
            "max_tokens": max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
            **kwargs,
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        response = self.llm_client.messages.create(**request_kwargs)
        self.last_usage = self._get_usage_from_response(response)
        return response.content[0].text


import logging as _logging

_gemini_rotation_logger = _logging.getLogger("desktopenv.gemini_rotation")


class _GeminiKeyRotationMixin:
    """Shared key-rotation logic for Gemini engine classes."""

    @staticmethod
    def _resolve_keys(api_keys, api_key):
        """Normalize api_keys / api_key into a list of key strings."""
        if api_keys:
            if isinstance(api_keys, str):
                return [api_keys]
            return list(api_keys)
        if api_key:
            return [api_key]
        env_key = os.getenv("GEMINI_API_KEY")
        if env_key:
            return [env_key]
        return []

    def _next_key(self):
        """Return the next API key via round-robin."""
        if not self._api_keys:
            raise ValueError(
                "No Gemini API keys available. Provide api_keys, api_key, "
                "or set the GEMINI_API_KEY environment variable."
            )
        key = self._api_keys[self._key_rr_index % len(self._api_keys)]
        self._key_rr_index = (self._key_rr_index + 1) % len(self._api_keys)
        return key

    def _get_or_create_genai_client(self, key):
        """Return a cached ``genai.Client`` for *key*."""
        if key not in self._genai_clients:
            self._genai_clients[key] = genai.Client(api_key=key)
        return self._genai_clients[key]

    @staticmethod
    def _is_rate_limit_error(exc):
        """Return True if *exc* is a Gemini rate-limit / quota error."""
        # openai SDK RateLimitError (used by Gemini OpenAI-compat endpoint)
        if isinstance(exc, RateLimitError):
            return True
        # google.api_core.exceptions.ResourceExhausted
        try:
            from google.api_core.exceptions import ResourceExhausted

            if isinstance(exc, ResourceExhausted):
                return True
        except ImportError:
            pass
        # Catch generic google-genai 429 errors by message
        exc_str = str(exc).lower()
        if (
            "429" in exc_str
            or "resource_exhausted" in exc_str
            or "resourceexhausted" in exc_str
        ):
            return True
        return False

    def _invoke_with_rotation(self, fn):
        """Try *fn(api_key)* with each key; raise ``GeminiRateLimitError`` when all exhausted."""
        if not self._api_keys:
            raise ValueError(
                "No Gemini API keys available. Provide api_keys, api_key, "
                "or set the GEMINI_API_KEY environment variable."
            )
        last_exc = None
        for _ in range(len(self._api_keys)):
            key = self._next_key()
            try:
                return fn(key)
            except Exception as exc:
                if self._is_rate_limit_error(exc):
                    _gemini_rotation_logger.warning(
                        "Gemini key ...%s hit rate limit (%s), rotating.",
                        key[-6:],
                        type(exc).__name__,
                    )
                    last_exc = exc
                    continue
                raise
        raise GeminiRateLimitError(
            f"All {len(self._api_keys)} Gemini API key(s) exhausted. "
            f"Last error: {last_exc}"
        )


class LMMEngineGemini(_GeminiKeyRotationMixin, LMMEngine):
    def __init__(
        self,
        base_url=None,
        api_key=None,
        api_keys=None,
        model=None,
        rate_limit=-1,
        thinking=False,
        thinking_budget=None,
        thinking_level=None,
        include_thoughts=False,
        **kwargs,
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self._api_keys = self._resolve_keys(api_keys, api_key)
        self._key_rr_index = 0
        self._genai_clients = {}
        self.thinking = thinking
        self.thinking_budget = thinking_budget
        self.thinking_level = thinking_level
        self.include_thoughts = include_thoughts
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.genai_client = None
        self.last_usage = None

    def _get_usage_from_response(self, response) -> dict:
        """Extract token usage from API response."""
        try:
            usage = response.usage
            return {
                "prompt_tokens": usage.prompt_tokens,
                "completion_tokens": usage.completion_tokens,
                "total_tokens": usage.total_tokens,
            }
        except AttributeError:
            pass

        try:
            usage = response.usage_metadata
            prompt_tokens = usage.prompt_token_count or 0
            completion_tokens = usage.candidates_token_count or 0
            total_tokens = usage.total_token_count or (
                prompt_tokens + completion_tokens
            )
            return {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": total_tokens,
            }
        except AttributeError:
            return None

    def _ensure_genai_client(self):
        if self.genai_client:
            return
        api_key = self.api_key or os.getenv("GEMINI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named GEMINI_API_KEY"
            )
        self.genai_client = genai.Client(api_key=api_key)

    def _resolve_thinking_level(self):
        if not isinstance(self.thinking_level, str):
            return None

        thinking_level = self.thinking_level.strip().lower()
        return thinking_level or None

    def _build_thinking_config(
        self, max_new_tokens=None, include_thoughts=None, force=False
    ):
        if not self.thinking and not force:
            return None

        thinking_kwargs = {}
        if include_thoughts is not None:
            thinking_kwargs["includeThoughts"] = include_thoughts

        thinking_level = self._resolve_thinking_level()
        if thinking_level is not None:
            return types.ThinkingConfig(**thinking_kwargs, thinkingLevel=thinking_level)

        budget = (
            self.thinking_budget
            if self.thinking_budget is not None
            else (
                max_new_tokens
                if max_new_tokens is not None
                else DEFAULT_THINKING_BUDGET
            )
        )
        return types.ThinkingConfig(**thinking_kwargs, thinkingBudget=budget)

    def _prepare_genai_contents(self, messages):
        contents = []
        system_instruction_parts = []

        for message in messages:
            if isinstance(message, types.Content):
                if message.role == "system":
                    system_instruction_parts.extend(message.parts or [])
                else:
                    contents.append(message)
                continue

            if isinstance(message, str):
                contents.append(
                    types.Content(
                        role="user",
                        parts=[types.Part.from_text(text=message)],
                    )
                )
                continue

            if not isinstance(message, dict):
                continue

            role = message.get("role")
            parts = self._parts_from_openai_content(message.get("content"))
            if role == "system":
                system_instruction_parts.extend(parts)
                continue

            if role == "assistant":
                role = "model"
            if role is None:
                role = "user"

            contents.append(types.Content(role=role, parts=parts))

        system_instruction = None
        if system_instruction_parts:
            instruction_text = "".join(
                part.text
                for part in system_instruction_parts
                if getattr(part, "text", None)
            )
            system_instruction = (
                instruction_text if instruction_text else system_instruction_parts
            )

        return contents, system_instruction

    def _parts_from_openai_content(self, content):
        parts = []
        if content is None:
            return parts

        if isinstance(content, str):
            return [types.Part.from_text(text=content)]

        if isinstance(content, dict):
            content = [content]

        if not isinstance(content, list):
            return parts

        for item in content:
            if isinstance(item, str):
                parts.append(types.Part.from_text(text=item))
                continue
            if not isinstance(item, dict):
                continue

            item_type = item.get("type")
            if item_type == "text":
                text = item.get("text", "")
                parts.append(types.Part.from_text(text=text))
            elif item_type == "image_url":
                image_url = item.get("image_url", {})
                url = image_url.get("url", "")
                mime_type, data = self._decode_data_url(url)
                if data:
                    parts.append(
                        types.Part.from_bytes(
                            data=data,
                            mime_type=mime_type or "image/png",
                        )
                    )
            elif item_type == "image":
                source = item.get("source", {})
                data = source.get("data")
                if data:
                    parts.append(
                        types.Part.from_bytes(
                            data=base64.b64decode(data),
                            mime_type=source.get("media_type", "image/png"),
                        )
                    )

        return parts

    def _decode_data_url(self, url):
        if not url or not url.startswith("data:"):
            return None, None
        try:
            header, b64_data = url.split(",", 1)
        except ValueError:
            return None, None
        mime_type = header.split(";")[0].replace("data:", "")
        try:
            return mime_type, base64.b64decode(b64_data)
        except (ValueError, base64.binascii.Error):
            return None, None

    def _extract_thoughts_and_answer(self, response):
        thoughts = []
        answer_parts = []

        for candidate in getattr(response, "candidates", []) or []:
            parts = getattr(getattr(candidate, "content", None), "parts", []) or []
            for part in parts:
                text = getattr(part, "text", None)
                if not text:
                    continue
                if getattr(part, "thought", False):
                    thoughts.append(text)
                else:
                    answer_parts.append(text)
            if thoughts or answer_parts:
                break

        if not answer_parts and getattr(response, "text", None):
            answer_parts.append(response.text)

        return "\n".join(thoughts).strip(), "\n".join(answer_parts).strip()

    def _generate_with_genai(
        self,
        messages,
        temperature=0.0,
        max_new_tokens=None,
        include_thoughts=None,
        force_thinking=False,
        genai_client=None,
        **kwargs,
    ):
        local_kwargs = dict(kwargs)
        client = genai_client
        if client is None:
            self._ensure_genai_client()
            client = self.genai_client
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(
            self.top_p, local_kwargs.pop("top_p", None)
        )
        contents, system_instruction = self._prepare_genai_contents(messages)

        config_kwargs = {
            "temperature": resolved_temperature,
            "max_output_tokens": max_new_tokens
            if max_new_tokens
            else DEFAULT_MAX_TOKENS,
            "media_resolution": types.MediaResolution.MEDIA_RESOLUTION_HIGH,
        }
        if resolved_top_p is not None:
            config_kwargs["top_p"] = resolved_top_p
        if system_instruction:
            config_kwargs["system_instruction"] = system_instruction

        thinking_config = self._build_thinking_config(
            max_new_tokens=max_new_tokens,
            include_thoughts=include_thoughts,
            force=force_thinking,
        )
        if thinking_config is not None:
            config_kwargs["thinking_config"] = thinking_config

        # Internal control kwargs not accepted by Google API.
        local_kwargs.pop("system_prompt", None)
        local_kwargs.pop("image_content", None)

        response = client.models.generate_content(
            model=self.model,
            contents=contents,
            config=types.GenerateContentConfig(**config_kwargs),
            **local_kwargs,
        )
        self.last_usage = self._get_usage_from_response(response)
        return response

    @backoff.on_exception(
        backoff.expo,
        GEMINI_BACKOFF_EXCEPTIONS,
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        def _do_generate(api_key):
            response = self._generate_with_genai(
                messages,
                temperature=temperature,
                max_new_tokens=max_new_tokens,
                include_thoughts=self.include_thoughts if self.thinking else None,
                force_thinking=self.thinking,
                genai_client=self._get_or_create_genai_client(api_key),
                **kwargs,
            )
            if self.thinking:
                _, answer = self._extract_thoughts_and_answer(response)
                if answer:
                    return answer
            return response.text

        return self._invoke_with_rotation(_do_generate)

    @backoff.on_exception(
        backoff.expo,
        GEMINI_BACKOFF_EXCEPTIONS,
        max_time=DEFAULT_MAX_TIME,
    )
    def generate_with_thinking(
        self, messages, temperature=0.0, max_new_tokens=None, **kwargs
    ):
        def _do_generate(api_key):
            gc = self._get_or_create_genai_client(api_key)
            response = self._generate_with_genai(
                messages,
                temperature=temperature,
                max_new_tokens=max_new_tokens,
                include_thoughts=True,
                force_thinking=True,
                genai_client=gc,
                **kwargs,
            )
            thoughts, answer = self._extract_thoughts_and_answer(response)
            return f"<thoughts>\n{thoughts}\n</thoughts>\n\n<answer>\n{answer}\n</answer>\n"

        return self._invoke_with_rotation(_do_generate)


class LMMEngineAnthropicBedrock(LMMEngine):
    """Bedrock engine using cross-region inference.

    Cross-region inference profiles (e.g. ``global.anthropic.claude-sonnet-4-6-v1:0``)
    let AWS automatically route requests to regions with available capacity.
    The client only needs to connect to a single source region.

    Multiple AWS credential sets (``aws_keys``) are still supported for quota
    pooling — the engine round-robins through them on throttling errors.
    """

    def __init__(self, model=None, region=None, aws_keys=None, thinking=True, **kwargs):
        assert model is not None, "model must be provided"
        self.model = model
        self.thinking = thinking
        self.region = (
            region or os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION")
        )
        self._keys = self._resolve_keys(aws_keys)
        self._clients = {}
        self._key_rr_index = 0
        self.bedrock_client = None
        self.last_usage = None
        self.prompt_cache_ttl = resolve_bedrock_prompt_cache_ttl(
            self.model,
            kwargs.get("prompt_cache_ttl", kwargs.get("bedrock_prompt_cache_ttl")),
        )

    def _resolve_keys(self, aws_keys):
        """Resolve AWS credential sets.

        Each entry can have ``aws_access_key_id``, ``aws_secret_access_key``,
        and optionally ``aws_profile``.  A ``None`` entry means "use the
        default credential chain".

        Returns a list of dicts (or ``[None]`` when no explicit keys are
        configured).
        """
        if not aws_keys:
            return [None]

        resolved = []
        for entry in aws_keys:
            if not isinstance(entry, dict):
                continue
            key_id = entry.get("aws_access_key_id", "").strip()
            secret = entry.get("aws_secret_access_key", "").strip()
            profile = entry.get("aws_profile", "").strip()
            if key_id and secret:
                resolved.append(
                    {"aws_access_key_id": key_id, "aws_secret_access_key": secret}
                )
            elif profile:
                resolved.append({"aws_profile": profile})

        return resolved if resolved else [None]

    def _get_or_create_client(self, key_entry):
        """Get or create a boto3 bedrock-runtime client for the given credentials."""
        if key_entry is None:
            cache_key = "default"
        elif "aws_profile" in key_entry:
            cache_key = f"profile:{key_entry['aws_profile']}"
        else:
            cache_key = f"key:{key_entry['aws_access_key_id']}"

        if cache_key not in self._clients:
            kwargs = {
                "service_name": "bedrock-runtime",
                "config": BotoConfig(
                    read_timeout=300,
                    connect_timeout=10,
                    retries={"max_attempts": 3},
                ),
            }
            if self.region:
                kwargs["region_name"] = self.region
            if key_entry is not None:
                if "aws_profile" in key_entry:
                    session = boto3.Session(profile_name=key_entry["aws_profile"])
                    self._clients[cache_key] = session.client(**kwargs)
                    return self._clients[cache_key]
                else:
                    kwargs["aws_access_key_id"] = key_entry["aws_access_key_id"]
                    kwargs["aws_secret_access_key"] = key_entry["aws_secret_access_key"]
            self._clients[cache_key] = boto3.client(**kwargs)
        return self._clients[cache_key]

    def _next_client(self):
        """Round-robin through credential sets."""
        key_entry = self._keys[self._key_rr_index % len(self._keys)]
        self._key_rr_index += 1
        client = self._get_or_create_client(key_entry)
        self.bedrock_client = client
        return client

    def _should_rotate_on_error(self, err):
        error_obj = err.response.get("Error", {}) if hasattr(err, "response") else {}
        error_code = str(error_obj.get("Code", "")).lower()
        error_message = str(error_obj.get("Message", "")).lower()

        throttling_codes = {
            "throttlingexception",
            "toomanyrequestsexception",
            "servicequotaexceededexception",
            "limitexceededexception",
            "requestlimitexceeded",
        }
        throttling_keywords = (
            "too many tokens",
            "too many requests",
            "rate exceeded",
            "rate limit",
            "throttle",
            "quota",
        )

        if error_code in throttling_codes:
            return True
        return any(keyword in error_message for keyword in throttling_keywords)

    def _invoke_model(self, body):
        """Invoke the model, rotating through credential sets on throttling.

        Raises ``BedrockRateLimitError`` immediately when a throttling /
        rate-limit error is detected so that the caller can abandon the
        task without wasting time on retries.
        """
        attempts = max(1, len(self._keys))
        last_error = None

        for _ in range(attempts):
            client = self._next_client()
            try:
                return client.invoke_model(body=body, modelId=self.model)
            except ClientError as err:
                last_error = err
                if self._should_rotate_on_error(err):
                    # Rate-limit / throttling error
                    if attempts == 1:
                        # Only one credential set — no point rotating, bail out
                        raise BedrockRateLimitError(
                            f"Bedrock rate limit error (no rotation available): {err}"
                        ) from err
                    # Multiple keys — try the next one before giving up
                    continue
                # Non-throttling error — raise immediately
                raise

        # All credential sets exhausted with throttling errors
        if last_error is not None:
            raise BedrockRateLimitError(
                f"Bedrock rate limit error (all {attempts} credential sets exhausted): {last_error}"
            ) from last_error

    def _get_usage_from_response(self, response_body) -> dict:
        """Extract token usage from API response."""
        try:
            usage = response_body.get("usage", {})
            input_tokens = usage.get("input_tokens", 0)
            output_tokens = usage.get("output_tokens", 0)
            return {
                "prompt_tokens": input_tokens,
                "completion_tokens": output_tokens,
                "total_tokens": input_tokens + output_tokens,
            }
        except (AttributeError, KeyError):
            return None

    def _is_adaptive_thinking_model(self):
        """Check if the model supports adaptive thinking (Sonnet 4.6, Opus 4.6)."""
        model_lower = self.model.lower()
        return "sonnet-4-6" in model_lower or "opus-4-6" in model_lower

    def _build_thinking_config(self):
        """Build the thinking configuration based on model type.

        Adaptive thinking (Sonnet 4.6, Opus 4.6) uses {"type": "adaptive"}
        with no budget_tokens — the model decides how much to think.
        Older models use {"type": "enabled", "budget_tokens": N}.
        """
        if self._is_adaptive_thinking_model():
            return {"type": "adaptive"}
        return {
            "type": "enabled",
            "budget_tokens": DEFAULT_MAX_TOKENS,
        }

    def _build_system_prompt(self, system_prompt):
        if not system_prompt:
            return None
        if not self.prompt_cache_ttl:
            return system_prompt
        return [
            {
                "type": "text",
                "text": system_prompt,
                "cache_control": build_bedrock_cache_control(self.prompt_cache_ttl),
            }
        ]

    def _convert_messages_to_bedrock_format(self, messages):
        """Convert messages from OpenAI/Anthropic format to Bedrock format."""
        system_prompt = None
        bedrock_messages = []

        for message in messages:
            role = message.get("role", "user")
            content = message.get("content", [])

            if role == "system":
                # Extract system prompt
                if isinstance(content, str):
                    system_prompt = content
                elif isinstance(content, list) and len(content) > 0:
                    # Handle list of content blocks
                    text_parts = []
                    for item in content:
                        if isinstance(item, dict) and item.get("type") == "text":
                            text_parts.append(item.get("text", ""))
                        elif isinstance(item, str):
                            text_parts.append(item)
                    system_prompt = "\n".join(text_parts)
                continue

            # Convert content to Bedrock format
            bedrock_content = []
            if isinstance(content, str):
                bedrock_content.append({"type": "text", "text": content})
            elif isinstance(content, list):
                for item in content:
                    if isinstance(item, str):
                        bedrock_content.append({"type": "text", "text": item})
                    elif isinstance(item, dict):
                        item_type = item.get("type")
                        if item_type == "text":
                            bedrock_content.append(
                                {"type": "text", "text": item.get("text", "")}
                            )
                        elif item_type == "image":
                            # Anthropic native format
                            source = item.get("source", {})
                            bedrock_content.append(
                                {
                                    "type": "image",
                                    "source": {
                                        "type": "base64",
                                        "media_type": source.get(
                                            "media_type", "image/png"
                                        ),
                                        "data": source.get("data", ""),
                                    },
                                }
                            )
                        elif item_type == "image_url":
                            # OpenAI format - convert base64 data URL
                            image_url = item.get("image_url", {})
                            url = image_url.get("url", "")
                            if url.startswith("data:"):
                                # Parse data URL: data:image/png;base64,<data>
                                try:
                                    header, b64_data = url.split(",", 1)
                                    media_type = header.split(";")[0].replace(
                                        "data:", ""
                                    )
                                    bedrock_content.append(
                                        {
                                            "type": "image",
                                            "source": {
                                                "type": "base64",
                                                "media_type": media_type,
                                                "data": b64_data,
                                            },
                                        }
                                    )
                                except ValueError:
                                    pass

            bedrock_messages.append({"role": role, "content": bedrock_content})

        return system_prompt, bedrock_messages

    @backoff.on_exception(backoff.expo, (ClientError,), max_time=DEFAULT_MAX_TIME)
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        system_prompt, bedrock_messages = self._convert_messages_to_bedrock_format(
            messages
        )
        apply_bedrock_cache_control_to_recent_messages(
            bedrock_messages,
            ttl=self.prompt_cache_ttl,
        )

        body_dict = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
            "messages": bedrock_messages,
        }

        if system_prompt:
            body_dict["system"] = self._build_system_prompt(system_prompt)

        if temperature > 0 and not self.thinking:
            body_dict["temperature"] = temperature

        if self.thinking:
            body_dict["thinking"] = self._build_thinking_config()
            if not self._is_adaptive_thinking_model():
                body_dict["max_tokens"] = DEFAULT_MAX_TOKENS + 4096

        body = json.dumps(body_dict)

        try:
            response = self._invoke_model(body=body)
            response_body = json.loads(response.get("body").read())
            self.last_usage = self._get_usage_from_response(response_body)

            if self.thinking:
                # Extract thinking and answer from response
                # Adaptive thinking may produce interleaved thinking/text blocks
                content = response_body.get("content", [])
                answer_parts = []
                for block in content:
                    if block.get("type") == "thinking":
                        pass
                    elif block.get("type") == "text":
                        answer_parts.append(block.get("text", ""))
                return "\n".join(answer_parts) if answer_parts else ""

            # Standard response
            content = response_body.get("content", [])
            if content and len(content) > 0:
                return content[0].get("text", "")
            return ""

        except ClientError as err:
            raise err

    @backoff.on_exception(backoff.expo, (ClientError,), max_time=DEFAULT_MAX_TIME)
    def generate_with_thinking(
        self, messages, temperature=0.0, max_new_tokens=None, **kwargs
    ):
        """Generate the next message based on previous messages, and keeps the thinking tokens."""
        system_prompt, bedrock_messages = self._convert_messages_to_bedrock_format(
            messages
        )
        apply_bedrock_cache_control_to_recent_messages(
            bedrock_messages,
            ttl=self.prompt_cache_ttl,
        )

        adaptive = self._is_adaptive_thinking_model()
        default_max = DEFAULT_MAX_TOKENS if adaptive else DEFAULT_MAX_TOKENS + 4096
        body_dict = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": max_new_tokens if max_new_tokens else default_max,
            "messages": bedrock_messages,
            "thinking": self._build_thinking_config(),
        }

        if system_prompt:
            body_dict["system"] = self._build_system_prompt(system_prompt)

        body = json.dumps(body_dict)

        try:
            response = self._invoke_model(body=body)
            response_body = json.loads(response.get("body").read())
            self.last_usage = self._get_usage_from_response(response_body)

            content = response_body.get("content", [])
            thoughts_parts = []
            answer_parts = []

            for block in content:
                if block.get("type") == "thinking":
                    thoughts_parts.append(block.get("thinking", ""))
                elif block.get("type") == "text":
                    answer_parts.append(block.get("text", ""))

            thoughts = "\n".join(thoughts_parts)
            answer = "\n".join(answer_parts)

            return f"<thoughts>\n{thoughts}\n</thoughts>\n\n<answer>\n{answer}\n</answer>\n"

        except ClientError as err:
            raise err


class LMMEngineQwen(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.thinking = kwargs.get("thinking", False)
        self.thinking_budget = kwargs.get("thinking_budget")
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")

    def _get_api_key(self):
        api_key = (
            self.api_key
            or os.getenv("DASHSCOPE_API_KEY")
            or os.getenv("OPENAI_API_KEY")
        )
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named DASHSCOPE_API_KEY"
            )
        return api_key

    def _get_client(self, api_key):
        if not self.llm_client:
            base_url = (
                self.base_url
                or os.getenv("DASHSCOPE_ENDPOINT_URL")
                or "https://dashscope.aliyuncs.com/compatible-mode/v1"
            )
            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
        return self.llm_client

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", None))
        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
        api_key = self._get_api_key()
        self._get_client(api_key)
        caller_extra = kwargs.pop("extra_body", {})
        merged_extra = {"vl_high_resolution_images": True}
        merged_extra.update(caller_extra)
        if self.thinking:
            merged_extra["enable_thinking"] = True
            if self.thinking_budget is not None:
                merged_extra["thinking_budget"] = self.thinking_budget
        else:
            merged_extra.pop("enable_thinking", None)
            merged_extra.pop("thinking_budget", None)
        request_kwargs = {
            "model": self.model,
            "messages": messages,
            "extra_body": merged_extra,
            **kwargs,
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        return (
            self.llm_client.chat.completions.create(**request_kwargs)
            .choices[0]
            .message.content
        )

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    # Compatible with Claude-3.7 Sonnet thinking mode
    def generate_with_thinking(
        self, messages, temperature=0.0, max_new_tokens=None, **kwargs
    ):
        """Generate the next message based on previous messages, and keeps the thinking tokens"""
        if not self.thinking:
            return self.generate(
                messages,
                temperature=temperature,
                max_new_tokens=max_new_tokens,
                **kwargs,
            ), None
        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
        api_key = self._get_api_key()
        self._get_client(api_key)
        caller_extra = kwargs.pop("extra_body", {})
        merged_extra = {"vl_high_resolution_images": True, "enable_thinking": True}
        merged_extra.update(caller_extra)
        merged_extra["enable_thinking"] = True
        if self.thinking_budget is not None:
            merged_extra["thinking_budget"] = self.thinking_budget

        completion_message = (
            self.llm_client.chat.completions.create(
                model=self.model,
                messages=messages,
                extra_body=merged_extra,
                **kwargs,
            )
            .choices[0]
            .message
        )

        thinking = getattr(completion_message, "reasoning_content", None)
        return completion_message.content, thinking


class LMMEngineOpenRouter(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", None))
        api_key = self.api_key or os.getenv("OPENROUTER_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENROUTER_API_KEY"
            )
        base_url = self.base_url or os.getenv("OPEN_ROUTER_ENDPOINT_URL")
        if base_url is None:
            raise ValueError(
                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named OPEN_ROUTER_ENDPOINT_URL"
            )
        if not self.llm_client:
            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
        request_kwargs = {
            "model": self.model,
            "messages": messages,
            "max_tokens": max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
            **kwargs,
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        return (
            self.llm_client.chat.completions.create(**request_kwargs)
            .choices[0]
            .message.content
        )


class LMMEngineAzureOpenAI(LMMEngine):
    def __init__(
        self,
        base_url=None,
        api_key=None,
        model=None,
        api_version=None,
        rate_limit=-1,
        **kwargs,
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.api_version = api_version
        self.api_key = api_key
        self.azure_endpoint = base_url
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.cost = 0.0
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", None))
        api_key = self.api_key or os.getenv("AZURE_OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named AZURE_OPENAI_API_KEY"
            )
        api_version = self.api_version or os.getenv("OPENAI_API_VERSION")
        if api_version is None:
            raise ValueError(
                "api_version must be provided either as a parameter or as an environment variable named OPENAI_API_VERSION"
            )
        azure_endpoint = self.azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
        if azure_endpoint is None:
            raise ValueError(
                "An Azure API endpoint needs to be provided in either the azure_endpoint parameter or as an environment variable named AZURE_OPENAI_ENDPOINT"
            )
        if not self.llm_client:
            self.llm_client = AzureOpenAI(
                azure_endpoint=azure_endpoint,
                api_key=api_key,
                api_version=api_version,
            )
        request_kwargs = {
            "model": self.model,
            "messages": messages,
            "max_tokens": max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
            **kwargs,
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        completion = self.llm_client.chat.completions.create(**request_kwargs)
        total_tokens = completion.usage.total_tokens
        self.cost += 0.02 * ((total_tokens + 500) / 1000)
        return completion.choices[0].message.content


class LMMEnginevLLM(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.api_key = api_key
        self.base_url = base_url
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(
        self,
        messages,
        temperature=0.0,
        top_p=0.8,
        repetition_penalty=1.05,
        max_new_tokens=512,
        **kwargs,
    ):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", top_p))
        api_key = self.api_key or os.getenv("vLLM_API_KEY")
        if api_key is None:
            raise ValueError(
                "A vLLM API key needs to be provided in either the api_key parameter or as an environment variable named vLLM_API_KEY"
            )
        base_url = self.base_url or os.getenv("vLLM_ENDPOINT_URL")
        if base_url is None:
            raise ValueError(
                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named vLLM_ENDPOINT_URL"
            )
        if not self.llm_client:
            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
        request_kwargs = {
            "model": self.model,
            "messages": messages,
            "max_tokens": max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
            "extra_body": {"repetition_penalty": repetition_penalty},
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        completion = self.llm_client.chat.completions.create(**request_kwargs)
        return completion.choices[0].message.content


class LMMEngineHuggingFace(LMMEngine):
    def __init__(self, base_url=None, api_key=None, rate_limit=-1, **kwargs):
        self.base_url = base_url
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", None))
        api_key = self.api_key or os.getenv("HF_TOKEN")
        if api_key is None:
            raise ValueError(
                "A HuggingFace token needs to be provided in either the api_key parameter or as an environment variable named HF_TOKEN"
            )
        base_url = self.base_url
        if base_url is None:
            raise ValueError(
                "HuggingFace endpoint must be provided as base_url parameter."
            )
        if not self.llm_client:
            self.llm_client = OpenAI(base_url=base_url, api_key=api_key)
        request_kwargs = {
            "model": "tgi",
            "messages": messages,
            "max_tokens": max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
            **kwargs,
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        return (
            self.llm_client.chat.completions.create(**request_kwargs)
            .choices[0]
            .message.content
        )


class LMMEngineParasail(LMMEngine):
    def __init__(self, api_key=None, model=None, rate_limit=-1, **kwargs):
        assert model is not None, "Parasail model id must be provided"
        self.model = model
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", None))
        api_key = self.api_key or os.getenv("PARASAIL_API_KEY")
        if api_key is None:
            raise ValueError(
                "A Parasail API key needs to be provided in either the api_key parameter or as an environment variable named PARASAIL_API_KEY"
            )
        if not self.llm_client:
            self.llm_client = OpenAI(
                base_url="https://api.parasail.io/v1", api_key=api_key
            )
        request_kwargs = {
            "model": self.model,
            "messages": messages,
            "max_tokens": max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
            **kwargs,
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        return (
            self.llm_client.chat.completions.create(**request_kwargs)
            .choices[0]
            .message.content
        )


class LMMEngineArk(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.temperature = kwargs.get("temperature")
        self.top_p = kwargs.get("top_p")

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        resolved_temperature = _resolve_sampling_value(self.temperature, temperature)
        resolved_top_p = _resolve_sampling_value(self.top_p, kwargs.pop("top_p", None))
        api_key = self.api_key or os.getenv("ARK_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ARK_API_KEY"
            )
        if not self.llm_client:
            if not self.base_url:
                self.llm_client = Ark(api_key=api_key)
            else:
                self.llm_client = Ark(base_url=self.base_url, api_key=api_key)
        request_kwargs = {
            "model": self.model,
            "messages": messages,
            "max_completion_tokens": max_new_tokens
            if max_new_tokens
            else DEFAULT_MAX_TOKENS,
            **kwargs,
        }
        if resolved_temperature is not None:
            request_kwargs["temperature"] = resolved_temperature
        if resolved_top_p is not None:
            request_kwargs["top_p"] = resolved_top_p
        return (
            self.llm_client.chat.completions.create(**request_kwargs)
            .choices[0]
            .message.content
        )


# class LMMEngineLiteLLM(LMMEngine):
#     def __init__(self, model=None, api_key=None, rate_limit=-1, **kwargs):
#         assert model is not None, "model must be provided"
#         self.model = model
#         self.api_key = api_key
#         self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

#     @backoff.on_exception(
#         backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=DEFAULT_MAX_TIME
#     )
#     def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
#         api_key = self.api_key
#         if api_key is None:
#             raise ValueError(
#                 "An API key needs to be provided in either the api_key parameter or as an environment variable named LITELLM_API_KEY"
#             )

#         response = completion(
#             model=self.model,
#             api_key=api_key,
#             max_tokens=max_new_tokens if max_new_tokens else DEFAULT_MAX_TOKENS,
#             temperature=temperature,
#             **kwargs,
#         )
#         return response.choices[0].message.content


class LMMEngineGeminiSearch(_GeminiKeyRotationMixin, LMMEngine):
    def __init__(
        self,
        base_url=None,
        api_key=None,
        api_keys=None,
        model=None,
        rate_limit=-1,
        **kwargs,
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self._api_keys = self._resolve_keys(api_keys, api_key)
        self._key_rr_index = 0
        self._genai_clients = {}
        self._openai_clients = {}
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit
        self.llm_client = None
        self.grounding_tool = types.Tool(google_search=types.GoogleSearch())

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        def _do_generate(api_key):
            client = self._get_or_create_genai_client(api_key)

            contents, system_instruction = self._prepare_genai_contents(messages)

            config_kwargs = {
                "temperature": temperature,
                "max_output_tokens": max_new_tokens
                if max_new_tokens
                else DEFAULT_MAX_TOKENS,
                "tools": [self.grounding_tool],
            }
            if system_instruction:
                config_kwargs["system_instruction"] = system_instruction

            config = types.GenerateContentConfig(**config_kwargs)

            response = client.models.generate_content(
                model=self.model,
                contents=contents,
                config=config,
                **kwargs,
            )

            return response.text

        return self._invoke_with_rotation(_do_generate)

    def _prepare_genai_contents(self, messages):
        contents = []
        system_instruction_parts = []

        for message in messages:
            if isinstance(message, types.Content):
                if message.role == "system":
                    system_instruction_parts.extend(message.parts or [])
                else:
                    contents.append(message)
                continue

            if isinstance(message, str):
                contents.append(
                    types.Content(
                        role="user",
                        parts=[types.Part.from_text(text=message)],
                    )
                )
                continue

            if not isinstance(message, dict):
                continue

            role = message.get("role")
            parts = self._parts_from_openai_content(message.get("content"))
            if role == "system":
                system_instruction_parts.extend(parts)
                continue

            if role == "assistant":
                role = "model"
            if role is None:
                role = "user"

            contents.append(types.Content(role=role, parts=parts))

        system_instruction = None
        if system_instruction_parts:
            instruction_text = "".join(
                part.text
                for part in system_instruction_parts
                if getattr(part, "text", None)
            )
            system_instruction = (
                instruction_text if instruction_text else system_instruction_parts
            )

        return contents, system_instruction

    def _parts_from_openai_content(self, content):
        parts = []
        if content is None:
            return parts

        if isinstance(content, str):
            return [types.Part.from_text(text=content)]

        if isinstance(content, dict):
            content = [content]

        if not isinstance(content, list):
            return parts

        for item in content:
            if isinstance(item, str):
                parts.append(types.Part.from_text(text=item))
                continue
            if not isinstance(item, dict):
                continue

            item_type = item.get("type")
            if item_type == "text":
                text = item.get("text", "")
                parts.append(types.Part.from_text(text=text))
            elif item_type == "image_url":
                image_url = item.get("image_url", {})
                url = image_url.get("url", "")
                mime_type, data = self._decode_data_url(url)
                if data:
                    parts.append(
                        types.Part.from_bytes(
                            data=data,
                            mime_type=mime_type or "image/png",
                        )
                    )
            elif item_type == "image":
                source = item.get("source", {})
                data = source.get("data")
                if data:
                    parts.append(
                        types.Part.from_bytes(
                            data=base64.b64decode(data),
                            mime_type=source.get("media_type", "image/png"),
                        )
                    )

        return parts

    def _decode_data_url(self, url):
        if not url or not url.startswith("data:"):
            return None, None
        try:
            header, b64_data = url.split(",", 1)
        except ValueError:
            return None, None
        mime_type = header.split(";")[0].replace("data:", "")
        try:
            return mime_type, base64.b64decode(b64_data)
        except (ValueError, base64.binascii.Error):
            return None, None


class LMMEngineOpenAISearch(LMMEngineOpenAI):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        super().__init__(
            base_url=base_url,
            api_key=api_key,
            model=model,
            rate_limit=rate_limit,
            **kwargs,
        )

    @backoff.on_exception(
        backoff.expo,
        (APIConnectionError, APIError, RateLimitError),
        max_time=DEFAULT_MAX_TIME,
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        api_key = self.api_key or os.getenv("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
            )
        client = self._get_client(api_key)
        response_kwargs, input_items, instructions = self._build_response_request(
            messages=messages,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            tools=[{"type": "web_search"}],
            **kwargs,
        )
        response = client.responses.create(**response_kwargs)
        self.last_usage = self._get_usage_from_response(response)
        self._store_response_state(input_items, instructions, response)
        return self._extract_response_text(response)
