Merge branch 'main' into images-in-telegram-transport

Author: dkolas

src/steamship/data/block.py

@@ -49,6 +49,8 @@ class Block(CamelModel):
     tags: Optional[List[Tag]] = []
     index_in_file: Optional[int] = Field(alias="index")
     mime_type: Optional[MimeTypes]
+    public_data: bool = False
+
     url: Optional[
         str
     ] = None  # Only for creation of blocks; used to fetch content from a public URL.
@@ -94,6 +96,7 @@ def create(
         content: Union[str, bytes] = None,
         url: Optional[str] = None,
         mime_type: Optional[MimeTypes] = None,
+        public_data: bool = False,
     ) -> Block:
         """
         Create a new Block within a File specified by file_id.
@@ -122,6 +125,7 @@ def create(
             "url": url,
             "mimeType": mime_type,
             "uploadType": upload_type,
+            "publicData": public_data,
         }
 
         file_data = (
@@ -199,6 +203,16 @@ def is_video(self):
         """Return whether this is a video Block."""
         return self.mime_type in [MimeTypes.MP4_VIDEO, MimeTypes.WEBM_VIDEO]
 
+    @property
+    def raw_data_url(self) -> Optional[str]:
+        """Return a URL at which the data content of this Block can be accessed.  If public_data is True,
+        this content can be accessed without an API key.
+        """
+        if self.client is not None:
+            return f"{self.client.config.api_base}block/{self.id}/raw"
+        else:
+            return None
+
     @property
     def chat_role(self) -> Optional[RoleTag]:
         return get_tag_value_key(

src/steamship/data/file.py

@@ -55,6 +55,7 @@ class File(CamelModel):
     blocks: List[Block] = []
     tags: List[Tag] = []
     filename: str = None
+    public_data: bool = False
 
     class CreateResponse(Response):
         data_: Any = None
@@ -115,6 +116,7 @@ def create(
         handle: str = None,
         blocks: List[Block] = None,
         tags: List[Tag] = None,
+        public_data: bool = False,
     ) -> Any:
 
         if content is None and blocks is None:
@@ -145,6 +147,7 @@ def create(
             "tags": [
                 tag.dict(by_alias=True, exclude_unset=True, exclude_none=True) for tag in tags or []
             ],
+            "publicData": public_data,
         }
 
         file_data = (
@@ -209,6 +212,16 @@ def raw(self):
             raw_response=True,
         )
 
+    @property
+    def raw_data_url(self) -> Optional[str]:
+        """Return a URL at which the data content of this File can be accessed.  If public_data is True,
+        this content can be accessed without an API key.
+        """
+        if self.client is not None:
+            return f"{self.client.config.api_base}file/{self.id}/raw"
+        else:
+            return None
+
     def blockify(self, plugin_instance: str = None, wait_on_tasks: List[Task] = None) -> Task:
         from steamship.data.operations.blockifier import BlockifyRequest
         from steamship.plugin.outputs.block_and_tag_plugin_output import BlockAndTagPluginOutput

src/steamship/data/operations/generator.py

@@ -51,6 +51,12 @@ class GenerateRequest(Request):
     # create a new file.
     output_file_id: Optional[str] = None
 
+    # May specify that the output blocks' bytes content should be made public-readable.
+    # Useful for generating images / audio / etc that will be shared.
+    # Defaults to False (private) content.
+    # Requires append_output_to_file to be True.
+    make_output_public: Optional[bool] = None
+
     # Arbitrary runtime options which may be passed to a generator
     options: Optional[dict]
 

src/steamship/data/plugin/plugin_instance.py

@@ -128,6 +128,7 @@ def generate(
         # url: Optional[str] = None, [Not yet implemented]
         append_output_to_file: bool = False,
         output_file_id: Optional[str] = None,
+        make_output_public: Optional[bool] = None,
         options: Optional[dict] = None,
     ) -> Task[GenerateResponse]:
         """See GenerateRequest for description of parameter options"""
@@ -143,6 +144,7 @@ def generate(
             # url=url,
             append_output_to_file=append_output_to_file,
             output_file_id=output_file_id,
+            make_output_public=make_output_public,
             options=options,
         )
         return self.client.post("plugin/instance/generate", req, expect=GenerateResponse)

src/steamship/utils/repl.py

@@ -118,12 +118,17 @@ class AgentREPL(SteamshipREPL):
     config = None
 
     def __init__(
-        self, agent_class: Type[AgentService], method: str, client: Optional[Steamship] = None
+        self,
+        agent_class: Type[AgentService],
+        method: str,
+        agent_package_config: Optional[Dict[str, Any]],
+        client: Optional[Steamship] = None,
     ):
         super().__init__()
         self.agent_class = agent_class
         self.method = method
         self.client = client or Steamship()
+        self.config = agent_package_config
 
     def run_with_client(self, client: Steamship):
         try:
@@ -136,7 +141,7 @@ def colored(text: str, color: str):
         print("Starting REPL for Agent...")
         print("If you make code changes, restart this REPL. Press CTRL+C to exit at any time.\n")
 
-        agent_service = self.agent_class(client=client)
+        agent_service = self.agent_class(client=client, config=self.config)
 
         while True:
             input_text = input(colored(text="Input: ", color="blue"))  # noqa: F821

tests/steamship_tests/data/test_block.py

@@ -1,4 +1,5 @@
 import pytest
+import requests
 from steamship_tests import TEST_ASSETS_PATH
 
 from steamship import MimeTypes
@@ -175,3 +176,40 @@ def test_append_is_present(client: Steamship):
     assert len(my_file.blocks) == 0
     my_file.append_block(text="first")
     assert len(my_file.blocks) == 1
+
+
+@pytest.mark.usefixtures("client")
+def test_append_block_public_data(client: Steamship):
+    file = File.create(client, blocks=[])
+
+    appended_block = Block.create(
+        client, file_id=file.id, content=bytes("This is a test", "utf-8"), public_data=True
+    )
+    assert appended_block.public_data
+
+    # Intentionally no API key
+    response = requests.get(appended_block.raw_data_url)
+
+    assert response.text == "This is a test"
+    assert response.headers["content-type"] == MimeTypes.TXT
+
+
+@pytest.mark.usefixtures("client")
+def test_append_block_private_data(client: Steamship):
+    file = File.create(client, blocks=[])
+
+    appended_block = Block.create(client, file_id=file.id, content=bytes("This is a test", "utf-8"))
+    assert not appended_block.public_data
+
+    # Intentionally no API key
+    failed_response = requests.get(appended_block.raw_data_url)
+    assert not failed_response.ok
+
+    # Should still be able to get with API key
+    response = requests.get(
+        appended_block.raw_data_url,
+        headers={"Authorization": f"Bearer {client.config.api_key.get_secret_value()}"},
+    )
+
+    assert response.text == "This is a test"
+    assert response.headers["content-type"] == MimeTypes.TXT

tests/steamship_tests/data/test_file.py

@@ -351,3 +351,34 @@ def test_append_with_tag(client: Steamship):
     assert file.blocks[1].text == "second"
     assert len(file.blocks[1].tags) == 1
     assert file.blocks[1].tags[0].kind == TagKind.DOCUMENT
+
+
+@pytest.mark.usefixtures("client")
+def test_file_public_data(client: Steamship):
+    file = File.create(client, content=bytes("This is a test", "utf-8"), public_data=True)
+    assert file.public_data
+
+    # Intentionally no API key
+    response = requests.get(file.raw_data_url)
+
+    assert response.text == "This is a test"
+    assert response.headers["content-type"] == MimeTypes.TXT
+
+
+@pytest.mark.usefixtures("client")
+def test_append_block_private_data(client: Steamship):
+    file = File.create(client, content=bytes("This is a test", "utf-8"))
+    assert not file.public_data
+
+    # Intentionally no API key
+    failed_response = requests.get(file.raw_data_url)
+    assert not failed_response.ok
+
+    # Should still be able to get with API key
+    response = requests.get(
+        file.raw_data_url,
+        headers={"Authorization": f"Bearer {client.config.api_key.get_secret_value()}"},
+    )
+
+    assert response.text == "This is a test"
+    assert response.headers["content-type"] == MimeTypes.TXT

tests/steamship_tests/plugin/integration/test_e2e_generator.py

@@ -1,9 +1,11 @@
+import pytest
+import requests
 from assets.plugins.generators.test_generator_returns_bytes import TEST_BYTES_STRING
 from steamship_tests import PLUGINS_PATH, TEST_ASSETS_PATH
 from steamship_tests.utils.deployables import deploy_plugin
 from steamship_tests.utils.fixtures import get_steamship_client
 
-from steamship import Block, File, MimeTypes, PluginInstance
+from steamship import Block, File, MimeTypes, PluginInstance, Steamship
 
 
 def test_e2e_generator():
@@ -173,3 +175,46 @@ def test_e2e_generate_returning_bytes():
         assert persisted_block.file_id is not None
         persisted_block_content = persisted_block.raw()
         assert persisted_block_content == TEST_BYTES_STRING.encode("utf-8")
+
+
+@pytest.mark.usefixtures("client")
+def test_generate_block_public_data(client: Steamship):
+    client = get_steamship_client()
+    plugin_instance = PluginInstance.create(client, plugin_handle="test-image-generator")
+    generate_task = plugin_instance.generate(
+        text="This won't be used", append_output_to_file=True, make_output_public=True
+    )
+
+    blocks = generate_task.wait().blocks
+    assert blocks is not None
+    assert blocks[0].public_data
+
+    # Intentionally no API key
+    response = requests.get(blocks[0].raw_data_url)
+
+    assert response.text == "PRETEND THIS IS THE DATA OF AN IMAGE"
+    assert response.headers["content-type"] == MimeTypes.PNG
+
+
+@pytest.mark.usefixtures("client")
+def test_generate_block_private_data(client: Steamship):
+    client = get_steamship_client()
+    plugin_instance = PluginInstance.create(client, plugin_handle="test-image-generator")
+    generate_task = plugin_instance.generate(text="This won't be used", append_output_to_file=True)
+
+    blocks = generate_task.wait().blocks
+    assert blocks is not None
+    assert not blocks[0].public_data
+
+    # Intentionally no API key
+    failed_response = requests.get(blocks[0].raw_data_url)
+    assert not failed_response.ok
+
+    # Should still be able to get with API key
+    response = requests.get(
+        blocks[0].raw_data_url,
+        headers={"Authorization": f"Bearer {client.config.api_key.get_secret_value()}"},
+    )
+
+    assert response.text == "PRETEND THIS IS THE DATA OF AN IMAGE"
+    assert response.headers["content-type"] == MimeTypes.PNG