Merge pull request #26 from e2b-dev/upload-artifacts

Allow to upload artifacts

agent/python/agent_protocol/agent.py

@@ -1,7 +1,8 @@
 import asyncio
 import os
 
-from fastapi import APIRouter
+import aiofiles
+from fastapi import APIRouter, UploadFile
 from fastapi.responses import FileResponse
 from hypercorn.asyncio import serve
 from hypercorn.config import Config
@@ -100,10 +101,6 @@ async def execute_agent_task_step(
     step = await _step_handler(step)
 
     step.status = Status.completed
-
-    if step.artifacts:
-        task.artifacts.extend(step.artifacts)
-
     return step
 
 
@@ -132,6 +129,31 @@ async def list_agent_task_artifacts(task_id: str) -> List[Artifact]:
     return task.artifacts
 
 
+@base_router.post(
+    "/agent/tasks/{task_id}/artifacts",
+    response_model=Artifact,
+    tags=["agent"],
+)
+async def upload_agent_task_artifacts(
+    task_id: str, file: UploadFile, relative_path: Optional[str] = None
+) -> Artifact:
+    """
+    Upload an artifact for the specified task.
+    """
+    await Agent.db.get_task(task_id)
+    artifact = await Agent.db.create_artifact(task_id, file.filename, 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:
+        while content := await file.read(1024 * 1024):  # async read chunk ~1MiB
+            await f.write(content)
+
+    return artifact
+
+
 @base_router.get(
     "/agent/tasks/{task_id}/artifacts/{artifact_id}",
     tags=["agent"],
@@ -172,13 +194,20 @@ def get_workspace(task_id: str) -> str:
         return os.path.join(os.getcwd(), Agent.workspace, task_id)
 
     @staticmethod
-    def get_artifact_path(task_id: str, artifact: Artifact) -> str:
+    def get_artifact_folder(task_id: str, artifact: Artifact) -> str:
         """
         Get the artifact path for the specified task and artifact.
         """
         workspace_path = Agent.get_workspace(task_id)
         relative_path = artifact.relative_path or ""
-        return os.path.join(workspace_path, relative_path, artifact.file_name)
+        return os.path.join(workspace_path, relative_path)
+
+    @staticmethod
+    def get_artifact_path(task_id: str, artifact: Artifact) -> str:
+        """
+        Get the artifact path for the specified task and artifact.
+        """
+        return os.path.join(Agent.get_artifact_folder(task_id, artifact), artifact.file_name)
 
     @staticmethod
     def start(port: int = 8000, router: APIRouter = base_router):

agent/python/agent_protocol/db.py

@@ -31,6 +31,15 @@ async def create_step(
     ) -> Step:
         raise NotImplementedError
 
+    async def create_artifact(
+        self,
+        task_id: str,
+        file_name: str,
+        relative_path: Optional[str] = None,
+        step_id: Optional[str] = None,
+    ) -> Artifact:
+        raise NotImplementedError
+
     async def get_task(self, task_id: str) -> Task:
         raise NotImplementedError
 
@@ -114,6 +123,26 @@ async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact:
             raise Exception(f"Artifact with id {artifact_id} not found")
         return artifact
 
+    async def create_artifact(
+        self,
+        task_id: str,
+        file_name: str,
+        relative_path: Optional[str] = None,
+        step_id: Optional[str] = None,
+    ) -> Artifact:
+        artifact_id = str(uuid.uuid4())
+        artifact = Artifact(
+            artifact_id=artifact_id, file_name=file_name, relative_path=relative_path
+        )
+        task = await self.get_task(task_id)
+        task.artifacts.append(artifact)
+
+        if step_id:
+            step = await self.get_step(task_id, step_id)
+            step.artifacts.append(artifact)
+
+        return artifact
+
     async def list_tasks(self) -> List[Task]:
         return [task for task in self._tasks.values()]
 

agent/python/agent_protocol/models.py

@@ -24,6 +24,13 @@ class Artifact(BaseModel):
     )
 
 
+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."
+    )
+
+
 class StepInput(BaseModel):
     __root__: Any = Field(
         ..., description="Input parameters for the task step. Any value is allowed."

agent/python/examples/smol_developer.py

@@ -61,16 +61,17 @@ async def _generate_code(task: Task, step: Step) -> Step:
     file_path = step.additional_properties["file_path"]
 
     code = await generate_code(task.input, shared_deps, file_path)
+    step.output = code
+
     write_file(os.path.join(Agent.get_workspace(task.task_id), file_path), code)
     path = Path("./" + file_path)
-    artifact = Artifact(
-        artifact_id=str(uuid.uuid4()),
-        file_name=path.name,
+    await Agent.db.create_artifact(
+        task_id=task.task_id,
+        step_id=step.step_id,
         relative_path=str(path.parent),
+        file_name=path.name,
     )
 
-    step.output = code
-    step.artifacts.append(artifact)
     return step
 
 

agent/python/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agent-protocol"
-version = "0.2.2"
+version = "0.2.3"
 description = "API for interacting with Agent"
 authors = ["e2b <hello@e2b.dev>"]
 license = "MIT"
@@ -17,6 +17,8 @@ pytest = "^7.0.0"
 pydantic = "^1.10.5, <2"
 click = "^8.1.6"
 requests = "^2.31.0"
+python-multipart = "^0.0.6"
+aiofiles = "^23.1.0"
 
 [tool.poetry.group.dev.dependencies]
 fastapi-code-generator = "^0.4.2"