chore(agents-api): misc validations

Author: Vedantsahai18

Diff for: agents-api/agents_api/activities/container.py

@@ -1,5 +1,19 @@
+from typing import Any
+
+from aiobotocore.client import AioBaseClient
+from asyncpg.pool import Pool
+
+
 class State:
-    pass
+    postgres_pool: Pool | None
+    s3_client: AioBaseClient | None
+
+    def __init__(self):
+        self.postgres_pool = None
+        self.s3_client = None
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        super().__setattr__(name, value)
 
 
 class Container:

Diff for: agents-api/agents_api/activities/execute_integration.py

@@ -9,7 +9,7 @@
 from ..common.exceptions.tools import IntegrationExecutionException
 from ..common.protocol.tasks import ExecutionInput, StepContext
 from ..env import testing
-from ..queries.tools import get_tool_args_from_metadata
+from ..queries import tools
 from .container import container
 
 
@@ -28,17 +28,22 @@ async def execute_integration(
 
     developer_id = context.execution_input.developer_id
     agent_id = context.execution_input.agent.id
+
+    if context.execution_input.task is None:
+        msg = "Task cannot be None in execution_input"
+        raise ValueError(msg)
+
     task_id = context.execution_input.task.id
 
-    merged_tool_args = await get_tool_args_from_metadata(
+    merged_tool_args = await tools.get_tool_args_from_metadata(
         developer_id=developer_id,
         agent_id=agent_id,
         task_id=task_id,
         arg_type="args",
         connection_pool=container.state.postgres_pool,
     )
 
-    merged_tool_setup = await get_tool_args_from_metadata(
+    merged_tool_setup = await tools.get_tool_args_from_metadata(
         developer_id=developer_id,
         agent_id=agent_id,
         task_id=task_id,

Diff for: agents-api/agents_api/activities/execute_system.py

@@ -23,7 +23,7 @@
 )
 from ..common.protocol.tasks import ExecutionInput, StepContext
 from ..env import testing
-from ..queries.developers import get_developer
+from ..queries import developers
 from .container import container
 from .utils import get_handler
 
@@ -95,7 +95,7 @@ async def execute_system(
 
         # Handle chat operations
         if system.operation == "chat" and system.resource == "session":
-            developer = await get_developer(
+            developer = await developers.get_developer(
                 developer_id=arguments["developer_id"],
                 connection_pool=container.state.postgres_pool,
             )

Diff for: agents-api/agents_api/activities/task_steps/transition_step.py

@@ -39,6 +39,10 @@ async def transition_step(
         msg = "Expected ExecutionInput type for context.execution_input"
         raise TypeError(msg)
 
+    if not context.execution_input.execution:
+        msg = "Execution is required in execution_input"
+        raise ValueError(msg)
+
     # Create transition
     try:
         transition = await create_execution_transition(

Diff for: agents-api/agents_api/activities/task_steps/yield_step.py

@@ -16,10 +16,20 @@ async def yield_step(context: StepContext) -> StepOutcome:
             msg = "Expected ExecutionInput type for context.execution_input"
             raise TypeError(msg)
 
+        # Add validation for task
+        if not context.execution_input.task:
+            msg = "Task is required in execution_input"
+            raise ValueError(msg)
+
         all_workflows = context.execution_input.task.workflows
         workflow = context.current_step.workflow
         exprs = context.current_step.arguments
 
+        # Validate workflows exists
+        if not all_workflows:
+            msg = "No workflows found in task"
+            raise ValueError(msg)
+
         assert workflow in [wf.name for wf in all_workflows], (
             f"Workflow {workflow} not found in task"
         )

Diff for: agents-api/agents_api/activities/utils.py

@@ -41,21 +41,60 @@ def safe_range(*args):
     return result
 
 
-def safe_json_loads(s: str):
+@beartype
+def safe_json_loads(s: str) -> Any:
+    """
+    Safely load a JSON string with size limits.
+
+    Args:
+        s: JSON string to parse
+
+    Returns:
+        Parsed JSON data
+
+    Raises:
+        ValueError: If string exceeds size limit
+    """
     if len(s) > MAX_STRING_LENGTH:
         msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}"
         raise ValueError(msg)
     return json.loads(s)
 
 
-def safe_yaml_load(s: str):
+@beartype
+def safe_yaml_load(s: str) -> Any:
+    """
+    Safely load a YAML string with size limits.
+
+    Args:
+        s: YAML string to parse
+
+    Returns:
+        Parsed YAML data
+
+    Raises:
+        ValueError: If string exceeds size limit
+    """
     if len(s) > MAX_STRING_LENGTH:
         msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}"
         raise ValueError(msg)
     return yaml.load(s)
 
 
+@beartype
 def safe_base64_decode(s: str) -> str:
+    """
+    Safely decode a base64 string with size limits.
+
+    Args:
+        s: Base64 string to decode
+
+    Returns:
+        Decoded UTF-8 string
+
+    Raises:
+        ValueError: If string exceeds size limit or is invalid base64
+    """
     if len(s) > MAX_STRING_LENGTH:
         msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}"
         raise ValueError(msg)
@@ -66,21 +105,66 @@ def safe_base64_decode(s: str) -> str:
         raise ValueError(msg)
 
 
+@beartype
 def safe_base64_encode(s: str) -> str:
+    """
+    Safely encode a string to base64 with size limits.
+
+    Args:
+        s: String to encode
+
+    Returns:
+        Base64 encoded string
+
+    Raises:
+        ValueError: If string exceeds size limit
+    """
     if len(s) > MAX_STRING_LENGTH:
         msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}"
         raise ValueError(msg)
     return base64.b64encode(s.encode("utf-8")).decode("utf-8")
 
 
-def safe_random_choice(seq):
+@beartype
+def safe_random_choice(seq: list[Any] | tuple[Any, ...] | str) -> Any:
+    """
+    Safely choose a random element from a sequence with size limits.
+
+    Args:
+        seq: A sequence (list, tuple, or string) to choose from
+
+    Returns:
+        A randomly selected element
+
+    Raises:
+        ValueError: If sequence exceeds size limit
+        TypeError: If input is not a valid sequence type
+    """
     if len(seq) > MAX_COLLECTION_SIZE:
         msg = f"Sequence exceeds maximum size of {MAX_COLLECTION_SIZE}"
         raise ValueError(msg)
     return random.choice(seq)
 
 
-def safe_random_sample(population, k):
+@beartype
+def safe_random_sample(population: list[T] | tuple[T, ...] | str, k: int) -> list[T]:
+    """
+    Safely sample k elements from a population with size limits.
+
+    Args:
+        population: A sequence to sample from
+        k: Number of elements to sample
+
+    Returns:
+        A list containing k randomly selected elements
+
+    Raises:
+        ValueError: If population/sample size exceeds limits
+        TypeError: If input is not a valid sequence type
+    """
+    if not isinstance(population, list | tuple | str):
+        msg = "Expected a sequence (list, tuple, or string)"
+        raise TypeError(msg)
     if len(population) > MAX_COLLECTION_SIZE:
         msg = f"Population exceeds maximum size of {MAX_COLLECTION_SIZE}"
         raise ValueError(msg)
@@ -93,9 +177,19 @@ def safe_random_sample(population, k):
     return random.sample(population, k)
 
 
+@beartype
 def chunk_doc(string: str) -> list[str]:
     """
     Chunk a string into sentences.
+
+    Args:
+        string: The text to chunk into sentences
+
+    Returns:
+        A list of sentence chunks
+
+    Raises:
+        ValueError: If string exceeds size limit
     """
     if len(string) > MAX_STRING_LENGTH:
         msg = f"String exceeds maximum length of {MAX_STRING_LENGTH}"

Diff for: agents-api/agents_api/clients/temporal.py

@@ -90,6 +90,10 @@ async def run_task_execution_workflow(
 ):
     from ..workflows.task_execution import TaskExecutionWorkflow
 
+    if execution_input.execution is None:
+        msg = "execution_input.execution cannot be None"
+        raise ValueError(msg)
+
     start: TransitionTarget = start or TransitionTarget(workflow="main", step=0)
 
     client = client or (await get_client())

Diff for: agents-api/agents_api/workflows/task_execution/__init__.py

@@ -143,6 +143,10 @@ async def run(
         start: TransitionTarget,
         previous_inputs: list,
     ) -> Any:
+        if not execution_input.task:
+            msg = "execution_input.task cannot be None"
+            raise ApplicationError(msg)
+
         workflow.logger.info(
             f"TaskExecutionWorkflow for task {execution_input.task.id}"
             f" [LOC {start.workflow}.{start.step}]"