(fix): AgentOps is only instantiated if an api key is passed (#22)

Author: areibman

src/multion/client.py

@@ -15,6 +15,8 @@
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
 
+record_function = lambda *args, **kwargs: lambda fn: fn
+
 
 class MultiOn(BaseMultiOn):
     """
@@ -52,25 +54,37 @@ def __init__(
         agentops_api_key: typing.Optional[str] = os.getenv("AGENTOPS_API_KEY"),
         **kwargs
     ):
+
         super().__init__(*args, **kwargs)
-        self.sessions = WrappedSessionsClient(client_wrapper=self._client_wrapper)
-        if agentops_api_key is not None:
+        self._agentops_api_key = agentops_api_key
+
+        self.sessions = WrappedSessionsClient(
+            client_wrapper=self._client_wrapper, use_agentops=self._agentops_api_key is not None)
+
+        if self._agentops_api_key is not None:
             agentops.init(
                 api_key=agentops_api_key,
                 parent_key=os.getenv("AGENTOPS_PARENT_KEY"),
                 auto_start_session=False,
             )
+            # Redefine the record_function globally
+            global record_function
+            record_function = agentops.record_function
+            # Apply the decorator at runtime
+            self.browse = record_function(event_name="browse")(self.browse)
+            self.retrieve = record_function(
+                event_name="retrieve")(self.retrieve)
 
-    @agentops.record_function(event_name="browse")  # type: ignore
     @wraps_function(BaseMultiOn.browse)
     def browse(self, *args, **kwargs) -> BrowseOutput:
-        agentops.start_session(tags=["multion-sdk"])
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
         return super().browse(*args, **kwargs)
 
-    @agentops.record_function(event_name="retrieve")  # type: ignore
     @wraps_function(BaseMultiOn.retrieve)
     def retrieve(self, *args, **kwargs) -> RetrieveOutput:
-        agentops.start_session(tags=["multion-sdk"])
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
         return super().retrieve(*args, **kwargs)
 
 
@@ -111,22 +125,31 @@ def __init__(
         **kwargs
     ):
         super().__init__(*args, **kwargs)
-        self.sessions = WrappedAsyncSessionsClient(client_wrapper=self._client_wrapper)
+        self._agentops_api_key = agentops_api_key
+        self.sessions = WrappedAsyncSessionsClient(client_wrapper=self._client_wrapper,
+                                                   use_agentops=self._agentops_api_key is not None)
         if agentops_api_key is not None:
             agentops.init(
                 api_key=agentops_api_key,
                 parent_key=os.getenv("AGENTOPS_PARENT_KEY"),
                 auto_start_session=False,
             )
+            # If agentops is set, we need to wrap the record function
+            global record_function
+            record_function = agentops.record_function
+
+        record_function = agentops.record_function
 
-    @agentops.record_function(event_name="browse")  # type: ignore
+    @record_function(event_name="browse")  # type: ignore
     @wraps_function(AsyncBaseMultiOn.browse)
     async def browse(self, *args, **kwargs) -> BrowseOutput:
-        agentops.start_session(tags=["multion-sdk"])
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
         return await super().browse(*args, **kwargs)
 
-    @agentops.record_function(event_name="retrieve")  # type: ignore
+    @record_function(event_name="retrieve")  # type: ignore
     @wraps_function(BaseMultiOn.retrieve)
     async def retrieve(self, *args, **kwargs) -> RetrieveOutput:
-        agentops.start_session(tags=["multion-sdk"])
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
         return await super().retrieve(*args, **kwargs)

src/multion/sessions/wrapped_client.py

@@ -1,3 +1,4 @@
+from ..wrappers import wraps_function
 from multion.sessions.client import AsyncSessionsClient, SessionsClient
 import agentops  # type: ignore
 from agentops import ActionEvent, LLMEvent, ErrorEvent  # type: ignore
@@ -11,13 +12,16 @@
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
 
-from ..wrappers import wraps_function
-
 
 class WrappedSessionsClient(SessionsClient):
+    def __init__(self, use_agentops: bool = False, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.use_agentops = use_agentops
+
     @wraps_function(SessionsClient.create)  # type: ignore
     def create(self, *args, **kwargs) -> SessionCreated:
-        agentops.start_session(tags=["multion-sdk"])
+        if self.use_agentops:
+            agentops.start_session(tags=["multion-sdk"])
         try:
             return super().create(*args, **kwargs)
         except Exception as e:
@@ -68,9 +72,14 @@ def close(self, *args, **kwargs) -> SessionsCloseResponse:
 
 
 class WrappedAsyncSessionsClient(AsyncSessionsClient):
+    def __init__(self, use_agentops: bool = False, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.use_agentops = use_agentops
+
     @wraps_function(AsyncSessionsClient.create)  # type: ignore
     async def create(self, *args, **kwargs) -> SessionCreated:
-        agentops.start_session(tags=["multion-sdk"])
+        if self.use_agentops:
+            agentops.start_session(tags=["multion-sdk"])
         try:
             return await super().create(*args, **kwargs)
         except Exception as e: