(fix): start session no longer throws on agent ops

Author: areibman

src/multion/client.py

@@ -1,5 +1,4 @@
 import typing
-
 from .base_client import AsyncBaseMultiOn, BaseMultiOn
 from .sessions.wrapped_client import (
     WrappedAsyncSessionsClient,
@@ -52,27 +51,48 @@ 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,
             )
 
-    @agentops.record_function(event_name="browse")  # type: ignore
-    @wraps_function(BaseMultiOn.browse)
-    def browse(self, *args, **kwargs) -> BrowseOutput:
-        agentops.start_session(tags=["multion-sdk"])
+    @wraps_function(BaseMultiOn.browse)  # type: ignore
+    def _browse(self, *args, **kwargs) -> BrowseOutput:
         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"])
+    @wraps_function(BaseMultiOn.retrieve)  # type: ignore
+    def _retrieve(self, *args, **kwargs) -> RetrieveOutput:
         return super().retrieve(*args, **kwargs)
 
+    def browse(self, *args, **kwargs) -> BrowseOutput:
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
+
+            @agentops.record_function(event_name="browse")
+            def wrapped_browse(*args, **kwargs):
+                return self._browse(*args, **kwargs)
+            return wrapped_browse(*args, **kwargs)
+        return self._browse(*args, **kwargs)
+
+    def retrieve(self, *args, **kwargs) -> RetrieveOutput:
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
+
+            @agentops.record_function(event_name="retrieve")
+            def wrapped_retrieve(*args, **kwargs):
+                return self._retrieve(*args, **kwargs)
+            return wrapped_retrieve(*args, **kwargs)
+        return self._retrieve(*args, **kwargs)
+
 
 class AsyncMultiOn(AsyncBaseMultiOn):
     """
@@ -111,22 +131,40 @@ 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,
             )
 
-    @agentops.record_function(event_name="browse")  # type: ignore
-    @wraps_function(AsyncBaseMultiOn.browse)
-    async def browse(self, *args, **kwargs) -> BrowseOutput:
-        agentops.start_session(tags=["multion-sdk"])
+    @wraps_function(AsyncBaseMultiOn.browse)  # type: ignore
+    async def _browse(self, *args, **kwargs) -> BrowseOutput:
         return await super().browse(*args, **kwargs)
 
-    @agentops.record_function(event_name="retrieve")  # type: ignore
-    @wraps_function(BaseMultiOn.retrieve)
-    async def retrieve(self, *args, **kwargs) -> RetrieveOutput:
-        agentops.start_session(tags=["multion-sdk"])
+    @wraps_function(BaseMultiOn.retrieve)  # type: ignore
+    async def _retrieve(self, *args, **kwargs) -> RetrieveOutput:
         return await super().retrieve(*args, **kwargs)
+
+    async def browse(self, *args, **kwargs) -> BrowseOutput:
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
+
+            @agentops.record_function(event_name="browse")
+            async def wrapped_browse(*args, **kwargs):
+                return await self._browse(*args, **kwargs)
+            return await wrapped_browse(*args, **kwargs)
+        return await self._browse(*args, **kwargs)
+
+    async def retrieve(self, *args, **kwargs) -> RetrieveOutput:
+        if self._agentops_api_key is not None:
+            agentops.start_session(tags=["multion-sdk"])
+
+            @agentops.record_function(event_name="retrieve")
+            async def wrapped_retrieve(*args, **kwargs):
+                return await self._retrieve(*args, **kwargs)  # type: ignore
+            return await wrapped_retrieve(*args, **kwargs)
+        return await self._retrieve(*args, **kwargs)  # type: ignore

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: