Update python sdk to v0.3

Author: jakubno

sdk/python/Makefile

@@ -7,3 +7,4 @@ generate:
 	mv agent_protocol/main.py agent_protocol/server.py
 	rm -rf agent_protocol/routers
 	rm agent_protocol/dependencies.py
+	black .

sdk/python/agent_protocol/__init__.py

@@ -1,16 +1,17 @@
 from .agent import Agent, StepHandler, TaskHandler, base_router as router
-from .models import Artifact, StepRequestBody, TaskRequestBody
+from .models import Artifact, Status, StepRequestBody, TaskRequestBody
 from .db import Step, Task, TaskDB
 
 
 __all__ = [
     "Agent",
     "Artifact",
+    "Status",
     "Step",
     "StepHandler",
+    "StepRequestBody",
     "Task",
     "TaskDB",
-    "StepRequestBody",
     "TaskHandler",
     "TaskRequestBody",
     "router",

sdk/python/agent_protocol/agent.py

@@ -1,27 +1,26 @@
 import asyncio
 import os
+from uuid import uuid4
 
 import aiofiles
 from fastapi import APIRouter, UploadFile, Form, File
 from fastapi.responses import FileResponse
 from hypercorn.asyncio import serve
 from hypercorn.config import Config
-from typing import Awaitable, Callable, List, Optional, Annotated
+from typing import Callable, List, Optional, Annotated, Coroutine, Any
 
-from .db import InMemoryTaskDB, TaskDB
+from .db import InMemoryTaskDB, Task, TaskDB, Step
 from .server import app
 from .models import (
     TaskRequestBody,
-    Step,
     StepRequestBody,
     Artifact,
-    Task,
     Status,
 )
 
 
-StepHandler = Callable[[Step], Awaitable[Step]]
-TaskHandler = Callable[[Task], Awaitable[None]]
+StepHandler = Callable[[Step], Coroutine[Any, Any, Step]]
+TaskHandler = Callable[[Task], Coroutine[Any, Any, None]]
 
 
 _task_handler: Optional[TaskHandler]
@@ -89,12 +88,17 @@ async def execute_agent_task_step(
     """
     Execute a step in the specified agent task.
     """
+    if not _step_handler:
+        raise Exception("Step handler not defined")
+
     task = await Agent.db.get_task(task_id)
     step = next(filter(lambda x: x.status == Status.created, task.steps), None)
 
     if not step:
         raise Exception("No steps to execute")
 
+    step.status = Status.running
+
     step.input = body.input if body else None
     step.additional_input = body.additional_input if body else None
 
@@ -109,7 +113,7 @@ async def execute_agent_task_step(
     response_model=Step,
     tags=["agent"],
 )
-async def get_agent_task_step(task_id: str, step_id: str = ...) -> Step:
+async def get_agent_task_step(task_id: str, step_id: str) -> Step:
     """
     Get details about a specified task step.
     """
@@ -142,14 +146,15 @@ async def upload_agent_task_artifacts(
     """
     Upload an artifact for the specified task.
     """
+    file_name = file.filename or str(uuid4())
     await Agent.db.get_task(task_id)
-    artifact = await Agent.db.create_artifact(task_id, file.filename, relative_path)
+    artifact = await Agent.db.create_artifact(task_id, file_name, relative_path)
 
     path = Agent.get_artifact_folder(task_id, artifact)
     if not os.path.exists(path):
         os.makedirs(path)
 
-    async with aiofiles.open(os.path.join(path, file.filename), "wb") as f:
+    async with aiofiles.open(os.path.join(path, file_name), "wb") as f:
         while content := await file.read(1024 * 1024):  # async read chunk ~1MiB
             await f.write(content)
 

sdk/python/agent_protocol/db.py

@@ -12,20 +12,32 @@ class Task(APITask):
     steps: List[Step] = []
 
 
+class NotFoundException(Exception):
+    """
+    Exception raised when a resource is not found.
+    """
+
+    def __init__(self, item_name: str, item_id: str):
+        self.item_name = item_name
+        self.item_id = item_id
+        super().__init__(f"{item_name} with {item_id} not found.")
+
+
 class TaskDB(ABC):
     async def create_task(
         self,
         input: Optional[str],
-        additional_input: Optional[str] = None,
-        artifacts: List[Artifact] = None,
-        steps: List[Step] = None,
+        additional_input: Any = None,
+        artifacts: Optional[List[Artifact]] = None,
+        steps: Optional[List[Step]] = None,
     ) -> Task:
         raise NotImplementedError
 
     async def create_step(
         self,
         task_id: str,
         name: Optional[str] = None,
+        input: Optional[str] = None,
         is_last: bool = False,
         additional_properties: Optional[Dict[str, str]] = None,
     ) -> Step:
@@ -52,7 +64,9 @@ async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact:
     async def list_tasks(self) -> List[Task]:
         raise NotImplementedError
 
-    async def list_steps(self, task_id: str) -> List[Step]:
+    async def list_steps(
+        self, task_id: str, status: Optional[Status] = None
+    ) -> List[Step]:
         raise NotImplementedError
 
 
@@ -62,9 +76,9 @@ class InMemoryTaskDB(TaskDB):
     async def create_task(
         self,
         input: Optional[str],
-        additional_input: Optional[str] = None,
-        artifacts: List[Artifact] = None,
-        steps: List[Step] = None,
+        additional_input: Any = None,
+        artifacts: Optional[List[Artifact]] = None,
+        steps: Optional[List[Step]] = None,
     ) -> Task:
         if not steps:
             steps = []
@@ -85,14 +99,16 @@ async def create_step(
         self,
         task_id: str,
         name: Optional[str] = None,
+        input: Optional[str] = None,
         is_last=False,
-        additional_properties: Dict[str, Any] = None,
+        additional_properties: Optional[Dict[str, Any]] = None,
     ) -> Step:
         step_id = str(uuid.uuid4())
         step = Step(
             task_id=task_id,
             step_id=step_id,
             name=name,
+            input=input,
             status=Status.created,
             is_last=is_last,
             additional_properties=additional_properties,
@@ -104,14 +120,14 @@ async def create_step(
     async def get_task(self, task_id: str) -> Task:
         task = self._tasks.get(task_id, None)
         if not task:
-            raise Exception(f"Task with id {task_id} not found")
+            raise NotFoundException("Task", task_id)
         return task
 
     async def get_step(self, task_id: str, step_id: str) -> Step:
         task = await self.get_task(task_id)
         step = next(filter(lambda s: s.task_id == task_id, task.steps), None)
         if not step:
-            raise Exception(f"Step with id {step_id} not found")
+            raise NotFoundException("Step", step_id)
         return step
 
     async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact:
@@ -120,7 +136,7 @@ async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact:
             filter(lambda a: a.artifact_id == artifact_id, task.artifacts), None
         )
         if not artifact:
-            raise Exception(f"Artifact with id {artifact_id} not found")
+            raise NotFoundException("Artifact", artifact_id)
         return artifact
 
     async def create_artifact(
@@ -146,6 +162,11 @@ async def create_artifact(
     async def list_tasks(self) -> List[Task]:
         return [task for task in self._tasks.values()]
 
-    async def list_steps(self, task_id: str) -> List[Step]:
+    async def list_steps(
+        self, task_id: str, status: Optional[Status] = None
+    ) -> List[Step]:
         task = await self.get_task(task_id)
-        return [step for step in task.steps]
+        steps = task.steps
+        if status:
+            steps = list(filter(lambda s: s.status == status, steps))
+        return steps

sdk/python/agent_protocol/middlewares.py

@@ -0,0 +1,13 @@
+from fastapi import Request
+from fastapi.responses import PlainTextResponse
+
+from agent_protocol.db import NotFoundException
+
+
+async def not_found_exception_handler(
+    request: Request, exc: NotFoundException
+) -> PlainTextResponse:
+    return PlainTextResponse(
+        str(exc),
+        status_code=404,
+    )

sdk/python/agent_protocol/models.py

@@ -1,6 +1,6 @@
 # generated by fastapi-codegen:
 #   filename:  ../../openapi.yml
-#   timestamp: 2023-08-07T12:14:43+00:00
+#   timestamp: 2023-08-11T14:24:22+00:00
 
 from __future__ import annotations
 
@@ -12,65 +12,111 @@
 
 class TaskInput(BaseModel):
     __root__: Any = Field(
-        ..., description="Input parameters for the task. Any value is allowed."
+        ...,
+        description="Input parameters for the task. Any value is allowed.",
+        example='{\n"debug": false,\n"mode": "benchmarks"\n}',
     )
 
 
 class Artifact(BaseModel):
-    artifact_id: str = Field(..., description="ID of the artifact.")
-    file_name: str = Field(..., description="Filename of the artifact.")
+    artifact_id: str = Field(
+        ...,
+        description="ID of the artifact.",
+        example="b225e278-8b4c-4f99-a696-8facf19f0e56",
+    )
+    file_name: str = Field(
+        ..., description="Filename of the artifact.", example="main.py"
+    )
     relative_path: Optional[str] = Field(
-        None, description="Relative path of the artifact in the agent's workspace."
+        None,
+        description="Relative path of the artifact in the agent's workspace.",
+        example="python/code/",
     )
 
 
 class ArtifactUpload(BaseModel):
     file: bytes = Field(..., description="File to upload.")
     relative_path: Optional[str] = Field(
-        None, description="Relative path of the artifact in the agent's workspace."
+        None,
+        description="Relative path of the artifact in the agent's workspace.",
+        example="python/code",
     )
 
 
 class StepInput(BaseModel):
     __root__: Any = Field(
-        ..., description="Input parameters for the task step. Any value is allowed."
+        ...,
+        description="Input parameters for the task step. Any value is allowed.",
+        example='{\n"file_to_refactor": "models.py"\n}',
     )
 
 
 class StepOutput(BaseModel):
     __root__: Any = Field(
-        ..., description="Output that the task step has produced. Any value is allowed."
+        ...,
+        description="Output that the task step has produced. Any value is allowed.",
+        example='{\n"tokens": 7894,\n"estimated_cost": "0,24$"\n}',
     )
 
 
 class TaskRequestBody(BaseModel):
-    input: Optional[str] = Field(None, description="Input prompt for the task.")
+    input: Optional[str] = Field(
+        None,
+        description="Input prompt for the task.",
+        example="Write the words you receive to the file 'output.txt'.",
+    )
     additional_input: Optional[TaskInput] = None
 
 
 class Task(TaskRequestBody):
-    task_id: str = Field(..., description="The ID of the task.")
+    task_id: str = Field(
+        ...,
+        description="The ID of the task.",
+        example="50da533e-3904-4401-8a07-c49adf88b5eb",
+    )
     artifacts: List[Artifact] = Field(
-        [], description="A list of artifacts that the task has produced."
+        [],
+        description="A list of artifacts that the task has produced.",
+        example=[
+            "7a49f31c-f9c6-4346-a22c-e32bc5af4d8e",
+            "ab7b4091-2560-4692-a4fe-d831ea3ca7d6",
+        ],
     )
 
 
 class StepRequestBody(BaseModel):
-    input: Optional[str] = Field(None, description="Input prompt for the step.")
+    input: Optional[str] = Field(
+        None, description="Input prompt for the step.", example="Washington"
+    )
     additional_input: Optional[StepInput] = None
 
 
 class Status(Enum):
     created = "created"
+    running = "running"
     completed = "completed"
 
 
 class Step(StepRequestBody):
-    task_id: str = Field(..., description="The ID of the task this step belongs to.")
-    step_id: str = Field(..., description="The ID of the task step.")
-    name: Optional[str] = Field(None, description="The name of the task step.")
+    task_id: str = Field(
+        ...,
+        description="The ID of the task this step belongs to.",
+        example="50da533e-3904-4401-8a07-c49adf88b5eb",
+    )
+    step_id: str = Field(
+        ...,
+        description="The ID of the task step.",
+        example="6bb1801a-fd80-45e8-899a-4dd723cc602e",
+    )
+    name: Optional[str] = Field(
+        None, description="The name of the task step.", example="Write to file"
+    )
     status: Status = Field(..., description="The status of the task step.")
-    output: Optional[str] = Field(None, description="Output of the task step.")
+    output: Optional[str] = Field(
+        None,
+        description="Output of the task step.",
+        example="I am going to use the write_to_file command and write Washington to a file called output.txt <write_to_file('output.txt', 'Washington')",
+    )
     additional_output: Optional[StepOutput] = None
     artifacts: List[Artifact] = Field(
         [], description="A list of artifacts that the step has produced."

sdk/python/agent_protocol/server.py

@@ -6,8 +6,13 @@
 
 from fastapi import FastAPI
 
+from agent_protocol.db import NotFoundException
+from agent_protocol.middlewares import not_found_exception_handler
+
 app = FastAPI(
     title="Agent Communication Protocol",
     description="Specification of the API protocol for communication with an agent.",
-    version="v0.2",
+    version="v0.3",
 )
+
+app.add_exception_handler(NotFoundException, not_found_exception_handler)

sdk/python/examples/smol_developer.py

@@ -80,7 +80,7 @@ async def task_handler(task: Task) -> None:
     await Agent.db.create_step(task.task_id, StepTypes.PLAN)
 
 
-async def step_handler(step: Step):
+async def step_handler(step: Step) -> Step:
     task = await Agent.db.get_task(step.task_id)
     if step.name == StepTypes.PLAN:
         return await _generate_shared_deps(step)