fix: Fix workflows serialization

julep-ai · Jul 24, 2024 · 72f6782 · 72f6782
1 parent 5ad3254
commit 72f6782
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 5 deletions.
diff --git a/agents-api/agents_api/activities/task_steps/__init__.py b/agents-api/agents_api/activities/task_steps/__init__.py
@@ -119,7 +119,7 @@ async def if_else_step(context: StepContext) -> dict:
         else context.definition.else_
     )
 
-    return {"workflow": next_workflow}
+    return {"goto_workflow": next_workflow}
 
 
 @activity.defn

diff --git a/agents-api/agents_api/common/utils/json.py b/agents-api/agents_api/common/utils/json.py
@@ -3,6 +3,7 @@
 import json
 from uuid import UUID
 from typing import Any
+from pydantic import BaseModel
 
 
 class CustomJSONEncoder(json.JSONEncoder):
@@ -37,6 +38,9 @@ def default(self, obj):
         if isinstance(obj, UUID):
             return str(obj)
 
+        if isinstance(obj, BaseModel):
+            return obj.model_dump()
+
         return obj
 
 

diff --git a/agents-api/agents_api/routers/agents/create_agent.py b/agents-api/agents_api/routers/agents/create_agent.py
@@ -8,7 +8,6 @@
 from ...dependencies.developer_id import get_developer_id
 from ...models.agent.create_agent import create_agent_query
 from ...autogen.openapi_model import CreateAgentRequest, ResourceCreatedResponse
-from ...common.utils.datetime import utcnow
 
 from .router import router
 

diff --git a/agents-api/agents_api/worker/__main__.py b/agents-api/agents_api/worker/__main__.py
@@ -24,7 +24,7 @@
     yield_step,
     # tool_call_step,
     # error_step,
-    # if_else_step,
+    if_else_step,
     transition_step,
 )
 
@@ -80,7 +80,7 @@ async def main():
         yield_step,
         # tool_call_step,
         # error_step,
-        # if_else_step,
+        if_else_step,
         transition_step,
     ]
 

diff --git a/agents-api/agents_api/workflows/task_execution.py b/agents-api/agents_api/workflows/task_execution.py
@@ -97,7 +97,7 @@ async def run(
                     context,
                     schedule_to_close_timeout=timedelta(seconds=600),
                 )
-                workflow_step: YieldWorkflowStep = outputs["workflow"]
+                workflow_step = YieldWorkflowStep(**outputs["goto_workflow"])
 
                 outputs = await workflow.execute_child_workflow(
                     TaskExecutionWorkflow.run,