Add Initial Support for Instrumenting OpenAI Python Library - Chat Completion Create

instrumentation/opentelemetry-instrumentation-openai/README.rst

@@ -1,5 +1,5 @@
 OpenTelemetry OpenAI Instrumentation
-===================================
+====================================    
 
 |pypi|
 

instrumentation/opentelemetry-instrumentation-openai/pyproject.toml

@@ -26,12 +26,15 @@ classifiers = [
 ]
 dependencies = [
   "opentelemetry-api ~= 1.12",
-  "opentelemetry-instrumentation == 0.48b0.dev",
+  "opentelemetry-instrumentation == 0.47b0",
+  "tiktoken>=0.1.1",
+  "pydantic>=1.8"
+
 ]
 
 [project.optional-dependencies]
 instruments = [
-  "openai ~= 1.37.1",
+  "openai >= 0.27.0",
 ]
 
 [project.entry-points.opentelemetry_instrumentor]

instrumentation/opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/__init__.py

@@ -47,9 +47,7 @@
 from opentelemetry.instrumentation.openai.package import _instruments
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
-from langtrace_python_sdk.instrumentation.openai.patch import (
-    chat_completions_create
-)
+from .patch import chat_completions_create
 
 
 class OpenAIInstrumentor(BaseInstrumentor):
@@ -58,16 +56,16 @@ def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
-        """Enable OpenAI instrumentation.
-        """
+        """Enable OpenAI instrumentation."""
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, "", tracer_provider)
         version = importlib.metadata.version("openai")
-
         wrap_function_wrapper(
             "openai.resources.chat.completions",
             "Completions.create",
-            chat_completions_create("openai.chat.completions.create", version, tracer),
+            chat_completions_create(
+                "openai.chat.completions.create", version, tracer
+            ),
         )
 
     def _uninstrument(self, **kwargs):

instrumentation/opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/package.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 
-_instruments = ("openai ~= 1.37.1",)
+_instruments = ("openai >= 0.27.0",)

instrumentation/opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/patch.py

@@ -13,16 +13,19 @@
 # limitations under the License.
 
 import json
+from typing import Optional, Union
 from opentelemetry import trace
 from opentelemetry.trace import SpanKind, Span
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.trace.propagation import set_span_in_context
-from openai._types import NOT_GIVEN
-from span_attributes import SpanAttributes, LLMSpanAttributes, Event
-from utils import estimate_tokens, silently_fail, extract_content, calculate_prompt_tokens
+from openai import NOT_GIVEN
+from .span_attributes import LLMSpanAttributes, SpanAttributes
 
+from .utils import silently_fail, extract_content
+from opentelemetry.trace import Tracer
 
-def chat_completions_create(original_method, version, tracer):
+
+def chat_completions_create(original_method, version, tracer: Tracer):
     """Wrap the `create` method of the `ChatCompletion` class to trace it."""
 
     def traced_method(wrapped, instance, args, kwargs):
@@ -34,7 +37,11 @@ def traced_method(wrapped, instance, args, kwargs):
                 for tool_call in tools:
                     tool_call_dict = {
                         "id": tool_call.id if hasattr(tool_call, "id") else "",
-                        "type": tool_call.type if hasattr(tool_call, "type") else "",
+                        "type": (
+                            tool_call.type
+                            if hasattr(tool_call, "type")
+                            else ""
+                        ),
                     }
                     if hasattr(tool_call, "function"):
                         tool_call_dict["function"] = {
@@ -60,8 +67,10 @@ def traced_method(wrapped, instance, args, kwargs):
 
         attributes = LLMSpanAttributes(**span_attributes)
 
+        span_name = f"{attributes.gen_ai_operation_name} {attributes.gen_ai_request_model}"
+
         span = tracer.start_span(
-            "openai.completion",
+            name=span_name,
             kind=SpanKind.CLIENT,
             context=set_span_in_context(trace.get_current_span()),
         )
@@ -70,36 +79,18 @@ def traced_method(wrapped, instance, args, kwargs):
         try:
             result = wrapped(*args, **kwargs)
             if is_streaming(kwargs):
-                prompt_tokens = 0
-                for message in kwargs.get("messages", {}):
-                    prompt_tokens += calculate_prompt_tokens(
-                        json.dumps(str(message)), kwargs.get("model")
-                    )
-
-                if (
-                    kwargs.get("functions") is not None
-                    and kwargs.get("functions") != NOT_GIVEN
-                ):
-                    for function in kwargs.get("functions"):
-                        prompt_tokens += calculate_prompt_tokens(
-                            json.dumps(function), kwargs.get("model")
-                        )
-
                 return StreamWrapper(
                     result,
                     span,
-                    prompt_tokens,
                     function_call=kwargs.get("functions") is not None,
                     tool_calls=kwargs.get("tools") is not None,
                 )
             else:
                 _set_response_attributes(span, kwargs, result)
-                span.set_status(StatusCode.OK)
                 span.end()
                 return result
 
         except Exception as error:
-            span.record_exception(error)
             span.set_status(Status(StatusCode.ERROR, str(error)))
             span.end()
             raise
@@ -109,32 +100,33 @@ def traced_method(wrapped, instance, args, kwargs):
 
 def get_tool_calls(item):
     if isinstance(item, dict):
-        if "tool_calls" in item and item["tool_calls"] is not None:
-            return item["tool_calls"]
-        return None
-
+        return item.get("tool_calls")
     else:
-        if hasattr(item, "tool_calls") and item.tool_calls is not None:
-            return item.tool_calls
-        return None
+        return getattr(item, "tool_calls", None)
 
 
 @silently_fail
-def _set_input_attributes(span, kwargs, attributes):
+def _set_input_attributes(span, kwargs, attributes: LLMSpanAttributes):
     tools = []
-    for field, value in attributes.model_dump(by_alias=True).items():
-        set_span_attribute(span, field, value)
 
-    if kwargs.get("functions") is not None and kwargs.get("functions") != NOT_GIVEN:
+    if (
+        kwargs.get("functions") is not None
+        and kwargs.get("functions") != NOT_GIVEN
+    ):
         for function in kwargs.get("functions"):
-            tools.append(json.dumps({"type": "function", "function": function}))
+            tools.append(
+                json.dumps({"type": "function", "function": function})
+            )
 
     if kwargs.get("tools") is not None and kwargs.get("tools") != NOT_GIVEN:
         tools.append(json.dumps(kwargs.get("tools")))
 
     if tools:
         set_span_attribute(span, SpanAttributes.LLM_TOOLS, json.dumps(tools))
 
+    for field, value in attributes.model_dump(by_alias=True).items():
+        set_span_attribute(span, field, value)
+
 
 @silently_fail
 def _set_response_attributes(span, kwargs, result):
@@ -149,7 +141,11 @@ def _set_response_attributes(span, kwargs, result):
                 ),
                 "content": extract_content(choice),
                 **(
-                    {"content_filter_results": choice["content_filter_results"]}
+                    {
+                        "content_filter_results": choice[
+                            "content_filter_results"
+                        ]
+                    }
                     if "content_filter_results" in choice
                     else {}
                 ),
@@ -212,15 +208,6 @@ def set_event_completion(span: Span, result_content):
     )
 
 
-def set_event_completion_chunk(span: Span, chunk):
-    span.add_event(
-        name=SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK,
-        attributes={
-            SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK: json.dumps(chunk),
-        },
-    )
-
-
 def set_span_attribute(span: Span, name, value):
     if value is not None:
         if value != "" or value != NOT_GIVEN:
@@ -232,31 +219,33 @@ def set_span_attribute(span: Span, name, value):
 
 
 def is_streaming(kwargs):
-    return not (
-        kwargs.get("stream") is False
-        or kwargs.get("stream") is None
-        or kwargs.get("stream") == NOT_GIVEN
-    )
+    return non_numerical_value_is_set(kwargs.get("stream"))
+
+
+def non_numerical_value_is_set(value: Optional[Union[bool, str]]):
+    return bool(value) and value != NOT_GIVEN
 
 
-def get_llm_request_attributes(kwargs, prompts=None, model=None, operation_name="chat"):
+def get_llm_request_attributes(
+    kwargs, prompts=None, model=None, operation_name="chat"
+):
 
-    user = kwargs.get("user", None)
+    user = kwargs.get("user")
     if prompts is None:
         prompts = (
             [{"role": user or "user", "content": kwargs.get("prompt")}]
             if "prompt" in kwargs
             else None
         )
     top_k = (
-        kwargs.get("n", None)
-        or kwargs.get("k", None)
-        or kwargs.get("top_k", None)
-        or kwargs.get("top_n", None)
+        kwargs.get("n")
+        or kwargs.get("k")
+        or kwargs.get("top_k")
+        or kwargs.get("top_n")
     )
 
-    top_p = kwargs.get("p", None) or kwargs.get("top_p", None)
-    tools = kwargs.get("tools", None)
+    top_p = kwargs.get("p") or kwargs.get("top_p")
+    tools = kwargs.get("tools")
     return {
         SpanAttributes.LLM_OPERATION_NAME: operation_name,
         SpanAttributes.LLM_REQUEST_MODEL: model or kwargs.get("model"),
@@ -267,7 +256,9 @@ def get_llm_request_attributes(kwargs, prompts=None, model=None, operation_name=
         SpanAttributes.LLM_USER: user,
         SpanAttributes.LLM_REQUEST_TOP_P: top_p,
         SpanAttributes.LLM_REQUEST_MAX_TOKENS: kwargs.get("max_tokens"),
-        SpanAttributes.LLM_SYSTEM_FINGERPRINT: kwargs.get("system_fingerprint"),
+        SpanAttributes.LLM_SYSTEM_FINGERPRINT: kwargs.get(
+            "system_fingerprint"
+        ),
         SpanAttributes.LLM_PRESENCE_PENALTY: kwargs.get("presence_penalty"),
         SpanAttributes.LLM_FREQUENCY_PENALTY: kwargs.get("frequency_penalty"),
         SpanAttributes.LLM_REQUEST_SEED: kwargs.get("seed"),
@@ -283,7 +274,12 @@ class StreamWrapper:
     span: Span
 
     def __init__(
-        self, stream, span, prompt_tokens, function_call=False, tool_calls=False
+        self,
+        stream,
+        span,
+        prompt_tokens=None,
+        function_call=False,
+        tool_calls=False,
     ):
         self.stream = stream
         self.span = span
@@ -297,12 +293,10 @@ def __init__(
 
     def setup(self):
         if not self._span_started:
-            self.span.add_event(Event.STREAM_START.value)
             self._span_started = True
 
     def cleanup(self):
         if self._span_started:
-            self.span.add_event(Event.STREAM_END.value)
             set_span_attribute(
                 self.span,
                 SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
@@ -339,13 +333,6 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.cleanup()
 
-    async def __aenter__(self):
-        self.setup()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        self.cleanup()
-
     def __iter__(self):
         return self
 
@@ -358,18 +345,6 @@ def __next__(self):
             self.cleanup()
             raise
 
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self):
-        try:
-            chunk = await self.stream.__anext__()
-            self.process_chunk(chunk)
-            return chunk
-        except StopAsyncIteration:
-            self.cleanup()
-            raise StopAsyncIteration
-
     def process_chunk(self, chunk):
         if hasattr(chunk, "model") and chunk.model is not None:
             set_span_attribute(
@@ -383,8 +358,6 @@ def process_chunk(self, chunk):
             if not self.function_call and not self.tool_calls:
                 for choice in chunk.choices:
                     if choice.delta and choice.delta.content is not None:
-                        token_counts = estimate_tokens(choice.delta.content)
-                        self.completion_tokens += token_counts
                         content = [choice.delta.content]
             elif self.function_call:
                 for choice in chunk.choices:
@@ -393,10 +366,6 @@ def process_chunk(self, chunk):
                         and choice.delta.function_call is not None
                         and choice.delta.function_call.arguments is not None
                     ):
-                        token_counts = estimate_tokens(
-                            choice.delta.function_call.arguments
-                        )
-                        self.completion_tokens += token_counts
                         content = [choice.delta.function_call.arguments]
             elif self.tool_calls:
                 for choice in chunk.choices:
@@ -409,30 +378,17 @@ def process_chunk(self, chunk):
                                 and tool_call.function is not None
                                 and tool_call.function.arguments is not None
                             ):
-                                token_counts = estimate_tokens(
-                                    tool_call.function.arguments
-                                )
-                                self.completion_tokens += token_counts
                                 content.append(tool_call.function.arguments)
-            set_event_completion_chunk(
-                self.span,
-                "".join(content) if len(content) > 0 and content[0] is not None else "",
-            )
+
             if content:
                 self.result_content.append(content[0])
 
         if hasattr(chunk, "text"):
-            token_counts = estimate_tokens(chunk.text)
-            self.completion_tokens += token_counts
             content = [chunk.text]
-            set_event_completion_chunk(
-                self.span,
-                "".join(content) if len(content) > 0 and content[0] is not None else "",
-            )
 
             if content:
                 self.result_content.append(content[0])
 
-        if hasattr(chunk, "usage_metadata"):
-            self.completion_tokens = chunk.usage_metadata.candidates_token_count
-            self.prompt_tokens = chunk.usage_metadata.prompt_token_count
+        if getattr(chunk, "usage"):
+            self.completion_tokens = chunk.usage.completion_tokens
+            self.prompt_tokens = chunk.usage.prompt_tokens