Add Support for Async openai instrumentation (open-telemetry#2984)

Author: alizenhom

Diff for: .github/component_owners.yml

@@ -73,4 +73,5 @@ components:
     - lzchen
     - gyliu513
     - nirga
+    - alizenhom
     - codefromthecrypt

Diff for: instrumentation-genai/opentelemetry-instrumentation-openai-v2/CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Support for `AsyncOpenAI/AsyncCompletions` ([#2984](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2984))
+
 ## Version 2.0b0 (2024-11-08)
 
 - Use generic `OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT` environment variable

Diff for: instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/__init__.py

@@ -52,7 +52,7 @@
 from opentelemetry.semconv.schemas import Schemas
 from opentelemetry.trace import get_tracer
 
-from .patch import chat_completions_create
+from .patch import async_chat_completions_create, chat_completions_create
 
 
 class OpenAIInstrumentor(BaseInstrumentor):
@@ -84,7 +84,16 @@ def _instrument(self, **kwargs):
             ),
         )
 
+        wrap_function_wrapper(
+            module="openai.resources.chat.completions",
+            name="AsyncCompletions.create",
+            wrapper=async_chat_completions_create(
+                tracer, event_logger, is_content_enabled()
+            ),
+        )
+
     def _uninstrument(self, **kwargs):
         import openai  # pylint: disable=import-outside-toplevel
 
         unwrap(openai.resources.chat.completions.Completions, "create")
+        unwrap(openai.resources.chat.completions.AsyncCompletions, "create")

Diff for: instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/patch.py

@@ -21,15 +21,12 @@
 from opentelemetry.semconv._incubating.attributes import (
     gen_ai_attributes as GenAIAttributes,
 )
-from opentelemetry.semconv.attributes import (
-    error_attributes as ErrorAttributes,
-)
 from opentelemetry.trace import Span, SpanKind, Tracer
-from opentelemetry.trace.status import Status, StatusCode
 
 from .utils import (
     choice_to_event,
     get_llm_request_attributes,
+    handle_span_exception,
     is_streaming,
     message_to_event,
     set_span_attribute,
@@ -72,12 +69,49 @@ def traced_method(wrapped, instance, args, kwargs):
                 return result
 
             except Exception as error:
-                span.set_status(Status(StatusCode.ERROR, str(error)))
+                handle_span_exception(span, error)
+                raise
+
+    return traced_method
+
+
+def async_chat_completions_create(
+    tracer: Tracer, event_logger: EventLogger, capture_content: bool
+):
+    """Wrap the `create` method of the `AsyncChatCompletion` class to trace it."""
+
+    async def traced_method(wrapped, instance, args, kwargs):
+        span_attributes = {**get_llm_request_attributes(kwargs, instance)}
+
+        span_name = f"{span_attributes[GenAIAttributes.GEN_AI_OPERATION_NAME]} {span_attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL]}"
+        with tracer.start_as_current_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+            attributes=span_attributes,
+            end_on_exit=False,
+        ) as span:
+            if span.is_recording():
+                for message in kwargs.get("messages", []):
+                    event_logger.emit(
+                        message_to_event(message, capture_content)
+                    )
+
+            try:
+                result = await wrapped(*args, **kwargs)
+                if is_streaming(kwargs):
+                    return StreamWrapper(
+                        result, span, event_logger, capture_content
+                    )
+
                 if span.is_recording():
-                    span.set_attribute(
-                        ErrorAttributes.ERROR_TYPE, type(error).__qualname__
+                    _set_response_attributes(
+                        span, result, event_logger, capture_content
                     )
                 span.end()
+                return result
+
+            except Exception as error:
+                handle_span_exception(span, error)
                 raise
 
     return traced_method
@@ -286,10 +320,19 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         try:
             if exc_type is not None:
-                self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
-                self.span.set_attribute(
-                    ErrorAttributes.ERROR_TYPE, exc_type.__qualname__
-                )
+                handle_span_exception(self.span, exc_val)
+        finally:
+            self.cleanup()
+        return False  # Propagate the exception
+
+    async def __aenter__(self):
+        self.setup()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        try:
+            if exc_type is not None:
+                handle_span_exception(self.span, exc_val)
         finally:
             self.cleanup()
         return False  # Propagate the exception
@@ -301,6 +344,9 @@ def close(self):
     def __iter__(self):
         return self
 
+    def __aiter__(self):
+        return self
+
     def __next__(self):
         try:
             chunk = next(self.stream)
@@ -310,10 +356,20 @@ def __next__(self):
             self.cleanup()
             raise
         except Exception as error:
-            self.span.set_status(Status(StatusCode.ERROR, str(error)))
-            self.span.set_attribute(
-                ErrorAttributes.ERROR_TYPE, type(error).__qualname__
-            )
+            handle_span_exception(self.span, error)
+            self.cleanup()
+            raise
+
+    async def __anext__(self):
+        try:
+            chunk = await self.stream.__anext__()
+            self.process_chunk(chunk)
+            return chunk
+        except StopAsyncIteration:
+            self.cleanup()
+            raise
+        except Exception as error:
+            handle_span_exception(self.span, error)
             self.cleanup()
             raise
 

Diff for: instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/utils.py

@@ -26,6 +26,10 @@
 from opentelemetry.semconv._incubating.attributes import (
     server_attributes as ServerAttributes,
 )
+from opentelemetry.semconv.attributes import (
+    error_attributes as ErrorAttributes,
+)
+from opentelemetry.trace.status import Status, StatusCode
 
 OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT = (
     "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
@@ -138,9 +142,11 @@ def choice_to_event(choice, capture_content):
 
     if choice.message:
         message = {
-            "role": choice.message.role
-            if choice.message and choice.message.role
-            else None
+            "role": (
+                choice.message.role
+                if choice.message and choice.message.role
+                else None
+            )
         }
         tool_calls = extract_tool_calls(choice.message, capture_content)
         if tool_calls:
@@ -210,3 +216,12 @@ def get_llm_request_attributes(
 
     # filter out None values
     return {k: v for k, v in attributes.items() if v is not None}
+
+
+def handle_span_exception(span, error):
+    span.set_status(Status(StatusCode.ERROR, str(error)))
+    if span.is_recording():
+        span.set_attribute(
+            ErrorAttributes.ERROR_TYPE, type(error).__qualname__
+        )
+    span.end()

Diff for: instrumentation-genai/opentelemetry-instrumentation-openai-v2/test-requirements-0.txt

@@ -5,6 +5,7 @@ importlib-metadata==6.11.0
 packaging==24.0
 pytest==7.4.4
 pytest-vcr==1.0.2
+pytest-asyncio==0.21.0
 wrapt==1.16.0
 opentelemetry-api==1.28  # when updating, also update in pyproject.toml
 opentelemetry-sdk==1.28  # when updating, also update in pyproject.toml

Diff for: instrumentation-genai/opentelemetry-instrumentation-openai-v2/test-requirements-1.txt

@@ -5,6 +5,7 @@ importlib-metadata==6.11.0
 packaging==24.0
 pytest==7.4.4
 pytest-vcr==1.0.2
+pytest-asyncio==0.21.0
 wrapt==1.16.0
 # test with the latest version of opentelemetry-api, sdk, and semantic conventions
 

Diff for: instrumentation-genai/opentelemetry-instrumentation-openai-v2/tests/cassettes/test_async_chat_completion_404.yaml

@@ -0,0 +1,89 @@
+interactions:
+- request:
+    body: |-
+      {
+        "messages": [
+          {
+            "role": "user",
+            "content": "Say this is a test"
+          }
+        ],
+        "model": "this-model-does-not-exist"
+      }
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      authorization:
+      - Bearer test_openai_api_key
+      connection:
+      - keep-alive
+      content-length:
+      - '103'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+      user-agent:
+      - AsyncOpenAI/Python 1.26.0
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - async:asyncio
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.26.0
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.12.5
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: |-
+        {
+          "error": {
+            "message": "The model `this-model-does-not-exist` does not exist or you do not have access to it.",
+            "type": "invalid_request_error",
+            "param": null,
+            "code": "model_not_found"
+          }
+        }
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 8e1a80827a861852-MRS
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json; charset=utf-8
+      Date:
+      - Wed, 13 Nov 2024 00:04:01 GMT
+      Server:
+      - cloudflare
+      Set-Cookie: test_set_cookie
+      Transfer-Encoding:
+      - chunked
+      X-Content-Type-Options:
+      - nosniff
+      alt-svc:
+      - h3=":443"; ma=86400
+      content-length:
+      - '231'
+      openai-organization: test_openai_org_id
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Origin
+      x-request-id:
+      - req_5cf06a7fabd45ebe21ee38c14c5b2f76
+    status:
+      code: 404
+      message: Not Found
+version: 1