fix(openai): track client attributes for v1 SDK of OpenAI (traceloop#522

)

Co-authored-by: Nir Gazit <nirga@users.noreply.github.com>

packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/__init__.py

@@ -32,10 +32,33 @@ def _set_span_attribute(span, name, value):
     return
 
 
+def _set_client_attributes(span, instance):
+    if not span.is_recording():
+        return
+
+    if not is_openai_v1():
+        return
+
+    try:
+        client = instance._client  # pylint: disable=protected-access
+        if isinstance(client, (openai.AsyncOpenAI, openai.OpenAI)):
+            _set_span_attribute(span, OPENAI_API_BASE, str(client.base_url))
+        if isinstance(client, (openai.AsyncAzureOpenAI, openai.AzureOpenAI)):
+            _set_span_attribute(span, OPENAI_API_VERSION, client._api_version)  # pylint: disable=protected-access
+
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.warning(
+            "Failed to set api attributes for openai v1 span, error: %s", str(ex)
+        )
+
+
 def _set_api_attributes(span):
     if not span.is_recording():
         return
 
+    if is_openai_v1():
+        return
+
     try:
         base_url = openai.base_url if hasattr(openai, "base_url") else openai.api_base
 

packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py

@@ -8,6 +8,7 @@
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.instrumentation.openai.utils import _with_tracer_wrapper
 from opentelemetry.instrumentation.openai.shared import (
+    _set_client_attributes,
     _set_request_attributes,
     _set_span_attribute,
     _set_functions_attributes,
@@ -39,7 +40,7 @@ def chat_wrapper(tracer, wrapped, instance, args, kwargs):
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     )
 
-    _handle_request(span, kwargs)
+    _handle_request(span, kwargs, instance)
     response = wrapped(*args, **kwargs)
 
     if is_streaming_response(response):
@@ -62,7 +63,7 @@ async def achat_wrapper(tracer, wrapped, instance, args, kwargs):
         kind=SpanKind.CLIENT,
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     )
-    _handle_request(span, kwargs)
+    _handle_request(span, kwargs, instance)
     response = await wrapped(*args, **kwargs)
 
     if is_streaming_response(response):
@@ -75,8 +76,9 @@ async def achat_wrapper(tracer, wrapped, instance, args, kwargs):
     return response
 
 
-def _handle_request(span, kwargs):
+def _handle_request(span, kwargs, instance):
     _set_request_attributes(span, kwargs)
+    _set_client_attributes(span, instance)
     if should_send_prompts():
         _set_prompts(span, kwargs.get("messages"))
         _set_functions_attributes(span, kwargs.get("functions"))

packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py

@@ -7,6 +7,7 @@
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.instrumentation.openai.utils import _with_tracer_wrapper
 from opentelemetry.instrumentation.openai.shared import (
+    _set_client_attributes,
     _set_request_attributes,
     _set_span_attribute,
     _set_functions_attributes,
@@ -39,7 +40,7 @@ def completion_wrapper(tracer, wrapped, instance, args, kwargs):
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     )
 
-    _handle_request(span, kwargs)
+    _handle_request(span, kwargs, instance)
     response = wrapped(*args, **kwargs)
 
     if is_streaming_response(response):
@@ -63,7 +64,7 @@ async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     )
 
-    _handle_request(span, kwargs)
+    _handle_request(span, kwargs, instance)
     response = await wrapped(*args, **kwargs)
 
     if is_streaming_response(response):
@@ -76,11 +77,12 @@ async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
     return response
 
 
-def _handle_request(span, kwargs):
+def _handle_request(span, kwargs, instance):
     _set_request_attributes(span, kwargs)
     if should_send_prompts():
         _set_prompts(span, kwargs.get("prompt"))
         _set_functions_attributes(span, kwargs.get("functions"))
+    _set_client_attributes(span, instance)
 
 
 def _handle_response(response, span):

packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/embeddings_wrappers.py

@@ -10,6 +10,7 @@
     start_as_current_span_async,
 )
 from opentelemetry.instrumentation.openai.shared import (
+    _set_client_attributes,
     _set_request_attributes,
     _set_span_attribute,
     _set_response_attributes,
@@ -37,7 +38,7 @@ def embeddings_wrapper(tracer, wrapped, instance, args, kwargs):
         kind=SpanKind.CLIENT,
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     ) as span:
-        _handle_request(span, kwargs)
+        _handle_request(span, kwargs, instance)
         response = wrapped(*args, **kwargs)
         _handle_response(response, span)
 
@@ -55,17 +56,18 @@ async def aembeddings_wrapper(tracer, wrapped, instance, args, kwargs):
         kind=SpanKind.CLIENT,
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     ) as span:
-        _handle_request(span, kwargs)
+        _handle_request(span, kwargs, instance)
         response = await wrapped(*args, **kwargs)
         _handle_response(response, span)
 
         return response
 
 
-def _handle_request(span, kwargs):
+def _handle_request(span, kwargs, instance):
     _set_request_attributes(span, kwargs)
     if should_send_prompts():
         _set_prompts(span, kwargs.get("input"))
+    _set_client_attributes(span, instance)
 
 
 def _handle_response(response, span):