Auto-GPT/autogpt/agent.py

import asyncio
import os
import typing

from fastapi import APIRouter, FastAPI, Response, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from hypercorn.asyncio import serve
from hypercorn.config import Config

from .db import AgentDB
from .middlewares import AgentMiddleware
from .routes.agent_protocol import base_router
from .schema import Artifact, Status, Step, StepRequestBody, Task, TaskRequestBody
from .utils import run
from .workspace import Workspace


class Agent:
    def __init__(self, database: AgentDB, workspace: Workspace):
        self.db = database
        self.workspace = workspace

    def start(self, port: int = 8000, router: APIRouter = base_router):
        """
        Start the agent server.
        """
        config = Config()
        config.bind = [f"localhost:{port}"]
        app = FastAPI(
            title="Auto-GPT Forge",
            description="Modified version of The Agent Protocol.",
            version="v0.4",
        )
        app.include_router(router)
        app.add_middleware(AgentMiddleware, agent=self)
        asyncio.run(serve(app, config))

    async def create_task(self, task_request: TaskRequestBody) -> Task:
        """
        Create a task for the agent.
        """
        try:
            task = await self.db.create_task(
                input=task_request.input if task_request.input else None,
                additional_input=task_request.additional_input
                if task_request.additional_input
                else None,
            )
            print(task)
        except Exception as e:
            return Response(status_code=500, content=str(e))
        print(task)
        return task

    async def list_tasks(self) -> typing.List[str]:
        """
        List the IDs of all tasks that the agent has created.
        """
        try:
            task_ids = [task.task_id for task in await self.db.list_tasks()]
        except Exception as e:
            return Response(status_code=500, content=str(e))
        return task_ids

    async def get_task(self, task_id: str) -> Task:
        """
        Get a task by ID.
        """
        if not task_id:
            return Response(status_code=400, content="Task ID is required.")
        if not isinstance(task_id, str):
            return Response(status_code=400, content="Task ID must be a string.")
        try:
            task = await self.db.get_task(task_id)
        except Exception as e:
            return Response(status_code=500, content=str(e))
        return task

    async def list_steps(self, task_id: str) -> typing.List[str]:
        """
        List the IDs of all steps that the task has created.
        """
        if not task_id:
            return Response(status_code=400, content="Task ID is required.")
        if not isinstance(task_id, str):
            return Response(status_code=400, content="Task ID must be a string.")
        try:
            steps_ids = [step.step_id for step in await self.db.list_steps(task_id)]
        except Exception as e:
            return Response(status_code=500, content=str(e))
        return steps_ids

    async def create_and_execute_step(
        self, task_id: str, step_request: StepRequestBody
    ) -> Step:
        """
        Create a step for the task.
        """
        if step_request.input != "y":
            step = await self.db.create_step(
                task_id=task_id,
                input=step_request.input if step_request else None,
                additional_properties=step_request.additional_input
                if step_request
                else None,
            )
            # utils.run
            artifacts = run(step.input)
            for artifact in artifacts:
                art = await self.db.create_artifact(
                    task_id=step.task_id,
                    file_name=artifact["file_name"],
                    uri=artifact["uri"],
                    agent_created=True,
                    step_id=step.step_id,
                )
                assert isinstance(
                    art, Artifact
                ), f"Artifact not instance of Artifact {type(art)}"
                step.artifacts.append(art)
            step.status = "completed"
        else:
            steps = await self.db.list_steps(task_id)
            artifacts = await self.db.list_artifacts(task_id)
            step = steps[-1]
            step.artifacts = artifacts
            step.output = "No more steps to run."
            step.is_last = True
        if isinstance(step.status, Status):
            step.status = step.status.value
        step.output = "Done some work"
        return JSONResponse(content=step.dict(), status_code=200)

    async def get_step(self, task_id: str, step_id: str) -> Step:
        """
        Get a step by ID.
        """
        if not task_id or not step_id:
            return Response(
                status_code=400, content="Task ID and step ID are required."
            )
        if not isinstance(task_id, str) or not isinstance(step_id, str):
            return Response(
                status_code=400, content="Task ID and step ID must be strings."
            )
        try:
            step = await self.db.get_step(task_id, step_id)
        except Exception as e:
            return Response(status_code=500, content=str(e))
        return step

    async def list_artifacts(self, task_id: str) -> typing.List[Artifact]:
        """
        List the artifacts that the task has created.
        """
        if not task_id:
            return Response(status_code=400, content="Task ID is required.")
        if not isinstance(task_id, str):
            return Response(status_code=400, content="Task ID must be a string.")
        try:
            artifacts = await self.db.list_artifacts(task_id)
        except Exception as e:
            return Response(status_code=500, content=str(e))
        return artifacts

    async def create_artifact(
        self,
        task_id: str,
        file: UploadFile | None = None,
        uri: str | None = None,
    ) -> Artifact:
        """
        Create an artifact for the task.
        """
        if not file and not uri:
            return Response(status_code=400, content="No file or uri provided")
        data = None
        if not uri:
            file_name = file.filename or str(uuid4())
            try:
                data = b""
                while contents := file.file.read(1024 * 1024):
                    data += contents
            except Exception as e:
                return Response(status_code=500, content=str(e))

            file_path = os.path.join(task_id / file_name)
            self.write(file_path, data)
            self.db.save_artifact(task_id, artifact)

        artifact = await self.create_artifact(
            task_id=task_id,
            file_name=file_name,
            uri=f"file://{file_path}",
            agent_created=False,
        )

        return artifact

    async def get_artifact(self, task_id: str, artifact_id: str) -> Artifact:
        """
        Get an artifact by ID.
        """
        artifact = await self.db.get_artifact(task_id, artifact_id)
        if not artifact.uri.startswith("file://"):
            return Response(
                status_code=500, content="Loading from none file uri not implemented"
            )
        file_path = artifact.uri.split("file://")[1]
        if not self.workspace.exists(file_path):
            return Response(status_code=500, content="File not found")
        retrieved_artifact = self.workspace.read(file_path)
        path = artifact.file_name
        with open(path, "wb") as f:
            f.write(retrieved_artifact)
        return FileResponse(
            # Note: mimetype is guessed in the FileResponse constructor
            path=path,
            filename=artifact.file_name,
        )