feat: add workspace API

Signed-off-by: Donnie Adams <[email protected]>

Author: thedadams

gptscript/gptscript.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import os
 import platform
@@ -291,6 +292,101 @@ async def get_dataset_element(self, workspace: str, datasetID: str, elementName:
         )
         return DatasetElement.model_validate_json(res)
 
+    async def create_workspace(self, provider_type: str, from_workspaces: list[str] = None) -> str:
+        return await self._run_basic_command(
+            "workspaces/create",
+            {
+                "providerType": provider_type,
+                "fromWorkspaces": from_workspaces,
+                "workspaceTool": self.opts.WorkspaceTool,
+                "env": self.opts.Env,
+            }
+        )
+
+    async def delete_workspace(self, workspace_id: str = ""):
+        if workspace_id == "":
+            workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
+
+        await self._run_basic_command(
+            "workspaces/delete",
+            {
+                "id": workspace_id,
+                "workspaceTool": self.opts.WorkspaceTool,
+                "env": self.opts.Env,
+            }
+        )
+
+    async def list_files_in_workspace(self, workspace_id: str = "", prefix: str = "") -> List[str]:
+        if workspace_id == "":
+            workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
+
+        return json.loads(await self._run_basic_command(
+            "workspaces/list",
+            {
+                "id": workspace_id,
+                "prefix": prefix,
+                "workspaceTool": self.opts.WorkspaceTool,
+                "env": self.opts.Env,
+            }
+        ))
+
+    async def remove_all(self, workspace_id: str = "", with_prefix: str = ""):
+        if workspace_id == "":
+            workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
+
+        await self._run_basic_command(
+            "workspaces/remove-all-with-prefix",
+            {
+                "id": workspace_id,
+                "prefix": with_prefix,
+                "workspaceTool": self.opts.WorkspaceTool,
+                "env": self.opts.Env,
+            }
+        )
+
+    async def write_file_in_workspace(self, file_path: str, contents: bytes, workspace_id: str = ""):
+        if workspace_id == "":
+            workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
+
+        await self._run_basic_command(
+            "workspaces/write-file",
+            {
+                "id": workspace_id,
+                "filePath": file_path,
+                "contents": base64.b64encode(contents).decode("utf-8") if contents is not None else None,
+                "workspaceTool": self.opts.WorkspaceTool,
+                "env": self.opts.Env,
+            }
+        )
+
+    async def delete_file_in_workspace(self, file_path: str, workspace_id: str = ""):
+        if workspace_id == "":
+            workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
+
+        await self._run_basic_command(
+            "workspaces/delete-file",
+            {
+                "id": workspace_id,
+                "filePath": file_path,
+                "workspaceTool": self.opts.WorkspaceTool,
+                "env": self.opts.Env,
+            }
+        )
+
+    async def read_file_in_workspace(self, file_path: str, workspace_id: str = "") -> bytes:
+        if workspace_id == "":
+            workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
+
+        return base64.b64decode(await self._run_basic_command(
+            "workspaces/read-file",
+            {
+                "id": workspace_id,
+                "filePath": file_path,
+                "workspaceTool": self.opts.WorkspaceTool,
+                "env": self.opts.Env,
+            }
+        ))
+
 
 def _get_command():
     if os.getenv("GPTSCRIPT_BIN") is not None:

gptscript/opts.py

@@ -13,6 +13,7 @@ def __init__(
             defaultModel: str = "",
             cacheDir: str = "",
             datasetToolRepo: str = "",
+            workspaceTool: str = "",
             env: list[str] = None,
     ):
         self.URL = url
@@ -23,6 +24,7 @@ def __init__(
         self.DefaultModelProvider = defaultModelProvider
         self.CacheDir = cacheDir
         self.DatasetToolRepo = datasetToolRepo
+        self.WorkspaceTool = workspaceTool
         if env is None:
             env = [f"{k}={v}" for k, v in os.environ.items()]
         elif isinstance(env, dict):
@@ -41,6 +43,7 @@ def merge(self, other: Self) -> Self:
         cp.DefaultModelProvider = other.DefaultModelProvider if other.DefaultModelProvider != "" else self.DefaultModelProvider
         cp.CacheDir = other.CacheDir if other.CacheDir != "" else self.CacheDir
         cp.DatasetToolRepo = other.DatasetToolRepo if other.DatasetToolRepo != "" else self.DatasetToolRepo
+        cp.WorkspaceTool = other.WorkspaceTool if other.WorkspaceTool != "" else self.WorkspaceTool
         cp.Env = (other.Env or [])
         cp.Env.extend(self.Env or [])
         return cp
@@ -81,8 +84,10 @@ def __init__(self,
                  defaultModel: str = "",
                  cacheDir: str = "",
                  datasetToolDir: str = "",
+                 workspaceTool: str = "",
                  ):
-        super().__init__(url, token, apiKey, baseURL, defaultModelProvider, defaultModel, cacheDir, datasetToolDir, env)
+        super().__init__(url, token, apiKey, baseURL, defaultModelProvider, defaultModel, cacheDir, datasetToolDir,
+                         workspaceTool, env)
         self.input = input
         self.disableCache = disableCache
         self.subTool = subTool

tests/test_gptscript.py

@@ -757,6 +757,7 @@ async def test_credentials(gptscript):
     res = await gptscript.delete_credential(name=name)
     assert not res.startswith("an error occurred"), "Unexpected error deleting credential: " + res
 
+
 @pytest.mark.asyncio
 async def test_datasets(gptscript):
     with tempfile.TemporaryDirectory(prefix="py-gptscript_") as tempdir:
@@ -770,7 +771,8 @@ async def test_datasets(gptscript):
         assert len(dataset.elements) == 0, "Expected dataset elements to be empty"
 
         # Add an element
-        element_meta = await gptscript.add_dataset_element(tempdir, dataset.id, "element1", "element1 contents", "element1 description")
+        element_meta = await gptscript.add_dataset_element(tempdir, dataset.id, "element1", "element1 contents",
+                                                           "element1 description")
         assert element_meta.name == "element1", "Expected element name to match"
         assert element_meta.description == "element1 description", "Expected element description to match"
 
@@ -792,3 +794,93 @@ async def test_datasets(gptscript):
         assert datasets[0].id == dataset.id, "Expected dataset id to match"
         assert datasets[0].name == dataset_name, "Expected dataset name to match"
         assert datasets[0].description == "this is a test dataset", "Expected dataset description to match"
+
+
+@pytest.mark.asyncio
+async def test_create_and_delete_workspace(gptscript):
+    workspace_id = await gptscript.create_workspace("directory")
+    assert workspace_id != "" and workspace_id.startswith("directory://"), "Expected workspace id to be set"
+    await gptscript.delete_workspace(workspace_id)
+
+
+@pytest.mark.asyncio
+async def test_create_read_and_delete_file_in_workspace(gptscript):
+    workspace_id = await gptscript.create_workspace("directory")
+    await gptscript.write_file_in_workspace("test.txt", b"test", workspace_id)
+    contents = await gptscript.read_file_in_workspace("test.txt", workspace_id)
+    assert contents == b"test"
+    await gptscript.delete_file_in_workspace("test.txt", workspace_id)
+    await gptscript.delete_workspace(workspace_id)
+
+
+@pytest.mark.asyncio
+async def test_ls_complex_workspace(gptscript):
+    workspace_id = await gptscript.create_workspace("directory")
+    await gptscript.write_file_in_workspace("test/test1.txt", b"hello1", workspace_id)
+    await gptscript.write_file_in_workspace("test1/test2.txt", b"hello2", workspace_id)
+    await gptscript.write_file_in_workspace("test1/test3.txt", b"hello3", workspace_id)
+    await gptscript.write_file_in_workspace(".hidden.txt", b"hidden", workspace_id)
+
+    files = await gptscript.list_files_in_workspace(workspace_id)
+    assert len(files) == 4
+
+    files = await gptscript.list_files_in_workspace(workspace_id, prefix="test1")
+    assert len(files) == 2
+
+    await gptscript.remove_all(workspace_id, with_prefix="test1")
+
+    files = await gptscript.list_files_in_workspace(workspace_id)
+    assert len(files) == 2
+
+    await gptscript.delete_workspace(workspace_id)
+
+
+@pytest.mark.skipif(
+    os.environ.get("AWS_ACCESS_KEY_ID") is None or os.environ.get("AWS_SECRET_ACCESS_KEY") is None,
+    reason="AWS credentials not set",
+)
+@pytest.mark.asyncio
+async def test_create_and_delete_workspace_s3(gptscript):
+    workspace_id = await gptscript.create_workspace("s3")
+    assert workspace_id != "" and workspace_id.startswith("s3://"), "Expected workspace id to be set"
+    await gptscript.delete_workspace(workspace_id)
+
+
+@pytest.mark.skipif(
+    os.environ.get("AWS_ACCESS_KEY_ID") is None or os.environ.get("AWS_SECRET_ACCESS_KEY") is None,
+    reason="AWS credentials not set",
+)
+@pytest.mark.asyncio
+async def test_create_read_and_delete_file_in_workspaces3(gptscript):
+    workspace_id = await gptscript.create_workspace("s3")
+    await gptscript.write_file_in_workspace("test.txt", b"test", workspace_id)
+    contents = await gptscript.read_file_in_workspace("test.txt", workspace_id)
+    assert contents == b"test"
+    await gptscript.delete_file_in_workspace("test.txt", workspace_id)
+    await gptscript.delete_workspace(workspace_id)
+
+
+@pytest.mark.skipif(
+    os.environ.get("AWS_ACCESS_KEY_ID") is None or os.environ.get("AWS_SECRET_ACCESS_KEY") is None,
+    reason="AWS credentials not set",
+)
+@pytest.mark.asyncio
+async def test_ls_complex_workspace_s3(gptscript):
+    workspace_id = await gptscript.create_workspace("s3")
+    await gptscript.write_file_in_workspace("test/test1.txt", b"hello1", workspace_id)
+    await gptscript.write_file_in_workspace("test1/test2.txt", b"hello2", workspace_id)
+    await gptscript.write_file_in_workspace("test1/test3.txt", b"hello3", workspace_id)
+    await gptscript.write_file_in_workspace(".hidden.txt", b"hidden", workspace_id)
+
+    files = await gptscript.list_files_in_workspace(workspace_id)
+    assert len(files) == 4
+
+    files = await gptscript.list_files_in_workspace(workspace_id, prefix="test1")
+    assert len(files) == 2
+
+    await gptscript.remove_all(workspace_id, with_prefix="test1")
+
+    files = await gptscript.list_files_in_workspace(workspace_id)
+    assert len(files) == 2
+
+    await gptscript.delete_workspace(workspace_id)

tox.ini

@@ -11,6 +11,11 @@ deps =
 passenv =
 	OPENAI_API_KEY
 	ANTHROPIC_API_KEY
+	AWS_ACCESS_KEY_ID
+	AWS_SECRET_ACCESS_KEY
+	AWS_REGION
+	WORKSPACE_PROVIDER_S3_BUCKET
+	WORKSPACE_PROVIDER_S3_BASE_ENDPOINT
 	GPTSCRIPT_BIN
 	GPTSCRIPT_URL
 	GPTSCRIPT_CONFIG_FILE