Moved retrieve due to sdk change

Author: HowieG

src/multion/client.py

@@ -9,6 +9,7 @@
 import os
 
 from .wrappers import wraps_function
+from .types.retrieve_output import RetrieveOutput
 
 
 # this is used as the default value for optional parameters
@@ -66,6 +67,12 @@ def browse(self, *args, **kwargs):
         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"])
+        return super().retrieve(*args, **kwargs)
+
 
 class AsyncMultiOn(AsyncBaseMultiOn):
     """
@@ -117,3 +124,9 @@ def __init__(
     async def browse(self, *args, **kwargs):
         agentops.start_session(tags=["multion-sdk"])
         return 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"])
+        return super().retrieve(*args, **kwargs)

src/multion/sessions/wrapped_client.py

@@ -3,7 +3,6 @@
 from agentops import ActionEvent, LLMEvent, ErrorEvent  # type: ignore
 
 import typing
-from ..types.retrieve_output import RetrieveOutput
 from ..types.session_created import SessionCreated
 from ..types.session_step_stream_chunk import SessionStepStreamChunk
 from ..types.session_step_success import SessionStepSuccess
@@ -14,6 +13,7 @@
 
 from ..wrappers import wraps_function
 
+
 class WrappedSessionsClient(SessionsClient):
     @wraps_function(SessionsClient.create)  # type: ignore
     def create(self, *args, **kwargs) -> SessionCreated:
@@ -50,7 +50,7 @@ def generator():
                 yield chunk
 
         return generator()
-    
+
     @agentops.record_function(event_name="step")  # type: ignore
     @wraps_function(SessionsClient.step)
     def step(self, *args, **kwargs) -> SessionStepSuccess:
@@ -66,10 +66,6 @@ def close(self, *args, **kwargs) -> SessionsCloseResponse:
         agentops.end_session("Success")
         return close_response
 
-    @agentops.record_function(event_name="retrieve")  # type: ignore
-    @wraps_function(SessionsClient.retrieve)
-    def retrieve(self, *args, **kwargs) -> RetrieveOutput:
-        return super().retrieve(*args, **kwargs)
 
 class WrappedAsyncSessionsClient(AsyncSessionsClient):
     @wraps_function(AsyncSessionsClient.create)  # type: ignore