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

Author: fern-support

src/multion/client.py

@@ -15,8 +15,6 @@
 # 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):
     """
@@ -54,37 +52,25 @@ def __init__(
         agentops_api_key: typing.Optional[str] = os.getenv("AGENTOPS_API_KEY"),
         **kwargs
     ):
-
         super().__init__(*args, **kwargs)
-        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:
+        self.sessions = WrappedSessionsClient(client_wrapper=self._client_wrapper)
+        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,
             )
-            # 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:
-        if self._agentops_api_key is not None:
-            agentops.start_session(tags=["multion-sdk"])
+        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:
-        if self._agentops_api_key is not None:
-            agentops.start_session(tags=["multion-sdk"])
+        agentops.start_session(tags=["multion-sdk"])
         return super().retrieve(*args, **kwargs)
 
 
@@ -125,31 +111,22 @@ def __init__(
         **kwargs
     ):
         super().__init__(*args, **kwargs)
-        self._agentops_api_key = agentops_api_key
-        self.sessions = WrappedAsyncSessionsClient(client_wrapper=self._client_wrapper,
-                                                   use_agentops=self._agentops_api_key is not None)
+        self.sessions = WrappedAsyncSessionsClient(client_wrapper=self._client_wrapper)
         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
 
-    @record_function(event_name="browse")  # type: ignore
+    @agentops.record_function(event_name="browse")  # type: ignore
     @wraps_function(AsyncBaseMultiOn.browse)
     async def browse(self, *args, **kwargs) -> BrowseOutput:
-        if self._agentops_api_key is not None:
-            agentops.start_session(tags=["multion-sdk"])
+        agentops.start_session(tags=["multion-sdk"])
         return await super().browse(*args, **kwargs)
 
-    @record_function(event_name="retrieve")  # type: ignore
+    @agentops.record_function(event_name="retrieve")  # type: ignore
     @wraps_function(BaseMultiOn.retrieve)
     async def retrieve(self, *args, **kwargs) -> RetrieveOutput:
-        if self._agentops_api_key is not None:
-            agentops.start_session(tags=["multion-sdk"])
+        agentops.start_session(tags=["multion-sdk"])
         return await super().retrieve(*args, **kwargs)

src/multion/sessions/wrapped_client.py

@@ -1,4 +1,3 @@
-from ..wrappers import wraps_function
 from multion.sessions.client import AsyncSessionsClient, SessionsClient
 import agentops  # type: ignore
 from agentops import ActionEvent, LLMEvent, ErrorEvent  # type: ignore
@@ -12,16 +11,13 @@
 # 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
 
+class WrappedSessionsClient(SessionsClient):
     @wraps_function(SessionsClient.create)  # type: ignore
     def create(self, *args, **kwargs) -> SessionCreated:
-        if self.use_agentops:
-            agentops.start_session(tags=["multion-sdk"])
+        agentops.start_session(tags=["multion-sdk"])
         try:
             return super().create(*args, **kwargs)
         except Exception as e:
@@ -72,14 +68,9 @@ 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:
-        if self.use_agentops:
-            agentops.start_session(tags=["multion-sdk"])
+        agentops.start_session(tags=["multion-sdk"])
         try:
             return await super().create(*args, **kwargs)
         except Exception as e: