fix(agents-api): set include remote to true in prepare_for_step

julep-ai · Nov 29, 2024 · b82d159 · b82d159
1 parent 5f370b3
commit b82d159
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/agents-api/agents_api/common/protocol/tasks.py b/agents-api/agents_api/common/protocol/tasks.py
@@ -243,7 +243,7 @@ def model_dump(self, *args, **kwargs) -> dict[str, Any]:
         return dump | execution_input
 
     async def prepare_for_step(
-        self, *args, include_remote: bool = False, **kwargs
+        self, *args, include_remote: bool = True, **kwargs
     ) -> dict[str, Any]:
         current_input = self.current_input
         inputs = self.inputs

diff --git a/agents-api/agents_api/common/storage_handler.py b/agents-api/agents_api/common/storage_handler.py
@@ -97,7 +97,7 @@ async def load_args(
                             }
                         )
                     elif isinstance(arg, BaseRemoteModel):
-                        new_args.append(arg.unload_all())
+                        new_args.append(await arg.unload_all())
 
                     elif isinstance(arg, BaseModel):
                         for field in arg.model_fields.keys():
@@ -119,7 +119,11 @@ async def load_args(
                                     ],
                                 )
                             elif isinstance(getattr(arg, field), BaseRemoteModel):
-                                setattr(arg, field, getattr(arg, field).unload_all())
+                                setattr(
+                                    arg,
+                                    field,
+                                    await getattr(arg, field).unload_all(),
+                                )
 
                         new_args.append(arg)
 
@@ -163,7 +167,11 @@ async def load_args(
                                     ],
                                 )
                             elif isinstance(getattr(v, field), BaseRemoteModel):
-                                setattr(v, field, getattr(v, field).unload_all())
+                                setattr(
+                                    v,
+                                    field,
+                                    await getattr(v, field).unload_all(),
+                                )
                         new_kwargs[k] = v
 
                     else: