From c17617f8bb73f3a374a2b8a1d2b06222bb1cf672 Mon Sep 17 00:00:00 2001
From: Dmitry Paramonov <asmatic075@gmail.com>
Date: Mon, 15 Jul 2024 22:34:36 +0300
Subject: [PATCH] WIP

---
 .../activities/task_steps/__init__.py         | 26 ++++++++--------
 .../agents_api/workflows/task_execution.py    | 30 ++++++++++---------
 2 files changed, 29 insertions(+), 27 deletions(-)

diff --git a/agents-api/agents_api/activities/task_steps/__init__.py b/agents-api/agents_api/activities/task_steps/__init__.py
index e376c4782..25fbabce2 100644
--- a/agents-api/agents_api/activities/task_steps/__init__.py
+++ b/agents-api/agents_api/activities/task_steps/__init__.py
@@ -9,9 +9,9 @@
     PromptWorkflowStep,
     # EvaluateWorkflowStep,
     # YieldWorkflowStep,
-    # ToolCallWorkflowStep,
+    ToolCallWorkflowStep,
     # ErrorWorkflowStep,
-    # IfElseWorkflowStep,
+    IfElseWorkflowStep,
     InputChatMLMessage,
 )
 
@@ -89,14 +89,14 @@ async def prompt_step(context: StepContext) -> dict:
 #     return {"test": "result"}
 
 
-# @activity.defn
-# async def tool_call_step(context: StepContext) -> dict:
-#     if not isinstance(context.definition, ToolCallWorkflowStep):
-#         return {}
+@activity.defn
+async def tool_call_step(context: StepContext) -> dict:
+    if not isinstance(context.definition, ToolCallWorkflowStep):
+        return {}
 
-#     # TODO: implement
+    # TODO: implement
 
-#     return {"test": "result"}
+    return {"test": "result"}
 
 
 # @activity.defn
@@ -107,12 +107,12 @@ async def prompt_step(context: StepContext) -> dict:
 #     return {"error": context.definition.error}
 
 
-# @activity.defn
-# async def if_else_step(context: StepContext) -> dict:
-#     if not isinstance(context.definition, IfElseWorkflowStep):
-#         return {}
+@activity.defn
+async def if_else_step(context: StepContext) -> dict:
+    if not isinstance(context.definition, IfElseWorkflowStep):
+        return {}
 
-#     return {"test": "result"}
+    return {"test": "result"}
 
 
 @activity.defn
diff --git a/agents-api/agents_api/workflows/task_execution.py b/agents-api/agents_api/workflows/task_execution.py
index 5daa53d95..dd2925ddb 100644
--- a/agents-api/agents_api/workflows/task_execution.py
+++ b/agents-api/agents_api/workflows/task_execution.py
@@ -9,6 +9,8 @@
     from ..activities.task_steps import (
         prompt_step,
         transition_step,
+        tool_call_step,
+        if_else_step,
     )
 
     from ..common.protocol.tasks import (
@@ -16,9 +18,9 @@
         PromptWorkflowStep,
         # EvaluateWorkflowStep,
         # YieldWorkflowStep,
-        # ToolCallWorkflowStep,
+        ToolCallWorkflowStep,
         # ErrorWorkflowStep,
-        # IfElseWorkflowStep,
+        IfElseWorkflowStep,
         StepContext,
         TransitionInfo,
     )
@@ -81,24 +83,24 @@ async def run(
             #         context,
             #         schedule_to_close_timeout=timedelta(seconds=600),
             #     )
-            # case ToolCallWorkflowStep():
-            #     result = await workflow.execute_activity(
-            #         tool_call_step,
-            #         context,
-            #         schedule_to_close_timeout=timedelta(seconds=600),
-            #     )
+            case ToolCallWorkflowStep():
+                outputs = await workflow.execute_activity(
+                    tool_call_step,
+                    context,
+                    schedule_to_close_timeout=timedelta(seconds=600),
+                )
             # case ErrorWorkflowStep():
             #     result = await workflow.execute_activity(
             #         error_step,
             #         context,
             #         schedule_to_close_timeout=timedelta(seconds=600),
             #     )
-            # case IfElseWorkflowStep():
-            #     result = await workflow.execute_activity(
-            #         if_else_step,
-            #         context,
-            #         schedule_to_close_timeout=timedelta(seconds=600),
-            #     )
+            case IfElseWorkflowStep():
+                outputs = await workflow.execute_activity(
+                    if_else_step,
+                    context,
+                    schedule_to_close_timeout=timedelta(seconds=600),
+                )
 
         # Transition type
         transition_type = (