[Feat] OpenAI Assistants (Beta) Integration

agentops/llms/providers/openai.py

@@ -1,4 +1,3 @@
-import inspect
 import pprint
 from typing import Optional
 
@@ -16,6 +15,7 @@
 class OpenAiProvider(InstrumentedProvider):
     original_create = None
     original_create_async = None
+    original_assistant_methods = None
 
     def __init__(self, client):
         super().__init__(client)
@@ -138,6 +138,7 @@ async def async_generator():
     def override(self):
         self._override_openai_v1_completion()
         self._override_openai_v1_async_completion()
+        self._override_openai_assistants_beta()
 
     def _override_openai_v1_completion(self):
         from openai.resources.chat import completions
@@ -228,9 +229,84 @@ async def patched_function(*args, **kwargs):
         # Override the original method with the patched one
         completions.AsyncCompletions.create = patched_function
 
+    def _override_openai_assistants_beta(self):
+        """Override OpenAI Assistants API methods"""
+        from openai._legacy_response import LegacyAPIResponse
+        from openai.resources import beta
+        from openai.pagination import BasePage
+
+        def handle_response(response, kwargs, init_timestamp, session: Optional[Session] = None) -> dict:
+            """Handle response based on return type"""
+            action_event = ActionEvent(init_timestamp=init_timestamp, params=kwargs)
+
+            try:
+                if session is not None:
+                    action_event.session_id = session.session_id
+
+                # Base ActionEvent for all API calls
+                if isinstance(response, BasePage):
+                    action_event.action_type = response.__class__.__name__.split("[")[1][:-1]
+                else:
+                    action_event.action_type = response.__class__.__name__
+
+                action_event.end_timestamp = get_ISO_time()
+                action_event.returns = response.model_dump()
+
+                self._safe_record(session, action_event)
+            except Exception as e:
+                self._safe_record(session, ErrorEvent(trigger_event=action_event, exception=e))
+
+                kwargs_str = pprint.pformat(kwargs)
+                response = pprint.pformat(response)
+                logger.warning(
+                    f"Unable to parse response for Assistants API. Skipping upload to AgentOps\n"
+                    f"response:\n {response}\n"
+                    f"kwargs:\n {kwargs_str}\n"
+                )
+
+            return response
+
+        def create_patched_function(original_func):
+            def patched_function(*args, **kwargs):
+                init_timestamp = get_ISO_time()
+
+                session = kwargs.get("session", None)
+                if "session" in kwargs.keys():
+                    del kwargs["session"]
+
+                response = original_func(*args, **kwargs)
+                if isinstance(response, LegacyAPIResponse):
+                    return response
+
+                return handle_response(response, kwargs, init_timestamp, session=session)
+
+            return patched_function
+
+        # Store and patch Assistant API methods
+        assistant_api_methods = {
+            beta.Assistants: ["create", "retrieve", "update", "delete", "list"],
+            beta.Threads: ["create", "retrieve", "update", "delete"],
+            beta.threads.Messages: ["create", "retrieve", "update", "list"],
+            beta.threads.Runs: ["create", "retrieve", "update", "list", "submit_tool_outputs", "cancel"],
+            beta.threads.runs.steps.Steps: ["retrieve", "list"],
+        }
+
+        self.original_assistant_methods = {
+            (cls, method): getattr(cls, method) for cls, methods in assistant_api_methods.items() for method in methods
+        }
+
+        # Override methods and verify
+        for (cls, method), original_func in self.original_assistant_methods.items():
+            patched_function = create_patched_function(original_func)
+            setattr(cls, method, patched_function)
+
     def undo_override(self):
         if self.original_create is not None and self.original_create_async is not None:
             from openai.resources.chat import completions
 
             completions.AsyncCompletions.create = self.original_create_async
             completions.Completions.create = self.original_create
+
+        if self.original_assistant_methods is not None:
+            for (cls, method), original in self.original_assistant_methods.items():
+                setattr(cls, method, original)

agentops/llms/tracker.py

@@ -27,7 +27,41 @@ class LlmTracker:
     SUPPORTED_APIS = {
         "litellm": {"1.3.1": ("openai_chat_completions.completion",)},
         "openai": {
-            "1.0.0": ("chat.completions.create",),
+            "1.0.0": (
+                "chat.completions.create",
+                # Assistants
+                "beta.assistants.create",
+                "beta.assistants.retrieve",
+                "beta.assistants.update",
+                "beta.assistants.delete",
+                "beta.assistants.list",
+                "beta.assistants.files.create",
+                "beta.assistants.files.retrieve",
+                "beta.assistants.files.delete",
+                "beta.assistants.files.list",
+                # Threads
+                "beta.threads.create",
+                "beta.threads.retrieve",
+                "beta.threads.update",
+                "beta.threads.delete",
+                # Messages
+                "beta.threads.messages.create",
+                "beta.threads.messages.retrieve",
+                "beta.threads.messages.update",
+                "beta.threads.messages.list",
+                "beta.threads.messages.files.retrieve",
+                "beta.threads.messages.files.list",
+                # Runs
+                "beta.threads.runs.create",
+                "beta.threads.runs.retrieve",
+                "beta.threads.runs.update",
+                "beta.threads.runs.list",
+                "beta.threads.runs.cancel",
+                "beta.threads.runs.submit_tool_outputs",
+                # Run Steps
+                "beta.threads.runs.steps.Steps.retrieve",
+                "beta.threads.runs.steps.Steps.list",
+            ),
             "0.0.0": (
                 "ChatCompletion.create",
                 "ChatCompletion.acreate",

docs/v1/examples/examples.mdx

@@ -20,6 +20,10 @@ mode: "wide"
   <Card title="Multi Session" icon="computer" href="/v1/examples/multi_session">
     Manage multiple sessions at the same time
   </Card>
+
+  <Card title="OpenAI Assistants" icon={<img src="https://www.github.com/agentops-ai/agentops/blob/main/docs/images/external/openai/openai-logomark.png?raw=true" alt="OpenAI Assistants" />} iconType="image" href="/v1/integrations/openai" href="/v1/examples/openai_assistants">
+    Observe OpenAI Assistants
+  </Card>
 </CardGroup>
 
 ### Integration Examples