[SHIP-834] [SHIP-818] Streaming Generators (#548)

This PR includes:

- Streaming Generator interface
- Additional methods and status on `Block` for streaming

Depends on: https://github.com/nludb/nludb/pull/631

(tests will fail until deploy of above)

---------

Co-authored-by: Ted Benson <edward.benson@gmail.com>

Author: dkolas

src/steamship/base/client.py

@@ -315,13 +315,15 @@ def _headers(  # noqa: C901
         return headers
 
     @staticmethod
-    def _prepare_data(payload: Union[Request, dict]):
+    def _prepare_data(payload: Union[Request, dict, bytes]):
         if payload is None:
             data = {}
         elif isinstance(payload, dict):
             data = payload
         elif isinstance(payload, BaseModel):
             data = payload.dict(by_alias=True)
+        elif isinstance(payload, bytes):
+            data = payload
         else:
             raise RuntimeError(f"Unable to parse payload of type {type(payload)}")
 
@@ -407,7 +409,7 @@ def call(  # noqa: C901
         self,
         verb: Verb,
         operation: str,
-        payload: Union[Request, dict] = None,
+        payload: Union[Request, dict, bytes] = None,
         file: Any = None,
         expect: Type[T] = None,
         debug: bool = False,
@@ -464,7 +466,10 @@ def call(  # noqa: C901
                 files = self._prepare_multipart_data(data, file)
                 resp = self._session.post(url, files=files, headers=headers, timeout=timeout_s)
             else:
-                resp = self._session.post(url, json=data, headers=headers, timeout=timeout_s)
+                if isinstance(data, bytes):
+                    resp = self._session.post(url, data=data, headers=headers, timeout=timeout_s)
+                else:
+                    resp = self._session.post(url, json=data, headers=headers, timeout=timeout_s)
         elif verb == Verb.GET:
             resp = self._session.get(url, params=data, headers=headers, timeout=timeout_s)
         else:
@@ -555,7 +560,7 @@ def call(  # noqa: C901
     def post(
         self,
         operation: str,
-        payload: Union[Request, dict, BaseModel] = None,
+        payload: Union[Request, dict, BaseModel, bytes] = None,
         file: Any = None,
         expect: Any = None,
         debug: bool = False,

src/steamship/cli/deploy.py

@@ -241,6 +241,7 @@ def _create_version(self, client: Steamship, manifest: Manifest, thing_id: str):
             handle=manifest.version,
             filename=f"build/archives/{manifest.handle}_v{manifest.version}.zip",
             plugin_id=thing_id,
+            streaming=manifest.plugin.streaming,
         )
 
     def create_object(self, client: Steamship, manifest: Manifest):

src/steamship/cli/manifest_init_wizard.py

@@ -58,7 +58,11 @@ def manifest_init_wizard(client: Steamship):
             trainable = click.confirm("Is the plugin trainable?", default=False)
         else:
             trainable = False
-        plugin_detail = PluginConfig(isTrainable=trainable, type=plugin_type)
+        if plugin_type == "generator":
+            streaming = click.confirm("Will the plugin stream its results?", default=False)
+        else:
+            streaming = False
+        plugin_detail = PluginConfig(isTrainable=trainable, type=plugin_type, streaming=streaming)
 
     public = click.confirm(f"Do you want this {deployable_type} to be public?", default=True)
 
@@ -71,7 +75,7 @@ def manifest_init_wizard(client: Steamship):
     if public:
         tagline = click.prompt(f"Want to give the {deployable_type} a tagline?", default="")
         author_github = click.prompt(
-            "If you'd like this associated with your github account, please your github username",
+            "If you'd like this associated with your github account, please input your github username",
             default="",
         )
 

src/steamship/data/block.py

@@ -27,6 +27,12 @@ class BlockUploadType(str, Enum):
     NONE = "none"  # No upload; plain text only.
 
 
+class StreamState(str, Enum):
+    STARTED = "started"  # A producer has begun streaming to this block.
+    COMPLETE = "complete"  # The producer has finished streaming to this block successfully.
+    ABORTED = "aborted"  # The producer finished streaming to the block, but there was an error.
+
+
 class Block(CamelModel):
     """A Block is a chunk of content within a File. It can be plain text content, image content,
     video content, etc. If the content is not text, the text value may be the empty string
@@ -41,6 +47,7 @@ class Block(CamelModel):
     index_in_file: Optional[int] = Field(alias="index")
     mime_type: Optional[MimeTypes]
     public_data: bool = False
+    stream_state: Optional[StreamState] = None
 
     url: Optional[
         str
@@ -88,6 +95,7 @@ def create(
         url: Optional[str] = None,
         mime_type: Optional[MimeTypes] = None,
         public_data: bool = False,
+        streaming: Optional[bool] = None,
     ) -> Block:
         """
         Create a new Block within a File specified by file_id.
@@ -117,6 +125,7 @@ def create(
             "mimeType": mime_type,
             "uploadType": upload_type,
             "publicData": public_data,
+            "streaming": streaming,
         }
 
         file_data = (
@@ -340,6 +349,24 @@ def as_llm_input(self, exclude_block_wrapper: Optional[bool] = False) -> str:
                 return f"{identifier}"
             return f"Block({identifier})"
 
+    def finish_stream(self):
+        self.client.post(
+            f"block/{self.id}/finishStream",
+            payload={},
+        )
+
+    def append_stream(self, bytes: bytes):
+        self.client.post(
+            f"block/{self.id}/appendStream",
+            payload=bytes,
+        )
+
+    def abort_stream(self):
+        self.client.post(
+            f"block/{self.id}/abortStream",
+            payload={},
+        )
+
 
 class BlockQueryResponse(Response):
     blocks: List[Block]

src/steamship/data/file.py

@@ -282,6 +282,7 @@ def generate(
         options: Optional[dict] = None,
         wait_on_tasks: List[Task] = None,
         make_output_public: bool = False,
+        streaming: Optional[bool] = False,
     ) -> Task[GenerateResponse]:
         """Generate new content from this file. Assumes this file as context for input and output.  May specify start and end blocks."""
         from steamship.data.operations.generator import GenerateRequest, GenerateResponse
@@ -301,6 +302,7 @@ def generate(
             output_file_id=output_file_id,
             options=options,
             make_output_public=make_output_public,
+            streaming=streaming,
         )
         return self.client.post(
             "plugin/instance/generate", req, expect=GenerateResponse, wait_on_tasks=wait_on_tasks

src/steamship/data/manifest.py

@@ -59,6 +59,7 @@ class PluginConfig(BaseModel):
     isTrainable: Optional[bool] = False  # noqa: N815
     transport: str = "jsonOverHttp"
     type: str  # Does not use PluginType due to circular import
+    streaming: bool = False
 
 
 class Manifest(BaseModel):

src/steamship/data/operations/generator.py

@@ -60,6 +60,13 @@ class GenerateRequest(Request):
     # Arbitrary runtime options which may be passed to a generator
     options: Optional[dict]
 
+    # Whether we wish to have the output blocks streamed back to us.
+    # If the blocks are streamed, they will be returned with a streamState=started,
+    # and the content can be fetched in a streaming manner by
+    # fetching Block.raw()
+    # Default behavior if not provided is streaming=false
+    streaming: Optional[bool] = None
+
 
 class GenerateResponse(Response):
     blocks: List[Block]

src/steamship/data/plugin/plugin_instance.py

@@ -130,6 +130,7 @@ def generate(
         output_file_id: Optional[str] = None,
         make_output_public: Optional[bool] = None,
         options: Optional[dict] = None,
+        streaming: Optional[bool] = None,
     ) -> Task[GenerateResponse]:
         """See GenerateRequest for description of parameter options"""
         req = GenerateRequest(
@@ -146,6 +147,7 @@ def generate(
             output_file_id=output_file_id,
             make_output_public=make_output_public,
             options=options,
+            streaming=streaming,
         )
         return self.client.post("plugin/instance/generate", req, expect=GenerateResponse)
 

src/steamship/data/plugin/plugin_version.py

@@ -23,6 +23,7 @@ class CreatePluginVersionRequest(Request):
     type: str = "file"
     # Note: this is a Dict[str, Any] but should be transmitted to the Engine as a JSON string
     config_template: str = None
+    streaming: Optional[bool] = None
 
 
 class ListPluginVersionsRequest(Request):
@@ -65,6 +66,7 @@ def create(
         is_public: Optional[bool] = None,
         is_default: Optional[bool] = None,
         config_template: Optional[Dict[str, Any]] = None,
+        streaming: Optional[bool] = None,
     ) -> PluginVersion:
 
         if filename is None and filebytes is None:
@@ -85,6 +87,7 @@ def create(
             is_public=is_public,
             is_default=is_default,
             config_template=json.dumps(config_template or {}),
+            streaming=streaming,
         )
 
         task = client.post(

src/steamship/plugin/inputs/raw_block_and_tag_plugin_input_with_preallocated_blocks.py

@@ -0,0 +1,8 @@
+from typing import List
+
+from steamship import Block
+from steamship.plugin.inputs.raw_block_and_tag_plugin_input import RawBlockAndTagPluginInput
+
+
+class RawBlockAndTagPluginInputWithPreallocatedBlocks(RawBlockAndTagPluginInput):
+    output_blocks: List[Block]