-
Notifications
You must be signed in to change notification settings - Fork 152
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
2cc293c
commit 54ef8fb
Showing
31 changed files
with
164 additions
and
94 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,47 +1,23 @@ | ||
from __future__ import annotations | ||
|
||
import base64 | ||
import string | ||
import time | ||
import random | ||
from typing import Optional | ||
from attr import define, field, Factory | ||
from griptape.artifacts import BlobArtifact | ||
from attr import define, field | ||
|
||
from griptape.artifacts import MediaArtifact | ||
|
||
|
||
@define | ||
class ImageArtifact(BlobArtifact): | ||
"""ImageArtifact is a type of BlobArtifact that represents an image. | ||
class ImageArtifact(MediaArtifact): | ||
"""ImageArtifact is a type of MediaArtifact representing an image. | ||
Attributes: | ||
value: Raw bytes representing the image. | ||
value: Raw bytes representing media data. | ||
media_type: The type of media, defaults to "image". | ||
format: The format of the media, like png, jpeg, or gif. | ||
name: Artifact name, generated using creation time and a random string. | ||
mime_type: The mime type of the image, like image/png or image/jpeg. | ||
width: The width of the image in pixels. | ||
height: The height of the image in pixels. | ||
model: Optionally specify the model used to generate the image. | ||
prompt: Optionally specify the prompt used to generate the image. | ||
model: Optionally specify the model used to generate the media. | ||
prompt: Optionally specify the prompt used to generate the media. | ||
""" | ||
|
||
mime_type: str = field(kw_only=True, default="image/png", metadata={"serializable": True}) | ||
media_type: str = "image" | ||
width: int = field(kw_only=True, metadata={"serializable": True}) | ||
height: int = field(kw_only=True, metadata={"serializable": True}) | ||
model: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True}) | ||
prompt: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True}) | ||
name: str = field( | ||
default=Factory(lambda self: self.make_name(), takes_self=True), kw_only=True, metadata={"serializable": True} | ||
) | ||
|
||
@property | ||
def base64(self) -> str: | ||
return base64.b64encode(self.value).decode("utf-8") | ||
|
||
def make_name(self) -> str: | ||
entropy = "".join(random.choices(string.ascii_lowercase + string.digits, k=4)) | ||
fmt_time = time.strftime("%y%m%d%H%M%S", time.localtime()) | ||
extension = self.mime_type.split("/")[1].split("+")[0] | ||
|
||
return f"image_artifact_{fmt_time}_{entropy}.{extension}" | ||
|
||
def to_text(self) -> str: | ||
return f"Image, dimensions: {self.width}x{self.height}, type: {self.mime_type}, size: {len(self.value)} bytes" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,53 @@ | ||
from __future__ import annotations | ||
|
||
import string | ||
import time | ||
import random | ||
from typing import Optional | ||
|
||
from attr import define, field | ||
from griptape.artifacts import BlobArtifact | ||
import base64 | ||
|
||
|
||
@define | ||
class MediaArtifact(BlobArtifact): | ||
"""MediaArtifact is a type of BlobArtifact that represents media (image, audio, video, etc.) | ||
and can be extended to support a specific media type. | ||
Attributes: | ||
value: Raw bytes representing media data. | ||
media_type: The type of media, like image, audio, or video. | ||
format: The format of the media, like png, wav, or mp4. | ||
name: Artifact name, generated using creation time and a random string. | ||
model: Optionally specify the model used to generate the media. | ||
prompt: Optionally specify the prompt used to generate the media. | ||
""" | ||
|
||
media_type: str = field(default="media", kw_only=True, metadata={"serializable": True}) | ||
format: str = field(kw_only=True, metadata={"serializable": True}) | ||
model: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True}) | ||
prompt: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True}) | ||
|
||
def __attrs_post_init__(self): | ||
# Generating the name string requires attributes set by child classes. | ||
# This waits until all attributes are available before generating a name. | ||
if self.name == self.id: | ||
self.name = self.make_name() | ||
|
||
@property | ||
def mime_type(self) -> str: | ||
return f"{self.media_type}/{self.format}" | ||
|
||
@property | ||
def base64(self) -> str: | ||
return base64.b64encode(self.value).decode("utf-8") | ||
|
||
def to_text(self) -> str: | ||
return f"Media, type: {self.mime_type}, size: {len(self.value)} bytes" | ||
|
||
def make_name(self) -> str: | ||
entropy = "".join(random.choices(string.ascii_lowercase + string.digits, k=4)) | ||
fmt_time = time.strftime("%y%m%d%H%M%S", time.localtime()) | ||
|
||
return f"{self.media_type}_artifact_{fmt_time}_{entropy}.{self.format}" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
import pytest | ||
|
||
from attr import define | ||
|
||
from griptape.artifacts import MediaArtifact | ||
|
||
|
||
class TestMediaArtifact: | ||
@define | ||
class ImaginaryMediaArtifact(MediaArtifact): | ||
media_type: str = "imagination" | ||
|
||
@pytest.fixture | ||
def media_artifact(self): | ||
return self.ImaginaryMediaArtifact(value=b"some binary dream data", format="dream") | ||
|
||
def test_to_dict(self, media_artifact): | ||
image_dict = media_artifact.to_dict() | ||
|
||
assert image_dict["format"] == "dream" | ||
assert image_dict["value"] == "c29tZSBiaW5hcnkgZHJlYW0gZGF0YQ==" | ||
|
||
def test_name(self, media_artifact): | ||
assert media_artifact.name.startswith("imagination_artifact") | ||
assert media_artifact.name.endswith(".dream") | ||
|
||
def test_mime_type(self, media_artifact): | ||
assert media_artifact.mime_type == "imagination/dream" | ||
|
||
def test_to_text(self, media_artifact): | ||
assert media_artifact.to_text() == "Media, type: imagination/dream, size: 22 bytes" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.