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

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

@@ -0,0 +1,368 @@
+from typing import Collection
+
+from opentelemetry._events import get_event_logger
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+
+from lmnr.sdk.log import get_default_logger
+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 ..shared.image_gen_wrappers import (
+    image_gen_metrics_wrapper,
+)
+from ..utils import is_metrics_enabled
+from .assistant_wrappers import (
+    assistants_create_wrapper,
+    messages_list_wrapper,
+    runs_create_and_stream_wrapper,
+    runs_create_wrapper,
+    runs_retrieve_wrapper,
+)
+
+from .responses_wrappers import (
+    async_responses_cancel_wrapper,
+    async_responses_get_or_create_wrapper,
+    responses_cancel_wrapper,
+    responses_get_or_create_wrapper,
+)
+
+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 >= 1.0.0",)
+logger = get_default_logger(__name__)
+
+
+class OpenAIV1Instrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _try_wrap(self, module, function, wrapper):
+        """
+        Wrap a function if it exists, otherwise do nothing.
+        This is useful for handling cases where the function is not available in
+        the older versions of the library.
+
+        Args:
+            module (str): The module to wrap, e.g. "openai.resources.chat.completions"
+            function (str): "Object.function" to wrap, e.g. "Completions.parse"
+            wrapper (callable): The wrapper to apply to the function.
+        """
+        try:
+            wrap_function_wrapper(module, function, wrapper)
+        except (AttributeError, ModuleNotFoundError, ImportError):
+            logger.debug(f"Failed to wrap {module}.{function}")
+            pass
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        # meter and counters are inited here
+        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)
+
+        wrap_function_wrapper(
+            "openai.resources.chat.completions",
+            "Completions.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.resources.completions",
+            "Completions.create",
+            completion_wrapper(tracer),
+        )
+
+        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.resources.embeddings",
+            "Embeddings.create",
+            embeddings_wrapper(
+                tracer,
+                tokens_histogram,
+                embeddings_vector_size_counter,
+                duration_histogram,
+                embeddings_exception_counter,
+            ),
+        )
+
+        wrap_function_wrapper(
+            "openai.resources.chat.completions",
+            "AsyncCompletions.create",
+            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.resources.completions",
+            "AsyncCompletions.create",
+            acompletion_wrapper(tracer),
+        )
+        wrap_function_wrapper(
+            "openai.resources.embeddings",
+            "AsyncEmbeddings.create",
+            aembeddings_wrapper(
+                tracer,
+                tokens_histogram,
+                embeddings_vector_size_counter,
+                duration_histogram,
+                embeddings_exception_counter,
+            ),
+        )
+        # in newer versions, Completions.parse are out of beta
+        self._try_wrap(
+            "openai.resources.chat.completions",
+            "Completions.parse",
+            chat_wrapper(
+                tracer,
+                tokens_histogram,
+                chat_choice_counter,
+                duration_histogram,
+                chat_exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+            ),
+        )
+        self._try_wrap(
+            "openai.resources.chat.completions",
+            "AsyncCompletions.parse",
+            achat_wrapper(
+                tracer,
+                tokens_histogram,
+                chat_choice_counter,
+                duration_histogram,
+                chat_exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+            ),
+        )
+
+        if is_metrics_enabled():
+            image_gen_exception_counter = meter.create_counter(
+                name=Meters.LLM_IMAGE_GENERATIONS_EXCEPTIONS,
+                unit="time",
+                description="Number of exceptions occurred during image generations operation",
+            )
+        else:
+            image_gen_exception_counter = None
+
+        wrap_function_wrapper(
+            "openai.resources.images",
+            "Images.generate",
+            image_gen_metrics_wrapper(duration_histogram, image_gen_exception_counter),
+        )
+
+        # Beta APIs may not be available consistently in all versions
+        self._try_wrap(
+            "openai.resources.beta.assistants",
+            "Assistants.create",
+            assistants_create_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.beta.chat.completions",
+            "Completions.parse",
+            chat_wrapper(
+                tracer,
+                tokens_histogram,
+                chat_choice_counter,
+                duration_histogram,
+                chat_exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+            ),
+        )
+        self._try_wrap(
+            "openai.resources.beta.chat.completions",
+            "AsyncCompletions.parse",
+            achat_wrapper(
+                tracer,
+                tokens_histogram,
+                chat_choice_counter,
+                duration_histogram,
+                chat_exception_counter,
+                streaming_time_to_first_token,
+                streaming_time_to_generate,
+            ),
+        )
+        self._try_wrap(
+            "openai.resources.beta.threads.runs",
+            "Runs.create",
+            runs_create_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.beta.threads.runs",
+            "Runs.retrieve",
+            runs_retrieve_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.beta.threads.runs",
+            "Runs.create_and_stream",
+            runs_create_and_stream_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.beta.threads.messages",
+            "Messages.list",
+            messages_list_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.responses",
+            "Responses.create",
+            responses_get_or_create_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.responses",
+            "Responses.retrieve",
+            responses_get_or_create_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.responses",
+            "Responses.cancel",
+            responses_cancel_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.responses",
+            "AsyncResponses.create",
+            async_responses_get_or_create_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.responses",
+            "AsyncResponses.retrieve",
+            async_responses_get_or_create_wrapper(tracer),
+        )
+        self._try_wrap(
+            "openai.resources.responses",
+            "AsyncResponses.cancel",
+            async_responses_cancel_wrapper(tracer),
+        )
+
+    def _uninstrument(self, **kwargs):
+        self.try_unwrap("openai.resources.chat.completions.Completions", "create")
+        self.try_unwrap("openai.resources.completions.Completions", "create")
+        self.try_unwrap("openai.resources.embeddings.Embeddings", "create")
+        self.try_unwrap("openai.resources.chat.completions.AsyncCompletions", "create")
+        self.try_unwrap("openai.resources.completions.AsyncCompletions", "create")
+        self.try_unwrap("openai.resources.embeddings.AsyncEmbeddings", "create")
+        self.try_unwrap("openai.resources.images.Images", "generate")
+        self.try_unwrap("openai.resources.chat.completions.Completions", "parse")
+        self.try_unwrap("openai.resources.chat.completions.AsyncCompletions", "parse")
+        self.try_unwrap("openai.resources.beta.assistants.Assistants", "create")
+        self.try_unwrap("openai.resources.beta.chat.completions.Completions", "parse")
+        self.try_unwrap(
+            "openai.resources.beta.chat.completions.AsyncCompletions", "parse"
+        )
+        self.try_unwrap("openai.resources.beta.threads.runs.Runs", "create")
+        self.try_unwrap("openai.resources.beta.threads.runs.Runs", "retrieve")
+        self.try_unwrap("openai.resources.beta.threads.runs.Runs", "create_and_stream")
+        self.try_unwrap("openai.resources.beta.threads.messages.Messages", "list")
+        self.try_unwrap("openai.resources.responses.Responses", "create")
+        self.try_unwrap("openai.resources.responses.Responses", "retrieve")
+        self.try_unwrap("openai.resources.responses.Responses", "cancel")
+        self.try_unwrap("openai.resources.responses.AsyncResponses", "create")
+        self.try_unwrap("openai.resources.responses.AsyncResponses", "retrieve")
+        self.try_unwrap("openai.resources.responses.AsyncResponses", "cancel")
+
+    def try_unwrap(self, module, function):
+        try:
+            unwrap(module, function)
+        except (AttributeError, ModuleNotFoundError, ImportError):
+            logger.debug(f"Failed to unwrap {module}.{function}")
+            pass

lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/assistant_wrappers.py

@@ -0,0 +1,325 @@
+import logging
+import time
+
+from opentelemetry import context as context_api
+from ..shared import (
+    _set_span_attribute,
+    model_as_dict,
+)
+from ..shared.config import Config
+from ..shared.event_emitter import emit_event
+from ..shared.event_models import (
+    ChoiceEvent,
+    MessageEvent,
+)
+from ..utils import (
+    _with_tracer_wrapper,
+    dont_throw,
+    should_emit_events,
+)
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    get_event_attributes_from_context,
+)
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
+from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes
+from opentelemetry.trace import SpanKind, Status, StatusCode
+
+from openai._legacy_response import LegacyAPIResponse
+from openai.types.beta.threads.run import Run
+
+logger = logging.getLogger(__name__)
+
+assistants = {}
+runs = {}
+
+
+@_with_tracer_wrapper
+def assistants_create_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    response = wrapped(*args, **kwargs)
+
+    assistants[response.id] = {
+        "model": kwargs.get("model"),
+        "instructions": kwargs.get("instructions"),
+    }
+
+    return response
+
+
+@_with_tracer_wrapper
+def runs_create_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    thread_id = kwargs.get("thread_id")
+    instructions = kwargs.get("instructions")
+
+    try:
+        response = wrapped(*args, **kwargs)
+        response_dict = model_as_dict(response)
+
+        runs[thread_id] = {
+            "start_time": time.time_ns(),
+            "assistant_id": kwargs.get("assistant_id"),
+            "instructions": instructions,
+            "run_id": response_dict.get("id"),
+        }
+
+        return response
+    except Exception as e:
+        runs[thread_id] = {
+            "exception": e,
+            "end_time": time.time_ns(),
+        }
+        raise
+
+
+@_with_tracer_wrapper
+def runs_retrieve_wrapper(tracer, wrapped, instance, args, kwargs):
+    @dont_throw
+    def process_response(response):
+        if type(response) is LegacyAPIResponse:
+            parsed_response = response.parse()
+        else:
+            parsed_response = response
+        assert type(parsed_response) is Run
+
+        if parsed_response.thread_id in runs:
+            thread_id = parsed_response.thread_id
+            runs[thread_id]["end_time"] = time.time_ns()
+            if parsed_response.usage:
+                runs[thread_id]["usage"] = parsed_response.usage
+
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    try:
+        response = wrapped(*args, **kwargs)
+        process_response(response)
+        return response
+    except Exception as e:
+        thread_id = kwargs.get("thread_id")
+        if thread_id in runs:
+            runs[thread_id]["exception"] = e
+            runs[thread_id]["end_time"] = time.time_ns()
+        raise
+
+
+@_with_tracer_wrapper
+def messages_list_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    id = kwargs.get("thread_id")
+
+    response = wrapped(*args, **kwargs)
+
+    response_dict = model_as_dict(response)
+    if id not in runs:
+        return response
+
+    run = runs[id]
+    messages = sorted(response_dict["data"], key=lambda x: x["created_at"])
+
+    span = tracer.start_span(
+        "openai.assistant.run",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+        start_time=run.get("start_time"),
+        context=get_current_context(),
+    )
+
+    if exception := run.get("exception"):
+        span.set_attribute(ERROR_TYPE, exception.__class__.__name__)
+        span.record_exception(exception, attributes=get_event_attributes_from_context())
+        span.set_status(Status(StatusCode.ERROR, str(exception)))
+        span.end(run.get("end_time"))
+
+    prompt_index = 0
+    if assistants.get(run["assistant_id"]) is not None or Config.enrich_assistant:
+        if Config.enrich_assistant:
+            assistant = model_as_dict(
+                instance._client.beta.assistants.retrieve(run["assistant_id"])
+            )
+            assistants[run["assistant_id"]] = assistant
+        else:
+            assistant = assistants[run["assistant_id"]]
+
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_SYSTEM,
+            "openai",
+        )
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_REQUEST_MODEL,
+            assistant["model"],
+        )
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_RESPONSE_MODEL,
+            assistant["model"],
+        )
+        if should_emit_events():
+            emit_event(MessageEvent(content=assistant["instructions"], role="system"))
+        else:
+            _set_span_attribute(
+                span, f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.role", "system"
+            )
+            _set_span_attribute(
+                span,
+                f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.content",
+                assistant["instructions"],
+            )
+        prompt_index += 1
+    _set_span_attribute(
+        span, f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.role", "system"
+    )
+    _set_span_attribute(
+        span,
+        f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.content",
+        run["instructions"],
+    )
+    emit_event(MessageEvent(content=run["instructions"], role="system"))
+    prompt_index += 1
+
+    completion_index = 0
+    for msg in messages:
+        prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{completion_index}"
+        content = msg.get("content")
+
+        message_content = content[0].get("text").get("value")
+        message_role = msg.get("role")
+        if message_role in ["user", "system"]:
+            if should_emit_events():
+                emit_event(MessageEvent(content=message_content, role=message_role))
+            else:
+                _set_span_attribute(
+                    span,
+                    f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.role",
+                    message_role,
+                )
+                _set_span_attribute(
+                    span,
+                    f"{SpanAttributes.LLM_PROMPTS}.{prompt_index}.content",
+                    message_content,
+                )
+            prompt_index += 1
+        else:
+            if should_emit_events():
+                emit_event(
+                    ChoiceEvent(
+                        index=completion_index,
+                        message={"content": message_content, "role": message_role},
+                    )
+                )
+            else:
+                _set_span_attribute(span, f"{prefix}.role", msg.get("role"))
+                _set_span_attribute(span, f"{prefix}.content", message_content)
+                _set_span_attribute(
+                    span, f"gen_ai.response.{completion_index}.id", msg.get("id")
+                )
+            completion_index += 1
+
+    if run.get("usage"):
+        usage_dict = model_as_dict(run.get("usage"))
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+            usage_dict.get("completion_tokens"),
+        )
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+            usage_dict.get("prompt_tokens"),
+        )
+
+    span.end(run.get("end_time"))
+
+    return response
+
+
+@_with_tracer_wrapper
+def runs_create_and_stream_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    assistant_id = kwargs.get("assistant_id")
+    instructions = kwargs.get("instructions")
+
+    span = tracer.start_span(
+        "openai.assistant.run_stream",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+        context=get_current_context(),
+    )
+
+    i = 0
+    if assistants.get(assistant_id) is not None or Config.enrich_assistant:
+        if Config.enrich_assistant:
+            assistant = model_as_dict(
+                instance._client.beta.assistants.retrieve(assistant_id)
+            )
+            assistants[assistant_id] = assistant
+        else:
+            assistant = assistants[assistant_id]
+
+        _set_span_attribute(
+            span, SpanAttributes.LLM_REQUEST_MODEL, assistants[assistant_id]["model"]
+        )
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_SYSTEM,
+            "openai",
+        )
+        _set_span_attribute(
+            span,
+            SpanAttributes.LLM_RESPONSE_MODEL,
+            assistants[assistant_id]["model"],
+        )
+        if should_emit_events():
+            emit_event(
+                MessageEvent(
+                    content=assistants[assistant_id]["instructions"], role="system"
+                )
+            )
+        else:
+            _set_span_attribute(
+                span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system"
+            )
+            _set_span_attribute(
+                span,
+                f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
+                assistants[assistant_id]["instructions"],
+            )
+        i += 1
+    if should_emit_events():
+        emit_event(MessageEvent(content=instructions, role="system"))
+    else:
+        _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system")
+        _set_span_attribute(
+            span, f"{SpanAttributes.LLM_PROMPTS}.{i}.content", instructions
+        )
+
+    from ..v1.event_handler_wrapper import (
+        EventHandlerWrapper,
+    )
+
+    kwargs["event_handler"] = EventHandlerWrapper(
+        original_handler=kwargs["event_handler"],
+        span=span,
+    )
+
+    try:
+        response = wrapped(*args, **kwargs)
+        return response
+    except Exception as e:
+        span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+        span.record_exception(e, attributes=get_event_attributes_from_context())
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        span.end()
+        raise