# pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument, too-many-branches
"""
Module for monitoring Azure OpenAI API calls.
"""

import logging
from opentelemetry.trace import SpanKind, Status, StatusCode
from opentelemetry.sdk.resources import TELEMETRY_SDK_NAME
from murnitur.lit.__helpers import get_chat_model_cost, get_embed_model_cost
from murnitur.lit.__helpers import get_image_model_cost, openai_tokens, handle_exception
from murnitur.lit.semcov import SemanticConvetion

# Initialize logger for logging potential issues and operations
logger = logging.getLogger(__name__)


def azure_chat_completions(
    gen_ai_endpoint,
    version,
    environment,
    application_name,
    tracer,
    pricing_info,
    trace_content,
    metrics,
    disable_metrics,
):
    """
    Generates a telemetry wrapper for chat completions to collect metrics.

    Args:
        gen_ai_endpoint: Endpoint identifier for logging and tracing.
        version: Version of the monitoring package.
        environment: Deployment environment (e.g., production, staging).
        application_name: Name of the application using the OpenAI API.
        tracer: OpenTelemetry tracer for creating spans.
        pricing_info: Information used for calculating the cost of OpenAI usage.
        trace_content: Flag indicating whether to trace the actual content.

    Returns:
        A function that wraps the chat completions method to add telemetry.
    """

    def wrapper(wrapped, instance, args, kwargs):
        """
        Wraps the 'chat.completions' API call to add telemetry.

        This collects metrics such as execution time, cost, and token usage, and handles errors
        gracefully, adding details to the trace for observability.

        Args:
            wrapped: The original 'chat.completions' method to be wrapped.
            instance: The instance of the class where the original method is defined.
            args: Positional arguments for the 'chat.completions' method.
            kwargs: Keyword arguments for the 'chat.completions' method.

        Returns:
            The response from the original 'chat.completions' method.
        """

        # Check if streaming is enabled for the API call
        streaming = kwargs.get("stream", False)

        # pylint: disable=no-else-return
        if streaming:
            # Special handling for streaming response to accommodate the nature of data flow
            def stream_generator():
                # pylint: disable=line-too-long
                with tracer.start_as_current_span(
                    gen_ai_endpoint, kind=SpanKind.CLIENT
                ) as span:
                    # Placeholder for aggregating streaming response
                    llmresponse = ""

                    # Loop through streaming events capturing relevant details
                    for chunk in wrapped(*args, **kwargs):
                        # Collect message IDs and aggregated response from events
                        if len(chunk.choices) > 0:
                            # pylint: disable=line-too-long
                            if hasattr(chunk.choices[0], "delta") and hasattr(
                                chunk.choices[0].delta, "content"
                            ):
                                content = chunk.choices[0].delta.content
                                if content:
                                    llmresponse += content
                        yield chunk
                        response_id = chunk.id
                        model = "azure_" + chunk.model

                    # Handling exception ensure observability without disrupting operation
                    try:
                        # Format 'messages' into a single string
                        message_prompt = kwargs.get("messages", "")
                        formatted_messages = []
                        for message in message_prompt:
                            role = message["role"]
                            content = message["content"]

                            if isinstance(content, list):
                                content_str = ", ".join(
                                    # pylint: disable=line-too-long
                                    (
                                        f'{item["type"]}: {item["text"] if "text" in item else item["image_url"]}'
                                        if "type" in item
                                        else f'text: {item["text"]}'
                                    )
                                    for item in content
                                )
                                formatted_messages.append(f"{role}: {content_str}")
                            else:
                                formatted_messages.append(f"{role}: {content}")
                        prompt = "\n".join(formatted_messages)

                        # Calculate tokens using input prompt and aggregated response
                        prompt_tokens = openai_tokens(
                            prompt, kwargs.get("model", "gpt-3.5-turbo")
                        )
                        completion_tokens = openai_tokens(
                            llmresponse, kwargs.get("model", "gpt-3.5-turbo")
                        )

                        # Calculate cost of the operation
                        cost = get_chat_model_cost(
                            model, pricing_info, prompt_tokens, completion_tokens
                        )

                        # Set Span attributes
                        span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_SYSTEM,
                            SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_TYPE,
                            SemanticConvetion.GEN_AI_TYPE_CHAT,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_ENDPOINT, gen_ai_endpoint
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_RESPONSE_ID, response_id
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_ENVIRONMENT, environment
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_APPLICATION_NAME, application_name
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_MODEL, model
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_USER,
                            kwargs.get("user", ""),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_TOOL_CHOICE,
                            kwargs.get("tool_choice", ""),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
                            kwargs.get("temperature", 1.0),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_PRESENCE_PENALTY,
                            kwargs.get("presence_penalty", 0.0),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_FREQUENCY_PENALTY,
                            kwargs.get("frequency_penalty", 0.0),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_SEED,
                            kwargs.get("seed", ""),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_IS_STREAM, True
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS, prompt_tokens
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
                            completion_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
                            prompt_tokens + completion_tokens,
                        )
                        span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)
                        if trace_content:
                            span.set_attributes(
                                attributes={
                                    SemanticConvetion.GEN_AI_CONTENT_PROMPT: prompt,
                                },
                            )
                            span.set_attributes(
                                attributes={
                                    SemanticConvetion.GEN_AI_CONTENT_COMPLETION: llmresponse,
                                },
                            )

                        span.set_status(Status(StatusCode.OK))

                        if disable_metrics is False:
                            attributes = {
                                TELEMETRY_SDK_NAME: "openlit",
                                SemanticConvetion.GEN_AI_APPLICATION_NAME: application_name,
                                SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                                SemanticConvetion.GEN_AI_ENVIRONMENT: environment,
                                SemanticConvetion.GEN_AI_TYPE: SemanticConvetion.GEN_AI_TYPE_CHAT,
                                SemanticConvetion.GEN_AI_REQUEST_MODEL: model,
                            }

                            metrics["genai_requests"].add(1, attributes)
                            metrics["genai_total_tokens"].add(
                                prompt_tokens + completion_tokens, attributes
                            )
                            metrics["genai_completion_tokens"].add(
                                completion_tokens, attributes
                            )
                            metrics["genai_prompt_tokens"].add(
                                prompt_tokens, attributes
                            )
                            metrics["genai_cost"].record(cost, attributes)

                    except Exception as e:
                        handle_exception(span, e)
                        logger.error("Error in trace creation: %s", e)

            return stream_generator()

        # Handling for non-streaming responses
        else:
            # pylint: disable=line-too-long
            with tracer.start_as_current_span(
                gen_ai_endpoint, kind=SpanKind.CLIENT
            ) as span:
                response = wrapped(*args, **kwargs)

                try:
                    # Find base model from response
                    model = "azure_" + response.model

                    # Format 'messages' into a single string
                    message_prompt = kwargs.get("messages", "")
                    formatted_messages = []
                    for message in message_prompt:
                        role = message["role"]
                        content = message["content"]

                        if isinstance(content, list):
                            content_str = ", ".join(
                                # pylint: disable=line-too-long
                                (
                                    f'{item["type"]}: {item["text"] if "text" in item else item["image_url"]}'
                                    if "type" in item
                                    else f'text: {item["text"]}'
                                )
                                for item in content
                            )
                            formatted_messages.append(f"{role}: {content_str}")
                        else:
                            formatted_messages.append(f"{role}: {content}")
                    prompt = "\n".join(formatted_messages)

                    # Set base span attribues
                    span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_SYSTEM,
                        SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_TYPE,
                        SemanticConvetion.GEN_AI_TYPE_CHAT,
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_ENDPOINT, gen_ai_endpoint
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_RESPONSE_ID, response.id
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_ENVIRONMENT, environment
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_APPLICATION_NAME, application_name
                    )
                    span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL, model)
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_USER, kwargs.get("user", "")
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_TOOL_CHOICE,
                        kwargs.get("tool_choice", ""),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
                        kwargs.get("temperature", 1.0),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_PRESENCE_PENALTY,
                        kwargs.get("presence_penalty", 0.0),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_FREQUENCY_PENALTY,
                        kwargs.get("frequency_penalty", 0.0),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_SEED, kwargs.get("seed", "")
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_IS_STREAM, False
                    )
                    if trace_content:
                        span.set_attributes(
                            attributes={
                                SemanticConvetion.GEN_AI_CONTENT_PROMPT: prompt,
                            },
                        )
                    if trace_content:
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_CONTENT_PROMPT, prompt
                        )

                    # Set span attributes when tools is not passed to the function call
                    if "tools" not in kwargs:
                        # Calculate cost of the operation
                        cost = get_chat_model_cost(
                            model,
                            pricing_info,
                            response.usage.prompt_tokens,
                            response.usage.completion_tokens,
                        )

                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
                            response.usage.prompt_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
                            response.usage.completion_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
                            response.usage.total_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_RESPONSE_FINISH_REASON,
                            [response.choices[0].finish_reason],
                        )
                        span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)

                        # Set span attributes for when n = 1 (default)
                        if "n" not in kwargs or kwargs["n"] == 1:
                            if trace_content:
                                span.set_attributes(
                                    attributes={
                                        SemanticConvetion.GEN_AI_CONTENT_COMPLETION: response.choices[
                                            0
                                        ].message.content,
                                    },
                                )

                        # Set span attributes for when n > 0
                        else:
                            i = 0
                            while i < kwargs["n"] and trace_content is True:
                                attribute_name = f"gen_ai.content.completion.{i}"
                                span.set_attributes(
                                    attributes={
                                        SemanticConvetion.GEN_AI_CONTENT_COMPLETION: response.choices[
                                            i
                                        ].message.content,
                                    },
                                )
                                i += 1

                            # Return original response
                            return response

                    # Set span attributes when tools is passed to the function call
                    elif "tools" in kwargs:
                        # Calculate cost of the operation
                        cost = get_chat_model_cost(
                            model,
                            pricing_info,
                            response.usage.prompt_tokens,
                            response.usage.completion_tokens,
                        )

                        span.set_attributes(
                            attributes={
                                SemanticConvetion.GEN_AI_CONTENT_COMPLETION: "Function called with tools",
                            },
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
                            response.usage.prompt_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
                            response.usage.completion_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
                            response.usage.total_tokens,
                        )
                        span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)

                    span.set_status(Status(StatusCode.OK))

                    if disable_metrics is False:
                        attributes = {
                            TELEMETRY_SDK_NAME: "openlit",
                            SemanticConvetion.GEN_AI_APPLICATION_NAME: application_name,
                            SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                            SemanticConvetion.GEN_AI_ENVIRONMENT: environment,
                            SemanticConvetion.GEN_AI_TYPE: SemanticConvetion.GEN_AI_TYPE_CHAT,
                            SemanticConvetion.GEN_AI_REQUEST_MODEL: model,
                        }

                        metrics["genai_requests"].add(1, attributes)
                        metrics["genai_total_tokens"].add(
                            response.usage.total_tokens, attributes
                        )
                        metrics["genai_completion_tokens"].add(
                            response.usage.completion_tokens, attributes
                        )
                        metrics["genai_prompt_tokens"].add(
                            response.usage.prompt_tokens, attributes
                        )
                        metrics["genai_cost"].record(cost, attributes)

                    # Return original response
                    return response

                except Exception as e:
                    handle_exception(span, e)
                    logger.error("Error in trace creation: %s", e)

                    # Return original response
                    return response

    return wrapper


def azure_completions(
    gen_ai_endpoint,
    version,
    environment,
    application_name,
    tracer,
    pricing_info,
    trace_content,
    metrics,
    disable_metrics,
):
    """
    Generates a telemetry wrapper for completions to collect metrics.

    Args:
        gen_ai_endpoint: Endpoint identifier for logging and tracing.
        version: Version of the monitoring package.
        environment: Deployment environment (e.g., production, staging).
        application_name: Name of the application using the OpenAI API.
        tracer: OpenTelemetry tracer for creating spans.
        pricing_info: Information used for calculating the cost of OpenAI usage.
        trace_content: Flag indicating whether to trace the actual content.

    Returns:
        A function that wraps the chat completions method to add telemetry.
    """

    def wrapper(wrapped, instance, args, kwargs):
        """
        Wraps the 'completions' API call to add telemetry.

        This collects metrics such as execution time, cost, and token usage, and handles errors
        gracefully, adding details to the trace for observability.

        Args:
            wrapped: The original 'completions' method to be wrapped.
            instance: The instance of the class where the original method is defined.
            args: Positional arguments for the 'completions' method.
            kwargs: Keyword arguments for the 'completions' method.

        Returns:
            The response from the original 'chat.completions' method.
        """

        # Check if streaming is enabled for the API call
        streaming = kwargs.get("stream", False)

        # pylint: disable=no-else-return
        if streaming:
            # Special handling for streaming response to accommodate the nature of data flow
            def stream_generator():
                # pylint: disable=line-too-long
                with tracer.start_as_current_span(
                    gen_ai_endpoint, kind=SpanKind.CLIENT
                ) as span:
                    # Placeholder for aggregating streaming response
                    llmresponse = ""

                    # Loop through streaming events capturing relevant details
                    for chunk in wrapped(*args, **kwargs):
                        # Collect message IDs and aggregated response from events
                        if len(chunk.choices) > 0:
                            if hasattr(chunk.choices[0], "text"):
                                content = chunk.choices[0].text
                                if content:
                                    llmresponse += content
                        yield chunk
                        response_id = chunk.id
                        model = "azure_" + chunk.model

                    # Handling exception ensure observability without disrupting operation
                    try:
                        prompt = kwargs.get("prompt", "")

                        # Calculate tokens using input prompt and aggregated response
                        prompt_tokens = openai_tokens(prompt, "gpt-3.5-turbo")
                        completion_tokens = openai_tokens(llmresponse, "gpt-3.5-turbo")

                        # Calculate cost of the operation
                        cost = get_chat_model_cost(
                            model, pricing_info, prompt_tokens, completion_tokens
                        )

                        # Set Span attributes
                        span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_SYSTEM,
                            SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_TYPE,
                            SemanticConvetion.GEN_AI_TYPE_CHAT,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_ENDPOINT, gen_ai_endpoint
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_RESPONSE_ID, response_id
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_ENVIRONMENT, environment
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_APPLICATION_NAME, application_name
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_MODEL, model
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_USER,
                            kwargs.get("user", ""),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_TOOL_CHOICE,
                            kwargs.get("tool_choice", ""),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
                            kwargs.get("temperature", 1.0),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_PRESENCE_PENALTY,
                            kwargs.get("presence_penalty", 0.0),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_FREQUENCY_PENALTY,
                            kwargs.get("frequency_penalty", 0.0),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_SEED,
                            kwargs.get("seed", ""),
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_REQUEST_IS_STREAM, True
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS, prompt_tokens
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
                            completion_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
                            prompt_tokens + completion_tokens,
                        )
                        span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)
                        if trace_content:
                            span.set_attributes(
                                attributes={
                                    SemanticConvetion.GEN_AI_CONTENT_PROMPT: prompt,
                                },
                            )
                            span.set_attributes(
                                attributes={
                                    SemanticConvetion.GEN_AI_CONTENT_COMPLETION: llmresponse,
                                },
                            )

                        span.set_status(Status(StatusCode.OK))

                        if disable_metrics is False:
                            attributes = {
                                TELEMETRY_SDK_NAME: "openlit",
                                SemanticConvetion.GEN_AI_APPLICATION_NAME: application_name,
                                SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                                SemanticConvetion.GEN_AI_ENVIRONMENT: environment,
                                SemanticConvetion.GEN_AI_TYPE: SemanticConvetion.GEN_AI_TYPE_CHAT,
                                SemanticConvetion.GEN_AI_REQUEST_MODEL: model,
                            }

                            metrics["genai_requests"].add(1, attributes)
                            metrics["genai_total_tokens"].add(
                                prompt_tokens + completion_tokens, attributes
                            )
                            metrics["genai_completion_tokens"].add(
                                completion_tokens, attributes
                            )
                            metrics["genai_prompt_tokens"].add(
                                prompt_tokens, attributes
                            )
                            metrics["genai_cost"].record(cost, attributes)

                    except Exception as e:
                        handle_exception(span, e)
                        logger.error("Error in trace creation: %s", e)

            return stream_generator()

        # Handling for non-streaming responses
        else:
            # pylint: disable=line-too-long
            with tracer.start_as_current_span(
                gen_ai_endpoint, kind=SpanKind.CLIENT
            ) as span:
                response = wrapped(*args, **kwargs)

                try:
                    # Find base model from response
                    model = "azure_" + response.model

                    # Set base span attribues
                    span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_SYSTEM,
                        SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_TYPE,
                        SemanticConvetion.GEN_AI_TYPE_CHAT,
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_ENDPOINT, gen_ai_endpoint
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_RESPONSE_ID, response.id
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_ENVIRONMENT, environment
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_APPLICATION_NAME, application_name
                    )
                    span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL, model)
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_USER, kwargs.get("user", "")
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_TOOL_CHOICE,
                        kwargs.get("tool_choice", ""),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
                        kwargs.get("temperature", 1.0),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_PRESENCE_PENALTY,
                        kwargs.get("presence_penalty", 0.0),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_FREQUENCY_PENALTY,
                        kwargs.get("frequency_penalty", 0.0),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_SEED, kwargs.get("seed", "")
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_IS_STREAM, False
                    )
                    if trace_content:
                        span.set_attributes(
                            attributes={
                                SemanticConvetion.GEN_AI_CONTENT_PROMPT: kwargs.get(
                                    "prompt", ""
                                ),
                            },
                        )

                    # Set span attributes when tools is not passed to the function call
                    if "tools" not in kwargs:
                        # Calculate cost of the operation
                        cost = get_chat_model_cost(
                            model,
                            pricing_info,
                            response.usage.prompt_tokens,
                            response.usage.completion_tokens,
                        )

                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
                            response.usage.prompt_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
                            response.usage.completion_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
                            response.usage.total_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_RESPONSE_FINISH_REASON,
                            [response.choices[0].finish_reason],
                        )
                        span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)

                        # Set span attributes for when n = 1 (default)
                        if "n" not in kwargs or kwargs["n"] == 1:
                            if trace_content:
                                span.set_attributes(
                                    attributes={
                                        SemanticConvetion.GEN_AI_CONTENT_COMPLETION: response.choices[
                                            0
                                        ].text,
                                    },
                                )

                        # Set span attributes for when n > 0
                        else:
                            i = 0
                            while i < kwargs["n"] and trace_content is True:
                                attribute_name = f"gen_ai.content.completion.{i}"
                                span.set_attributes(
                                    attributes={
                                        SemanticConvetion.GEN_AI_CONTENT_COMPLETION: response.choices[
                                            i
                                        ].text,
                                    },
                                )
                                i += 1
                            return response

                    # Set span attributes when tools is passed to the function call
                    elif "tools" in kwargs:
                        # Calculate cost of the operation
                        cost = get_chat_model_cost(
                            model,
                            pricing_info,
                            response.usage.prompt_tokens,
                            response.usage.completion_tokens,
                        )

                        span.set_attribute(
                            SemanticConvetion.GEN_AI_CONTENT_COMPLETION,
                            "Function called with tools",
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
                            response.usage.prompt_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
                            response.usage.completion_tokens,
                        )
                        span.set_attribute(
                            SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
                            response.usage.total_tokens,
                        )
                        span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)

                    span.set_status(Status(StatusCode.OK))

                    if disable_metrics is False:
                        attributes = {
                            TELEMETRY_SDK_NAME: "openlit",
                            SemanticConvetion.GEN_AI_APPLICATION_NAME: application_name,
                            SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                            SemanticConvetion.GEN_AI_ENVIRONMENT: environment,
                            SemanticConvetion.GEN_AI_TYPE: SemanticConvetion.GEN_AI_TYPE_CHAT,
                            SemanticConvetion.GEN_AI_REQUEST_MODEL: model,
                        }

                        metrics["genai_requests"].add(1, attributes)
                        metrics["genai_total_tokens"].add(
                            response.usage.total_tokens, attributes
                        )
                        metrics["genai_completion_tokens"].add(
                            response.usage.completion_tokens, attributes
                        )
                        metrics["genai_prompt_tokens"].add(
                            response.usage.prompt_tokens, attributes
                        )
                        metrics["genai_cost"].record(cost, attributes)

                    # Return original response
                    return response

                except Exception as e:
                    handle_exception(span, e)
                    logger.error("Error in trace creation: %s", e)

                    # Return original response
                    return response

    return wrapper


def azure_embedding(
    gen_ai_endpoint,
    version,
    environment,
    application_name,
    tracer,
    pricing_info,
    trace_content,
    metrics,
    disable_metrics,
):
    """
    Generates a telemetry wrapper for embeddings to collect metrics.

    Args:
        gen_ai_endpoint: Endpoint identifier for logging and tracing.
        version: Version of the monitoring package.
        environment: Deployment environment (e.g., production, staging).
        application_name: Name of the application using the OpenAI API.
        tracer: OpenTelemetry tracer for creating spans.
        pricing_info: Information used for calculating the cost of OpenAI usage.
        trace_content: Flag indicating whether to trace the actual content.

    Returns:
        A function that wraps the embeddings method to add telemetry.
    """

    def wrapper(wrapped, instance, args, kwargs):
        """
        Wraps the 'embeddings' API call to add telemetry.

        This collects metrics such as execution time, cost, and token usage, and handles errors
        gracefully, adding details to the trace for observability.

        Args:
            wrapped: The original 'embeddings' method to be wrapped.
            instance: The instance of the class where the original method is defined.
            args: Positional arguments for the 'embeddings' method.
            kwargs: Keyword arguments for the 'embeddings' method.

        Returns:
            The response from the original 'embeddings' method.
        """

        with tracer.start_as_current_span(
            gen_ai_endpoint, kind=SpanKind.CLIENT
        ) as span:
            response = wrapped(*args, **kwargs)

            try:
                # Calculate cost of the operation
                cost = get_embed_model_cost(
                    "azure_" + response.model,
                    pricing_info,
                    response.usage.prompt_tokens,
                )

                # Set Span attributes
                span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
                span.set_attribute(
                    SemanticConvetion.GEN_AI_SYSTEM,
                    SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                )
                span.set_attribute(
                    SemanticConvetion.GEN_AI_TYPE,
                    SemanticConvetion.GEN_AI_TYPE_EMBEDDING,
                )
                span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT, gen_ai_endpoint)
                span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT, environment)
                span.set_attribute(
                    SemanticConvetion.GEN_AI_APPLICATION_NAME, application_name
                )
                span.set_attribute(
                    SemanticConvetion.GEN_AI_REQUEST_MODEL, "azure_" + response.model
                )
                span.set_attribute(
                    SemanticConvetion.GEN_AI_REQUEST_EMBEDDING_FORMAT,
                    kwargs.get("encoding_format", "float"),
                )
                span.set_attribute(
                    SemanticConvetion.GEN_AI_REQUEST_EMBEDDING_DIMENSION,
                    kwargs.get("dimensions", ""),
                )
                span.set_attribute(
                    SemanticConvetion.GEN_AI_REQUEST_USER, kwargs.get("user", "")
                )
                span.set_attribute(
                    SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
                    response.usage.prompt_tokens,
                )
                span.set_attribute(
                    SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
                    response.usage.total_tokens,
                )
                span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST, cost)
                if trace_content:
                    span.set_attributes(
                        attributes={
                            SemanticConvetion.GEN_AI_CONTENT_PROMPT: kwargs.get(
                                "input", ""
                            ),
                        },
                    )

                span.set_status(Status(StatusCode.OK))

                if disable_metrics is False:
                    attributes = {
                        TELEMETRY_SDK_NAME: "openlit",
                        SemanticConvetion.GEN_AI_APPLICATION_NAME: application_name,
                        SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                        SemanticConvetion.GEN_AI_ENVIRONMENT: environment,
                        SemanticConvetion.GEN_AI_TYPE: SemanticConvetion.GEN_AI_TYPE_EMBEDDING,
                        SemanticConvetion.GEN_AI_REQUEST_MODEL: "azure_"
                        + response.model,
                    }

                    metrics["genai_requests"].add(1, attributes)
                    metrics["genai_total_tokens"].add(
                        response.usage.total_tokens, attributes
                    )
                    metrics["genai_prompt_tokens"].add(
                        response.usage.prompt_tokens, attributes
                    )
                    metrics["genai_cost"].record(cost, attributes)

                # Return original response
                return response

            except Exception as e:
                handle_exception(span, e)
                logger.error("Error in trace creation: %s", e)

                # Return original response
                return response

    return wrapper


def azure_image_generate(
    gen_ai_endpoint,
    version,
    environment,
    application_name,
    tracer,
    pricing_info,
    trace_content,
    metrics,
    disable_metrics,
):
    """
    Generates a telemetry wrapper for image generation to collect metrics.

    Args:
        gen_ai_endpoint: Endpoint identifier for logging and tracing.
        version: Version of the monitoring package.
        environment: Deployment environment (e.g., production, staging).
        application_name: Name of the application using the OpenAI API.
        tracer: OpenTelemetry tracer for creating spans.
        pricing_info: Information used for calculating the cost of OpenAI image generation.
        trace_content: Flag indicating whether to trace the input prompt and generated images.

    Returns:
        A function that wraps the image generation method to add telemetry.
    """

    def wrapper(wrapped, instance, args, kwargs):
        """
        Wraps the 'images.generate' API call to add telemetry.

        This collects metrics such as execution time, cost, and handles errors
        gracefully, adding details to the trace for observability.

        Args:
            wrapped: The original 'images.generate' method to be wrapped.
            instance: The instance of the class where the original method is defined.
            args: Positional arguments for the 'images.generate' method.
            kwargs: Keyword arguments for the 'images.generate' method.

        Returns:
            The response from the original 'images.generate' method.
        """

        with tracer.start_as_current_span(
            gen_ai_endpoint, kind=SpanKind.CLIENT
        ) as span:
            response = wrapped(*args, **kwargs)
            images_count = 0

            try:
                # Find Image format
                if (
                    "response_format" in kwargs
                    and kwargs["response_format"] == "b64_json"
                ):
                    image = "b64_json"
                else:
                    image = "url"

                # Calculate cost of the operation
                cost = get_image_model_cost(
                    "azure_" + kwargs.get("model", "dall-e-3"),
                    pricing_info,
                    kwargs.get("size", "1024x1024"),
                    kwargs.get("quality", "standard"),
                )

                for items in response.data:
                    # Set Span attributes
                    span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_SYSTEM,
                        SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_TYPE,
                        SemanticConvetion.GEN_AI_TYPE_IMAGE,
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_ENDPOINT, gen_ai_endpoint
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_RESPONSE_ID, response.created
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_ENVIRONMENT, environment
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_APPLICATION_NAME, application_name
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_MODEL,
                        "azure_" + kwargs.get("model", "dall-e-3"),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_IMAGE_SIZE,
                        kwargs.get("size", "1024x1024"),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_IMAGE_QUALITY,
                        kwargs.get("quality", "standard"),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_IMAGE_STYLE,
                        kwargs.get("style", "vivid"),
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_CONTENT_REVISED_PROMPT,
                        items.revised_prompt if response.revised_prompt else "",
                    )
                    span.set_attribute(
                        SemanticConvetion.GEN_AI_REQUEST_USER, kwargs.get("user", "")
                    )
                    if trace_content:
                        span.set_attributes(
                            attributes={
                                SemanticConvetion.GEN_AI_CONTENT_PROMPT: kwargs.get(
                                    "prompt", ""
                                ),
                            },
                        )
                        attribute_name = f"gen_ai.response.image.{images_count}"
                        span.set_attributes(
                            attributes={
                                SemanticConvetion.GEN_AI_CONTENT_COMPLETION: getattr(
                                    items, image
                                ),
                            },
                        )

                    images_count += 1

                span.set_attribute(
                    SemanticConvetion.GEN_AI_USAGE_COST, len(response.data) * cost
                )
                span.set_status(Status(StatusCode.OK))

                if disable_metrics is False:
                    attributes = {
                        TELEMETRY_SDK_NAME: "openlit",
                        SemanticConvetion.GEN_AI_APPLICATION_NAME: application_name,
                        SemanticConvetion.GEN_AI_SYSTEM: SemanticConvetion.GEN_AI_SYSTEM_AZURE_OPENAI,
                        SemanticConvetion.GEN_AI_ENVIRONMENT: environment,
                        SemanticConvetion.GEN_AI_TYPE: SemanticConvetion.GEN_AI_TYPE_IMAGE,
                        SemanticConvetion.GEN_AI_REQUEST_MODEL: "azure_"
                        + kwargs.get("model", "dall-e-3"),
                    }

                    metrics["genai_requests"].add(1, attributes)
                    metrics["genai_cost"].record(cost, attributes)

                # Return original response
                return response

            except Exception as e:
                handle_exception(span, e)
                logger.error("Error in trace creation: %s", e)

                # Return original response
                return response

    return wrapper
