feat(agents-api): Add WaitForInputStep definition

Signed-off-by: Diwank Tomer <diwank@julep.ai>
julep-ai · Aug 2, 2024 · fdcfe57 · fdcfe57
1 parent a1da3e6
commit fdcfe57
Show file tree

Hide file tree

Showing 48 changed files with 2,356 additions and 735 deletions.
diff --git a/agents-api/agents_api/autogen/Tasks.py b/agents-api/agents_api/autogen/Tasks.py
@@ -14,6 +14,18 @@
 from .Tools import FunctionDef
 
 
+class BaseWorkflowStep(BaseModel):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    kind_: Literal[
+        "tool_call", "yield", "prompt", "evaluate", "if_else", "wait_for_input", "error"
+    ]
+    """
+    The kind of step
+    """
+
+
 class CreateTaskRequest(BaseModel):
     """
     Payload for creating a task
@@ -30,6 +42,7 @@ class CreateTaskRequest(BaseModel):
         | YieldStep
         | PromptStep
         | ErrorWorkflowStep
+        | WaitForInputStep
         | IfElseWorkflowStep
     ]
     """
@@ -50,40 +63,55 @@ class CreateTaskRequest(BaseModel):
     metadata: dict[str, Any] | None = None
 
 
-class ErrorWorkflowStep(BaseModel):
+class ErrorWorkflowStep(BaseWorkflowStep):
     model_config = ConfigDict(
         populate_by_name=True,
     )
+    kind_: Literal["error"] = "error"
     error: str
     """
     The error message
     """
 
 
-class EvaluateStep(BaseModel):
+class EvaluateStep(BaseWorkflowStep):
     model_config = ConfigDict(
         populate_by_name=True,
     )
+    kind_: Literal["evaluate"] = "evaluate"
     evaluate: dict[str, str]
     """
     The expression to evaluate
     """
 
 
-class IfElseWorkflowStep(BaseModel):
+class IfElseWorkflowStep(BaseWorkflowStep):
     model_config = ConfigDict(
         populate_by_name=True,
     )
+    kind_: Literal["if_else"] = "if_else"
     if_: Annotated[str, Field(alias="if")]
     """
     The condition to evaluate
     """
-    then: EvaluateStep | ToolCallStep | YieldStep | PromptStep | ErrorWorkflowStep
+    then: (
+        Any
+        | ToolCallStep
+        | YieldStep
+        | PromptStep
+        | ErrorWorkflowStep
+        | WaitForInputStep
+    )
     """
     The steps to run if the condition is true
     """
     else_: Annotated[
-        EvaluateStep | ToolCallStep | YieldStep | PromptStep | ErrorWorkflowStep,
+        Any
+        | ToolCallStep
+        | YieldStep
+        | PromptStep
+        | ErrorWorkflowStep
+        | WaitForInputStep,
         Field(alias="else"),
     ]
     """
@@ -107,6 +135,7 @@ class PatchTaskRequest(BaseModel):
             | YieldStep
             | PromptStep
             | ErrorWorkflowStep
+            | WaitForInputStep
             | IfElseWorkflowStep
         ]
         | None
@@ -129,10 +158,11 @@ class PatchTaskRequest(BaseModel):
     metadata: dict[str, Any] | None = None
 
 
-class PromptStep(BaseModel):
+class PromptStep(BaseWorkflowStep):
     model_config = ConfigDict(
         populate_by_name=True,
     )
+    kind_: Literal["prompt"] = "prompt"
     prompt: str | list[InputChatMLMessage]
     """
     The prompt to run
@@ -356,6 +386,7 @@ class Task(BaseModel):
         | YieldStep
         | PromptStep
         | ErrorWorkflowStep
+        | WaitForInputStep
         | IfElseWorkflowStep
     ]
     """
@@ -407,10 +438,11 @@ class TaskTool(BaseModel):
     api_call: Any | None = None
 
 
-class ToolCallStep(BaseModel):
+class ToolCallStep(BaseWorkflowStep):
     model_config = ConfigDict(
         populate_by_name=True,
     )
+    kind_: Literal["tool_call"] = "tool_call"
     tool: Annotated[
         str, Field(pattern="^(function|integration|system|api_call)\\.(\\w+)$")
     ]
@@ -438,6 +470,7 @@ class UpdateTaskRequest(BaseModel):
         | YieldStep
         | PromptStep
         | ErrorWorkflowStep
+        | WaitForInputStep
         | IfElseWorkflowStep
     ]
     """
@@ -458,10 +491,22 @@ class UpdateTaskRequest(BaseModel):
     metadata: dict[str, Any] | None = None
 
 
-class YieldStep(BaseModel):
+class WaitForInputStep(BaseWorkflowStep):
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    kind_: Literal["wait_for_input"] = "wait_for_input"
+    info: str | dict[str, Any]
+    """
+    Any additional info or data
+    """
+
+
+class YieldStep(BaseWorkflowStep):
     model_config = ConfigDict(
         populate_by_name=True,
     )
+    kind_: Literal["yield"] = "yield"
     workflow: str
     """
     The subworkflow to run

diff --git a/agents-api/poetry.lock b/agents-api/poetry.lock
diff --git a/sdks/python/julep/api/__init__.py b/sdks/python/julep/api/__init__.py
@@ -139,14 +139,29 @@
     TaskExecutionsRouteListRequestDirection,
     TaskExecutionsRouteListRequestSortBy,
     TaskExecutionsRouteListResponse,
+    TasksBaseWorkflowStep,
     TasksCreateTaskRequest,
     TasksCreateTaskRequestMainItem,
+    TasksCreateTaskRequestMainItem_Error,
+    TasksCreateTaskRequestMainItem_Evaluate,
+    TasksCreateTaskRequestMainItem_IfElse,
+    TasksCreateTaskRequestMainItem_Prompt,
+    TasksCreateTaskRequestMainItem_ToolCall,
+    TasksCreateTaskRequestMainItem_WaitForInput,
+    TasksCreateTaskRequestMainItem_Yield,
     TasksErrorWorkflowStep,
     TasksEvaluateStep,
     TasksIfElseWorkflowStep,
     TasksIfElseWorkflowStepElse,
     TasksIfElseWorkflowStepThen,
     TasksPatchTaskRequestMainItem,
+    TasksPatchTaskRequestMainItem_Error,
+    TasksPatchTaskRequestMainItem_Evaluate,
+    TasksPatchTaskRequestMainItem_IfElse,
+    TasksPatchTaskRequestMainItem_Prompt,
+    TasksPatchTaskRequestMainItem_ToolCall,
+    TasksPatchTaskRequestMainItem_WaitForInput,
+    TasksPatchTaskRequestMainItem_Yield,
     TasksPromptStep,
     TasksPromptStepPrompt,
     TasksPromptStepSettings,
@@ -158,9 +173,25 @@
     TasksRouteListResponse,
     TasksTask,
     TasksTaskMainItem,
+    TasksTaskMainItem_Error,
+    TasksTaskMainItem_Evaluate,
+    TasksTaskMainItem_IfElse,
+    TasksTaskMainItem_Prompt,
+    TasksTaskMainItem_ToolCall,
+    TasksTaskMainItem_WaitForInput,
+    TasksTaskMainItem_Yield,
     TasksTaskTool,
     TasksToolCallStep,
     TasksUpdateTaskRequestMainItem,
+    TasksUpdateTaskRequestMainItem_Error,
+    TasksUpdateTaskRequestMainItem_Evaluate,
+    TasksUpdateTaskRequestMainItem_IfElse,
+    TasksUpdateTaskRequestMainItem_Prompt,
+    TasksUpdateTaskRequestMainItem_ToolCall,
+    TasksUpdateTaskRequestMainItem_WaitForInput,
+    TasksUpdateTaskRequestMainItem_Yield,
+    TasksWaitForInputStep,
+    TasksWaitForInputStepInfo,
     TasksYieldStep,
     ToolsChosenFunctionCall,
     ToolsChosenToolCall,
@@ -331,14 +362,29 @@
     "TaskExecutionsRouteListRequestDirection",
     "TaskExecutionsRouteListRequestSortBy",
     "TaskExecutionsRouteListResponse",
+    "TasksBaseWorkflowStep",
     "TasksCreateTaskRequest",
     "TasksCreateTaskRequestMainItem",
+    "TasksCreateTaskRequestMainItem_Error",
+    "TasksCreateTaskRequestMainItem_Evaluate",
+    "TasksCreateTaskRequestMainItem_IfElse",
+    "TasksCreateTaskRequestMainItem_Prompt",
+    "TasksCreateTaskRequestMainItem_ToolCall",
+    "TasksCreateTaskRequestMainItem_WaitForInput",
+    "TasksCreateTaskRequestMainItem_Yield",
     "TasksErrorWorkflowStep",
     "TasksEvaluateStep",
     "TasksIfElseWorkflowStep",
     "TasksIfElseWorkflowStepElse",
     "TasksIfElseWorkflowStepThen",
     "TasksPatchTaskRequestMainItem",
+    "TasksPatchTaskRequestMainItem_Error",
+    "TasksPatchTaskRequestMainItem_Evaluate",
+    "TasksPatchTaskRequestMainItem_IfElse",
+    "TasksPatchTaskRequestMainItem_Prompt",
+    "TasksPatchTaskRequestMainItem_ToolCall",
+    "TasksPatchTaskRequestMainItem_WaitForInput",
+    "TasksPatchTaskRequestMainItem_Yield",
     "TasksPromptStep",
     "TasksPromptStepPrompt",
     "TasksPromptStepSettings",
@@ -350,9 +396,25 @@
     "TasksRouteListResponse",
     "TasksTask",
     "TasksTaskMainItem",
+    "TasksTaskMainItem_Error",
+    "TasksTaskMainItem_Evaluate",
+    "TasksTaskMainItem_IfElse",
+    "TasksTaskMainItem_Prompt",
+    "TasksTaskMainItem_ToolCall",
+    "TasksTaskMainItem_WaitForInput",
+    "TasksTaskMainItem_Yield",
     "TasksTaskTool",
     "TasksToolCallStep",
     "TasksUpdateTaskRequestMainItem",
+    "TasksUpdateTaskRequestMainItem_Error",
+    "TasksUpdateTaskRequestMainItem_Evaluate",
+    "TasksUpdateTaskRequestMainItem_IfElse",
+    "TasksUpdateTaskRequestMainItem_Prompt",
+    "TasksUpdateTaskRequestMainItem_ToolCall",
+    "TasksUpdateTaskRequestMainItem_WaitForInput",
+    "TasksUpdateTaskRequestMainItem_Yield",
+    "TasksWaitForInputStep",
+    "TasksWaitForInputStepInfo",
     "TasksYieldStep",
     "ToolsChosenFunctionCall",
     "ToolsChosenToolCall",