Generator Request Tags

src/steamship/data/operations/generator.py

@@ -5,6 +5,7 @@
 from steamship.base.request import Request
 from steamship.base.response import Response
 from steamship.data.block import Block
+from steamship.data.tags.tag import Tag
 
 
 class GenerateRequest(Request):
@@ -67,6 +68,10 @@ class GenerateRequest(Request):
     # Default behavior if not provided is streaming=false
     streaming: Optional[bool] = None
 
+    # Tags which will be applied to all Blocks that are generated as part of
+    # this request.
+    tags: Optional[List[Tag]] = None
+
 
 class GenerateResponse(Response):
     blocks: List[Block]

src/steamship/data/plugin/plugin_instance.py

@@ -22,6 +22,7 @@
     HostingTimeout,
     HostingType,
 )
+from steamship.data.tags.tag import Tag
 from steamship.plugin.inputs.export_plugin_input import ExportPluginInput
 from steamship.plugin.inputs.training_parameter_plugin_input import TrainingParameterPluginInput
 from steamship.plugin.outputs.train_plugin_output import TrainPluginOutput
@@ -131,6 +132,7 @@ def generate(
         make_output_public: Optional[bool] = None,
         options: Optional[dict] = None,
         streaming: Optional[bool] = None,
+        tags: Optional[List[Tag]] = None,
     ) -> Task[GenerateResponse]:
         """See GenerateRequest for description of parameter options"""
         req = GenerateRequest(
@@ -148,6 +150,7 @@ def generate(
             make_output_public=make_output_public,
             options=options,
             streaming=streaming,
+            tags=tags,
         )
         return self.client.post("plugin/instance/generate", req, expect=GenerateResponse)
 

tests/steamship_tests/plugin/integration/test_e2e_generator.py

@@ -5,7 +5,7 @@
 from steamship_tests.utils.deployables import deploy_plugin
 from steamship_tests.utils.fixtures import get_steamship_client
 
-from steamship import Block, File, MimeTypes, PluginInstance, Steamship
+from steamship import Block, File, MimeTypes, PluginInstance, Steamship, Tag
 
 
 def test_e2e_generator():
@@ -218,3 +218,28 @@
 
     assert response.text == "PRETEND THIS IS THE DATA OF AN IMAGE"
     assert response.headers["content-type"] == MimeTypes.PNG
+
+
+@pytest.mark.usefixtures("client")
+def test_generation_with_tags(client: Steamship):
+    plugin_instance = PluginInstance.create(client, plugin_handle="test-generator")
+    file = File.create(client, blocks=[Block(text="One block"), Block(text="two blocks")])
+    tags = [
+        Tag(kind="test_kind_0", name="test_name_0", value={"test_value": "test_value_0"}),
+        Tag(kind="test_kind_1", name="test_name_1", value={"test_value": "test_value_1"}),
+    ]
+    generate_task = plugin_instance.generate(
+        input_file_id=file.id, append_output_to_file=True, tags=tags
+    )
+    blocks = generate_task.wait().blocks
+    assert blocks is not None
+    assert len(blocks) == 2
+    for block in blocks:
+        result_tags = block.tags
+        assert result_tags is not None
+        assert len(result_tags) == 2
+        for i in range(2):
+            result_tag = result_tags[i]
+            assert result_tag.kind == f"test_kind_{i}"
+            assert result_tag.name == f"test_name_{i}"
+            assert result_tag.value == {"test_value": f"test_value_{i}"}