Async Agent Interface Refactor (#2467)

Signed-off-by: Jan Fiedler <jan@union.ai>
flyteorg · Jul 25, 2024 · 7764cbf · 7764cbf
1 parent 04cbf01
commit 7764cbf
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py
@@ -152,8 +152,8 @@ def metadata_type(self) -> ResourceMeta:
     def create(
         self,
         task_template: TaskTemplate,
+        output_prefix: str,
         inputs: Optional[LiteralMap],
-        output_prefix: Optional[str],
         task_execution_metadata: Optional[TaskExecutionMetadata],
         **kwargs,
     ) -> ResourceMeta:
@@ -297,7 +297,9 @@ def execute(self: PythonTask, **kwargs) -> LiteralMap:
         task_template = get_serializable(OrderedDict(), ss, self).template
         self._agent = AgentRegistry.get_agent(task_template.type, task_template.task_type_version)
 
-        resource_mata = asyncio.run(self._create(task_template, output_prefix, kwargs))
+        resource_mata = asyncio.run(
+            self._create(task_template=task_template, output_prefix=output_prefix, inputs=kwargs)
+        )
         resource = asyncio.run(self._get(resource_meta=resource_mata))
 
         if resource.phase != TaskExecution.SUCCEEDED: