from __future__ import annotations

import inspect
from enum import Enum
from types import SimpleNamespace
from typing import (
    Any,
    Callable,
    Collection,
    Dict,
    Iterable,
    Iterator,
    List,
    Mapping,
    Optional,
    Tuple,
    TypeVar,
    Union,
)

from opentelemetry import context as context_api
from opentelemetry import trace as trace_api
from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
from opentelemetry.util.types import AttributeValue

from openinference.instrumentation import (
    OITracer,
    Tool,
    TraceConfig,
    get_attributes_from_context,
    get_input_attributes,
    get_llm_attributes,
    get_llm_invocation_parameter_attributes,
    get_llm_provider_attributes,
    get_llm_tool_attributes,
    get_output_attributes,
    safe_json_dumps,
)
from openinference.instrumentation.litellm._anthropic_messages_attributes import (
    ANTHROPIC_KEYS_TO_REDACT,
    AnthropicMessagesStreamAccumulator,
    _get_attributes_from_anthropic_input_messages,
    _get_attributes_from_anthropic_output_message,
    _get_output_text_from_anthropic_response,
)
from openinference.instrumentation.litellm._responses_attributes import (
    _get_attributes_from_response_input,
    _get_attributes_from_response_output,
)
from openinference.instrumentation.litellm.package import _instruments
from openinference.instrumentation.litellm.version import __version__
from openinference.semconv.trace import (
    EmbeddingAttributes,
    ImageAttributes,
    MessageAttributes,
    MessageContentAttributes,
    OpenInferenceLLMProviderValues,
    OpenInferenceLLMSystemValues,
    OpenInferenceMimeTypeValues,
    OpenInferenceSpanKindValues,
    SpanAttributes,
    ToolAttributes,
    ToolCallAttributes,
)

# Skip capture
KEYS_TO_REDACT = ["api_key", "messages"]

RESPONSES_KEYS_TO_REDACT = KEYS_TO_REDACT + ["input", "tools", "model"]


# Helper functions to set span attributes
def _set_span_attribute(span: trace_api.Span, name: str, value: AttributeValue) -> None:
    if value is not None and value != "":
        span.set_attribute(name, value)


_LITELLM_TO_OPENINFERENCE_PROVIDERS = {
    "openai": OpenInferenceLLMProviderValues.OPENAI,
    "anthropic": OpenInferenceLLMProviderValues.ANTHROPIC,
    "cohere": OpenInferenceLLMProviderValues.COHERE,
    "mistral": OpenInferenceLLMProviderValues.MISTRALAI,
    "vertex_ai": OpenInferenceLLMProviderValues.GOOGLE,
    "gemini": OpenInferenceLLMProviderValues.GOOGLE,
    "azure": OpenInferenceLLMProviderValues.AZURE,
    "bedrock": OpenInferenceLLMProviderValues.AWS,
    "xai": OpenInferenceLLMProviderValues.XAI,
    "deepseek": OpenInferenceLLMProviderValues.DEEPSEEK,
}


def _get_oi_provider_from_litellm_model_name(
    model_name: str,
) -> Optional[OpenInferenceLLMProviderValues]:
    import litellm

    try:
        get_llm_provider = getattr(litellm, "get_llm_provider", None)
        if get_llm_provider is None:
            return None
        _, litellm_provider, _, _ = get_llm_provider(model=model_name)
    except Exception:
        return None
    return _LITELLM_TO_OPENINFERENCE_PROVIDERS.get(litellm_provider)


T = TypeVar("T", bound=type)


def is_iterable_of(lst: Iterable[object], tp: T) -> bool:
    return isinstance(lst, Iterable) and all(isinstance(x, tp) for x in lst)


def _set_output_message_value(span: trace_api.Span, result: Any) -> Any:
    from litellm.types.utils import Choices

    if (
        result.choices
        and isinstance(result.choices[-1], Choices)
        and (output_value := result.choices[-1].message.content)
    ):
        _set_span_attribute(span, SpanAttributes.OUTPUT_VALUE, output_value)
    else:
        _set_span_attribute(span, SpanAttributes.OUTPUT_VALUE, result.model_dump_json())
        _set_span_attribute(
            span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.JSON.value
        )


def _get_attributes_from_message_param(
    message: Any,
) -> Iterator[Tuple[str, AttributeValue]]:
    if not hasattr(message, "get"):
        return
    if role := message.get("role"):
        yield (
            MessageAttributes.MESSAGE_ROLE,
            role.value if isinstance(role, Enum) else role,
        )

    reasoning_blocks = _get_reasoning_content_blocks(message)
    content_index = 0
    for block in reasoning_blocks:
        for key, value in _get_attributes_from_reasoning_block(block):
            yield f"{MessageAttributes.MESSAGE_CONTENTS}.{content_index}.{key}", value
        content_index += 1

    if content := message.get("content"):
        if isinstance(content, str):
            if reasoning_blocks:
                prefix = f"{MessageAttributes.MESSAGE_CONTENTS}.{content_index}"
                yield f"{prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TYPE}", "text"
                yield f"{prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TEXT}", content
                content_index += 1
            else:
                yield MessageAttributes.MESSAGE_CONTENT, content
        elif is_iterable_of(content, dict):
            for c in content:
                for key, value in _get_attributes_from_message_content(c):
                    yield f"{MessageAttributes.MESSAGE_CONTENTS}.{content_index}.{key}", value
                content_index += 1

    if tool_calls := message.get("tool_calls"):
        if isinstance(tool_calls, Iterable):
            for tool_call_index, tool_call in enumerate(tool_calls):
                tool_call_prefix = f"{MessageAttributes.MESSAGE_TOOL_CALLS}.{tool_call_index}"
                tool_call_id = tool_call.get("id")
                reasoning_signature = None
                if tool_call_id:
                    tool_call_id, reasoning_signature = _split_litellm_tool_call_id(tool_call_id)
                    yield (
                        f"{tool_call_prefix}.{ToolCallAttributes.TOOL_CALL_ID}",
                        tool_call_id,
                    )
                    if reasoning_signature:
                        yield (
                            f"{tool_call_prefix}.{ToolCallAttributes.TOOL_CALL_REASONING_SIGNATURE}",
                            reasoning_signature,
                        )
                function_name = None
                function_arguments = None
                if function := tool_call.get("function"):
                    if function_name := function.get("name"):
                        yield (
                            f"{tool_call_prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}",
                            function_name,
                        )
                    if function_arguments := function.get("arguments"):
                        yield (
                            f"{tool_call_prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            function_arguments,
                        )
                if reasoning_blocks:
                    prefix = f"{MessageAttributes.MESSAGE_CONTENTS}.{content_index}"
                    yield f"{prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TYPE}", "tool_use"
                    if tool_call_id:
                        yield f"{prefix}.{ToolCallAttributes.TOOL_CALL_ID}", tool_call_id
                    if reasoning_signature:
                        yield (
                            f"{prefix}.{ToolCallAttributes.TOOL_CALL_REASONING_SIGNATURE}",
                            reasoning_signature,
                        )
                    if function_name:
                        yield (
                            f"{prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}",
                            function_name,
                        )
                    if function_arguments:
                        yield (
                            f"{prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
                            function_arguments,
                        )
                    content_index += 1


_THOUGHT_SIGNATURE_ID_SEPARATOR = "__thought__"


def _split_litellm_tool_call_id(tool_call_id: str) -> Tuple[str, Optional[str]]:
    """
    LiteLLM encodes a Gemini functionCall's thoughtSignature into the tool call id
    as ``<id>__thought__<signature>`` for OpenAI-client compatibility. Split it back
    into the clean id and the signature, if present.
    """
    call_id, separator, signature = tool_call_id.partition(_THOUGHT_SIGNATURE_ID_SEPARATOR)
    return (call_id, signature) if separator else (tool_call_id, None)


def _normalize_thinking_block(block: Mapping[str, Any]) -> Dict[str, Any]:
    normalized: Dict[str, Any] = {}
    if thinking := block.get("thinking"):
        normalized["text"] = thinking
    if data := block.get("data"):  # redacted_thinking payload
        normalized["data"] = data
    if signature := block.get("signature"):
        normalized["signature"] = signature
    return normalized


def _normalize_reasoning_item(item: Mapping[str, Any]) -> Dict[str, Any]:
    normalized: Dict[str, Any] = {}
    if item_id := item.get("id"):
        normalized["id"] = item_id
    summary = item.get("summary") or []
    if is_iterable_of(summary, dict):
        texts = [
            text
            for part in summary
            if part.get("type") == "summary_text" and (text := part.get("text"))
        ]
        if texts:
            normalized["text"] = "\n".join(texts)
    if encrypted_content := item.get("encrypted_content"):
        normalized["encrypted_content"] = encrypted_content
    return normalized


def _get_reasoning_content_blocks(message: Any) -> List[Mapping[str, Any]]:
    thinking_blocks = message.get("thinking_blocks")
    if thinking_blocks and is_iterable_of(thinking_blocks, dict):
        normalized_blocks: List[Mapping[str, Any]] = [
            normalized
            for block in thinking_blocks
            if (normalized := _normalize_thinking_block(block))
        ]
        if normalized_blocks:
            return normalized_blocks

    reasoning_items = message.get("reasoning_items")
    if reasoning_items and is_iterable_of(reasoning_items, dict):
        blocks: List[Mapping[str, Any]] = [
            normalized
            for item in reasoning_items
            if (normalized := _normalize_reasoning_item(item))
        ]
        if blocks:
            return blocks

    reasoning_content = message.get("reasoning_content")
    if isinstance(reasoning_content, str) and reasoning_content:
        return [{"text": reasoning_content}]
    return []


def _get_attributes_from_reasoning_block(
    block: Mapping[str, Any],
) -> Iterator[Tuple[str, AttributeValue]]:
    yield MessageContentAttributes.MESSAGE_CONTENT_TYPE, "reasoning"
    if block_id := block.get("id"):
        yield MessageContentAttributes.MESSAGE_CONTENT_ID, block_id
    if text := block.get("text"):
        yield MessageContentAttributes.MESSAGE_CONTENT_TEXT, text
    if data := block.get("data"):
        yield MessageContentAttributes.MESSAGE_CONTENT_DATA, data
    if signature := block.get("signature"):
        yield MessageContentAttributes.MESSAGE_CONTENT_SIGNATURE, signature
    if encrypted_content := block.get("encrypted_content"):
        yield MessageContentAttributes.MESSAGE_CONTENT_ENCRYPTED_CONTENT, encrypted_content


def _get_attributes_from_message_content(
    content: Mapping[str, Any],
) -> Iterator[Tuple[str, AttributeValue]]:
    content = dict(content)
    type_ = content.pop("type")
    if type_ == "text":
        yield f"{MessageContentAttributes.MESSAGE_CONTENT_TYPE}", "text"
        if text := content.pop("text"):
            yield f"{MessageContentAttributes.MESSAGE_CONTENT_TEXT}", text
    elif type_ == "image_url":
        yield f"{MessageContentAttributes.MESSAGE_CONTENT_TYPE}", "image"
        if image := content.pop("image_url"):
            for key, value in _get_attributes_from_image(image):
                yield f"{MessageContentAttributes.MESSAGE_CONTENT_IMAGE}.{key}", value
    elif type_ in ("thinking", "redacted_thinking"):
        yield from _get_attributes_from_reasoning_block(_normalize_thinking_block(content))


def _get_attributes_from_image(
    image: Mapping[str, Any],
) -> Iterator[Tuple[str, AttributeValue]]:
    image = dict(image)
    if url := image.pop("url"):
        yield f"{ImageAttributes.IMAGE_URL}", url


def _instrument_func_type_responses(span: trace_api.Span, kwargs: Dict[str, Any]) -> None:
    """
    Currently instruments the functions:
        litellm.responses()
        litellm.aresponses() (async version of responses)
    """
    _set_span_attribute(
        span, SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.LLM.value
    )
    model = kwargs.get("model")
    if model:
        span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
        provider = _get_oi_provider_from_litellm_model_name(model)
        span.set_attributes(get_llm_provider_attributes(provider))
    if message_objs := _get_attributes_from_response_input(kwargs):
        span.set_attributes(get_llm_attributes(input_messages=message_objs))
        span.set_attributes(get_input_attributes(kwargs.get("input")))
    invocation_params = {k: v for k, v in kwargs.items() if k not in RESPONSES_KEYS_TO_REDACT}
    span.set_attributes(get_llm_invocation_parameter_attributes(invocation_params))
    span.set_attributes(
        get_llm_tool_attributes([Tool(json_schema=tool) for tool in kwargs.get("tools", [])])
    )


def _instrument_func_type_completion(span: trace_api.Span, kwargs: Dict[str, Any]) -> None:
    """
    Currently instruments the functions:
        litellm.completion()
        litellm.acompletion() (async version of completion)
        litellm.completion_with_retries()
        litellm.acompletion_with_retries() (async version of completion_with_retries)
    """
    from litellm.types.utils import Message as LitellmMessage

    _set_span_attribute(
        span, SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.LLM.value
    )
    model = kwargs.get("model")
    if model:
        span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
        provider = _get_oi_provider_from_litellm_model_name(model)
        span.set_attributes(get_llm_provider_attributes(provider))

    if messages := kwargs.get("messages"):
        messages_as_dicts = []
        for input_message in messages:
            if isinstance(input_message, LitellmMessage):
                messages_as_dicts.append(input_message.json())  # type: ignore[no-untyped-call]
            else:
                messages_as_dicts.append(input_message)

        for index, input_message in enumerate(messages):
            for key, value in _get_attributes_from_message_param(input_message):
                _set_span_attribute(
                    span, f"{SpanAttributes.LLM_INPUT_MESSAGES}.{index}.{key}", value
                )

        if messages_as_dicts:
            _set_span_attribute(
                span, SpanAttributes.INPUT_VALUE, safe_json_dumps({"messages": messages_as_dicts})
            )
            _set_span_attribute(span, SpanAttributes.INPUT_MIME_TYPE, "application/json")

    invocation_params = {k: v for k, v in kwargs.items() if k not in KEYS_TO_REDACT}
    _set_span_attribute(
        span, SpanAttributes.LLM_INVOCATION_PARAMETERS, safe_json_dumps(invocation_params)
    )

    # Capture tool schemas
    if tools := kwargs.get("tools"):
        if isinstance(tools, list):
            for idx, tool in enumerate(tools):
                _set_span_attribute(
                    span,
                    f"{SpanAttributes.LLM_TOOLS}.{idx}.{ToolAttributes.TOOL_JSON_SCHEMA}",
                    safe_json_dumps(tool),
                )


def _instrument_func_type_embedding(span: trace_api.Span, kwargs: Dict[str, Any]) -> None:
    """
    Currently instruments the functions:
        litellm.embedding()
        litellm.aembedding() (async version of embedding)
    """
    _set_span_attribute(
        span,
        SpanAttributes.OPENINFERENCE_SPAN_KIND,
        OpenInferenceSpanKindValues.EMBEDDING.value,
    )
    _set_span_attribute(
        span, SpanAttributes.EMBEDDING_MODEL_NAME, kwargs.get("model", "unknown_model")
    )

    # Extract invocation parameters (exclude sensitive keys and input)
    invocation_params = {
        k: v for k, v in kwargs.items() if k not in KEYS_TO_REDACT and k != "input"
    }
    if invocation_params:
        _set_span_attribute(
            span, SpanAttributes.EMBEDDING_INVOCATION_PARAMETERS, safe_json_dumps(invocation_params)
        )

    # Extract text from embedding input - only records text, not token IDs
    embedding_input = kwargs.get("input")
    if embedding_input is not None:
        if isinstance(embedding_input, str):
            # Single string input
            _set_span_attribute(
                span,
                f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.0.{EmbeddingAttributes.EMBEDDING_TEXT}",
                embedding_input,
            )
        elif isinstance(embedding_input, list) and embedding_input:
            # Check if it's a list of strings (not tokens)
            if all(isinstance(item, str) for item in embedding_input):
                # List of strings
                for index, text in enumerate(embedding_input):
                    _set_span_attribute(
                        span,
                        f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.{index}.{EmbeddingAttributes.EMBEDDING_TEXT}",
                        text,
                    )

    _set_span_attribute(span, SpanAttributes.INPUT_VALUE, str(kwargs.get("input")))


def _instrument_func_type_image_generation(span: trace_api.Span, kwargs: Dict[str, Any]) -> None:
    """
    Currently instruments the functions:
        litellm.image_generation()
        litellm.aimage_generation() (async version of image_generation)
    """
    _set_span_attribute(
        span, SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.LLM.value
    )
    if model := kwargs.get("model"):
        _set_span_attribute(span, SpanAttributes.LLM_MODEL_NAME, model)
    if prompt := kwargs.get("prompt"):
        _set_span_attribute(span, SpanAttributes.INPUT_VALUE, str(prompt))


_ANTHROPIC_MESSAGES_POSITIONAL_PARAMETERS = (
    "max_tokens",
    "messages",
    "model",
    "metadata",
    "stop_sequences",
    "stream",
    "system",
    "temperature",
    "thinking",
    "tool_choice",
    "tools",
    "top_k",
    "top_p",
    "container",
)


def _bind_anthropic_messages_arguments(
    signature: Optional[inspect.Signature],
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
) -> Dict[str, Any]:
    if signature is not None:
        try:
            bound = signature.bind_partial(*args, **kwargs)
        except TypeError:
            pass
        else:
            arguments: Dict[str, Any] = {}
            for name, value in bound.arguments.items():
                parameter = signature.parameters[name]
                if parameter.kind is inspect.Parameter.VAR_KEYWORD and isinstance(value, Mapping):
                    arguments.update(value)
                elif parameter.kind is not inspect.Parameter.VAR_POSITIONAL:
                    arguments[name] = value
            return arguments

    arguments = dict(zip(_ANTHROPIC_MESSAGES_POSITIONAL_PARAMETERS, args))
    arguments.update(kwargs)
    return arguments


def _instrument_func_type_anthropic_messages(span: trace_api.Span, kwargs: Dict[str, Any]) -> None:
    """
    Instruments:
        litellm.anthropic.create() / litellm.anthropic.messages.create()
        litellm.anthropic.acreate() / litellm.anthropic.messages.acreate()
    """
    _set_span_attribute(
        span, SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.LLM.value
    )
    _set_span_attribute(
        span, SpanAttributes.LLM_SYSTEM, OpenInferenceLLMSystemValues.ANTHROPIC.value
    )
    model = kwargs.get("model")
    if model:
        span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
        provider = _get_oi_provider_from_litellm_model_name(model)
        span.set_attributes(get_llm_provider_attributes(provider))

    messages = kwargs.get("messages") or []
    system = kwargs.get("system")
    for key, value in _get_attributes_from_anthropic_input_messages(messages, system):
        _set_span_attribute(span, key, value)

    input_payload: Dict[str, Any] = {"messages": list(messages)}
    if system:
        input_payload["system"] = system
    _set_span_attribute(span, SpanAttributes.INPUT_VALUE, safe_json_dumps(input_payload))
    _set_span_attribute(
        span, SpanAttributes.INPUT_MIME_TYPE, OpenInferenceMimeTypeValues.JSON.value
    )

    invocation_params = {k: v for k, v in kwargs.items() if k not in ANTHROPIC_KEYS_TO_REDACT}
    _set_span_attribute(
        span, SpanAttributes.LLM_INVOCATION_PARAMETERS, safe_json_dumps(invocation_params)
    )

    if tools := kwargs.get("tools"):
        if isinstance(tools, list):
            for idx, tool in enumerate(tools):
                _set_span_attribute(
                    span,
                    f"{SpanAttributes.LLM_TOOLS}.{idx}.{ToolAttributes.TOOL_JSON_SCHEMA}",
                    safe_json_dumps(tool),
                )


def _as_anthropic_messages_mapping(result: Any) -> Optional[Mapping[str, Any]]:
    if isinstance(result, Mapping):
        return result
    if hasattr(result, "model_dump"):
        try:
            dumped = result.model_dump()
            if isinstance(dumped, Mapping):
                return dumped
        except Exception:
            pass
    if hasattr(result, "__dict__"):
        data = vars(result)
        return data if isinstance(data, Mapping) else None
    return None


def _finalize_anthropic_messages_span(span: trace_api.Span, result: Any) -> None:
    response = _as_anthropic_messages_mapping(result)
    if response is None:
        _set_span_status(span, result)
        return

    if model := response.get("model"):
        _set_span_attribute(span, SpanAttributes.LLM_MODEL_NAME, model)

    for key, value in _get_attributes_from_anthropic_output_message(response):
        _set_span_attribute(span, key, value)

    if output_text := _get_output_text_from_anthropic_response(response):
        span.set_attributes(get_output_attributes(output_text))
    else:
        _set_span_attribute(span, SpanAttributes.OUTPUT_VALUE, safe_json_dumps(dict(response)))
        _set_span_attribute(
            span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.JSON.value
        )

    usage = response.get("usage")
    if usage is not None:
        if isinstance(usage, Mapping):
            usage_obj = SimpleNamespace(**dict(usage))
        else:
            usage_obj = usage
        _set_token_counts_from_usage(span, SimpleNamespace(usage=usage_obj))
        input_tokens = _get_value(usage_obj, "input_tokens")
        cache_creation_input_tokens = _get_value(usage_obj, "cache_creation_input_tokens")
        cache_read_input_tokens = _get_value(usage_obj, "cache_read_input_tokens")
        output_tokens = _get_value(usage_obj, "output_tokens")
        if any(
            value is not None
            for value in (
                input_tokens,
                cache_creation_input_tokens,
                cache_read_input_tokens,
                output_tokens,
            )
        ):
            # Re-derive the prompt count to include cache tokens (Anthropic reports
            # input_tokens exclusive of cache reads/writes) and set the total, which
            # Anthropic usage does not provide directly. Matches the anthropic instrumentor.
            prompt_tokens = (
                (input_tokens or 0)
                + (cache_creation_input_tokens or 0)
                + (cache_read_input_tokens or 0)
            )
            _set_span_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_PROMPT, prompt_tokens)
            _set_span_attribute(
                span,
                SpanAttributes.LLM_TOKEN_COUNT_TOTAL,
                prompt_tokens + (output_tokens or 0),
            )

    _set_cost_from_response(span, result)
    _set_span_status(span, result)


def _finalize_sync_anthropic_messages_streaming_span(span: trace_api.Span, stream: Any) -> Any:
    accumulator = AnthropicMessagesStreamAccumulator()
    try:
        with trace_api.use_span(
            span, end_on_exit=False, record_exception=False, set_status_on_exception=False
        ):
            for chunk in stream:
                accumulator.process_chunk(chunk)
                yield chunk
        accumulator.finish()
        _finalize_anthropic_messages_span(span, accumulator.to_response())
    except Exception as e:
        span.record_exception(e)
        span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(e)))
        raise
    finally:
        span.end()


async def _finalize_async_anthropic_messages_streaming_span(
    span: trace_api.Span, stream: Any
) -> Any:
    accumulator = AnthropicMessagesStreamAccumulator()
    try:
        with trace_api.use_span(
            span, end_on_exit=False, record_exception=False, set_status_on_exception=False
        ):
            async for chunk in stream:
                accumulator.process_chunk(chunk)
                yield chunk
        accumulator.finish()
        _finalize_anthropic_messages_span(span, accumulator.to_response())
    except Exception as e:
        span.record_exception(e)
        span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(e)))
        raise
    finally:
        span.end()


def _finalize_anthropic_messages_result(span: trace_api.Span, result: Any) -> Any:
    if hasattr(result, "__anext__"):
        return _finalize_async_anthropic_messages_streaming_span(span, result)
    if hasattr(result, "__next__"):
        return _finalize_sync_anthropic_messages_streaming_span(span, result)
    _finalize_anthropic_messages_span(span, result)
    span.end()
    return result


async def _finalize_anthropic_messages_awaitable(span: trace_api.Span, result: Any) -> Any:
    try:
        with trace_api.use_span(
            span, end_on_exit=False, record_exception=False, set_status_on_exception=False
        ):
            result = await result
    except Exception as e:
        span.record_exception(e)
        span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(e)))
        span.end()
        raise
    return _finalize_anthropic_messages_result(span, result)


def _finalize_span(span: trace_api.Span, result: Any) -> None:
    from openai.types.image import Image

    from litellm.types.llms.openai import ResponsesAPIResponse
    from litellm.types.utils import Choices, EmbeddingResponse, ImageResponse, ModelResponse

    if isinstance(result, ModelResponse):
        _set_output_message_value(span, result)
        for idx, choice in enumerate(result.choices):
            if not isinstance(choice, Choices):
                continue

            for key, value in _get_attributes_from_message_param(choice.message):
                _set_span_attribute(
                    span, f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}.{key}", value
                )

    elif isinstance(result, EmbeddingResponse):
        # Extract model name from response (may differ from request model name)
        if model_name := getattr(result, "model", None):
            _set_span_attribute(span, SpanAttributes.EMBEDDING_MODEL_NAME, model_name)

        if result_data := getattr(result, "data", None):
            # Extract embedding vectors directly
            for embedding_item in result_data:
                # LiteLLM may return embedding items as dicts or OpenAIObject instances
                if isinstance(embedding_item, dict):
                    raw_vector = embedding_item.get("embedding")
                    index = embedding_item.get("index")
                else:
                    raw_vector = getattr(embedding_item, "embedding", None)
                    index = getattr(embedding_item, "index", None)

                # Skip entries without a usable vector or explicit index
                if raw_vector is None or index is None:
                    continue

                # Skip empty embeddings to avoid recording invalid vectors
                if raw_vector == [] or raw_vector == "":
                    continue

                vector: Union[Tuple[Any, ...], str, None] = None
                # Record numeric vectors as tuples
                if isinstance(raw_vector, (list, tuple)):
                    if all(isinstance(x, (int, float)) for x in raw_vector):
                        vector = tuple(raw_vector)
                # Record base64-encoded vectors directly
                elif isinstance(raw_vector, str):
                    vector = raw_vector
                else:
                    continue

                if vector is not None:
                    _set_span_attribute(
                        span,
                        f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.{index}.{EmbeddingAttributes.EMBEDDING_VECTOR}",
                        vector,
                    )
    elif isinstance(result, ImageResponse):
        if result.data and len(result.data) > 0:
            if img_data := result.data[0]:
                if isinstance(img_data, Image) and (url := (img_data.url or img_data.b64_json)):
                    _set_span_attribute(span, ImageAttributes.IMAGE_URL, url)
                    _set_span_attribute(span, SpanAttributes.OUTPUT_VALUE, url)
                elif isinstance(img_data, dict) and (
                    url := (img_data.get("url") or img_data.get("b64_json"))
                ):
                    _set_span_attribute(span, ImageAttributes.IMAGE_URL, url)
                    _set_span_attribute(span, SpanAttributes.OUTPUT_VALUE, url)
    elif isinstance(result, ResponsesAPIResponse):
        span.set_attributes(_get_attributes_from_response_output(result))

    _set_token_counts_from_usage(span, result)
    _set_cost_from_response(span, result)
    _set_span_status(span, result)


# Gets values safely from an object
def _get_value(obj: object, key: str) -> Any:
    if hasattr(obj, key):
        return getattr(obj, key)
    return None


def _set_token_counts_from_usage(span: trace_api.Span, result: Any) -> None:
    """
    Sets token count attributes on a span based on the usage information in result.
    """
    # Return early if no usage information
    if not hasattr(result, "usage"):
        return

    usage = result.usage
    if not usage:
        return

    prompt_tokens = _get_value(usage, "prompt_tokens") or _get_value(usage, "input_tokens")
    if prompt_tokens is not None:
        _set_span_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_PROMPT, prompt_tokens)

    prompt_token_details = _get_value(usage, "prompt_tokens_details") or _get_value(
        usage, "input_tokens_details"
    )
    if prompt_token_details is not None:
        cached_tokens = _get_value(prompt_token_details, "cached_tokens")
        if cached_tokens is not None:
            _set_span_attribute(
                span, SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, cached_tokens
            )

        audio_tokens = _get_value(prompt_token_details, "audio_tokens")
        if audio_tokens is not None:
            _set_span_attribute(
                span, SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO, audio_tokens
            )
        text_tokens = _get_value(prompt_token_details, "text_tokens")
        if text_tokens is not None:
            _set_span_attribute(
                span, SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_INPUT, text_tokens
            )

    completion_tokens = _get_value(usage, "completion_tokens") or _get_value(usage, "output_tokens")
    if completion_tokens is not None:
        _set_span_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, completion_tokens)

    completion_tokens_details = _get_value(usage, "completion_tokens_details") or _get_value(
        usage, "output_tokens_details"
    )
    if completion_tokens_details is not None:
        reasoning_tokens = _get_value(completion_tokens_details, "reasoning_tokens")
        if reasoning_tokens is not None:
            _set_span_attribute(
                span, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, reasoning_tokens
            )

        text_tokens = _get_value(completion_tokens_details, "text_tokens")
        if text_tokens is not None:
            _set_span_attribute(
                span, SpanAttributes.LLM_COST_COMPLETION_DETAILS_OUTPUT, text_tokens
            )

        completion_audio_tokens = _get_value(completion_tokens_details, "audio_tokens")
        if completion_audio_tokens is not None:
            _set_span_attribute(
                span,
                SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO,
                completion_audio_tokens,
            )

    total_tokens = _get_value(usage, "total_tokens")
    if total_tokens is not None:
        _set_span_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_TOTAL, total_tokens)

    cache_creation_input_tokens = _get_value(usage, "cache_creation_input_tokens")
    if cache_creation_input_tokens is not None:
        _set_span_attribute(
            span,
            SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE,
            cache_creation_input_tokens,
        )

    cache_read_input_tokens = _get_value(usage, "cache_read_input_tokens")
    if cache_read_input_tokens is not None:
        _set_span_attribute(
            span, SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, cache_read_input_tokens
        )

    _set_cost_from_usage(span, usage)


def _set_cost_from_usage(span: trace_api.Span, usage: Any) -> None:
    """
    Sets cost attribute on a span from usage object (streaming responses).
    """
    if not usage:
        return

    cost = _get_value(usage, "cost")
    if cost is not None:
        _set_span_attribute(span, SpanAttributes.LLM_COST_TOTAL, cost)


def _set_cost_from_response(span: trace_api.Span, result: Any) -> None:
    """
    Sets cost attribute on a span from response's _hidden_params (non-streaming responses).
    """
    hidden_params = _get_value(result, "_hidden_params")
    if not hidden_params:
        return

    if isinstance(hidden_params, dict):
        response_cost = hidden_params.get("response_cost")
    else:
        response_cost = _get_value(hidden_params, "response_cost")
    if response_cost is not None:
        _set_span_attribute(span, SpanAttributes.LLM_COST_TOTAL, response_cost)


def _set_span_status(span: trace_api.Span, result: Any) -> None:
    """
    Sets the span status based on whether the result contains an error.
    """
    error = _get_value(result, "error")
    if error is None and isinstance(result, dict):
        error = result.get("error", None)

    if error is not None:
        span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(error)))
    else:
        span.set_status(trace_api.Status(trace_api.StatusCode.OK))


def _accumulate_tool_calls(
    tool_calls: Any,
    accumulated: Dict[int, Dict[str, Any]],
) -> None:
    if not tool_calls:
        return

    for tool_call in tool_calls:
        tool_call_index = getattr(tool_call, "index", None)
        if tool_call_index is None:
            continue

        if tool_call_index not in accumulated:
            accumulated[tool_call_index] = {
                "id": None,
                "type": "function",
                "function": {"name": None, "arguments": ""},
            }

        tool_call_entry = accumulated[tool_call_index]

        if tool_call_id := getattr(tool_call, "id", None):
            tool_call_entry["id"] = tool_call_id

        if function := getattr(tool_call, "function", None):
            if name := getattr(function, "name", None):
                tool_call_entry["function"]["name"] = name
            if arguments := getattr(function, "arguments", None):
                tool_call_entry["function"]["arguments"] += arguments


def _accumulate_thinking_blocks(
    thinking_blocks: Any,
    accumulated: List[Dict[str, Any]],
) -> None:
    if not thinking_blocks:
        return

    for block in thinking_blocks:
        if not hasattr(block, "get"):
            continue
        block_type = block.get("type")
        if block_type == "redacted_thinking":
            if data := block.get("data"):
                accumulated.append({"type": "redacted_thinking", "data": data})
        elif block_type == "thinking":
            last = accumulated[-1] if accumulated else None
            if last is None or last.get("type") != "thinking" or last.get("signature"):
                last = {"type": "thinking", "thinking": "", "signature": ""}
                accumulated.append(last)
            if thinking := block.get("thinking"):
                last["thinking"] += thinking
            if signature := block.get("signature"):
                last["signature"] = signature


def _build_message_from_accumulated(msg: Dict[str, Any]) -> Dict[str, Any]:
    message: Dict[str, Any] = {
        "role": msg.get("role"),
        "content": msg.get("content"),
    }

    if reasoning_content := msg.get("reasoning_content"):
        message["reasoning_content"] = reasoning_content

    if thinking_blocks := msg.get("thinking_blocks"):
        message["thinking_blocks"] = thinking_blocks

    if reasoning_items := msg.get("reasoning_items"):
        message["reasoning_items"] = reasoning_items

    if tool_calls_dict := msg.get("tool_calls"):
        message["tool_calls"] = [tool_calls_dict[idx] for idx in sorted(tool_calls_dict.keys())]

    return message


def _remove_redundant_reasoning_entries(
    output_messages: Dict[int, Dict[str, Any]],
    reasoning_items: List[Mapping[str, Any]],
) -> None:
    authoritative_texts = {
        text
        for item in reasoning_items
        for part in (item.get("summary") or [])
        if isinstance(part, Mapping)
        and part.get("type") == "summary_text"
        and (text := part.get("text"))
    }
    for index, message in list(output_messages.items()):
        if index == 0:
            continue
        if not message.get("content") and set(message) <= {
            "role",
            "content",
            "reasoning_content",
        }:
            if message.get("reasoning_content") in authoritative_texts:
                del output_messages[index]


def _finalize_sync_streaming_span(span: trace_api.Span, stream: Any) -> Any:
    output_messages: Dict[int, Dict[str, Any]] = {}
    usage_stats = None
    aggregated_output = None
    try:
        for token in stream:
            if token.choices:
                for choice in token.choices:
                    idx = choice.index
                    entry = output_messages.get(idx)
                    if entry is None:
                        entry = output_messages[idx] = {"role": None, "content": ""}
                    delta = choice.delta
                    if delta:
                        role = getattr(delta, "role", None)
                        content = getattr(delta, "content", None)
                        if role is not None and entry["role"] is None:
                            entry["role"] = role
                        if content is not None:
                            entry["content"] += content
                        if reasoning_content := getattr(delta, "reasoning_content", None):
                            entry["reasoning_content"] = (
                                entry.get("reasoning_content", "") + reasoning_content
                            )
                        if thinking_blocks := getattr(delta, "thinking_blocks", None):
                            _accumulate_thinking_blocks(
                                thinking_blocks, entry.setdefault("thinking_blocks", [])
                            )
                        if reasoning_items := getattr(delta, "reasoning_items", None):
                            # Emitted once, as a full list, on the final response.completed
                            # chunk (see litellm's litellm_responses_transformation), so this
                            # is a replace, not an incremental accumulation like the above.
                            entry["reasoning_items"] = reasoning_items
                        if tool_calls := getattr(delta, "tool_calls", None):
                            _accumulate_tool_calls(tool_calls, entry.setdefault("tool_calls", {}))
            usage_attrs = getattr(token, "usage", None)
            if usage_attrs:
                usage_stats = usage_attrs
            yield token
        if reasoning_items := output_messages.get(0, {}).get("reasoning_items"):
            _remove_redundant_reasoning_entries(output_messages, reasoning_items)
        aggregated_output = output_messages.get(0, {}).get("content", "")
        _set_span_attribute(span, SpanAttributes.OUTPUT_VALUE, aggregated_output)
        for idx, msg in output_messages.items():
            message = _build_message_from_accumulated(msg)
            for key, value in _get_attributes_from_message_param(message):
                _set_span_attribute(
                    span, f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}.{key}", value
                )

        if usage_stats:
            _set_token_counts_from_usage(span, SimpleNamespace(usage=usage_stats))
    except Exception as e:
        span.record_exception(e)
        raise
    else:
        _set_span_status(span, aggregated_output)
    finally:
        span.end()


async def _finalize_streaming_span(span: trace_api.Span, stream: Any) -> Any:
    output_messages: Dict[int, Dict[str, Any]] = {}
    usage_stats = None
    try:
        async for token in stream:
            if token.choices:
                for choice in token.choices:
                    idx = choice.index
                    entry = output_messages.get(idx)
                    if entry is None:
                        entry = output_messages[idx] = {"role": None, "content": ""}
                    delta = choice.delta
                    if delta:
                        role = getattr(delta, "role", None)
                        content = getattr(delta, "content", None)
                        if role is not None and entry["role"] is None:
                            entry["role"] = role
                        if content is not None:
                            entry["content"] += content
                        if reasoning_content := getattr(delta, "reasoning_content", None):
                            entry["reasoning_content"] = (
                                entry.get("reasoning_content", "") + reasoning_content
                            )
                        if thinking_blocks := getattr(delta, "thinking_blocks", None):
                            _accumulate_thinking_blocks(
                                thinking_blocks, entry.setdefault("thinking_blocks", [])
                            )
                        if reasoning_items := getattr(delta, "reasoning_items", None):
                            # Emitted once, as a full list, on the final response.completed
                            # chunk (see litellm's litellm_responses_transformation), so this
                            # is a replace, not an incremental accumulation like the above.
                            entry["reasoning_items"] = reasoning_items
                        if tool_calls := getattr(delta, "tool_calls", None):
                            _accumulate_tool_calls(tool_calls, entry.setdefault("tool_calls", {}))
            usage_attrs = getattr(token, "usage", None)
            if usage_attrs:
                usage_stats = usage_attrs
            yield token
        if reasoning_items := output_messages.get(0, {}).get("reasoning_items"):
            _remove_redundant_reasoning_entries(output_messages, reasoning_items)
        aggregated_output = output_messages.get(0, {}).get("content", "")
        _set_span_attribute(span, SpanAttributes.OUTPUT_VALUE, aggregated_output)
        for idx, msg in output_messages.items():
            message = _build_message_from_accumulated(msg)
            for key, value in _get_attributes_from_message_param(message):
                _set_span_attribute(
                    span, f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}.{key}", value
                )
        if usage_stats:
            _set_token_counts_from_usage(span, SimpleNamespace(usage=usage_stats))
    except Exception as e:
        span.record_exception(e)
        raise
    else:
        _set_span_status(span, aggregated_output)
    finally:
        span.end()


async def _finalize_aresponses_streaming_span(span: trace_api.Span, stream: Any) -> Any:
    from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
    from litellm.types.llms.openai import ResponsesAPIStreamEvents

    if isinstance(stream, ResponsesAPIStreamingIterator):
        async for token in stream:
            if token.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
                _finalize_span(span, token.response)
                span.end()
            yield token


def _finalize_responses_streaming_span(span: trace_api.Span, stream: Any) -> Any:
    from litellm.responses.streaming_iterator import SyncResponsesAPIStreamingIterator
    from litellm.types.llms.openai import ResponsesAPIStreamEvents

    if isinstance(stream, SyncResponsesAPIStreamingIterator):
        for token in stream:
            if token.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
                _finalize_span(span, token.response)
                span.end()
            yield token


class LiteLLMInstrumentor(BaseInstrumentor):  # type: ignore
    original_litellm_funcs: Dict[
        str, Callable[..., Any]
    ] = {}  # Dictionary for original uninstrumented liteLLM functions
    original_anthropic_funcs: Dict[
        str, Callable[..., Any]
    ] = {}  # Original litellm.anthropic create/acreate
    original_anthropic_signatures: Dict[str, inspect.Signature] = {}

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments

    def _instrument(self, **kwargs: Any) -> None:
        import litellm

        if not (tracer_provider := kwargs.get("tracer_provider")):
            tracer_provider = trace_api.get_tracer_provider()
        if not (config := kwargs.get("config")):
            config = TraceConfig()
        else:
            assert isinstance(config, TraceConfig)
        self._tracer = OITracer(
            trace_api.get_tracer(__name__, __version__, tracer_provider),
            config=config,
        )

        functions_to_instrument = {
            "completion": self._completion_wrapper,
            "acompletion": self._acompletion_wrapper,
            "responses": self._responses_wrapper,
            "aresponses": self._aresponses_wrapper,
            "completion_with_retries": self._completion_with_retries_wrapper,
            # Bug report filed on GitHub for acompletion_with_retries: https://github.com/BerriAI/litellm/issues/4908
            # "acompletion_with_retries": self._acompletion_with_retries_wrapper,
            "embedding": self._embedding_wrapper,
            "aembedding": self._aembedding_wrapper,
            "image_generation": self._image_generation_wrapper,
            "aimage_generation": self._aimage_generation_wrapper,
        }

        for func_name, func_wrapper in functions_to_instrument.items():
            if hasattr(litellm, func_name):
                original_func = getattr(litellm, func_name)
                self.original_litellm_funcs[func_name] = (
                    original_func  # Add original liteLLM function to dictionary
                )
                setattr(
                    litellm, func_name, func_wrapper
                )  # Monkey patch each function with their respective wrapper
                self._set_wrapper_attr(func_wrapper)

        self._instrument_anthropic_messages()

    def _instrument_anthropic_messages(self) -> None:
        try:
            import litellm.anthropic_interface as anthropic_module
            import litellm.anthropic_interface.messages as anthropic_messages_module
        except ImportError:
            return

        anthropic_functions = {
            "create": self._create_wrapper,
            "acreate": self._acreate_wrapper,
        }
        for func_name, func_wrapper in anthropic_functions.items():
            if not hasattr(anthropic_messages_module, func_name):
                continue
            original_func = getattr(anthropic_messages_module, func_name)
            self.original_anthropic_funcs[func_name] = original_func
            try:
                self.original_anthropic_signatures[func_name] = inspect.signature(original_func)
            except (TypeError, ValueError):
                pass
            setattr(anthropic_messages_module, func_name, func_wrapper)
            if hasattr(anthropic_module, func_name):
                setattr(anthropic_module, func_name, func_wrapper)
            self._set_wrapper_attr(func_wrapper)

    def _uninstrument(self, **kwargs: Any) -> None:
        import litellm

        for func_name, original_func in LiteLLMInstrumentor.original_litellm_funcs.items():
            setattr(litellm, func_name, original_func)
        self.original_litellm_funcs.clear()
        self._uninstrument_anthropic_messages()

    def _uninstrument_anthropic_messages(self) -> None:
        if not self.original_anthropic_funcs:
            return
        try:
            import litellm.anthropic_interface as anthropic_module
            import litellm.anthropic_interface.messages as anthropic_messages_module
        except ImportError:
            self.original_anthropic_funcs.clear()
            self.original_anthropic_signatures.clear()
            return

        for func_name, original_func in list(self.original_anthropic_funcs.items()):
            setattr(anthropic_messages_module, func_name, original_func)
            if hasattr(anthropic_module, func_name):
                setattr(anthropic_module, func_name, original_func)
        self.original_anthropic_funcs.clear()
        self.original_anthropic_signatures.clear()

    def _create_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_anthropic_funcs["create"](*args, **kwargs)

        arguments = _bind_anthropic_messages_arguments(
            self.original_anthropic_signatures.get("create"), args, kwargs
        )
        span = self._tracer.start_span(
            name="messages.create", attributes=dict(get_attributes_from_context())
        )
        _instrument_func_type_anthropic_messages(span, arguments)
        try:
            with trace_api.use_span(
                span, end_on_exit=False, record_exception=False, set_status_on_exception=False
            ):
                result = self.original_anthropic_funcs["create"](*args, **kwargs)
        except Exception as e:
            span.record_exception(e)
            span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(e)))
            span.end()
            raise
        if inspect.isawaitable(result):
            return _finalize_anthropic_messages_awaitable(span, result)
        return _finalize_anthropic_messages_result(span, result)

    async def _acreate_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            result = self.original_anthropic_funcs["acreate"](*args, **kwargs)
            if hasattr(result, "__await__"):
                return await result
            return result

        arguments = _bind_anthropic_messages_arguments(
            self.original_anthropic_signatures.get("acreate"), args, kwargs
        )
        span = self._tracer.start_span(
            name="messages.create", attributes=dict(get_attributes_from_context())
        )
        _instrument_func_type_anthropic_messages(span, arguments)
        try:
            with trace_api.use_span(
                span, end_on_exit=False, record_exception=False, set_status_on_exception=False
            ):
                result = self.original_anthropic_funcs["acreate"](*args, **kwargs)
                if inspect.isawaitable(result):
                    result = await result
        except Exception as e:
            span.record_exception(e)
            span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(e)))
            span.end()
            raise
        return _finalize_anthropic_messages_result(span, result)

    def _responses_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        from litellm.responses.streaming_iterator import SyncResponsesAPIStreamingIterator

        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["responses"](*args, **kwargs)

        if kwargs.get("stream", False):
            span = self._tracer.start_span(
                name="responses", attributes=dict(get_attributes_from_context())
            )
            _instrument_func_type_responses(span, kwargs)

            result = self.original_litellm_funcs["responses"](*args, **kwargs)
            if isinstance(result, SyncResponsesAPIStreamingIterator):
                return _finalize_responses_streaming_span(span, result)
            return result
        else:
            with self._tracer.start_as_current_span(
                name="responses", attributes=dict(get_attributes_from_context())
            ) as span:
                _instrument_func_type_responses(span, kwargs)
                result = self.original_litellm_funcs["responses"](*args, **kwargs)
                _finalize_span(span, result)
                return result

    async def _aresponses_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return await self.original_litellm_funcs["aresponses"](*args, **kwargs)

        if kwargs.get("stream", False):
            span = self._tracer.start_span(
                name="aresponses", attributes=dict(get_attributes_from_context())
            )
            _instrument_func_type_responses(span, kwargs)
            result = await self.original_litellm_funcs["aresponses"](*args, **kwargs)
            if hasattr(result, "__aiter__"):
                return _finalize_aresponses_streaming_span(span, result)
            return result
        else:
            with self._tracer.start_as_current_span(
                name="aresponses", attributes=dict(get_attributes_from_context())
            ) as span:
                _instrument_func_type_responses(span, kwargs)
                result = await self.original_litellm_funcs["aresponses"](*args, **kwargs)
                _finalize_span(span, result)
                return result

    def _completion_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper

        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["completion"](*args, **kwargs)

        if kwargs.get("stream", False):
            span = self._tracer.start_span(
                name="completion", attributes=dict(get_attributes_from_context())
            )
            _instrument_func_type_completion(span, kwargs)

            result = self.original_litellm_funcs["completion"](*args, **kwargs)

            if isinstance(result, CustomStreamWrapper):
                return _finalize_sync_streaming_span(span, result)

            _finalize_span(span, result)
            span.end()
            return result
        else:
            with self._tracer.start_as_current_span(
                name="completion", attributes=dict(get_attributes_from_context())
            ) as span:
                _instrument_func_type_completion(span, kwargs)
                result = self.original_litellm_funcs["completion"](*args, **kwargs)
                _finalize_span(span, result)
                return result

    async def _acompletion_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return await self.original_litellm_funcs["acompletion"](*args, **kwargs)

        if kwargs.get("stream", False):
            span = self._tracer.start_span(
                name="acompletion", attributes=dict(get_attributes_from_context())
            )
            _instrument_func_type_completion(span, kwargs)

            result = await self.original_litellm_funcs["acompletion"](*args, **kwargs)

            if hasattr(result, "__aiter__"):
                return _finalize_streaming_span(span, result)

            _finalize_span(span, result)
            span.end()
            return result
        else:
            with self._tracer.start_as_current_span(
                name="acompletion", attributes=dict(get_attributes_from_context())
            ) as span:
                _instrument_func_type_completion(span, kwargs)
                result = await self.original_litellm_funcs["acompletion"](*args, **kwargs)
                _finalize_span(span, result)
                return result

    def _completion_with_retries_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["completion_with_retries"](*args, **kwargs)
        with self._tracer.start_as_current_span(
            name="completion_with_retries", attributes=dict(get_attributes_from_context())
        ) as span:
            _instrument_func_type_completion(span, kwargs)
            result = self.original_litellm_funcs["completion_with_retries"](*args, **kwargs)
            _finalize_span(span, result)
        return result

    async def _acompletion_with_retries_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["acompletion_with_retries"](*args, **kwargs)
        with self._tracer.start_as_current_span(
            name="acompletion_with_retries", attributes=dict(get_attributes_from_context())
        ) as span:
            _instrument_func_type_completion(span, kwargs)
            result = await self.original_litellm_funcs["acompletion_with_retries"](*args, **kwargs)
            _finalize_span(span, result)
        return result

    def _embedding_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["embedding"](*args, **kwargs)
        with self._tracer.start_as_current_span(
            name="CreateEmbeddings", attributes=dict(get_attributes_from_context())
        ) as span:
            _instrument_func_type_embedding(span, kwargs)
            result = self.original_litellm_funcs["embedding"](*args, **kwargs)
            _finalize_span(span, result)
        return result

    async def _aembedding_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["aembedding"](*args, **kwargs)
        with self._tracer.start_as_current_span(
            name="CreateEmbeddings", attributes=dict(get_attributes_from_context())
        ) as span:
            _instrument_func_type_embedding(span, kwargs)
            result = await self.original_litellm_funcs["aembedding"](*args, **kwargs)
            _finalize_span(span, result)
        return result

    def _image_generation_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["image_generation"](*args, **kwargs)
        with self._tracer.start_as_current_span(
            name="image_generation", attributes=dict(get_attributes_from_context())
        ) as span:
            _instrument_func_type_image_generation(span, kwargs)
            result = self.original_litellm_funcs["image_generation"](*args, **kwargs)
            _finalize_span(span, result)
        return result

    async def _aimage_generation_wrapper(self, *args: Any, **kwargs: Any) -> Any:
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return self.original_litellm_funcs["aimage_generation"](*args, **kwargs)
        with self._tracer.start_as_current_span(
            name="aimage_generation", attributes=dict(get_attributes_from_context())
        ) as span:
            _instrument_func_type_image_generation(span, kwargs)
            result = await self.original_litellm_funcs["aimage_generation"](*args, **kwargs)
            _finalize_span(span, result)
        return result

    def _set_wrapper_attr(self, func_wrapper: Any) -> None:
        func_wrapper.__func__.is_wrapper = True
