chore: add add_dataset_elements function

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

gptscript/gptscript.py

@@ -270,6 +270,30 @@ async def add_dataset_element(self, workspace_id: str, datasetID: str, elementNa
         )
         return DatasetElementMeta.model_validate_json(res)
 
+    async def add_dataset_elements(self, workspace_id: str, datasetID: str, elements: List[DatasetElement]) -> str:
+        if workspace_id == "":
+            workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]
+
+        if datasetID == "":
+            raise ValueError("datasetID cannot be empty")
+        elif not elements:
+            raise ValueError("elements cannot be empty")
+
+        res = await self._run_basic_command(
+            "datasets/add-elements",
+            {
+                "input": json.dumps({
+                    "datasetID": datasetID,
+                    "elements": [element.model_dump() for element in elements],
+                }),
+                "workspaceID": workspace_id,
+                "datasetToolRepo": self.opts.DatasetToolRepo,
+                "env": self.opts.Env
+            }
+        )
+        return res
+
+
     async def list_dataset_elements(self, workspace_id: str, datasetID: str) -> List[DatasetElementMeta]:
         if workspace_id == "":
             workspace_id = os.environ["GPTSCRIPT_WORKSPACE_ID"]

tests/test_gptscript.py

@@ -11,6 +11,7 @@
 
 from gptscript.confirm import AuthResponse
 from gptscript.credentials import Credential
+from gptscript.datasets import DatasetElement
 from gptscript.exec_utils import get_env
 from gptscript.frame import RunEventType, CallFrame, RunFrame, RunState, PromptFrame
 from gptscript.gptscript import GPTScript
@@ -775,17 +776,35 @@ async def test_datasets(gptscript):
     assert element_meta.name == "element1", "Expected element name to match"
     assert element_meta.description == "element1 description", "Expected element description to match"
 
-    # Get the element
-    element = await gptscript.get_dataset_element(workspace_id, dataset.id, "element1")
-    assert element.name == "element1", "Expected element name to match"
-    assert element.contents == "element1 contents", "Expected element contents to match"
-    assert element.description == "element1 description", "Expected element description to match"
+    # Add two more elements
+    await gptscript.add_dataset_elements(workspace_id, dataset.id, [
+        DatasetElement(name="element2", contents="element2 contents", description="element2 description"),
+        DatasetElement(name="element3", contents="element3 contents", description="element3 description"),
+    ])
+
+    # Get the elements
+    e1 = await gptscript.get_dataset_element(workspace_id, dataset.id, "element1")
+    assert e1.name == "element1", "Expected element name to match"
+    assert e1.contents == "element1 contents", "Expected element contents to match"
+    assert e1.description == "element1 description", "Expected element description to match"
+    e2 = await gptscript.get_dataset_element(workspace_id, dataset.id, "element2")
+    assert e2.name == "element2", "Expected element name to match"
+    assert e2.contents == "element2 contents", "Expected element contents to match"
+    assert e2.description == "element2 description", "Expected element description to match"
+    e3 = await gptscript.get_dataset_element(workspace_id, dataset.id, "element3")
+    assert e3.name == "element3", "Expected element name to match"
+    assert e3.contents == "element3 contents", "Expected element contents to match"
+    assert e3.description == "element3 description", "Expected element description to match"
 
     # List elements in the dataset
     elements = await gptscript.list_dataset_elements(workspace_id, dataset.id)
-    assert len(elements) == 1, "Expected one element in the dataset"
+    assert len(elements) == 3, "Expected one element in the dataset"
     assert elements[0].name == "element1", "Expected element name to match"
     assert elements[0].description == "element1 description", "Expected element description to match"
+    assert elements[1].name == "element2", "Expected element name to match"
+    assert elements[1].description == "element2 description", "Expected element description to match"
+    assert elements[2].name == "element3", "Expected element name to match"
+    assert elements[2].description == "element3 description", "Expected element description to match"
 
     # List datasets
     datasets = await gptscript.list_datasets(workspace_id)