opentelemetry-instrumentation-openai 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

opentelemetry/instrumentation/openai/__init__.py

@@ -1,355 +1,13 @@
-import logging
-import os
-import json
-import types
-import pkg_resources
 from typing import Collection
-from wrapt import wrap_function_wrapper
-import openai
-
-from opentelemetry import context as context_api
-from opentelemetry.trace import get_tracer, SpanKind
-from opentelemetry.trace.status import Status, StatusCode
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
-from opentelemetry.instrumentation.utils import (
-    _SUPPRESS_INSTRUMENTATION_KEY,
-    unwrap,
-)
-
-from opentelemetry.semconv.ai import SpanAttributes, LLMRequestTypeValues
-from opentelemetry.instrumentation.openai.version import __version__
 
-logger = logging.getLogger(__name__)
+from opentelemetry.instrumentation.openai.utils import is_openai_v1
+from opentelemetry.instrumentation.openai.v0 import OpenAIV0Instrumentor
+from opentelemetry.instrumentation.openai.v1 import OpenAIV1Instrumentor
 
 _instruments = ("openai >= 0.27.0",)
 
-WRAPPED_METHODS_VERSION_0 = [
-    {
-        "module": "openai",
-        "object": "ChatCompletion",
-        "method": "create",
-        "span_name": "openai.chat",
-    },
-    {
-        "module": "openai",
-        "object": "Completion",
-        "method": "create",
-        "span_name": "openai.completion",
-    },
-]
-
-WRAPPED_METHODS_VERSION_1 = [
-    {
-        "module": "openai.resources.chat.completions",
-        "object": "Completions",
-        "method": "create",
-        "span_name": "openai.chat",
-    },
-    {
-        "module": "openai.resources.completions",
-        "object": "Completions",
-        "method": "create",
-        "span_name": "openai.completion",
-    },
-]
-
-
-def should_send_prompts():
-    return (
-        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
-    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
-
-
-def is_openai_v1():
-    return pkg_resources.get_distribution("openai").version >= "1.0.0"
-
-
-def _set_span_attribute(span, name, value):
-    if value is not None:
-        if value != "":
-            span.set_attribute(name, value)
-    return
-
-
-def _set_api_attributes(span):
-    _set_span_attribute(
-        span,
-        OpenAISpanAttributes.OPENAI_API_BASE,
-        openai.base_url if hasattr(openai, "base_url") else openai.api_base,
-    )
-    _set_span_attribute(span, OpenAISpanAttributes.OPENAI_API_TYPE, openai.api_type)
-    _set_span_attribute(
-        span, OpenAISpanAttributes.OPENAI_API_VERSION, openai.api_version
-    )
-
-    return
-
-
-def _set_span_prompts(span, messages):
-    if messages is None:
-        return
-
-    for i, msg in enumerate(messages):
-        prefix = f"{SpanAttributes.LLM_PROMPTS}.{i}"
-        if isinstance(msg.get("content"), str):
-            content = msg.get("content")
-        elif isinstance(msg.get("content"), list):
-            content = json.dumps(msg.get("content"))
-
-        _set_span_attribute(span, f"{prefix}.role", msg.get("role"))
-        _set_span_attribute(span, f"{prefix}.content", content)
-
-
-def _set_input_attributes(span, llm_request_type, kwargs):
-    _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("model"))
-    _set_span_attribute(
-        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_tokens")
-    )
-    _set_span_attribute(span, SpanAttributes.LLM_TEMPERATURE, kwargs.get("temperature"))
-    _set_span_attribute(span, SpanAttributes.LLM_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")))
-
-    if should_send_prompts():
-        if llm_request_type == LLMRequestTypeValues.CHAT:
-            _set_span_prompts(span, kwargs.get("messages"))
-        elif llm_request_type == LLMRequestTypeValues.COMPLETION:
-            prompt = kwargs.get("prompt")
-            _set_span_attribute(
-                span,
-                f"{SpanAttributes.LLM_PROMPTS}.0.user",
-                prompt[0] if isinstance(prompt, list) else prompt,
-            )
-
-        functions = kwargs.get("functions")
-        if functions:
-            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"))
-                )
-
-    return
-
-
-def _set_span_completions(span, llm_request_type, choices):
-    if choices is None:
-        return
-
-    for choice in choices:
-        if is_openai_v1() and not isinstance(choice, dict):
-            choice = choice.__dict__
-
-        index = choice.get("index")
-        prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
-        _set_span_attribute(
-            span, f"{prefix}.finish_reason", choice.get("finish_reason")
-        )
-
-        if llm_request_type == LLMRequestTypeValues.CHAT:
-            message = choice.get("message")
-            if message is not None:
-                if is_openai_v1() and not isinstance(message, dict):
-                    message = message.__dict__
-
-                _set_span_attribute(span, f"{prefix}.role", message.get("role"))
-                _set_span_attribute(span, f"{prefix}.content", message.get("content"))
-                function_call = message.get("function_call")
-                if function_call:
-                    if is_openai_v1() and not isinstance(function_call, dict):
-                        function_call = function_call.__dict__
-
-                    _set_span_attribute(
-                        span, f"{prefix}.function_call.name", function_call.get("name")
-                    )
-                    _set_span_attribute(
-                        span,
-                        f"{prefix}.function_call.arguments",
-                        function_call.get("arguments"),
-                    )
-        elif llm_request_type == LLMRequestTypeValues.COMPLETION:
-            _set_span_attribute(span, f"{prefix}.content", choice.get("text"))
-
-
-def _set_response_attributes(span, llm_request_type, response):
-    _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, response.get("model"))
-    if should_send_prompts():
-        _set_span_completions(span, llm_request_type, response.get("choices"))
-
-    usage = response.get("usage")
-    if usage is not None:
-        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")
-        )
-
-    return
-
-
-def _build_from_streaming_response(span, llm_request_type, response):
-    complete_response = {"choices": [], "model": ""}
-    for item in response:
-        item_to_yield = item
-        if is_openai_v1():
-            item = item.__dict__
-
-        for choice in item.get("choices"):
-            if is_openai_v1():
-                choice = choice.__dict__
-
-            index = choice.get("index")
-            if len(complete_response.get("choices")) <= index:
-                complete_response["choices"].append(
-                    {"index": index, "message": {"content": "", "role": ""}}
-                    if llm_request_type == LLMRequestTypeValues.CHAT
-                    else {"index": index, "text": ""}
-                )
-            complete_choice = complete_response.get("choices")[index]
-            if choice.get("finish_reason"):
-                complete_choice["finish_reason"] = choice.get("finish_reason")
-            if llm_request_type == LLMRequestTypeValues.CHAT:
-                delta = choice.get("delta")
-                if is_openai_v1():
-                    delta = delta.__dict__
-
-                if delta.get("content"):
-                    complete_choice["message"]["content"] += delta.get("content")
-                if delta.get("role"):
-                    complete_choice["message"]["role"] = delta.get("role")
-            else:
-                complete_choice["text"] += choice.get("text")
-
-        yield item_to_yield
-
-    _set_response_attributes(
-        span,
-        llm_request_type,
-        complete_response,
-    )
-    span.set_status(Status(StatusCode.OK))
-    span.end()
-
-
-def _with_tracer_wrapper(func):
-    """Helper for providing tracer for wrapper functions."""
-
-    def _with_tracer(tracer, to_wrap):
-        def wrapper(wrapped, instance, args, kwargs):
-            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
-
-        return wrapper
-
-    return _with_tracer
-
-
-def _llm_request_type_by_module_object(module_name, object_name):
-    if is_openai_v1():
-        if module_name == "openai.resources.chat.completions":
-            return LLMRequestTypeValues.CHAT
-        elif module_name == "openai.resources.completions":
-            return LLMRequestTypeValues.COMPLETION
-        else:
-            return LLMRequestTypeValues.UNKNOWN
-    else:
-        if object_name == "Completion":
-            return LLMRequestTypeValues.COMPLETION
-        elif object_name == "ChatCompletion":
-            return LLMRequestTypeValues.CHAT
-        else:
-            return LLMRequestTypeValues.UNKNOWN
-
-
-def is_streaming_response(response):
-    return isinstance(response, types.GeneratorType) or (
-        is_openai_v1() and isinstance(response, openai.Stream)
-    )
-
-
-@_with_tracer_wrapper
-def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
-    """Instruments and calls every function defined in TO_WRAP."""
-    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
-        return wrapped(*args, **kwargs)
-
-    name = to_wrap.get("span_name")
-    llm_request_type = _llm_request_type_by_module_object(
-        to_wrap.get("module"), to_wrap.get("object")
-    )
-
-    span = tracer.start_span(
-        name,
-        kind=SpanKind.CLIENT,
-        attributes={
-            SpanAttributes.LLM_VENDOR: "OpenAI",
-            SpanAttributes.LLM_REQUEST_TYPE: llm_request_type.value,
-        },
-    )
-
-    if span.is_recording():
-        _set_api_attributes(span)
-    try:
-        if span.is_recording():
-            _set_input_attributes(span, llm_request_type, kwargs)
-
-    except Exception as ex:  # pylint: disable=broad-except
-        logger.warning(
-            "Failed to set input attributes for openai span, error: %s", str(ex)
-        )
-
-    response = wrapped(*args, **kwargs)
-
-    if response:
-        try:
-            if span.is_recording():
-                if is_streaming_response(response):
-                    return _build_from_streaming_response(
-                        span, llm_request_type, response
-                    )
-                else:
-                    _set_response_attributes(
-                        span,
-                        llm_request_type,
-                        response.__dict__ if is_openai_v1() else response,
-                    )
-
-        except Exception as ex:  # pylint: disable=broad-except
-            logger.warning(
-                "Failed to set response attributes for openai span, error: %s",
-                str(ex),
-            )
-        if span.is_recording():
-            span.set_status(Status(StatusCode.OK))
-
-    span.end()
-    return response
-
-
-class OpenAISpanAttributes:
-    OPENAI_API_VERSION = "openai.api_version"
-    OPENAI_API_BASE = "openai.api_base"
-    OPENAI_API_TYPE = "openai.api_type"
-
 
 class OpenAIInstrumentor(BaseInstrumentor):
     """An instrumentor for OpenAI's client library."""
@@ -358,26 +16,13 @@ class OpenAIInstrumentor(BaseInstrumentor):
         return _instruments
 
     def _instrument(self, **kwargs):
-        tracer_provider = kwargs.get("tracer_provider")
-        tracer = get_tracer(__name__, __version__, tracer_provider)
-
-        wrapped_methods = (
-            WRAPPED_METHODS_VERSION_1 if is_openai_v1() else WRAPPED_METHODS_VERSION_0
-        )
-        for wrapped_method in wrapped_methods:
-            wrap_module = wrapped_method.get("module")
-            wrap_object = wrapped_method.get("object")
-            wrap_method = wrapped_method.get("method")
-            wrap_function_wrapper(
-                wrap_module,
-                f"{wrap_object}.{wrap_method}",
-                _wrap(tracer, wrapped_method),
-            )
+        if is_openai_v1():
+            OpenAIV1Instrumentor().instrument(**kwargs)
+        else:
+            OpenAIV0Instrumentor().instrument(**kwargs)
 
     def _uninstrument(self, **kwargs):
-        wrapped_methods = (
-            WRAPPED_METHODS_VERSION_1 if is_openai_v1() else WRAPPED_METHODS_VERSION_0
-        )
-        for wrapped_method in wrapped_methods:
-            wrap_object = wrapped_method.get("object")
-            unwrap(f"openai.{wrap_object}", wrapped_method.get("method"))
+        if is_openai_v1():
+            OpenAIV1Instrumentor().uninstrument(**kwargs)
+        else:
+            OpenAIV0Instrumentor().uninstrument(**kwargs)

opentelemetry/instrumentation/openai/shared/__init__.py

@@ -0,0 +1,136 @@
+import os
+import openai
+import json
+import types
+import logging
+
+from opentelemetry import context as context_api
+
+from opentelemetry.semconv.ai import SpanAttributes
+from opentelemetry.instrumentation.openai.utils import is_openai_v1
+
+
+OPENAI_API_VERSION = "openai.api_version"
+OPENAI_API_BASE = "openai.api_base"
+OPENAI_API_TYPE = "openai.api_type"
+
+logger = logging.getLogger(__name__)
+
+
+def should_send_prompts():
+    return (
+        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
+    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
+
+
+def _set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+def _set_api_attributes(span):
+    if not span.is_recording():
+        return
+
+    try:
+        base_url = openai.base_url if hasattr(openai, "base_url") else openai.api_base
+
+        _set_span_attribute(span, OPENAI_API_BASE, base_url)
+        _set_span_attribute(span, OPENAI_API_TYPE, openai.api_type)
+        _set_span_attribute(span, OPENAI_API_VERSION, openai.api_version)
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.warning(
+            "Failed to set api attributes for openai span, error: %s", str(ex)
+        )
+
+    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_request_attributes(span, llm_request_type, kwargs):
+    if not span.is_recording():
+        return
+
+    try:
+        _set_api_attributes(span)
+        _set_span_attribute(span, SpanAttributes.LLM_VENDOR, "OpenAI")
+        _set_span_attribute(
+            span, SpanAttributes.LLM_REQUEST_TYPE, llm_request_type.value
+        )
+        _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("model"))
+        _set_span_attribute(
+            span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_tokens")
+        )
+        _set_span_attribute(
+            span, SpanAttributes.LLM_TEMPERATURE, kwargs.get("temperature")
+        )
+        _set_span_attribute(span, SpanAttributes.LLM_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"))
+        )
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.warning(
+            "Failed to set input attributes for openai span, error: %s", str(ex)
+        )
+
+
+def _set_response_attributes(span, response):
+    if not span.is_recording():
+        return
+
+    try:
+        _set_span_attribute(
+            span, SpanAttributes.LLM_RESPONSE_MODEL, response.get("model")
+        )
+
+        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")
+        )
+
+        return
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.warning(
+            "Failed to set response attributes for openai span, error: %s", str(ex)
+        )
+
+
+def is_streaming_response(response):
+    return isinstance(response, types.GeneratorType) or (
+        is_openai_v1() and isinstance(response, openai.Stream)
+    )

opentelemetry/instrumentation/openai/shared/chat_wrappers.py

@@ -0,0 +1,170 @@
+import json
+import logging
+
+from opentelemetry import context as context_api
+
+from opentelemetry.semconv.ai import SpanAttributes, LLMRequestTypeValues
+
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.instrumentation.openai.utils import (
+    _with_tracer_wrapper,
+    start_as_current_span_async,
+)
+from opentelemetry.instrumentation.openai.shared import (
+    _set_request_attributes,
+    _set_span_attribute,
+    _set_functions_attributes,
+    _set_response_attributes,
+    is_streaming_response,
+    should_send_prompts,
+)
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+
+from opentelemetry.instrumentation.openai.utils import is_openai_v1
+
+SPAN_NAME = "openai.chat"
+LLM_REQUEST_TYPE = LLMRequestTypeValues.CHAT
+
+logger = logging.getLogger(__name__)
+
+
+@_with_tracer_wrapper
+def chat_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    # span needs to be opened and closed manually because the response is a generator
+    span = tracer.start_span(SPAN_NAME, kind=SpanKind.CLIENT)
+
+    _handle_request(span, kwargs)
+    response = wrapped(*args, **kwargs)
+
+    if is_streaming_response(response):
+        # span will be closed after the generator is done
+        return _build_from_streaming_response(span, response)
+    else:
+        _handle_response(response, span)
+
+    span.end()
+
+    return response
+
+
+@_with_tracer_wrapper
+async def achat_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    async with start_as_current_span_async(
+        tracer=tracer, name=SPAN_NAME, kind=SpanKind.CLIENT
+    ) as span:
+        _handle_request(span, kwargs)
+        response = await wrapped(*args, **kwargs)
+        _handle_response(response, span)
+
+        return response
+
+
+def _handle_request(span, kwargs):
+    _set_request_attributes(span, LLM_REQUEST_TYPE, kwargs)
+    if should_send_prompts():
+        _set_prompts(span, kwargs.get("messages"))
+        _set_functions_attributes(span, kwargs.get("functions"))
+
+
+def _handle_response(response, span):
+    if is_openai_v1():
+        response_dict = response.model_dump()
+    else:
+        response_dict = response
+
+    _set_response_attributes(span, response_dict)
+
+    if should_send_prompts():
+        _set_completions(span, response_dict.get("choices"))
+
+    return response
+
+
+def _set_prompts(span, messages):
+    if not span.is_recording() or messages is None:
+        return
+
+    try:
+        for i, msg in enumerate(messages):
+            prefix = f"{SpanAttributes.LLM_PROMPTS}.{i}"
+            if isinstance(msg.get("content"), str):
+                content = msg.get("content")
+            elif isinstance(msg.get("content"), list):
+                content = json.dumps(msg.get("content"))
+
+            _set_span_attribute(span, f"{prefix}.role", msg.get("role"))
+            _set_span_attribute(span, f"{prefix}.content", content)
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.warning("Failed to set prompts for openai span, error: %s", str(ex))
+
+
+def _set_completions(span, choices):
+    if choices is None:
+        return
+
+    for choice in choices:
+        index = choice.get("index")
+        prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+        _set_span_attribute(
+            span, f"{prefix}.finish_reason", choice.get("finish_reason")
+        )
+
+        message = choice.get("message")
+        if not message:
+            return
+
+        _set_span_attribute(span, f"{prefix}.role", message.get("role"))
+        _set_span_attribute(span, f"{prefix}.content", message.get("content"))
+
+        function_call = message.get("function_call")
+        if not function_call:
+            return
+
+        _set_span_attribute(
+            span, f"{prefix}.function_call.name", function_call.get("name")
+        )
+        _set_span_attribute(
+            span, f"{prefix}.function_call.arguments", function_call.get("arguments")
+        )
+
+
+def _build_from_streaming_response(span, response):
+    complete_response = {"choices": [], "model": ""}
+    for item in response:
+        item_to_yield = item
+        if is_openai_v1():
+            item = item.model_dump()
+
+        for choice in item.get("choices"):
+            index = choice.get("index")
+            if len(complete_response.get("choices")) <= index:
+                complete_response["choices"].append(
+                    {"index": index, "message": {"content": "", "role": ""}}
+                )
+            complete_choice = complete_response.get("choices")[index]
+            if choice.get("finish_reason"):
+                complete_choice["finish_reason"] = choice.get("finish_reason")
+
+            delta = choice.get("delta")
+
+            if delta.get("content"):
+                complete_choice["message"]["content"] += delta.get("content")
+            if delta.get("role"):
+                complete_choice["message"]["role"] = delta.get("role")
+
+        yield item_to_yield
+
+    _set_response_attributes(span, complete_response)
+
+    if should_send_prompts():
+        _set_completions(span, complete_response.get("choices"))
+
+    span.set_status(Status(StatusCode.OK))
+    span.end()

opentelemetry/instrumentation/openai/shared/completion_wrappers.py

@@ -0,0 +1,142 @@
+import logging
+
+from opentelemetry import context as context_api
+
+from opentelemetry.semconv.ai import SpanAttributes, LLMRequestTypeValues
+
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.instrumentation.openai.utils import (
+    _with_tracer_wrapper,
+    start_as_current_span_async,
+)
+from opentelemetry.instrumentation.openai.shared import (
+    _set_request_attributes,
+    _set_span_attribute,
+    _set_functions_attributes,
+    _set_response_attributes,
+    is_streaming_response,
+    should_send_prompts,
+)
+
+from opentelemetry.instrumentation.openai.utils import is_openai_v1
+
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+
+SPAN_NAME = "openai.completion"
+LLM_REQUEST_TYPE = LLMRequestTypeValues.COMPLETION
+
+logger = logging.getLogger(__name__)
+
+
+@_with_tracer_wrapper
+def completion_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    # span needs to be opened and closed manually because the response is a generator
+    span = tracer.start_span(SPAN_NAME, kind=SpanKind.CLIENT)
+
+    _handle_request(span, kwargs)
+    response = wrapped(*args, **kwargs)
+
+    if is_streaming_response(response):
+        # span will be closed after the generator is done
+        return _build_from_streaming_response(span, response)
+    else:
+        _handle_response(response, span)
+
+    span.end()
+    return response
+
+
+@_with_tracer_wrapper
+async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    async with start_as_current_span_async(
+        tracer=tracer, name=SPAN_NAME, kind=SpanKind.CLIENT
+    ) as span:
+        _handle_request(span, kwargs)
+        response = await wrapped(*args, **kwargs)
+        _handle_response(response, span)
+
+        return response
+
+
+def _handle_request(span, kwargs):
+    _set_request_attributes(span, LLM_REQUEST_TYPE, kwargs)
+    if should_send_prompts():
+        _set_prompts(span, kwargs.get("prompt"))
+        _set_functions_attributes(span, kwargs.get("functions"))
+
+
+def _handle_response(response, span):
+    if is_openai_v1():
+        response_dict = response.model_dump()
+    else:
+        response_dict = response
+
+    _set_response_attributes(span, response_dict)
+
+    if should_send_prompts():
+        _set_completions(span, response_dict.get("choices"))
+
+
+def _set_prompts(span, prompt):
+    if not span.is_recording() or not prompt:
+        return
+
+    try:
+        _set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_PROMPTS}.0.user",
+            prompt[0] if isinstance(prompt, list) else prompt,
+        )
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.warning("Failed to set prompts for openai span, error: %s", str(ex))
+
+
+def _set_completions(span, choices):
+    if not span.is_recording() or not choices:
+        return
+
+    try:
+        for choice in choices:
+            index = choice.get("index")
+            prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+            _set_span_attribute(
+                span, f"{prefix}.finish_reason", choice.get("finish_reason")
+            )
+            _set_span_attribute(span, f"{prefix}.content", choice.get("text"))
+    except Exception as e:
+        logger.warning("Failed to set completion attributes, error: %s", str(e))
+
+
+def _build_from_streaming_response(span, response):
+    complete_response = {"choices": [], "model": ""}
+    for item in response:
+        item_to_yield = item
+        if is_openai_v1():
+            item = item.model_dump()
+
+        for choice in item.get("choices"):
+            index = choice.get("index")
+            if len(complete_response.get("choices")) <= index:
+                complete_response["choices"].append({"index": index, "text": ""})
+            complete_choice = complete_response.get("choices")[index]
+            if choice.get("finish_reason"):
+                complete_choice["finish_reason"] = choice.get("finish_reason")
+
+            complete_choice["text"] += choice.get("text")
+
+        yield item_to_yield
+
+    _set_response_attributes(span, complete_response)
+
+    if should_send_prompts():
+        _set_completions(span, complete_response.get("choices"))
+
+    span.set_status(Status(StatusCode.OK))
+    span.end()

opentelemetry/instrumentation/openai/utils.py

@@ -0,0 +1,22 @@
+from importlib.metadata import version
+from contextlib import asynccontextmanager
+
+
+def is_openai_v1():
+    return version("openai") >= "1.0.0"
+
+
+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

opentelemetry/instrumentation/openai/v0/__init__.py

@@ -0,0 +1,28 @@
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.instrumentation.openai.shared.chat_wrappers import chat_wrapper, achat_wrapper
+from opentelemetry.instrumentation.openai.shared.completion_wrappers import completion_wrapper, acompletion_wrapper
+from opentelemetry.instrumentation.openai.version import __version__
+
+_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)
+
+        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))
+        wrap_function_wrapper("openai", "ChatCompletion.acreate", achat_wrapper(tracer))
+
+    def _uninstrument(self, **kwargs):
+        pass

opentelemetry/instrumentation/openai/v1/__init__.py

@@ -0,0 +1,28 @@
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.instrumentation.openai.shared.chat_wrappers import chat_wrapper, achat_wrapper
+from opentelemetry.instrumentation.openai.shared.completion_wrappers import completion_wrapper, acompletion_wrapper
+from opentelemetry.instrumentation.openai.version import __version__
+
+_instruments = ("openai >= 1.0.0",)
+
+
+class OpenAIV1Instrumentor(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)
+
+        wrap_function_wrapper("openai.resources.chat.completions", "Completions.create", chat_wrapper(tracer))
+        wrap_function_wrapper("openai.resources.completions", "Completions.create", completion_wrapper(tracer))
+        wrap_function_wrapper("openai.resources.chat.completions", "AsyncCompletions.create", achat_wrapper(tracer))
+        wrap_function_wrapper("openai.resources.completions", "AsyncCompletions.create", acompletion_wrapper(tracer))
+
+    def _uninstrument(self, **kwargs):
+        pass

opentelemetry/instrumentation/openai/version.py

@@ -1 +1 @@
-__version__ = "0.5.0"
+__version__ = "0.5.1"

{opentelemetry_instrumentation_openai-0.5.0.dist-info → opentelemetry_instrumentation_openai-0.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: opentelemetry-instrumentation-openai
-Version: 0.5.0
+Version: 0.5.1
 Summary: OpenTelemetry OpenAI instrumentation
 License: Apache-2.0
 Author: Gal Kleinman

opentelemetry_instrumentation_openai-0.5.1.dist-info/RECORD

@@ -0,0 +1,11 @@
+opentelemetry/instrumentation/openai/__init__.py,sha256=iiuJICVwzv0K2OeskonjZIYpPnxooSEsIkXOS7uZEf8,937
+opentelemetry/instrumentation/openai/shared/__init__.py,sha256=I3mHPmUPjb22ZsbgmbPVtuIxo8_FM6F87dyqBxfxTnA,4285
+opentelemetry/instrumentation/openai/shared/chat_wrappers.py,sha256=EcER4tO6Zwu4eC-Z-L2dmYbgTnsKmZG8hmVIhpEOURA,5354
+opentelemetry/instrumentation/openai/shared/completion_wrappers.py,sha256=M3K4IkQOS2bnBehLI_Mh_ApvD1trSMX2_PjSxbZO0yc,4449
+opentelemetry/instrumentation/openai/utils.py,sha256=m405jldmW6dsgQz_qF75WNM0fMTJUMD_IVMluu7s-nY,546
+opentelemetry/instrumentation/openai/v0/__init__.py,sha256=0olHqFUCAfVtD-TnOzIVDwK5YaieZU36RewZPLX8Cos,1236
+opentelemetry/instrumentation/openai/v1/__init__.py,sha256=FnDPr5UQcKnTBLt1dAbr_GYtE8pOCtR7U2PtY0thnFo,1320
+opentelemetry/instrumentation/openai/version.py,sha256=eZ1bOun1DDVV0YLOBW4wj2FP1ajReLjbIrGmzN7ASBw,22
+opentelemetry_instrumentation_openai-0.5.1.dist-info/METADATA,sha256=hbb4z20Hf0DJhbpyla--bVKQnwcqWTY6Sh4GFApfD9Q,1767
+opentelemetry_instrumentation_openai-0.5.1.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+opentelemetry_instrumentation_openai-0.5.1.dist-info/RECORD,,

opentelemetry_instrumentation_openai-0.5.0.dist-info/RECORD

@@ -1,5 +0,0 @@
-opentelemetry/instrumentation/openai/__init__.py,sha256=reFlctrambaZ2nVXNaeoZKqG_dS9k7XhnqruYexrBAk,12912
-opentelemetry/instrumentation/openai/version.py,sha256=LBK46heutvn3KmsCrKIYu8RQikbfnjZaj2xFrXaeCzQ,22
-opentelemetry_instrumentation_openai-0.5.0.dist-info/METADATA,sha256=FGjF7OQ6EEXPB855HoT22MMspZKhloDjKmQP_XB4qAw,1767
-opentelemetry_instrumentation_openai-0.5.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-opentelemetry_instrumentation_openai-0.5.0.dist-info/RECORD,,