diff --git a/src/banks/env.py b/src/banks/env.py
index 247dd34..7d495a8 100644
--- a/src/banks/env.py
+++ b/src/banks/env.py
@@ -4,7 +4,7 @@
from jinja2 import Environment, PackageLoader, select_autoescape
from .config import config
-from .filters import cache_control, lemmatize, tool
+from .filters import cache_control, image, lemmatize, tool
def _add_extensions(_env):
@@ -38,7 +38,9 @@ def _add_extensions(_env):
# Setup custom filters and defaults
-env.filters["lemmatize"] = lemmatize
env.filters["cache_control"] = cache_control
+env.filters["image"] = image
+env.filters["lemmatize"] = lemmatize
env.filters["tool"] = tool
+
_add_extensions(env)
diff --git a/src/banks/extensions/chat.py b/src/banks/extensions/chat.py
index ee75fa2..db6cb15 100644
--- a/src/banks/extensions/chat.py
+++ b/src/banks/extensions/chat.py
@@ -33,12 +33,12 @@ def content(self) -> ChatMessageContent:
return self._content_blocks
- def handle_starttag(self, tag, _):
- if tag == "content_block_txt":
+ def handle_starttag(self, tag, attrs): # noqa
+ if tag == "content_block":
self._parse_block_content = True
def handle_endtag(self, tag):
- if tag == "content_block_txt":
+ if tag == "content_block":
self._parse_block_content = False
def handle_data(self, data):
diff --git a/src/banks/filters/__init__.py b/src/banks/filters/__init__.py
index 27645ad..60ad899 100644
--- a/src/banks/filters/__init__.py
+++ b/src/banks/filters/__init__.py
@@ -2,11 +2,13 @@
#
# SPDX-License-Identifier: MIT
from .cache_control import cache_control
+from .image import image
from .lemmatize import lemmatize
from .tool import tool
__all__ = (
"cache_control",
+ "image",
"lemmatize",
"tool",
)
diff --git a/src/banks/filters/cache_control.py b/src/banks/filters/cache_control.py
index 65f4640..0e0284c 100644
--- a/src/banks/filters/cache_control.py
+++ b/src/banks/filters/cache_control.py
@@ -17,8 +17,8 @@ def cache_control(value: str, cache_type: str = "ephemeral") -> str:
```
Important:
- this filter marks the content to cache by surrounding it with `` and
- ``, so it's only useful when used within a `{% chat %}` block.
+ this filter marks the content to cache by surrounding it with `` and
+ ``, so it's only useful when used within a `{% chat %}` block.
"""
block = ContentBlock.model_validate({"type": "text", "text": value, "cache_control": {"type": cache_type}})
- return f"{block.model_dump_json()}"
+ return f"{block.model_dump_json()}"
diff --git a/src/banks/filters/image.py b/src/banks/filters/image.py
new file mode 100644
index 0000000..81ea7cd
--- /dev/null
+++ b/src/banks/filters/image.py
@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi
+#
+# SPDX-License-Identifier: MIT
+from pathlib import Path
+from urllib.parse import urlparse
+
+from banks.types import ContentBlock, ImageUrl
+
+
+def _is_url(string: str) -> bool:
+ try:
+ result = urlparse(string)
+ return all([result.scheme, result.netloc])
+ except ValueError:
+ return False
+
+
+def image(value: str) -> str:
+ """Wrap the filtered value into a ContentBlock with the proper cache_control field set.
+
+ The resulting ChatMessage will have the field `content` populated with a list of ContentBlock objects.
+
+ Example:
+ ```jinja
+ Describe what you see
+
+ {{ "path/to/image/file" | image }}
+ ```
+
+ Important:
+ this filter marks the content to cache by surrounding it with `` and
+ ``, so it's only useful when used within a `{% chat %}` block.
+ """
+ if _is_url(value):
+ image_url = ImageUrl(url=value)
+ else:
+ image_url = ImageUrl.from_path(Path(value))
+
+ block = ContentBlock.model_validate({"type": "image_url", "image_url": image_url})
+ return f"{block.model_dump_json()}"
diff --git a/src/banks/types.py b/src/banks/types.py
index fb011b9..7650034 100644
--- a/src/banks/types.py
+++ b/src/banks/types.py
@@ -1,8 +1,10 @@
# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi
#
# SPDX-License-Identifier: MIT
+import base64
from enum import Enum
from inspect import Parameter, getdoc, signature
+from pathlib import Path
from typing import Callable
from pydantic import BaseModel
@@ -26,9 +28,14 @@ class ImageUrl(BaseModel):
url: str
@classmethod
- def from_base64(cls, media_type: str, base64_str: str):
+ def from_base64(cls, media_type: str, base64_str: str) -> Self:
return cls(url=f"data:{media_type};base64,{base64_str}")
+ @classmethod
+ def from_path(cls, file_path: Path) -> Self:
+ with open(file_path, "rb") as image_file:
+ return cls.from_base64("image/jpeg", base64.b64encode(image_file.read()).decode("utf-8"))
+
class ContentBlock(BaseModel):
type: ContentBlockType
diff --git a/tests/test_cache_control.py b/tests/test_cache_control.py
index 100153b..7be8b3b 100644
--- a/tests/test_cache_control.py
+++ b/tests/test_cache_control.py
@@ -3,6 +3,6 @@
def test_cache_control():
res = cache_control("foo", "ephemeral")
- res = res.replace("", "")
- res = res.replace("", "")
+ res = res.replace("", "")
+ res = res.replace("", "")
assert res == '{"type":"text","cache_control":{"type":"ephemeral"},"text":"foo","image_url":null}'
diff --git a/tests/test_chat.py b/tests/test_chat.py
index 71e3be1..38a90c0 100644
--- a/tests/test_chat.py
+++ b/tests/test_chat.py
@@ -30,7 +30,7 @@ def test_content_block_parser_init():
def test_content_block_parser_single_with_cache_control():
p = _ContentBlockParser()
p.feed(
- '{"type":"text","cache_control":{"type":"ephemeral"},"text":"foo","source":null}'
+ '{"type":"text","cache_control":{"type":"ephemeral"},"text":"foo","source":null}'
)
assert p.content == [
ContentBlock(type=ContentBlockType.text, cache_control=CacheControl(type="ephemeral"), text="foo", source=None)
@@ -39,15 +39,15 @@ def test_content_block_parser_single_with_cache_control():
def test_content_block_parser_single_no_cache_control():
p = _ContentBlockParser()
- p.feed('{"type":"text","cache_control":null,"text":"foo","source":null}')
+ p.feed('{"type":"text","cache_control":null,"text":"foo","source":null}')
assert p.content == "foo"
def test_content_block_parser_multiple():
p = _ContentBlockParser()
p.feed(
- '{"type":"text","cache_control":null,"text":"foo","source":null}'
- '{"type":"text","cache_control":null,"text":"bar","source":null}'
+ '{"type":"text","cache_control":null,"text":"foo","source":null}'
+ '{"type":"text","cache_control":null,"text":"bar","source":null}'
)
assert p.content == [
ContentBlock(type=ContentBlockType.text, cache_control=None, text="foo", source=None),