Add Support for Async openai instrumentation

.github/component_owners.yml

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

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")

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
 

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()

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

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
 

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

@@ -0,0 +1,73 @@
+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
+      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: "{\n    \"error\": {\n        \"message\": \"The model `this-model-does-not-exist`
+        does not exist or you do not have access to it.\",\n        \"type\": \"invalid_request_error\",\n
+        \       \"param\": null,\n        \"code\": \"model_not_found\"\n    }\n}\n"
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 8df5e31d6e45e284-MRS
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json; charset=utf-8
+      Date:
+      - Fri, 08 Nov 2024 13:25:06 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_organization
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      vary:
+      - Origin
+      x-request-id:
+      - req_82c60ab486740fa07db9b846df2d532b
+    status:
+      code: 404
+      message: Not Found
+version: 1

instrumentation-genai/opentelemetry-instrumentation-openai-v2/tests/cassettes/test_async_chat_completion_extra_params.yaml

@@ -0,0 +1,98 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": "Say this is a test"}], "model":
+      "gpt-4o-mini", "max_tokens": 50, "seed": 42, "stream": false, "temperature":
+      0.5, "service_tier": "default"}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '183'
+      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: "{\n  \"id\": \"chatcmpl-ARJJHh1NYn4MAJ5vmO2u6MC0QWvS3\",\n  \"object\":
+        \"chat.completion\",\n  \"created\": 1731072447,\n  \"model\": \"gpt-4o-mini-2024-07-18\",\n
+        \ \"choices\": [\n    {\n      \"index\": 0,\n      \"message\": {\n        \"role\":
+        \"assistant\",\n        \"content\": \"This is a test. How can I assist you
+        further?\",\n        \"refusal\": null\n      },\n      \"logprobs\": null,\n
+        \     \"finish_reason\": \"stop\"\n    }\n  ],\n  \"usage\": {\n    \"prompt_tokens\":
+        12,\n    \"completion_tokens\": 12,\n    \"total_tokens\": 24,\n    \"prompt_tokens_details\":
+        {\n      \"cached_tokens\": 0,\n      \"audio_tokens\": 0\n    },\n    \"completion_tokens_details\":
+        {\n      \"reasoning_tokens\": 0,\n      \"audio_tokens\": 0,\n      \"accepted_prediction_tokens\":
+        0,\n      \"rejected_prediction_tokens\": 0\n    }\n  },\n  \"service_tier\":
+        \"default\",\n  \"system_fingerprint\": \"fp_0ba0d124f1\"\n}\n"
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 8df5e68e7df7e20d-MRS
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Fri, 08 Nov 2024 13:27:28 GMT
+      Server:
+      - cloudflare
+      Set-Cookie: test_set_cookie
+      Transfer-Encoding:
+      - chunked
+      X-Content-Type-Options:
+      - nosniff
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      content-length:
+      - '825'
+      openai-organization: test_organization
+      openai-processing-ms:
+      - '373'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      x-ratelimit-limit-requests:
+      - '30000'
+      x-ratelimit-limit-tokens:
+      - '150000000'
+      x-ratelimit-remaining-requests:
+      - '29999'
+      x-ratelimit-remaining-tokens:
+      - '149999943'
+      x-ratelimit-reset-requests:
+      - 2ms
+      x-ratelimit-reset-tokens:
+      - 0s
+      x-request-id:
+      - req_bea7292421fe218c0a9e158dbb7699d7
+    status:
+      code: 200
+      message: OK
+version: 1