Add Bria Image Generation Driver

README.md

@@ -1,142 +1,39 @@
-# Griptape Extension Template
+# Griptape Bria Extension
 
-A Github template repository for creating Griptape extensions.
+## Overview
+This extension provides an [Image Generation Driver](https://docs.griptape.ai/stable/griptape-framework/drivers/image-generation-drivers) for [Bria](https://bria.ai/).
 
-## Getting Started
+```python
+from griptape.bria.drivers.bria_image_generation_driver import BriaImageGenerationDriver
+from griptape.engines import PromptImageGenerationEngine
+from griptape.structures import Agent
+from griptape.tools import FileManagerTool, PromptImageGenerationTool
 
-Via github web page:
+agent = Agent(
+    tools=[
+        PromptImageGenerationTool(
+            engine=PromptImageGenerationEngine(
+                image_generation_driver=BriaImageGenerationDriver(model="2.3")
+            ),
+            off_prompt=True,
+        ),
+        FileManagerTool(),
+    ]
+)
 
-Click on `Use this template`
-
-![](https://docs.github.com/assets/cb-36544/images/help/repository/use-this-template-button.png)
-
-
-Via `gh`:
-
-```
-$ gh repo create griptape-extension-name -p griptape/griptape-extension-template
-```
-
-## What is a Griptape Extension?
-
-Griptape Extensions can add new functionality to the [Griptape framework](https://github.com/griptape-ai/griptape), such as new Tools, Drivers, Tasks, or Structures.
-
-This repository provides a recommended structure for organizing your extension code, as well as helpful tools for testing and development.
-
-## Extension Structure
-
-The template repository is structured as follows:
-
-```bash
-tree -I __init__.py -I __pycache__
-
-├── griptape
-│   └── extension_name # Name whatever you want
-│       └── tools
-│           └── reverse_string
-│               └── tool.py
-            ...more directories for other interfaces (drivers, tasks, structures, etc)...
-└── tests
-    └── unit
-        └── tools
-            └── test_reverse_string_tool.py
-├── examples
-    └── tools
-        └── example_agent.py # Example usage of the extension 
-├── LICENSE # Choose the appropriate license
-├── Makefile # Contains useful commands for development
-├── pyproject.toml # Contains the project's metadata
-├── README.md #  Describes the extension and how to use it
-```
-
-## Development
-
-### Poetry
-
-This project uses [Poetry](https://python-poetry.org/) for dependency management.
-It is recommended to configure Poetry to use [in-project](https://python-poetry.org/docs/configuration/#virtualenvsin-project) virtual environments:
-
-```bash
-poetry config virtualenvs.in-project true
-```
-
-This will create a `.venv` directory in the project root, where the virtual environment will be stored.
-This ensures that the virtual environment is always in the same location, regardless of where the project is cloned.
-
-### Useful Commands
-
-#### Installing Dependencies
-
-```bash
-make install
+agent.run(
+    "Save a picture of a watercolor painting of a dog riding a skateboard to the desktop."
+)
 ```
 
-#### Running Tests
+## Installation
 
+Poetry:
 ```bash
-make test
+poetry add https://github.com/griptape-ai/griptape-bria.git
 ```
 
-#### Running Checks (linting, formatting, etc)
-
+Pip:
 ```bash
-make check
+pip install git+https://github.com/griptape-ai/griptape-bria.git
 ```
-
-#### Running Formatter
-
-```bash
-make format
-```
-
-#### Running Example
-
-This template includes an [example](https://github.com/griptape-ai/tool-template/blob/main/examples/tools/example_agent.py) demonstrating how to use the extension. It shows how to import the `ReverseStringTool`, provide it to an Agent, and run it.
-
-1. Set the required environment variables. The example needs the `OPENAI_API_KEY` environment variable to be set.
-2. Run the example:
-
-```bash
-poetry run python examples/tools/example_agent.py
-```
-
-If successful, you should see:
-```
-[11/18/24 14:55:14] INFO     ToolkitTask 6bb7fa5581d147b2a39e801631c98005
-                             Input: Use the ReverseStringTool to reverse 'Griptape'
-[11/18/24 14:55:15] INFO     Subtask c3036471831144529b8d5300c6849203
-                             Actions: [
-                               {
-                                 "tag": "call_VE4tGBFL7iB7VDbkKaIFIkwY",
-                                 "name": "ReverseStringTool",
-                                 "path": "reverse_string",
-                                 "input": {
-                                   "values": {
-                                     "input": "Griptape"
-                                   }
-                                 }
-                               }
-                             ]
-                    INFO     Subtask c3036471831144529b8d5300c6849203
-                             Response: epatpirG
-[11/18/24 14:55:16] INFO     ToolkitTask 6bb7fa5581d147b2a39e801631c98005
-                             Output: The reversed string of "Griptape" is "epatpirG".
-```
-
-### Installing in Other Projects
-
-Extensions are designed to be shared. Extensions are made to easily install into existing Python projects.
-
-The easiest way to include your extension into an existing project is to install directly from the repository, like so:
-```bash
-poetry add git+https://github.com/{your-org}/{your-extension-name}.git
-```
-
-To install a local copy of the extension for development, run:
-```bash
-poetry add -e /path/to/your/extension
-```
-
-Any changes made to the extension will be automatically reflected in the project without needing to reinstall it.
-
-Advanced customers may seek to publish their extensions to PyPi. Those instructions are beyond the scope of this README.

examples/drivers/example_agent.py

@@ -0,0 +1,20 @@
+from griptape.bria.drivers.bria_image_generation_driver import BriaImageGenerationDriver
+from griptape.engines import PromptImageGenerationEngine
+from griptape.structures import Agent
+from griptape.tools import FileManagerTool, PromptImageGenerationTool
+
+agent = Agent(
+    tools=[
+        PromptImageGenerationTool(
+            engine=PromptImageGenerationEngine(
+                image_generation_driver=BriaImageGenerationDriver(model="2.3")
+            ),
+            off_prompt=True,
+        ),
+        FileManagerTool(),
+    ]
+)
+
+agent.run(
+    "Save a picture of a watercolor painting of a dog riding a skateboard to the desktop."
+)

griptape/bria/drivers/bria_image_generation_driver.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+import os
+from PIL import Image
+from io import BytesIO
+
+from urllib.parse import urljoin
+from griptape.artifacts import ImageArtifact
+import requests
+
+from attrs import define, field, Factory
+
+from griptape.drivers import BaseImageGenerationDriver
+
+
+@define
+class BriaImageGenerationDriver(BaseImageGenerationDriver):
+    base_url: str = field(
+        default="https://engine.prod.bria-api.com",
+        kw_only=True,
+        metadata={"serializable": False},
+    )
+    api_key: str | None = field(
+        default=Factory(lambda: os.environ["BRIA_API_KEY"]),
+        kw_only=True,
+        metadata={"serializable": False},
+    )
+    headers: dict[str, str] = field(
+        default=Factory(
+            lambda self: {
+                "Content-Type": "application/json",
+                "api_token": self.api_key,
+            },
+            takes_self=True,
+        ),
+        kw_only=True,
+        metadata={"serializable": False},
+    )
+    extra_params: dict[str, str] = field(
+        default=Factory(dict),
+        kw_only=True,
+        metadata={"serializable": False},
+    )
+
+    def try_text_to_image(
+        self, prompts: list[str], negative_prompts: list[str] | None = None
+    ) -> ImageArtifact:
+        url = urljoin(self.base_url, f"/v1/text-to-image/base/{self.model}")
+        prompt = " ".join(prompts)
+        negative_prompt = " ".join(negative_prompts) if negative_prompts else ""
+
+        payload = {
+            "prompt": prompt,
+            "num_results": 1,
+            "sync": True,
+            "negative_prompt": negative_prompt,
+            **self.extra_params,
+        }
+
+        response = requests.post(url, json=payload, headers=self.headers)
+        url = response.json()["result"][0]["urls"][0]
+
+        image_response = requests.get(url)
+        image_bytes = image_response.content
+
+        width, height = Image.open(BytesIO(image_bytes)).size
+
+        return ImageArtifact(
+            value=image_bytes, format="jpeg", width=width, height=height
+        )
+
+    def try_image_variation(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        negative_prompts: list[str] | None = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support image variation"
+        )
+
+    def try_image_inpainting(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        mask: ImageArtifact,
+        negative_prompts: list[str] | None = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support inpainting"
+        )
+
+    def try_image_outpainting(
+        self,
+        prompts: list[str],
+        image: ImageArtifact,
+        mask: ImageArtifact,
+        negative_prompts: list[str] | None = None,
+    ) -> ImageArtifact:
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support outpainting"
+        )