Skip to content

Commit 94fe795

Browse files
thedadamsthedadams
committed
feat: add workspace API
Signed-off-by: Donnie Adams <donnie@acorn.io>
1 parent ed2e33b commit 94fe795

File tree

4 files changed

+200
-2
lines changed

4 files changed

+200
-2
lines changed

‎gptscript/gptscript.py‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import base64
12
import json
23
import os
34
import platform
@@ -291,6 +292,101 @@ async def get_dataset_element(self, workspace: str, datasetID: str, elementName:
291292
)
292293
return DatasetElement.model_validate_json(res)
293294

295+
async def create_workspace(self, provider_type: str, from_workspaces: list[str] = None) -> str:
296+
return await self._run_basic_command(
297+
"workspaces/create",
298+
{
299+
"providerType": provider_type,
300+
"fromWorkspaces": from_workspaces,
301+
"workspaceTool": self.opts.WorkspaceTool,
302+
"env": self.opts.Env,
303+
}
304+
)
305+
306+
async def delete_workspace(self, workspace_id: str = ""):
307+
if workspace_id == "":
308+
workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
309+
310+
await self._run_basic_command(
311+
"workspaces/delete",
312+
{
313+
"id": workspace_id,
314+
"workspaceTool": self.opts.WorkspaceTool,
315+
"env": self.opts.Env,
316+
}
317+
)
318+
319+
async def list_files_in_workspace(self, workspace_id: str = "", prefix: str = "") -> List[str]:
320+
if workspace_id == "":
321+
workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
322+
323+
return json.loads(await self._run_basic_command(
324+
"workspaces/list",
325+
{
326+
"id": workspace_id,
327+
"prefix": prefix,
328+
"workspaceTool": self.opts.WorkspaceTool,
329+
"env": self.opts.Env,
330+
}
331+
))
332+
333+
async def remove_all(self, workspace_id: str = "", with_prefix: str = ""):
334+
if workspace_id == "":
335+
workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
336+
337+
await self._run_basic_command(
338+
"workspaces/remove-all-with-prefix",
339+
{
340+
"id": workspace_id,
341+
"prefix": with_prefix,
342+
"workspaceTool": self.opts.WorkspaceTool,
343+
"env": self.opts.Env,
344+
}
345+
)
346+
347+
async def write_file_in_workspace(self, file_path: str, contents: bytes, workspace_id: str = ""):
348+
if workspace_id == "":
349+
workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
350+
351+
await self._run_basic_command(
352+
"workspaces/write-file",
353+
{
354+
"id": workspace_id,
355+
"filePath": file_path,
356+
"contents": base64.b64encode(contents).decode("utf-8") if contents is not None else None,
357+
"workspaceTool": self.opts.WorkspaceTool,
358+
"env": self.opts.Env,
359+
}
360+
)
361+
362+
async def delete_file_in_workspace(self, file_path: str, workspace_id: str = ""):
363+
if workspace_id == "":
364+
workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
365+
366+
await self._run_basic_command(
367+
"workspaces/delete-file",
368+
{
369+
"id": workspace_id,
370+
"filePath": file_path,
371+
"workspaceTool": self.opts.WorkspaceTool,
372+
"env": self.opts.Env,
373+
}
374+
)
375+
376+
async def read_file_in_workspace(self, file_path: str, workspace_id: str = "") -> bytes:
377+
if workspace_id == "":
378+
workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
379+
380+
return base64.b64decode(await self._run_basic_command(
381+
"workspaces/read-file",
382+
{
383+
"id": workspace_id,
384+
"filePath": file_path,
385+
"workspaceTool": self.opts.WorkspaceTool,
386+
"env": self.opts.Env,
387+
}
388+
))
389+
294390

295391
def _get_command():
296392
if os.getenv("GPTSCRIPT_BIN") is not None:

‎gptscript/opts.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ def __init__(
1313
defaultModel: str = "",
1414
cacheDir: str = "",
1515
datasetToolRepo: str = "",
16+
workspaceTool: str = "",
1617
env: list[str] = None,
1718
):
1819
self.URL = url
@@ -23,6 +24,7 @@ def __init__(
2324
self.DefaultModelProvider = defaultModelProvider
2425
self.CacheDir = cacheDir
2526
self.DatasetToolRepo = datasetToolRepo
27+
self.WorkspaceTool = workspaceTool
2628
if env is None:
2729
env = [f"{k}={v}" for k, v in os.environ.items()]
2830
elif isinstance(env, dict):
@@ -41,6 +43,7 @@ def merge(self, other: Self) -> Self:
4143
cp.DefaultModelProvider = other.DefaultModelProvider if other.DefaultModelProvider != "" else self.DefaultModelProvider
4244
cp.CacheDir = other.CacheDir if other.CacheDir != "" else self.CacheDir
4345
cp.DatasetToolRepo = other.DatasetToolRepo if other.DatasetToolRepo != "" else self.DatasetToolRepo
46+
cp.WorkspaceTool = other.WorkspaceTool if other.WorkspaceTool != "" else self.WorkspaceTool
4447
cp.Env = (other.Env or [])
4548
cp.Env.extend(self.Env or [])
4649
return cp
@@ -81,8 +84,10 @@ def __init__(self,
8184
defaultModel: str = "",
8285
cacheDir: str = "",
8386
datasetToolDir: str = "",
87+
workspaceTool: str = "",
8488
):
85-
super().__init__(url, token, apiKey, baseURL, defaultModelProvider, defaultModel, cacheDir, datasetToolDir, env)
89+
super().__init__(url, token, apiKey, baseURL, defaultModelProvider, defaultModel, cacheDir, datasetToolDir,
90+
workspaceTool, env)
8691
self.input = input
8792
self.disableCache = disableCache
8893
self.subTool = subTool

‎tests/test_gptscript.py‎

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -757,6 +757,7 @@ async def test_credentials(gptscript):
757757
res = await gptscript.delete_credential(name=name)
758758
assert not res.startswith("an error occurred"), "Unexpected error deleting credential: " + res
759759

760+
760761
@pytest.mark.asyncio
761762
async def test_datasets(gptscript):
762763
with tempfile.TemporaryDirectory(prefix="py-gptscript_") as tempdir:
@@ -770,7 +771,8 @@ async def test_datasets(gptscript):
770771
assert len(dataset.elements) == 0, "Expected dataset elements to be empty"
771772

772773
# Add an element
773-
element_meta = await gptscript.add_dataset_element(tempdir, dataset.id, "element1", "element1 contents", "element1 description")
774+
element_meta = await gptscript.add_dataset_element(tempdir, dataset.id, "element1", "element1 contents",
775+
"element1 description")
774776
assert element_meta.name == "element1", "Expected element name to match"
775777
assert element_meta.description == "element1 description", "Expected element description to match"
776778

@@ -792,3 +794,93 @@ async def test_datasets(gptscript):
792794
assert datasets[0].id == dataset.id, "Expected dataset id to match"
793795
assert datasets[0].name == dataset_name, "Expected dataset name to match"
794796
assert datasets[0].description == "this is a test dataset", "Expected dataset description to match"
797+
798+
799+
@pytest.mark.asyncio
800+
async def test_create_and_delete_workspace(gptscript):
801+
workspace_id = await gptscript.create_workspace("directory")
802+
assert workspace_id != "" and workspace_id.startswith("directory://"), "Expected workspace id to be set"
803+
await gptscript.delete_workspace(workspace_id)
804+
805+
806+
@pytest.mark.asyncio
807+
async def test_create_read_and_delete_file_in_workspace(gptscript):
808+
workspace_id = await gptscript.create_workspace("directory")
809+
await gptscript.write_file_in_workspace("test.txt", b"test", workspace_id)
810+
contents = await gptscript.read_file_in_workspace("test.txt", workspace_id)
811+
assert contents == b"test"
812+
await gptscript.delete_file_in_workspace("test.txt", workspace_id)
813+
await gptscript.delete_workspace(workspace_id)
814+
815+
816+
@pytest.mark.asyncio
817+
async def test_ls_complex_workspace(gptscript):
818+
workspace_id = await gptscript.create_workspace("directory")
819+
await gptscript.write_file_in_workspace("test/test1.txt", b"hello1", workspace_id)
820+
await gptscript.write_file_in_workspace("test1/test2.txt", b"hello2", workspace_id)
821+
await gptscript.write_file_in_workspace("test1/test3.txt", b"hello3", workspace_id)
822+
await gptscript.write_file_in_workspace(".hidden.txt", b"hidden", workspace_id)
823+
824+
files = await gptscript.list_files_in_workspace(workspace_id)
825+
assert len(files) == 4
826+
827+
files = await gptscript.list_files_in_workspace(workspace_id, prefix="test1")
828+
assert len(files) == 2
829+
830+
await gptscript.remove_all(workspace_id, with_prefix="test1")
831+
832+
files = await gptscript.list_files_in_workspace(workspace_id)
833+
assert len(files) == 2
834+
835+
await gptscript.delete_workspace(workspace_id)
836+
837+
838+
@pytest.mark.skipif(
839+
os.environ.get("AWS_ACCESS_KEY_ID") is None or os.environ.get("AWS_SECRET_ACCESS_KEY") is None,
840+
reason="AWS credentials not set",
841+
)
842+
@pytest.mark.asyncio
843+
async def test_create_and_delete_workspace_s3(gptscript):
844+
workspace_id = await gptscript.create_workspace("s3")
845+
assert workspace_id != "" and workspace_id.startswith("s3://"), "Expected workspace id to be set"
846+
await gptscript.delete_workspace(workspace_id)
847+
848+
849+
@pytest.mark.skipif(
850+
os.environ.get("AWS_ACCESS_KEY_ID") is None or os.environ.get("AWS_SECRET_ACCESS_KEY") is None,
851+
reason="AWS credentials not set",
852+
)
853+
@pytest.mark.asyncio
854+
async def test_create_read_and_delete_file_in_workspaces3(gptscript):
855+
workspace_id = await gptscript.create_workspace("s3")
856+
await gptscript.write_file_in_workspace("test.txt", b"test", workspace_id)
857+
contents = await gptscript.read_file_in_workspace("test.txt", workspace_id)
858+
assert contents == b"test"
859+
await gptscript.delete_file_in_workspace("test.txt", workspace_id)
860+
await gptscript.delete_workspace(workspace_id)
861+
862+
863+
@pytest.mark.skipif(
864+
os.environ.get("AWS_ACCESS_KEY_ID") is None or os.environ.get("AWS_SECRET_ACCESS_KEY") is None,
865+
reason="AWS credentials not set",
866+
)
867+
@pytest.mark.asyncio
868+
async def test_ls_complex_workspace_s3(gptscript):
869+
workspace_id = await gptscript.create_workspace("s3")
870+
await gptscript.write_file_in_workspace("test/test1.txt", b"hello1", workspace_id)
871+
await gptscript.write_file_in_workspace("test1/test2.txt", b"hello2", workspace_id)
872+
await gptscript.write_file_in_workspace("test1/test3.txt", b"hello3", workspace_id)
873+
await gptscript.write_file_in_workspace(".hidden.txt", b"hidden", workspace_id)
874+
875+
files = await gptscript.list_files_in_workspace(workspace_id)
876+
assert len(files) == 4
877+
878+
files = await gptscript.list_files_in_workspace(workspace_id, prefix="test1")
879+
assert len(files) == 2
880+
881+
await gptscript.remove_all(workspace_id, with_prefix="test1")
882+
883+
files = await gptscript.list_files_in_workspace(workspace_id)
884+
assert len(files) == 2
885+
886+
await gptscript.delete_workspace(workspace_id)

‎tox.ini‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,11 @@ deps =
1111
passenv =
1212
OPENAI_API_KEY
1313
ANTHROPIC_API_KEY
14+
AWS_ACCESS_KEY_ID
15+
AWS_SECRET_ACCESS_KEY
16+
AWS_REGION
17+
WORKSPACE_PROVIDER_S3_BUCKET
18+
WORKSPACE_PROVIDER_S3_BASE_ENDPOINT
1419
GPTSCRIPT_BIN
1520
GPTSCRIPT_URL
1621
GPTSCRIPT_CONFIG_FILE

0 commit comments

Comments
 (0)