lmnr 0.4.53.dev0__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/__init__.py

@@ -0,0 +1,472 @@
+import json
+import logging
+import types
+from importlib.metadata import version
+
+from ..shared.config import Config
+from ..utils import (
+    dont_throw,
+    is_openai_v1,
+    should_record_stream_token_usage,
+)
+from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
+    GEN_AI_RESPONSE_ID,
+)
+from opentelemetry.semconv_ai import SpanAttributes
+from opentelemetry.trace.propagation import set_span_in_context
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+import openai
+import pydantic
+
+OPENAI_LLM_USAGE_TOKEN_TYPES = ["prompt_tokens", "completion_tokens"]
+PROMPT_FILTER_KEY = "prompt_filter_results"
+PROMPT_ERROR = "prompt_error"
+
+_PYDANTIC_VERSION = version("pydantic")
+
+# tiktoken encodings map for different model, key is model_name, value is tiktoken encoding
+tiktoken_encodings = {}
+
+logger = logging.getLogger(__name__)
+
+
+def _set_span_attribute(span, name, value):
+    if value is None or value == "":
+        return
+
+    if hasattr(openai, "NOT_GIVEN") and value == openai.NOT_GIVEN:
+        return
+
+    span.set_attribute(name, value)
+
+
+def _set_client_attributes(span, instance):
+    if not span.is_recording():
+        return
+
+    if not is_openai_v1():
+        return
+
+    client = instance._client  # pylint: disable=protected-access
+    if isinstance(client, (openai.AsyncOpenAI, openai.OpenAI)):
+        _set_span_attribute(
+            span, SpanAttributes.LLM_OPENAI_API_BASE, str(client.base_url)
+        )
+    if isinstance(client, (openai.AsyncAzureOpenAI, openai.AzureOpenAI)):
+        _set_span_attribute(
+            span, SpanAttributes.LLM_OPENAI_API_VERSION, client._api_version
+        )  # pylint: disable=protected-access
+
+
+def _set_api_attributes(span):
+    if not span.is_recording():
+        return
+
+    if is_openai_v1():
+        return
+
+    base_url = openai.base_url if hasattr(openai, "base_url") else openai.api_base
+
+    _set_span_attribute(span, SpanAttributes.LLM_OPENAI_API_BASE, base_url)
+    _set_span_attribute(span, SpanAttributes.LLM_OPENAI_API_TYPE, openai.api_type)
+    _set_span_attribute(span, SpanAttributes.LLM_OPENAI_API_VERSION, openai.api_version)
+
+    return
+
+
+def _set_functions_attributes(span, functions):
+    if not functions:
+        return
+
+    for i, function in enumerate(functions):
+        prefix = f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{i}"
+        _set_span_attribute(span, f"{prefix}.name", function.get("name"))
+        _set_span_attribute(span, f"{prefix}.description", function.get("description"))
+        _set_span_attribute(
+            span, f"{prefix}.parameters", json.dumps(function.get("parameters"))
+        )
+
+
+def set_tools_attributes(span, tools):
+    if not tools:
+        return
+
+    for i, tool in enumerate(tools):
+        function = tool.get("function")
+        if not function:
+            continue
+
+        prefix = f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{i}"
+        _set_span_attribute(span, f"{prefix}.name", function.get("name"))
+        _set_span_attribute(span, f"{prefix}.description", function.get("description"))
+        _set_span_attribute(
+            span, f"{prefix}.parameters", json.dumps(function.get("parameters"))
+        )
+
+
+def _set_request_attributes(span, kwargs, instance=None):
+    if not span.is_recording():
+        return
+
+    _set_api_attributes(span)
+
+    base_url = _get_openai_base_url(instance) if instance else ""
+    vendor = _get_vendor_from_url(base_url)
+    _set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
+
+    model = kwargs.get("model")
+    if vendor == "AWS" and model and "." in model:
+        model = _cross_region_check(model)
+    elif vendor == "OpenRouter":
+        model = _extract_model_name_from_provider_format(model)
+
+    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model)
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_tokens")
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, kwargs.get("temperature")
+    )
+    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, kwargs.get("top_p"))
+    _set_span_attribute(
+        span, SpanAttributes.LLM_FREQUENCY_PENALTY, kwargs.get("frequency_penalty")
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_PRESENCE_PENALTY, kwargs.get("presence_penalty")
+    )
+    _set_span_attribute(span, SpanAttributes.LLM_USER, kwargs.get("user"))
+    _set_span_attribute(span, SpanAttributes.LLM_HEADERS, str(kwargs.get("headers")))
+    # The new OpenAI SDK removed the `headers` and create new field called `extra_headers`
+    if kwargs.get("extra_headers") is not None:
+        _set_span_attribute(
+            span, SpanAttributes.LLM_HEADERS, str(kwargs.get("extra_headers"))
+        )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_IS_STREAMING, kwargs.get("stream") or False
+    )
+    _set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_REASONING_EFFORT,
+        kwargs.get("reasoning_effort"),
+    )
+    _set_span_attribute(
+        span,
+        "openai.request.service_tier",
+        kwargs.get("service_tier"),
+    )
+    if response_format := kwargs.get("response_format"):
+        # backward-compatible check for
+        # openai.types.shared_params.response_format_json_schema.ResponseFormatJSONSchema
+        if (
+            isinstance(response_format, dict)
+            and response_format.get("type") == "json_schema"
+            and response_format.get("json_schema")
+        ):
+            schema = dict(response_format.get("json_schema")).get("schema")
+            if schema:
+                _set_span_attribute(
+                    span,
+                    # TODO: change to SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA
+                    # when we upgrade to opentelemetry-semantic-conventions-ai>=0.4.10
+                    "gen_ai.request.structured_output_schema",
+                    json.dumps(schema),
+                )
+        else:
+            try:
+                from openai.lib._parsing._completions import (
+                    type_to_response_format_param,
+                )
+
+                response_format_param = type_to_response_format_param(response_format)
+                if response_format_param.get("type") == "json_schema":
+                    schema = response_format_param.get("json_schema").get("schema")
+                    if schema:
+                        _set_span_attribute(
+                            span,
+                            # TODO: change to SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA
+                            # when we upgrade to opentelemetry-semantic-conventions-ai>=0.4.10
+                            "gen_ai.request.structured_output_schema",
+                            json.dumps(schema),
+                        )
+            except (ImportError, TypeError, AttributeError):
+                # if we fail to import from openai.lib._parsing._completions,
+                # we fallback to the pydantic-based approach
+                if isinstance(response_format, pydantic.BaseModel) or (
+                    hasattr(response_format, "model_json_schema")
+                    and callable(response_format.model_json_schema)
+                ):
+                    _set_span_attribute(
+                        span,
+                        # TODO: change to SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA
+                        # when we upgrade to opentelemetry-semantic-conventions-ai>=0.4.10
+                        "gen_ai.request.structured_output_schema",
+                        json.dumps(response_format.model_json_schema()),
+                    )
+                else:
+                    schema = None
+                    try:
+                        schema = json.dumps(
+                            pydantic.TypeAdapter(response_format).json_schema()
+                        )
+                    except Exception:
+                        try:
+                            schema = json.dumps(response_format)
+                        except Exception:
+                            pass
+
+                    if schema:
+                        _set_span_attribute(
+                            span,
+                            # TODO: change to SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA
+                            # when we upgrade to opentelemetry-semantic-conventions-ai>=0.4.10
+                            "gen_ai.request.structured_output_schema",
+                            schema,
+                        )
+
+
+@dont_throw
+def _set_response_attributes(span, response):
+    if not span.is_recording():
+        return
+
+    if "error" in response:
+        _set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_PROMPTS}.{PROMPT_ERROR}",
+            json.dumps(response.get("error")),
+        )
+        return
+
+    response_model = response.get("model")
+    if response_model:
+        response_model = _extract_model_name_from_provider_format(response_model)
+    _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, response_model)
+    _set_span_attribute(span, GEN_AI_RESPONSE_ID, response.get("id"))
+
+    _set_span_attribute(
+        span,
+        SpanAttributes.LLM_OPENAI_RESPONSE_SYSTEM_FINGERPRINT,
+        response.get("system_fingerprint"),
+    )
+    _log_prompt_filter(span, response)
+    usage = response.get("usage")
+    _set_span_attribute(
+        span,
+        "openai.response.service_tier",
+        response.get("service_tier"),
+    )
+    if not usage:
+        return
+
+    if is_openai_v1() and not isinstance(usage, dict):
+        usage = usage.__dict__
+
+    _set_span_attribute(
+        span, SpanAttributes.LLM_USAGE_TOTAL_TOKENS, usage.get("total_tokens")
+    )
+    _set_span_attribute(
+        span,
+        SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+        usage.get("completion_tokens"),
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_USAGE_PROMPT_TOKENS, usage.get("prompt_tokens")
+    )
+    prompt_tokens_details = dict(usage.get("prompt_tokens_details", {}))
+    _set_span_attribute(
+        span,
+        SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
+        prompt_tokens_details.get("cached_tokens", 0),
+    )
+
+    if completion_token_details := dict(usage.get("completion_tokens_details", {})):
+        reasoning_tokens = completion_token_details.get("reasoning_tokens")
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_REASONING_TOKENS,
+            reasoning_tokens or 0,
+        )
+
+    return
+
+
+def _log_prompt_filter(span, response_dict):
+    if response_dict.get("prompt_filter_results"):
+        _set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_PROMPTS}.{PROMPT_FILTER_KEY}",
+            json.dumps(response_dict.get("prompt_filter_results")),
+        )
+
+
+@dont_throw
+def _set_span_stream_usage(span, prompt_tokens, completion_tokens):
+    if not span.is_recording():
+        return
+
+    if isinstance(completion_tokens, int) and completion_tokens >= 0:
+        _set_span_attribute(
+            span, SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens
+        )
+
+    if isinstance(prompt_tokens, int) and prompt_tokens >= 0:
+        _set_span_attribute(span, SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens)
+
+    if (
+        isinstance(prompt_tokens, int)
+        and isinstance(completion_tokens, int)
+        and completion_tokens + prompt_tokens >= 0
+    ):
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+            completion_tokens + prompt_tokens,
+        )
+
+
+def _get_openai_base_url(instance):
+    if hasattr(instance, "_client"):
+        client = instance._client  # pylint: disable=protected-access
+        if isinstance(client, (openai.AsyncOpenAI, openai.OpenAI)):
+            return str(client.base_url)
+
+    return ""
+
+
+def _get_vendor_from_url(base_url):
+    if not base_url:
+        return "openai"
+
+    if "openai.azure.com" in base_url:
+        return "Azure"
+    elif "amazonaws.com" in base_url or "bedrock" in base_url:
+        return "AWS"
+    elif "googleapis.com" in base_url or "vertex" in base_url:
+        return "Google"
+    elif "openrouter.ai" in base_url:
+        return "OpenRouter"
+
+    return "openai"
+
+
+def _cross_region_check(value):
+    if not value or "." not in value:
+        return value
+
+    prefixes = ["us", "us-gov", "eu", "apac"]
+    if any(value.startswith(prefix + ".") for prefix in prefixes):
+        parts = value.split(".")
+        if len(parts) > 2:
+            return parts[2]
+        else:
+            return value
+    else:
+        vendor, model = value.split(".", 1)
+        return model
+
+
+def _extract_model_name_from_provider_format(model_name):
+    """
+    Extract model name from provider/model format.
+    E.g., 'openai/gpt-4o' -> 'gpt-4o', 'anthropic/claude-3-sonnet' -> 'claude-3-sonnet'
+    """
+    if not model_name:
+        return model_name
+
+    if "/" in model_name:
+        parts = model_name.split("/")
+        return parts[-1]  # Return the last part (actual model name)
+
+    return model_name
+
+
+def is_streaming_response(response):
+    if is_openai_v1():
+        return isinstance(response, openai.Stream) or isinstance(
+            response, openai.AsyncStream
+        )
+
+    return isinstance(response, types.GeneratorType) or isinstance(
+        response, types.AsyncGeneratorType
+    )
+
+
+def model_as_dict(model):
+    if isinstance(model, dict):
+        return model
+    if _PYDANTIC_VERSION < "2.0.0":
+        return model.dict()
+    if hasattr(model, "model_dump"):
+        return model.model_dump()
+    elif hasattr(model, "parse"):  # Raw API response
+        return model_as_dict(model.parse())
+    else:
+        return model
+
+
+def get_token_count_from_string(string: str, model_name: str):
+    if not should_record_stream_token_usage():
+        return None
+
+    import tiktoken
+
+    if tiktoken_encodings.get(model_name) is None:
+        try:
+            encoding = tiktoken.encoding_for_model(model_name)
+        except KeyError as ex:
+            # no such model_name in tiktoken
+            logger.warning(
+                f"Failed to get tiktoken encoding for model_name {model_name}, error: {str(ex)}"
+            )
+            return None
+        except Exception as ex:
+            # Other exceptions in tiktoken
+            logger.warning(
+                f"Failed to get tiktoken encoding for model_name {model_name}, error: {str(ex)}"
+            )
+            return None
+
+        tiktoken_encodings[model_name] = encoding
+    else:
+        encoding = tiktoken_encodings.get(model_name)
+
+    token_count = len(encoding.encode(string))
+    return token_count
+
+
+def _token_type(token_type: str):
+    if token_type == "prompt_tokens":
+        return "input"
+    elif token_type == "completion_tokens":
+        return "output"
+
+    return None
+
+
+def metric_shared_attributes(
+    response_model: str, operation: str, server_address: str, is_streaming: bool = False
+):
+    attributes = Config.get_common_metrics_attributes()
+    vendor = _get_vendor_from_url(server_address)
+
+    return {
+        **attributes,
+        SpanAttributes.LLM_SYSTEM: vendor,
+        SpanAttributes.LLM_RESPONSE_MODEL: response_model,
+        "gen_ai.operation.name": operation,
+        "server.address": server_address,
+        "stream": is_streaming,
+    }
+
+
+def propagate_trace_context(span, kwargs):
+    if is_openai_v1():
+        extra_headers = kwargs.get("extra_headers", {})
+        ctx = set_span_in_context(span)
+        TraceContextTextMapPropagator().inject(extra_headers, context=ctx)
+        kwargs["extra_headers"] = extra_headers
+    else:
+        headers = kwargs.get("headers", {})
+        ctx = set_span_in_context(span)
+        TraceContextTextMapPropagator().inject(headers, context=ctx)
+        kwargs["headers"] = headers