From 62df429ffa7179ef761a65160950d9ce4bf87171 Mon Sep 17 00:00:00 2001 From: Reinier van der Leer Date: Tue, 17 Oct 2023 18:37:56 -0700 Subject: [PATCH] AutoGPT: Fix artifact creation & downloading --- .../autogpt/app/agent_protocol_server.py | 7 ++-- .../autogpt/commands/file_operations.py | 4 +-- .../autogpt/file_workspace/file_workspace.py | 7 ++-- .../tests/unit/test_file_operations.py | 32 +++++++++++-------- 4 files changed, 30 insertions(+), 20 deletions(-) diff --git a/autogpts/autogpt/autogpt/app/agent_protocol_server.py b/autogpts/autogpt/autogpt/app/agent_protocol_server.py index 279c2044..00c41640 100644 --- a/autogpts/autogpt/autogpt/app/agent_protocol_server.py +++ b/autogpts/autogpt/autogpt/app/agent_protocol_server.py @@ -6,7 +6,7 @@ from uuid import uuid4 from fastapi import APIRouter, FastAPI, UploadFile from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import FileResponse, RedirectResponse, StreamingResponse +from fastapi.responses import RedirectResponse, StreamingResponse from fastapi.staticfiles import StaticFiles from forge.sdk.db import AgentDB from forge.sdk.errors import NotFoundError @@ -317,7 +317,7 @@ class AgentProtocolServer: file_path = os.path.join(relative_path, file_name) workspace = get_task_agent_file_workspace(task_id, self.agent_manager) - workspace.write_file(file_path, data) + await workspace.write_file(file_path, data) artifact = await self.db.create_artifact( task_id=task_id, @@ -365,7 +365,8 @@ def get_task_agent_file_workspace( root=agent_manager.get_agent_dir( agent_id=task_agent_id(task_id), must_exist=True, - ), + ) + / "workspace", restrict_to_root=True, ) diff --git a/autogpts/autogpt/autogpt/commands/file_operations.py b/autogpts/autogpt/autogpt/commands/file_operations.py index 8e3b54ab..29935dc8 100644 --- a/autogpts/autogpt/autogpt/commands/file_operations.py +++ b/autogpts/autogpt/autogpt/commands/file_operations.py @@ -209,7 +209,7 @@ def ingest_file( aliases=["create_file"], ) @sanitize_path_arg("filename") -def write_to_file(filename: Path, contents: str, agent: Agent) -> str: +async def write_to_file(filename: Path, contents: str, agent: Agent) -> str: """Write contents to a file Args: @@ -225,7 +225,7 @@ def write_to_file(filename: Path, contents: str, agent: Agent) -> str: directory = os.path.dirname(filename) os.makedirs(directory, exist_ok=True) - agent.workspace.write_file(filename, contents) + await agent.workspace.write_file(filename, contents) log_operation("write", filename, agent, checksum) return f"File {filename.name} has been written successfully." diff --git a/autogpts/autogpt/autogpt/file_workspace/file_workspace.py b/autogpts/autogpt/autogpt/file_workspace/file_workspace.py index 37da359e..61b3c42f 100644 --- a/autogpts/autogpt/autogpt/file_workspace/file_workspace.py +++ b/autogpts/autogpt/autogpt/file_workspace/file_workspace.py @@ -3,6 +3,7 @@ The FileWorkspace class provides an interface for interacting with a file worksp """ from __future__ import annotations +import inspect import logging from pathlib import Path from typing import Any, Callable, Optional @@ -65,7 +66,7 @@ class FileWorkspace: with self.open_file(path, "rb" if binary else "r") as file: return file.read() - def write_file(self, path: str | Path, content: str | bytes): + async def write_file(self, path: str | Path, content: str | bytes): """Write to a file in the workspace.""" with self.open_file(path, "wb" if type(content) is bytes else "w") as file: file.write(content) @@ -74,7 +75,9 @@ class FileWorkspace: path = Path(path) if path.is_absolute(): path = path.relative_to(self.root) - self.on_write_file(path) + res = self.on_write_file(path) + if inspect.isawaitable(res): + await res def list_files(self, path: str | Path = "."): """List all files in a directory in the workspace.""" diff --git a/autogpts/autogpt/tests/unit/test_file_operations.py b/autogpts/autogpt/tests/unit/test_file_operations.py index 3636ce80..6d4c69ac 100644 --- a/autogpts/autogpt/tests/unit/test_file_operations.py +++ b/autogpts/autogpt/tests/unit/test_file_operations.py @@ -204,53 +204,59 @@ def test_read_file_not_found(agent: Agent): file_ops.read_file(filename, agent=agent) -def test_write_to_file_relative_path(test_file_name: Path, agent: Agent): +@pytest.mark.asyncio +async def test_write_to_file_relative_path(test_file_name: Path, agent: Agent): new_content = "This is new content.\n" - file_ops.write_to_file(str(test_file_name), new_content, agent=agent) + await file_ops.write_to_file(test_file_name, new_content, agent=agent) with open(agent.workspace.get_path(test_file_name), "r", encoding="utf-8") as f: content = f.read() assert content == new_content -def test_write_to_file_absolute_path(test_file_path: Path, agent: Agent): +@pytest.mark.asyncio +async def test_write_to_file_absolute_path(test_file_path: Path, agent: Agent): new_content = "This is new content.\n" - file_ops.write_to_file(str(test_file_path), new_content, agent=agent) + await file_ops.write_to_file(test_file_path, new_content, agent=agent) with open(test_file_path, "r", encoding="utf-8") as f: content = f.read() assert content == new_content -def test_write_file_logs_checksum(test_file_name: Path, agent: Agent): +@pytest.mark.asyncio +async def test_write_file_logs_checksum(test_file_name: Path, agent: Agent): new_content = "This is new content.\n" new_checksum = file_ops.text_checksum(new_content) - file_ops.write_to_file(str(test_file_name), new_content, agent=agent) + await file_ops.write_to_file(test_file_name, new_content, agent=agent) with open(agent.file_manager.file_ops_log_path, "r", encoding="utf-8") as f: log_entry = f.read() assert log_entry == f"write: {test_file_name} #{new_checksum}\n" -def test_write_file_fails_if_content_exists(test_file_name: Path, agent: Agent): +@pytest.mark.asyncio +async def test_write_file_fails_if_content_exists(test_file_name: Path, agent: Agent): new_content = "This is new content.\n" file_ops.log_operation( "write", - str(test_file_name), + test_file_name, agent=agent, checksum=file_ops.text_checksum(new_content), ) with pytest.raises(DuplicateOperationError): - file_ops.write_to_file(str(test_file_name), new_content, agent=agent) + await file_ops.write_to_file(test_file_name, new_content, agent=agent) -def test_write_file_succeeds_if_content_different( +@pytest.mark.asyncio +async def test_write_file_succeeds_if_content_different( test_file_with_content_path: Path, agent: Agent ): new_content = "This is different content.\n" - file_ops.write_to_file(str(test_file_with_content_path), new_content, agent=agent) + await file_ops.write_to_file(test_file_with_content_path, new_content, agent=agent) -def test_append_to_file(test_nested_file: Path, agent: Agent): +@pytest.mark.asyncio +async def test_append_to_file(test_nested_file: Path, agent: Agent): append_text = "This is appended text.\n" - file_ops.write_to_file(test_nested_file, append_text, agent=agent) + await file_ops.write_to_file(test_nested_file, append_text, agent=agent) file_ops.append_to_file(test_nested_file, append_text, agent=agent)