diff --git a/packages/sdk/python/agent_protocol/db.py b/packages/sdk/python/agent_protocol/db.py
index 0c52c15..24efb73 100644
--- a/packages/sdk/python/agent_protocol/db.py
+++ b/packages/sdk/python/agent_protocol/db.py
@@ -5,7 +5,7 @@
 
 
 class Step(APIStep):
-    additional_properties: Optional[Dict[str, str]] = None
+    additional_properties: Optional[Dict[str, Any]] = None
 
 
 class Task(APITask):
diff --git a/packages/sdk/python/agent_protocol/models.py b/packages/sdk/python/agent_protocol/models.py
index a8145ed..00ee368 100644
--- a/packages/sdk/python/agent_protocol/models.py
+++ b/packages/sdk/python/agent_protocol/models.py
@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Dict, Any
 
 from pydantic import BaseModel, Field
 
@@ -25,10 +25,6 @@ class Pagination(BaseModel):
     page_size: int = Field(..., description="Number of items per page.", example=25)
 
 
-class TaskInput(BaseModel):
-    pass
-
-
 class Artifact(BaseModel):
     artifact_id: str = Field(
         ...,
@@ -61,10 +57,6 @@ class ArtifactUpload(BaseModel):
     )
 
 
-class StepInput(BaseModel):
-    pass
-
-
 class StepOutput(BaseModel):
     pass
 
@@ -75,7 +67,7 @@ class TaskRequestBody(BaseModel):
         description="Input prompt for the task.",
         example="Write 'Washington' to the file 'output.txt'.",
     )
-    additional_input: Optional[TaskInput] = None
+    additional_input: Optional[Dict[str, Any]] = None
 
 
 class Task(TaskRequestBody):
@@ -100,7 +92,7 @@ class StepRequestBody(BaseModel):
         description="Input prompt for the step.",
         example="Write the words you receive to the file 'output.txt'.",
     )
-    additional_input: Optional[StepInput] = None
+    additional_input: Optional[Dict[str, Any]] = None
 
 
 class Status(Enum):