lmnr 0.6.18__py3-none-any.whl → 0.6.19__py3-none-any.whl

lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/__init__.py

@@ -11,6 +11,7 @@ from google.genai import types
 from .config import (
     Config,
 )
+from .schema_utils import SchemaJSONEncoder, process_schema
 from .utils import (
     dont_throw,
     get_content,
@@ -24,8 +25,9 @@ from opentelemetry.trace import Tracer
 from wrapt import wrap_function_wrapper
 
 from opentelemetry import context as context_api
-from opentelemetry.trace import get_tracer, SpanKind, Span
+from opentelemetry.trace import get_tracer, SpanKind, Span, Status, StatusCode
 from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
@@ -78,7 +80,7 @@ WRAPPED_METHODS = [
 
 def should_send_prompts():
     return (
-        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
+        os.getenv("LAMINAR_TRACE_CONTENT") or "true"
     ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
 
 
@@ -128,6 +130,29 @@ def _set_request_attributes(span, args, kwargs):
         span, gen_ai_attributes.GEN_AI_REQUEST_SEED, config_dict.get("seed")
     )
 
+    if schema := config_dict.get("response_schema"):
+        try:
+            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(process_schema(schema), cls=SchemaJSONEncoder),
+            )
+        except Exception:
+            pass
+    elif json_schema := config_dict.get("response_json_schema"):
+        try:
+            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(json_schema),
+            )
+        except Exception:
+            pass
+
     tools: list[types.FunctionDeclaration] = []
     arg_tools = config_dict.get("tools", kwargs.get("tools"))
     if arg_tools:
@@ -454,16 +479,20 @@ def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
     if span.is_recording():
         _set_request_attributes(span, args, kwargs)
 
-    if to_wrap.get("is_streaming"):
-        return _build_from_streaming_response(span, wrapped(*args, **kwargs))
-    else:
+    try:
         response = wrapped(*args, **kwargs)
-
-    if span.is_recording():
-        _set_response_attributes(span, response)
-
-    span.end()
-    return response
+        if to_wrap.get("is_streaming"):
+            return _build_from_streaming_response(span, response)
+        if span.is_recording():
+            _set_response_attributes(span, response)
+        span.end()
+        return response
+    except Exception as e:
+        span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+        span.record_exception(e)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        span.end()
+        raise e
 
 
 @with_tracer_wrapper
@@ -485,16 +514,22 @@ async def _awrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
     if span.is_recording():
         _set_request_attributes(span, args, kwargs)
 
-    if to_wrap.get("is_streaming"):
-        return _abuild_from_streaming_response(span, await wrapped(*args, **kwargs))
-    else:
+    try:
         response = await wrapped(*args, **kwargs)
-
-    if span.is_recording():
-        _set_response_attributes(span, response)
-
-    span.end()
-    return response
+        if to_wrap.get("is_streaming"):
+            return _abuild_from_streaming_response(span, response)
+        else:
+            if span.is_recording():
+                _set_response_attributes(span, response)
+
+            span.end()
+            return response
+    except Exception as e:
+        span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+        span.record_exception(e)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        span.end()
+        raise e
 
 
 class GoogleGenAiSdkInstrumentor(BaseInstrumentor):

lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/schema_utils.py

@@ -0,0 +1,23 @@
+from typing import Any
+from google.genai._api_client import BaseApiClient
+from google.genai._transformers import t_schema
+from google.genai.types import JSONSchemaType
+
+import json
+
+DUMMY_CLIENT = BaseApiClient(api_key="dummy")
+
+
+def process_schema(schema: Any) -> dict[str, Any]:
+    # The only thing we need from the client is the t_schema function
+    json_schema = t_schema(DUMMY_CLIENT, schema).json_schema.model_dump(
+        exclude_unset=True, exclude_none=True
+    )
+    return json_schema
+
+
+class SchemaJSONEncoder(json.JSONEncoder):
+    def default(self, o: Any) -> Any:
+        if isinstance(o, JSONSchemaType):
+            return o.value
+        return super().default(o)

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

@@ -0,0 +1,61 @@
+"""
+Initially copied over from openllmetry, commit
+b3a18c9f7e6ff2368c8fb0bc35fd9123f11121c4
+"""
+
+from typing import Callable, Collection, Optional
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from .shared.config import Config
+from .utils import is_openai_v1
+from typing_extensions import Coroutine
+
+_instruments = ("openai >= 0.27.0",)
+
+
+class OpenAIInstrumentor(BaseInstrumentor):
+    """An instrumentor for OpenAI's client library."""
+
+    def __init__(
+        self,
+        enrich_assistant: bool = False,
+        enrich_token_usage: bool = False,
+        exception_logger=None,
+        get_common_metrics_attributes: Callable[[], dict] = lambda: {},
+        upload_base64_image: Optional[
+            Callable[[str, str, str, str], Coroutine[None, None, str]]
+        ] = lambda *args: "",
+        enable_trace_context_propagation: bool = True,
+        use_legacy_attributes: bool = True,
+    ):
+        super().__init__()
+        Config.enrich_assistant = enrich_assistant
+        Config.enrich_token_usage = enrich_token_usage
+        Config.exception_logger = exception_logger
+        Config.get_common_metrics_attributes = get_common_metrics_attributes
+        Config.upload_base64_image = upload_base64_image
+        Config.enable_trace_context_propagation = enable_trace_context_propagation
+        Config.use_legacy_attributes = use_legacy_attributes
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        if is_openai_v1():
+            from .v1 import OpenAIV1Instrumentor
+
+            OpenAIV1Instrumentor().instrument(**kwargs)
+        else:
+            from .v0 import OpenAIV0Instrumentor
+
+            OpenAIV0Instrumentor().instrument(**kwargs)
+
+    def _uninstrument(self, **kwargs):
+        if is_openai_v1():
+            from .v1 import OpenAIV1Instrumentor
+
+            OpenAIV1Instrumentor().uninstrument(**kwargs)
+        else:
+            from .v0 import OpenAIV0Instrumentor
+
+            OpenAIV0Instrumentor().uninstrument(**kwargs)

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

@@ -0,0 +1,442 @@
+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
+    )
+    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")
+    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),
+    )
+    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
+
+        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