openai: apply instrumentation to beta.chat.Completions.parse (elastic#65)

Author: anuraaga

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

@@ -99,7 +99,9 @@ def _instrument(self, **kwargs):
 
         register_post_import_hook(self._patch, "openai")
 
-    def _patch(self, _module):
+    def _patch(self, module):
+        version = tuple([int(x) for x in getattr(getattr(module, "version"), "VERSION").split(".")])
+        self.beta_chat_available = version >= (1, 40, 0)
         wrap_function_wrapper(
             "openai.resources.chat.completions",
             "Completions.create",
@@ -110,6 +112,17 @@ def _patch(self, _module):
             "AsyncCompletions.create",
             self._async_chat_completion_wrapper,
         )
+        if self.beta_chat_available:
+            wrap_function_wrapper(
+                "openai.resources.beta.chat.completions",
+                "Completions.parse",
+                self._chat_completion_wrapper,
+            )
+            wrap_function_wrapper(
+                "openai.resources.beta.chat.completions",
+                "AsyncCompletions.parse",
+                self._async_chat_completion_wrapper,
+            )
         wrap_function_wrapper(
             "openai.resources.embeddings",
             "Embeddings.create",
@@ -128,11 +141,14 @@ def _uninstrument(self, **kwargs):
 
         unwrap(openai.resources.chat.completions.Completions, "create")
         unwrap(openai.resources.chat.completions.AsyncCompletions, "create")
+        if self.beta_chat_available:
+            unwrap(openai.resources.beta.chat.completions.Completions, "parse")
+            unwrap(openai.resources.beta.chat.completions.AsyncCompletions, "parse")
         unwrap(openai.resources.embeddings.Embeddings, "create")
         unwrap(openai.resources.embeddings.AsyncEmbeddings, "create")
 
     def _chat_completion_wrapper(self, wrapped, instance, args, kwargs):
-        logger.debug(f"openai.resources.chat.completions.Completions.create kwargs: {kwargs}")
+        logger.debug(f"{wrapped} kwargs: {kwargs}")
 
         span_attributes = _get_attributes_from_wrapper(instance, kwargs)
         event_attributes = _get_event_attributes()