artifacts

@define()
class ActionArtifact(BaseArtifact):
    """Represents the LLM taking an action to use a Tool.

    Attributes:
        value: The Action to take. Currently only supports ToolAction.
    """

    value: ToolAction = field(metadata={"serializable": True})

    def to_text(self) -> str:
        return str(self.value)

def to_text(self) -> str:
    return str(self.value)

@define
class AudioArtifact(BlobArtifact):
    """Stores audio data.

    Attributes:
        format: The audio format, e.g. "wav" or "mp3".
    """

    format: str = field(kw_only=True, metadata={"serializable": True})

    def __attrs_post_init__(self) -> None:
        # 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"audio/{self.format}"

    def to_bytes(self, **kwargs) -> bytes:
        return self.value

    def to_text(self) -> str:
        return f"Audio, format: {self.format}, 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"audio_artifact_{fmt_time}_{entropy}.{self.format}"

def __attrs_post_init__(self) -> None:
    # 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()

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"audio_artifact_{fmt_time}_{entropy}.{self.format}"

def to_bytes(self, **kwargs) -> bytes:
    return self.value

def to_text(self) -> str:
    return f"Audio, format: {self.format}, size: {len(self.value)} bytes"

@define
class BaseArtifact(SerializableMixin, ABC):
    """Serves as the base class for all Artifacts.

    Artifacts are used to encapsulate data and enhance it with metadata.

    Attributes:
        id: The unique identifier of the Artifact. Defaults to a random UUID.
        reference: The optional Reference to the Artifact.
        meta: The metadata associated with the Artifact. Defaults to an empty dictionary.
        name: The name of the Artifact. Defaults to the id.
        value: The value of the Artifact.
        encoding: The encoding to use when encoding/decoding the value.
        encoding_error_handler: The error on_event to use when encoding/decoding the value.
    """

    id: str = field(default=Factory(lambda: uuid.uuid4().hex), kw_only=True, metadata={"serializable": True})
    reference: Optional[Reference] = field(default=None, kw_only=True, metadata={"serializable": True})
    meta: dict[str, Any] = field(factory=dict, kw_only=True, metadata={"serializable": True})
    name: str = field(
        default=Factory(lambda self: self.id, takes_self=True),
        kw_only=True,
        metadata={"serializable": True},
    )
    value: Any = field()
    encoding_error_handler: str = field(default="strict", kw_only=True)
    encoding: str = field(default="utf-8", kw_only=True)

    def __str__(self) -> str:
        return self.to_text()

    def __bool__(self) -> bool:
        return bool(self.value)

    def __len__(self) -> int:
        return len(self.value)

    def to_bytes(self, **kwargs) -> bytes:
        return self.to_text().encode(encoding=self.encoding, errors=self.encoding_error_handler)

    @abstractmethod
    def to_text(self) -> str: ...

def __bool__(self) -> bool:
    return bool(self.value)

def __len__(self) -> int:
    return len(self.value)

def __str__(self) -> str:
    return self.to_text()

def to_bytes(self, **kwargs) -> bytes:
    return self.to_text().encode(encoding=self.encoding, errors=self.encoding_error_handler)

@abstractmethod
def to_text(self) -> str: ...

@define
class BlobArtifact(BaseArtifact):
    """Stores arbitrary binary data.

    Attributes:
        value: The binary data.
    """

    value: bytes = field(
        converter=lambda value: value if isinstance(value, bytes) else str(value).encode(),
        metadata={"serializable": True},
    )

    @property
    def base64(self) -> str:
        return base64.b64encode(self.value).decode(self.encoding)

    @property
    def mime_type(self) -> str:
        return "application/octet-stream"

    def to_bytes(self, **kwargs) -> bytes:
        return self.value

    def to_text(self) -> str:
        return self.value.decode(encoding=self.encoding, errors=self.encoding_error_handler)

def to_bytes(self, **kwargs) -> bytes:
    return self.value

def to_text(self) -> str:
    return self.value.decode(encoding=self.encoding, errors=self.encoding_error_handler)

@define
class BooleanArtifact(BaseArtifact):
    """Stores a boolean value.

    Attributes:
        value: The boolean value.
    """

    value: bool = field(converter=bool, metadata={"serializable": True})

    @classmethod
    def parse_bool(cls, value: Union[str, bool]) -> BooleanArtifact:
        """Convert a string literal or bool to a BooleanArtifact. The string must be either "true" or "false"."""
        if value is not None:
            if isinstance(value, str):
                if value.lower() == "true":
                    return BooleanArtifact(value=True)
                if value.lower() == "false":
                    return BooleanArtifact(value=False)
            elif isinstance(value, bool):
                return BooleanArtifact(value)
        raise ValueError(f"Cannot convert '{value}' to BooleanArtifact")

    def __add__(self, other: BaseArtifact) -> BooleanArtifact:
        raise ValueError("Cannot add BooleanArtifact with other artifacts")

    def __eq__(self, value: object) -> bool:
        return self.value == value

    def to_text(self) -> str:
        return str(self.value).lower()

def __add__(self, other: BaseArtifact) -> BooleanArtifact:
    raise ValueError("Cannot add BooleanArtifact with other artifacts")

def __eq__(self, value: object) -> bool:
    return self.value == value

@classmethod
def parse_bool(cls, value: Union[str, bool]) -> BooleanArtifact:
    """Convert a string literal or bool to a BooleanArtifact. The string must be either "true" or "false"."""
    if value is not None:
        if isinstance(value, str):
            if value.lower() == "true":
                return BooleanArtifact(value=True)
            if value.lower() == "false":
                return BooleanArtifact(value=False)
        elif isinstance(value, bool):
            return BooleanArtifact(value)
    raise ValueError(f"Cannot convert '{value}' to BooleanArtifact")

def to_text(self) -> str:
    return str(self.value).lower()

@define
class ErrorArtifact(BaseArtifact):
    """Represents an error that may want to be conveyed to the LLM.

    Attributes:
        value: The error message.
        exception: The exception that caused the error. Defaults to None.
    """

    value: str = field(converter=str, metadata={"serializable": True})
    exception: Optional[Exception] = field(default=None, kw_only=True, metadata={"serializable": False})

    def to_text(self) -> str:
        return self.value

def to_text(self) -> str:
    return self.value

@define
class GenericArtifact(BaseArtifact, Generic[T]):
    """Serves as an escape hatch for artifacts that don't fit into any other category.

    Attributes:
        value: The value of the Artifact.
    """

    value: T = field(metadata={"serializable": True})

    def to_text(self) -> str:
        return str(self.value)

def to_text(self) -> str:
    return str(self.value)

@define
class ImageArtifact(BlobArtifact):
    """Stores image data.

    Attributes:
        format: The format of the image data. Used when building the MIME type.
        width: The width of the image.
        height: The height of the image
    """

    format: str = field(kw_only=True, metadata={"serializable": True})
    width: int = field(kw_only=True, metadata={"serializable": True})
    height: int = field(kw_only=True, metadata={"serializable": True})

    def __attrs_post_init__(self) -> None:
        # 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"image/{self.format}"

    def to_bytes(self, **kwargs) -> bytes:
        return self.value

    def to_text(self) -> str:
        return f"Image, format: {self.format}, 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"image_artifact_{fmt_time}_{entropy}.{self.format}"

def __attrs_post_init__(self) -> None:
    # 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()

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"image_artifact_{fmt_time}_{entropy}.{self.format}"

def to_bytes(self, **kwargs) -> bytes:
    return self.value

def to_text(self) -> str:
    return f"Image, format: {self.format}, size: {len(self.value)} bytes"

class ImageUrlArtifact(UrlArtifact):
    """Stores a url to an image.

    Attributes:
        value: The url.
    """

@define
class InfoArtifact(BaseArtifact):
    """Represents helpful info that can be conveyed to the LLM.

    For example, "No results found" or "Please try again.".

    Attributes:
        value: The info to convey.
    """

    value: str = field(converter=str, metadata={"serializable": True})

    def to_text(self) -> str:
        return self.value

def to_text(self) -> str:
    return self.value

@define
class JsonArtifact(BaseArtifact):
    """Stores JSON data.

    Attributes:
        value: The JSON data. Values will automatically be converted to a JSON-compatible format.
    """

    value: Json = field(converter=lambda value: JsonArtifact.value_to_json(value), metadata={"serializable": True})

    @classmethod
    def value_to_json(cls, value: Any) -> Json:
        if isinstance(value, str):
            return json.loads(value)
        return json.loads(json.dumps(value))

    def to_text(self) -> str:
        return json.dumps(self.value)

def to_text(self) -> str:
    return json.dumps(self.value)

@classmethod
def value_to_json(cls, value: Any) -> Json:
    if isinstance(value, str):
        return json.loads(value)
    return json.loads(json.dumps(value))

@define
class ListArtifact(BaseArtifact, Generic[T_co]):
    value: Sequence[T_co] = field(factory=list, metadata={"serializable": True})
    item_separator: str = field(default="\n\n", kw_only=True, metadata={"serializable": True})
    validate_uniform_types: bool = field(default=False, kw_only=True, metadata={"serializable": True})

    @value.validator  # pyright: ignore[reportAttributeAccessIssue]
    def validate_value(self, _: Attribute, value: list[T_co]) -> None:
        if self.validate_uniform_types and len(value) > 0:
            first_type = type(value[0])

            if not all(isinstance(v, first_type) for v in value):
                raise ValueError("list elements in 'value' are not the same type")

    @property
    def child_type(self) -> Optional[type]:
        if self.value:
            return type(self.value[0])
        return None

    def __getitem__(self, key: int) -> T_co:
        return self.value[key]

    def __bool__(self) -> bool:
        return len(self) > 0

    def __add__(self, other: BaseArtifact) -> ListArtifact[T_co]:
        return ListArtifact(self.value + other.value)

    def __iter__(self) -> Iterator[T_co]:
        return iter(self.value)

    def to_text(self) -> str:
        return self.item_separator.join([v.to_text() for v in self.value])

    def is_type(self, target_type: type) -> bool:
        if self.value:
            return isinstance(self.value[0], target_type)
        return False

    def has_items(self) -> bool:
        return len(self) > 0

def __add__(self, other: BaseArtifact) -> ListArtifact[T_co]:
    return ListArtifact(self.value + other.value)

def __bool__(self) -> bool:
    return len(self) > 0

def __getitem__(self, key: int) -> T_co:
    return self.value[key]

def __iter__(self) -> Iterator[T_co]:
    return iter(self.value)

def has_items(self) -> bool:
    return len(self) > 0

def is_type(self, target_type: type) -> bool:
    if self.value:
        return isinstance(self.value[0], target_type)
    return False

def to_text(self) -> str:
    return self.item_separator.join([v.to_text() for v in self.value])

@value.validator  # pyright: ignore[reportAttributeAccessIssue]
def validate_value(self, _: Attribute, value: list[T_co]) -> None:
    if self.validate_uniform_types and len(value) > 0:
        first_type = type(value[0])

        if not all(isinstance(v, first_type) for v in value):
            raise ValueError("list elements in 'value' are not the same type")

@define
class ModelArtifact(GenericArtifact[BaseModel]):
    """Stores Pydantic models as Artifacts.

    Required since Pydantic models require a custom serialization method.

    Attributes:
        value: The pydantic model to store.
    """

    # We must explicitly define the type rather than rely on the parent T since
    # generic type information is lost at runtime.
    value: BaseModel = field(metadata={"serializable": True})

    def to_text(self) -> str:
        return self.value.model_dump_json()

def to_text(self) -> str:
    return self.value.model_dump_json()

@define
class TextArtifact(BaseArtifact):
    value: str = field(converter=str, metadata={"serializable": True})
    embedding: Optional[list[float]] = field(default=None, kw_only=True)

    def __add__(self, other: BaseArtifact) -> TextArtifact:
        return TextArtifact(self.value + other.value)

    def __bool__(self) -> bool:
        return bool(self.value.strip())

    def to_text(self) -> str:
        return self.value

    def generate_embedding(self, driver: BaseEmbeddingDriver) -> list[float]:
        embedding = driver.embed(str(self.value))

        if self.embedding is None:
            self.embedding = []
        self.embedding.clear()
        self.embedding.extend(embedding)

        return self.embedding

    def token_count(self, tokenizer: BaseTokenizer) -> int:
        return tokenizer.count_tokens(str(self.value))

def __add__(self, other: BaseArtifact) -> TextArtifact:
    return TextArtifact(self.value + other.value)

def __bool__(self) -> bool:
    return bool(self.value.strip())

def generate_embedding(self, driver: BaseEmbeddingDriver) -> list[float]:
    embedding = driver.embed(str(self.value))

    if self.embedding is None:
        self.embedding = []
    self.embedding.clear()
    self.embedding.extend(embedding)

    return self.embedding

def to_text(self) -> str:
    return self.value

def token_count(self, tokenizer: BaseTokenizer) -> int:
    return tokenizer.count_tokens(str(self.value))

@define
class UrlArtifact(BaseArtifact):
    """Stores a url.

    Attributes:
        value: The url.
    """

    value: str = field(metadata={"serializable": True})

    def to_bytes(self, *, headers: dict | None = None, **kwargs: dict) -> bytes:
        """Fetches the content of the URL and returns it as bytes.

        Args:
            headers: Optional headers to include in the request.
            **kwargs: Additional keyword arguments, not used.

        Returns:
            bytes: The content of the URL as bytes.

        """
        response = requests.get(self.value, headers=headers)
        response.raise_for_status()

        return response.content

    def to_text(self) -> str:
        """Returns the URL as is."""
        return self.value

def to_bytes(self, *, headers: dict | None = None, **kwargs: dict) -> bytes:
    """Fetches the content of the URL and returns it as bytes.

    Args:
        headers: Optional headers to include in the request.
        **kwargs: Additional keyword arguments, not used.

    Returns:
        bytes: The content of the URL as bytes.

    """
    response = requests.get(self.value, headers=headers)
    response.raise_for_status()

    return response.content

def to_text(self) -> str:
    """Returns the URL as is."""
    return self.value