OpenAIWrapper

def register_model_client(self, model_client_cls: ModelClient, **kwargs: Any):
    """Register a model client.

    Args:
        model_client_cls: A custom client class that follows the ModelClient interface
        kwargs: The kwargs for the custom client class to be initialized with
    """
    existing_client_class = False
    for i, client in enumerate(self._clients):
        if isinstance(client, PlaceHolderClient):
            placeholder_config = client.config

            if placeholder_config.get("model_client_cls") == model_client_cls.__name__:
                self._clients[i] = model_client_cls(placeholder_config, **kwargs)
                return
        elif isinstance(client, model_client_cls):
            existing_client_class = True

    if existing_client_class:
        logger.warning(
            f"Model client {model_client_cls.__name__} is already registered. Add more entries in the config_list to use multiple model clients."
        )
    else:
        raise ValueError(
            f'Model client "{model_client_cls.__name__}" is being registered but was not found in the config_list. '
            f'Please make sure to include an entry in the config_list with "model_client_cls": "{model_client_cls.__name__}"'
        )

@classmethod
def instantiate(
    cls,
    template: str | Callable[[dict[str, Any]], str] | None,
    context: dict[str, Any] | None = None,
    allow_format_str_template: bool | None = False,
) -> str | None:
    if not context or template is None:
        return template  # type: ignore [return-value]
    if isinstance(template, str):
        return template.format(**context) if allow_format_str_template else template
    return template(context)

def create(self, **config: Any) -> ModelClient.ModelClientResponseProtocol:
    """Make a completion for a given config using available clients.
    Besides the kwargs allowed in openai's [or other] client, we allow the following additional kwargs.
    The config in each client will be overridden by the config.

    Args:
        **config: The config for the completion.

    Raises:
        RuntimeError: If all declared custom model clients are not registered
        APIError: If any model client create call raises an APIError
    """
    # if ERROR:
    #     raise ERROR
    invocation_id = str(uuid.uuid4())
    last = len(self._clients) - 1
    # Check if all configs in config list are activated
    non_activated = [
        client.config["model_client_cls"] for client in self._clients if isinstance(client, PlaceHolderClient)
    ]
    if non_activated:
        raise RuntimeError(
            f"Model client(s) {non_activated} are not activated. Please register the custom model clients using `register_model_client` or filter them out form the config list."
        )

    ordered_clients_indices = list(range(len(self._clients)))
    if self.routing_method == "round_robin" and len(self._clients) > 0:
        ordered_clients_indices = (
            ordered_clients_indices[self._round_robin_index :] + ordered_clients_indices[: self._round_robin_index]
        )
        self._round_robin_index = (self._round_robin_index + 1) % len(self._clients)

    for i in ordered_clients_indices:
        # merge the input config with the i-th config in the config list
        client_config = self._config_list[i]
        full_config = merge_config_with_tools(config, client_config)

        # separate the config into create_config and extra_kwargs
        create_config, extra_kwargs = self._separate_create_config(full_config)
        # construct the create params
        params = self._construct_create_params(create_config, extra_kwargs)
        # get the cache_seed, filter_func and context
        cache_seed = extra_kwargs.get("cache_seed")
        cache = extra_kwargs.get("cache")
        filter_func = extra_kwargs.get("filter_func")
        context = extra_kwargs.get("context")
        agent = extra_kwargs.get("agent")
        price = extra_kwargs.get("price", None)
        if isinstance(price, list):
            price = tuple(price)
        elif isinstance(price, (float, int)):
            logger.warning(
                "Input price is a float/int. Using the same price for prompt and completion tokens. Use a list/tuple if prompt and completion token prices are different."
            )
            price = (price, price)

        total_usage = None
        actual_usage = None

        cache_client = None
        if cache is not None:
            # Use the cache object if provided.
            cache_client = cache
        elif cache_seed is not None:
            # Legacy cache behavior, if cache_seed is given, use DiskCache.
            cache_client = Cache.disk(cache_seed, LEGACY_CACHE_DIR)

        client = self._clients[i]
        log_cache_seed_value(cache if cache is not None else cache_seed, client=client)

        if cache_client is not None:
            with cache_client as cache:
                # Try to get the response from cache
                key = get_key(
                    {
                        **params,
                        **{"response_format": json.dumps(TypeAdapter(params["response_format"]).json_schema())},
                    }
                    if "response_format" in params and not isinstance(params["response_format"], dict)
                    else params
                )
                request_ts = get_current_ts()

                response: ChatCompletionExtended | None = cache.get(key, None)

                if response is not None:
                    # Backward compatibility: set message_retrieval_function for ChatCompletionExtended
                    if hasattr(response, "message_retrieval_function"):
                        response.message_retrieval_function = client.message_retrieval

                    try:
                        response.cost
                    except AttributeError:
                        # update attribute if cost is not calculated
                        response.cost = client.cost(response)
                        cache.set(key, response)
                    total_usage = client.get_usage(response)

                    if logging_enabled():
                        # Log the cache hit
                        # TODO: log the config_id and pass_filter etc.
                        log_chat_completion(
                            invocation_id=invocation_id,
                            client_id=id(client),
                            wrapper_id=id(self),
                            agent=agent,
                            request=params,
                            response=response,
                            is_cached=1,
                            cost=response.cost if response.cost is not None else 0.0,
                            start_time=request_ts,
                        )

                    # check the filter
                    pass_filter = filter_func is None or filter_func(context=context, response=response)
                    if pass_filter or i == last:
                        # Store metadata for serializable responses
                        if hasattr(response, "id"):
                            self._store_response_metadata(response.id, client, i, pass_filter)

                        # Backward compatibility: set attributes on ChatCompletionExtended
                        if hasattr(response, "config_id"):
                            response.config_id = i
                        if hasattr(response, "pass_filter"):
                            response.pass_filter = pass_filter
                        self._update_usage(actual_usage=actual_usage, total_usage=total_usage)
                        return response
                    continue  # filter is not passed; try the next config
        try:
            request_ts = get_current_ts()
            response = client.create(params)
        except Exception as e:
            if openai_result.is_successful:
                if APITimeoutError is not None and isinstance(e, APITimeoutError):
                    # logger.debug(f"config {i} timed out", exc_info=True)
                    if i == last:
                        raise TimeoutError(
                            "OpenAI API call timed out. This could be due to congestion or too small a timeout value. The timeout can be specified by setting the 'timeout' value (in seconds) in the llm_config (if you are using agents) or the OpenAIWrapper constructor (if you are using the OpenAIWrapper directly)."
                        ) from e
                elif APIError is not None and isinstance(e, APIError):
                    error_code = getattr(e, "code", None)
                    if logging_enabled():
                        log_chat_completion(
                            invocation_id=invocation_id,
                            client_id=id(client),
                            wrapper_id=id(self),
                            agent=agent,
                            request=params,
                            response=f"error_code:{error_code}, config {i} failed",
                            is_cached=0,
                            cost=0,
                            start_time=request_ts,
                        )

                    if error_code == "content_filter":
                        # raise the error for content_filter
                        raise
                    # logger.debug(f"config {i} failed", exc_info=True)
                    if i == last:
                        raise
                else:
                    raise
            else:
                raise
        except (
            gemini_InternalServerError,
            gemini_ResourceExhausted,
            anthorpic_InternalServerError,
            anthorpic_RateLimitError,
            mistral_SDKError,
            mistral_HTTPValidationError,
            together_TogetherException,
            groq_InternalServerError,
            groq_RateLimitError,
            groq_APIConnectionError,
            cohere_InternalServerError,
            cohere_TooManyRequestsError,
            cohere_ServiceUnavailableError,
            ollama_RequestError,
            ollama_ResponseError,
            bedrock_BotoCoreError,
            bedrock_ClientError,
            cerebras_AuthenticationError,
            cerebras_InternalServerError,
            cerebras_RateLimitError,
        ):
            # logger.debug(f"config {i} failed", exc_info=True)
            if i == last:
                raise
        else:
            # add cost calculation before caching no matter filter is passed or not
            if price is not None:
                response.cost = self._cost_with_customized_price(response, price)
            else:
                response.cost = client.cost(response)
            actual_usage = client.get_usage(response)
            total_usage = actual_usage.copy() if actual_usage is not None else total_usage
            self._update_usage(actual_usage=actual_usage, total_usage=total_usage)

            if cache_client is not None:
                # Cache the response
                with cache_client as cache:
                    cache.set(key, response)

            if logging_enabled():
                # TODO: log the config_id and pass_filter etc.
                log_chat_completion(
                    invocation_id=invocation_id,
                    client_id=id(client),
                    wrapper_id=id(self),
                    agent=agent,
                    request=params,
                    response=response,
                    is_cached=0,
                    cost=response.cost,
                    start_time=request_ts,
                )

            # Store metadata instead of mutating response
            # Keep backward compatibility by setting message_retrieval_function for now
            if hasattr(response, "message_retrieval_function"):
                response.message_retrieval_function = client.message_retrieval

            # check the filter
            pass_filter = filter_func is None or filter_func(context=context, response=response)
            if pass_filter or i == last:
                # Store metadata for serializable responses
                if hasattr(response, "id"):
                    self._store_response_metadata(response.id, client, i, pass_filter)

                # Backward compatibility: set attributes on ChatCompletionExtended
                if hasattr(response, "config_id"):
                    response.config_id = i
                if hasattr(response, "pass_filter"):
                    response.pass_filter = pass_filter

                # Return the response if it passes the filter or it is the last client
                return response
            continue  # filter is not passed; try the next config
    raise RuntimeError("Should not reach here.")

def print_usage_summary(self, mode: str | list[str] = ["actual", "total"]) -> None:
    """Print the usage summary."""
    iostream = IOStream.get_default()

    if isinstance(mode, list):
        if len(mode) == 0 or len(mode) > 2:
            raise ValueError(f'Invalid mode: {mode}, choose from "actual", "total", ["actual", "total"]')
        if "actual" in mode and "total" in mode:
            mode = "both"
        elif "actual" in mode:
            mode = "actual"
        elif "total" in mode:
            mode = "total"

    iostream.send(
        UsageSummaryEvent(
            actual_usage_summary=self.actual_usage_summary, total_usage_summary=self.total_usage_summary, mode=mode
        )
    )

def clear_usage_summary(self) -> None:
    """Clear the usage summary."""
    self.total_usage_summary = None
    self.actual_usage_summary = None

def extract_text_or_completion_object(self, response: Any) -> list[str] | list[dict[str, Any]]:
    """Extract the text or ChatCompletion objects from a completion or chat response.

    Supports both legacy responses (with message_retrieval_function) and new serializable responses.

    Args:
        response: The response from any client (ChatCompletion, UnifiedResponse, etc.)

    Returns:
        A list of text, or a list of message dicts if function_call/tool_calls are present.
    """
    # Option 1: Legacy path - response has message_retrieval_function attached
    if hasattr(response, "message_retrieval_function") and callable(response.message_retrieval_function):
        return response.message_retrieval_function(response)  # type: ignore [misc]

    # Option 2: Use stored metadata to find client
    if hasattr(response, "id") and response.id in self._response_metadata:
        metadata = self._response_metadata[response.id]
        client = metadata["client"]
        return client.message_retrieval(response)

    # Option 3: Fallback - try to extract from response structure directly
    # This handles cases where response is not in buffer
    if hasattr(response, "choices"):
        # OpenAI-style response
        return [
            choice.message
            if hasattr(choice.message, "tool_calls") and choice.message.tool_calls
            else getattr(choice.message, "content", "")
            for choice in response.choices
        ]

    # Last resort: return empty list
    warnings.warn(
        f"Could not extract messages from response type {type(response).__name__}. "
        "Response may not be in metadata buffer or may not support extraction.",
        UserWarning,
    )
    return []