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

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

@@ -0,0 +1,100 @@
+from dataclasses import asdict
+from enum import Enum
+from typing import Union
+
+from opentelemetry._events import Event
+from ..shared.event_models import (
+    ChoiceEvent,
+    MessageEvent,
+)
+from ..utils import (
+    should_emit_events,
+    should_send_prompts,
+)
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+
+from .config import Config
+
+
+class Roles(Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+    TOOL = "tool"
+
+
+VALID_MESSAGE_ROLES = {role.value for role in Roles}
+"""The valid roles for naming the message event."""
+
+EVENT_ATTRIBUTES = {
+    GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.OPENAI.value
+}
+"""The attributes to be used for the event."""
+
+
+def emit_event(event: Union[MessageEvent, ChoiceEvent]) -> None:
+    """
+    Emit an event to the OpenTelemetry SDK.
+
+    Args:
+        event: The event to emit.
+    """
+    if not should_emit_events():
+        return
+
+    if isinstance(event, MessageEvent):
+        _emit_message_event(event)
+    elif isinstance(event, ChoiceEvent):
+        _emit_choice_event(event)
+    else:
+        raise TypeError("Unsupported event type")
+
+
+def _emit_message_event(event: MessageEvent) -> None:
+    body = asdict(event)
+
+    if event.role in VALID_MESSAGE_ROLES:
+        name = "gen_ai.{}.message".format(event.role)
+        # According to the semantic conventions, the role is conditionally required if available
+        # and not equal to the "role" in the message name. So, remove the role from the body if
+        # it is the same as the in the event name.
+        body.pop("role", None)
+    else:
+        name = "gen_ai.user.message"
+
+    # According to the semantic conventions, only the assistant role has tool call
+    if event.role != Roles.ASSISTANT.value and event.tool_calls is not None:
+        del body["tool_calls"]
+    elif event.tool_calls is None:
+        del body["tool_calls"]
+
+    if not should_send_prompts():
+        del body["content"]
+        if body.get("tool_calls") is not None:
+            for tool_call in body["tool_calls"]:
+                tool_call["function"].pop("arguments", None)
+
+    Config.event_logger.emit(Event(name=name, body=body, attributes=EVENT_ATTRIBUTES))
+
+
+def _emit_choice_event(event: ChoiceEvent) -> None:
+    body = asdict(event)
+    if event.message["role"] == Roles.ASSISTANT.value:
+        # According to the semantic conventions, the role is conditionally required if available
+        # and not equal to "assistant", so remove the role from the body if it is "assistant".
+        body["message"].pop("role", None)
+
+    if event.tool_calls is None:
+        del body["tool_calls"]
+
+    if not should_send_prompts():
+        body["message"].pop("content", None)
+        if body.get("tool_calls") is not None:
+            for tool_call in body["tool_calls"]:
+                tool_call["function"].pop("arguments", None)
+
+    Config.event_logger.emit(
+        Event(name="gen_ai.choice", body=body, attributes=EVENT_ATTRIBUTES)
+    )

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

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, TypedDict
+
+
+class _FunctionToolCall(TypedDict):
+    function_name: str
+    arguments: Optional[dict[str, Any]]
+
+
+class ToolCall(TypedDict):
+    """Represents a tool call in the AI model."""
+
+    id: str
+    function: _FunctionToolCall
+    type: Literal["function"]
+
+
+class CompletionMessage(TypedDict):
+    """Represents a message in the AI model."""
+
+    content: Any
+    role: str = "assistant"
+
+
+@dataclass
+class MessageEvent:
+    """Represents an input event for the AI model."""
+
+    content: Any
+    role: str = "user"
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+@dataclass
+class ChoiceEvent:
+    """Represents a completion event for the AI model."""
+
+    index: int
+    message: CompletionMessage
+    finish_reason: str = "unknown"
+    tool_calls: Optional[List[ToolCall]] = None

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

@@ -0,0 +1,68 @@
+import time
+
+from opentelemetry import context as context_api
+from ..utils import is_openai_v1
+from ..shared import (
+    _get_openai_base_url,
+    metric_shared_attributes,
+    model_as_dict,
+)
+from ..utils import (
+    _with_image_gen_metric_wrapper,
+)
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.metrics import Counter, Histogram
+from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+
+
+@_with_image_gen_metric_wrapper
+def image_gen_metrics_wrapper(
+    duration_histogram: Histogram,
+    exception_counter: Counter,
+    wrapped,
+    instance,
+    args,
+    kwargs,
+):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
+        SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+    ):
+        return wrapped(*args, **kwargs)
+
+    try:
+        # record time for duration
+        start_time = time.time()
+        response = wrapped(*args, **kwargs)
+        end_time = time.time()
+    except Exception as e:  # pylint: disable=broad-except
+        end_time = time.time()
+        duration = end_time - start_time if "start_time" in locals() else 0
+
+        attributes = {
+            "error.type": e.__class__.__name__,
+        }
+
+        if duration > 0 and duration_histogram:
+            duration_histogram.record(duration, attributes=attributes)
+        if exception_counter:
+            exception_counter.add(1, attributes=attributes)
+
+        raise
+
+    if is_openai_v1():
+        response_dict = model_as_dict(response)
+    else:
+        response_dict = response
+
+    # not provide response.model in ImagesResponse response, use model in request kwargs
+    shared_attributes = metric_shared_attributes(
+        response_model=kwargs.get("model") or None,
+        operation="image_gen",
+        server_address=_get_openai_base_url(instance),
+    )
+
+    duration = end_time - start_time
+    if duration_histogram:
+        duration_histogram.record(duration, attributes=shared_attributes)
+
+    return response

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

@@ -0,0 +1,185 @@
+import asyncio
+import logging
+import os
+import threading
+import traceback
+from contextlib import asynccontextmanager
+from importlib.metadata import version
+
+from opentelemetry import context as context_api
+from opentelemetry._events import EventLogger
+from .shared.config import Config
+
+import openai
+
+_OPENAI_VERSION = version("openai")
+
+TRACELOOP_TRACE_CONTENT = "TRACELOOP_TRACE_CONTENT"
+
+
+def is_openai_v1():
+    return _OPENAI_VERSION >= "1.0.0"
+
+
+def is_azure_openai(instance):
+    return is_openai_v1() and isinstance(
+        instance._client, (openai.AsyncAzureOpenAI, openai.AzureOpenAI)
+    )
+
+
+def is_metrics_enabled() -> bool:
+    return (os.getenv("TRACELOOP_METRICS_ENABLED") or "true").lower() == "true"
+
+
+def should_record_stream_token_usage():
+    return Config.enrich_token_usage
+
+
+def _with_image_gen_metric_wrapper(func):
+    def _with_metric(duration_histogram, exception_counter):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(
+                duration_histogram,
+                exception_counter,
+                wrapped,
+                instance,
+                args,
+                kwargs,
+            )
+
+        return wrapper
+
+    return _with_metric
+
+
+def _with_embeddings_telemetry_wrapper(func):
+    def _with_embeddings_telemetry(
+        tracer,
+        token_counter,
+        vector_size_counter,
+        duration_histogram,
+        exception_counter,
+    ):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(
+                tracer,
+                token_counter,
+                vector_size_counter,
+                duration_histogram,
+                exception_counter,
+                wrapped,
+                instance,
+                args,
+                kwargs,
+            )
+
+        return wrapper
+
+    return _with_embeddings_telemetry
+
+
+def _with_chat_telemetry_wrapper(func):
+    def _with_chat_telemetry(
+        tracer,
+        token_counter,
+        choice_counter,
+        duration_histogram,
+        exception_counter,
+        streaming_time_to_first_token,
+        streaming_time_to_generate,
+    ):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(
+                tracer,
+                token_counter,
+                choice_counter,
+                duration_histogram,
+                exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+                wrapped,
+                instance,
+                args,
+                kwargs,
+            )
+
+        return wrapper
+
+    return _with_chat_telemetry
+
+
+def _with_tracer_wrapper(func):
+    def _with_tracer(tracer):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer
+
+
+@asynccontextmanager
+async def start_as_current_span_async(tracer, *args, **kwargs):
+    with tracer.start_as_current_span(*args, **kwargs) as span:
+        yield span
+
+
+def dont_throw(func):
+    """
+    A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+    Works for both synchronous and asynchronous functions.
+    """
+    logger = logging.getLogger(func.__module__)
+
+    async def async_wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            _handle_exception(e, func, logger)
+
+    def sync_wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            _handle_exception(e, func, logger)
+
+    def _handle_exception(e, func, logger):
+        logger.debug(
+            "OpenLLMetry failed to trace in %s, error: %s",
+            func.__name__,
+            traceback.format_exc(),
+        )
+        if Config.exception_logger:
+            Config.exception_logger(e)
+
+    return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
+
+
+def run_async(method):
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop and loop.is_running():
+        thread = threading.Thread(target=lambda: asyncio.run(method))
+        thread.start()
+        thread.join()
+    else:
+        asyncio.run(method)
+
+
+def should_send_prompts():
+    return (
+        os.getenv(TRACELOOP_TRACE_CONTENT) or "true"
+    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
+
+
+def should_emit_events() -> bool:
+    """
+    Checks if the instrumentation isn't using the legacy attributes
+    and if the event logger is not None.
+    """
+    return not Config.use_legacy_attributes and isinstance(
+        Config.event_logger, EventLogger
+    )

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

@@ -0,0 +1,176 @@
+from typing import Collection
+
+from opentelemetry._events import get_event_logger
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from ..shared.chat_wrappers import (
+    achat_wrapper,
+    chat_wrapper,
+)
+from ..shared.completion_wrappers import (
+    acompletion_wrapper,
+    completion_wrapper,
+)
+from ..shared.config import Config
+from ..shared.embeddings_wrappers import (
+    aembeddings_wrapper,
+    embeddings_wrapper,
+)
+from ..utils import is_metrics_enabled
+from ..version import __version__
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.metrics import get_meter
+from opentelemetry.semconv._incubating.metrics import gen_ai_metrics as GenAIMetrics
+from opentelemetry.semconv_ai import Meters
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+_instruments = ("openai >= 0.27.0", "openai < 1.0.0")
+
+
+class OpenAIV0Instrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        meter_provider = kwargs.get("meter_provider")
+        meter = get_meter(__name__, __version__, meter_provider)
+
+        if not Config.use_legacy_attributes:
+            event_logger_provider = kwargs.get("event_logger_provider")
+            Config.event_logger = get_event_logger(
+                __name__, __version__, event_logger_provider=event_logger_provider
+            )
+
+        if is_metrics_enabled():
+            tokens_histogram = meter.create_histogram(
+                name=Meters.LLM_TOKEN_USAGE,
+                unit="token",
+                description="Measures number of input and output tokens used",
+            )
+
+            chat_choice_counter = meter.create_counter(
+                name=Meters.LLM_GENERATION_CHOICES,
+                unit="choice",
+                description="Number of choices returned by chat completions call",
+            )
+
+            duration_histogram = meter.create_histogram(
+                name=Meters.LLM_OPERATION_DURATION,
+                unit="s",
+                description="GenAI operation duration",
+            )
+
+            chat_exception_counter = meter.create_counter(
+                name=Meters.LLM_COMPLETIONS_EXCEPTIONS,
+                unit="time",
+                description="Number of exceptions occurred during chat completions",
+            )
+
+            streaming_time_to_first_token = meter.create_histogram(
+                name=GenAIMetrics.GEN_AI_SERVER_TIME_TO_FIRST_TOKEN,
+                unit="s",
+                description="Time to first token in streaming chat completions",
+            )
+            streaming_time_to_generate = meter.create_histogram(
+                name=Meters.LLM_STREAMING_TIME_TO_GENERATE,
+                unit="s",
+                description="Time between first token and completion in streaming chat completions",
+            )
+        else:
+            (
+                tokens_histogram,
+                chat_choice_counter,
+                duration_histogram,
+                chat_exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+            ) = (None, None, None, None, None, None)
+
+        if is_metrics_enabled():
+            embeddings_vector_size_counter = meter.create_counter(
+                name=Meters.LLM_EMBEDDINGS_VECTOR_SIZE,
+                unit="element",
+                description="he size of returned vector",
+            )
+            embeddings_exception_counter = meter.create_counter(
+                name=Meters.LLM_EMBEDDINGS_EXCEPTIONS,
+                unit="time",
+                description="Number of exceptions occurred during embeddings operation",
+            )
+        else:
+            (
+                tokens_histogram,
+                embeddings_vector_size_counter,
+                embeddings_exception_counter,
+            ) = (None, None, None)
+
+        wrap_function_wrapper(
+            "openai",
+            "Completion.create",
+            completion_wrapper(tracer),
+        )
+
+        wrap_function_wrapper(
+            "openai",
+            "Completion.acreate",
+            acompletion_wrapper(tracer),
+        )
+        wrap_function_wrapper(
+            "openai",
+            "ChatCompletion.create",
+            chat_wrapper(
+                tracer,
+                tokens_histogram,
+                chat_choice_counter,
+                duration_histogram,
+                chat_exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+            ),
+        )
+        wrap_function_wrapper(
+            "openai",
+            "ChatCompletion.acreate",
+            achat_wrapper(
+                tracer,
+                tokens_histogram,
+                chat_choice_counter,
+                duration_histogram,
+                chat_exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+            ),
+        )
+        wrap_function_wrapper(
+            "openai",
+            "Embedding.create",
+            embeddings_wrapper(
+                tracer,
+                tokens_histogram,
+                embeddings_vector_size_counter,
+                duration_histogram,
+                embeddings_exception_counter,
+            ),
+        )
+        wrap_function_wrapper(
+            "openai",
+            "Embedding.acreate",
+            aembeddings_wrapper(
+                tracer,
+                tokens_histogram,
+                embeddings_vector_size_counter,
+                duration_histogram,
+                embeddings_exception_counter,
+            ),
+        )
+
+    def _uninstrument(self, **kwargs):
+        unwrap("openai", "Completion.create")
+        unwrap("openai", "Completion.acreate")
+        unwrap("openai", "ChatCompletion.create")
+        unwrap("openai", "ChatCompletion.acreate")
+        unwrap("openai", "Embedding.create")
+        unwrap("openai", "Embedding.acreate")