feat: Support documents (#66)

Author: bschloss

CLAUDE.md

@@ -4,7 +4,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 
 ## Project Overview
 
-Banks is a Python prompt programming language and templating system for LLM applications. It provides a Jinja2-based template engine with specialized extensions and filters for creating dynamic prompts, managing chat messages, handling multimodal content (images/audio), and integrating with various LLM providers through LiteLLM.
+Banks is a Python prompt programming language and templating system for LLM applications. It provides a Jinja2-based template engine with specialized extensions and filters for creating dynamic prompts, managing chat messages, handling multimodal content (images/audio/documents), and integrating with various LLM providers through LiteLLM.
 
 ## Development Commands
 
@@ -42,7 +42,7 @@ Banks is a Python prompt programming language and templating system for LLM appl
 
 **Type System** (`src/banks/types.py`):
 - `ChatMessage`: Core chat message structure with role and content
-- `ContentBlock`: Handles different content types (text, image_url, audio) with optional cache control
+- `ContentBlock`: Handles different content types (text, image_url, audio, document) with optional cache control
 - `Tool`: Function calling support with automatic schema generation from Python callables
 - `CacheControl`: Anthropic-style prompt caching metadata
 
@@ -67,6 +67,7 @@ Banks is a Python prompt programming language and templating system for LLM appl
 **Core Filters** (`src/banks/filters/`):
 - `image`: Convert file paths/URLs to base64-encoded image content blocks
 - `audio`: Convert audio files to base64-encoded audio content blocks  
+- `document`: Convert documents (PDF, TXT, HTML, CSS, XML, CSV, RTF, JS, JSON) to base64-encoded content blocks
 - `cache_control`: Add Anthropic cache control metadata to content blocks
 - `tool`: Convert Python callables to LLM function call schemas
 - `lemmatize`: Text lemmatization using simplemma
@@ -95,7 +96,7 @@ Banks is a Python prompt programming language and templating system for LLM appl
 4. Caching layer prevents re-rendering identical contexts
 
 ### Multimodal Content Handling
-- Images/audio converted to base64 during filter application
+- Images/audio/documents converted to base64 during filter application
 - Content blocks maintain type safety and metadata
 - Cache control integrated at content block level
 

README.md

@@ -94,11 +94,11 @@ print(p.chat_messages({"persona": "helpful assistant"}))
 # [
 #   ChatMessage(role='system', content=[
 #      ContentBlock(type=<ContentBlockType.text: 'text'>, cache_control=None, text='You are a helpful assistant.', 
-#                  image_url=None, input_audio=None)
+#                  image_url=None, input_audio=None, input_document=None)
 #   ], tool_call_id=None, name=None), 
 #   ChatMessage(role='user', content=[
 #      ContentBlock(type=<ContentBlockType.text: 'text'>, cache_control=None, text='Hello, how are you?', 
-#                  image_url=None, input_audio=None)
+#                  image_url=None, input_audio=None, input_document=None)
 #   ], tool_call_id=None, name=None)
 # ]
 ```

src/banks/env.py

@@ -4,7 +4,7 @@
 from jinja2 import Environment, select_autoescape
 
 from .config import config
-from .filters import audio, cache_control, image, lemmatize, tool, xml
+from .filters import audio, cache_control, document, image, lemmatize, tool, xml
 
 
 def _add_extensions(_env):
@@ -38,6 +38,7 @@ def _add_extensions(_env):
 env.filters["lemmatize"] = lemmatize
 env.filters["tool"] = tool
 env.filters["audio"] = audio
+env.filters["document"] = document
 env.filters["to_xml"] = xml
 
 _add_extensions(env)

src/banks/filters/__init__.py

@@ -3,9 +3,10 @@
 # SPDX-License-Identifier: MIT
 from .audio import audio
 from .cache_control import cache_control
+from .document import document
 from .image import image
 from .lemmatize import lemmatize
 from .tool import tool
 from .xml import xml
 
-__all__ = ("cache_control", "image", "lemmatize", "tool", "audio", "xml")
+__all__ = ("cache_control", "image", "lemmatize", "tool", "audio", "document", "xml")

src/banks/filters/document.py

@@ -0,0 +1,96 @@
+# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <mpippi@gmail.com>
+#
+# SPDX-License-Identifier: MIT
+import re
+from pathlib import Path
+from typing import cast
+from urllib.parse import urlparse
+
+from banks.types import ContentBlock, DocumentFormat, InputDocument
+
+BASE64_DOCUMENT_REGEX = re.compile(r"(text|application)\/.*;base64,.*")
+
+
+def _is_url(string: str) -> bool:
+    """Check if a string is a URL."""
+    result = urlparse(string)
+    if not result.scheme:
+        return False
+
+    if not result.netloc:
+        # The only valid format when netloc is empty is base64 data urls
+        return all([result.scheme == "data", BASE64_DOCUMENT_REGEX.match(result.path)])
+
+    return True
+
+
+def _get_document_format_from_url(url: str) -> DocumentFormat:
+    """Extract document format from URL.
+
+    Tries to determine format from URL path or defaults to pdf.
+    """
+    parsed = urlparse(url)
+    path = parsed.path.lower()
+    # Gemini supported file types https://ai.google.dev/gemini-api/docs/file-input-methods
+    # text/html
+    # text/css
+    # text/plain
+    # text/xml
+    # text/scv
+    # text/rtf
+    # text/javascript
+    # application/json
+    # application/pdf
+
+    # Claude supported file types
+    # application/pdf
+    # text/plain
+
+    # OpenAI supported file types
+    # application/pdf
+
+    for fmt in (
+        "pdf",
+        "html",
+        "htm",
+        "xhtml",
+        "css",
+        "txt",
+        "md",
+        "markdown",
+        "rst",
+        "xml",
+        "csv",
+        "rtf",
+        "js",
+        "mjs",
+        "cjs",
+        "javascript",
+        "json",
+    ):
+        if path.endswith(f".{fmt}"):
+            return cast(DocumentFormat, fmt)
+    # Default to pdf if format cannot be determined
+    return "pdf"
+
+
+def document(value: str) -> str:
+    """Wrap the filtered value into a ContentBlock of type document.
+
+    The resulting ChatMessage will have the field `content` populated with a list of ContentBlock objects.
+
+    Supports both file paths and URLs (including data URLs).
+
+    Example:
+        ```jinja
+        {{ "path/to/document/file.pdf" | document }}
+        {{ "https://example.com/document.pdf" | document }}
+        ```
+    """
+    if _is_url(value):
+        document_format = _get_document_format_from_url(value)
+        input_document = InputDocument.from_url(value, document_format)
+    else:
+        input_document = InputDocument.from_path(Path(value))
+    block = ContentBlock.model_validate({"type": "document", "input_document": input_document})
+    return f"<content_block>{block.model_dump_json()}</content_block>"

src/banks/types.py

@@ -23,6 +23,7 @@ class ContentBlockType(str, Enum):
     text = "text"
     image_url = "image_url"
     audio = "audio"
+    document = "document"
 
 
 class CacheControl(BaseModel):
@@ -43,6 +44,7 @@ def from_path(cls, file_path: Path) -> Self:
 
 
 AudioFormat = Literal["mp3", "wav", "m4a", "webm", "ogg", "flac"]
+DocumentFormat = Literal["pdf", "html", "css", "plain", "xml", "csv", "rtf", "javascript", "json"]
 
 
 class InputAudio(BaseModel):
@@ -70,12 +72,38 @@ def from_url(cls, url: str, audio_format: AudioFormat) -> Self:
         return cls(data=url, format=audio_format)
 
 
+class InputDocument(BaseModel):
+    data: str
+    format: DocumentFormat
+
+    @classmethod
+    def from_path(cls, file_path: Path) -> Self:
+        with open(file_path, "rb") as document_file:
+            encoded_str = base64.b64encode(document_file.read()).decode("utf-8")
+            file_format = cast(DocumentFormat, file_path.suffix[1:])
+            return cls(data=encoded_str, format=file_format)
+
+    @classmethod
+    def from_url(cls, url: str, document_format: DocumentFormat) -> Self:
+        """Create InputDocument from a URL.
+
+        Args:
+            url: The URL to the document file
+            document_format: The document format
+
+        Returns:
+            InputDocument instance with the URL as data
+        """
+        return cls(data=url, format=document_format)
+
+
 class ContentBlock(BaseModel):
     type: ContentBlockType
     cache_control: CacheControl | None = None
     text: str | None = None
     image_url: ImageUrl | None = None
     input_audio: InputAudio | None = None
+    input_document: InputDocument | None = None
 
 
 ChatMessageContent = Union[list[ContentBlock], str]

tests/data/1x1.pdf

@@ -5,6 +5,7 @@ def test_cache_control():
     res = cache_control("foo", "ephemeral")
     res = res.replace("<content_block>", "")
     res = res.replace("</content_block>", "")
-    assert (
-        res == '{"type":"text","cache_control":{"type":"ephemeral"},"text":"foo","image_url":null,"input_audio":null}'
+    assert res == (
+        '{"type":"text","cache_control":{"type":"ephemeral"},"text":"foo","image_url":null,"input_audio":null,'
+        '"input_document":null}'
     )

tests/test_cache_control.py

@@ -0,0 +1,74 @@
+import json
+from pathlib import Path
+
+import pytest
+
+from banks import Prompt
+from banks.filters.document import _get_document_format_from_url, _is_url, document
+
+
+@pytest.fixture
+def tiny_pdf():
+    here = Path(__file__).parent
+    return here / "data" / "1x1.pdf"
+
+
+def test_document_with_file_path(tiny_pdf):
+    """Test document filter with a file path input"""
+    result = document(str(tiny_pdf))
+
+    # Verify the content block wrapper
+    assert result.startswith("<content_block>")
+    assert result.endswith("</content_block>")
+
+    # Parse the JSON content
+    json_content = result[15:-16]  # Remove wrapper tags
+    content_block = json.loads(json_content)
+
+    assert content_block["type"] == "document"
+    assert content_block["input_document"]["format"].startswith("pdf")
+
+
+def test_document_with_nonexistent_file():
+    """Test document filter with a nonexistent file path"""
+    with pytest.raises(FileNotFoundError):
+        document("nonexistent/document.pdf")
+
+
+def test_document_with_url():
+    """Test document filter with a URL input (no filesystem access)."""
+    url = "https://example.com/document.css"
+    result = document(url)
+
+    assert result.startswith("<content_block>")
+    assert result.endswith("</content_block>")
+
+    json_content = result[15:-16]  # Remove wrapper tags
+    content_block = json.loads(json_content)
+
+    assert content_block["type"] == "document"
+    assert content_block["input_document"]["data"] == url
+    assert content_block["input_document"]["format"] == "css"
+
+
+def test_is_url_variants():
+    assert _is_url("relative/path.pdf") is False
+    assert _is_url("https://example.com/document.pdf") is True
+    assert _is_url("data:application/pdf;base64,AAAA") is True
+    assert _is_url("data:text/plain;base64,AAAA") is True
+    assert _is_url("data:audio/mp3;base64,AAAA") is False
+
+
+def test_get_document_format_from_url():
+    assert _get_document_format_from_url("https://example.com/document.WAV") == "pdf"
+    assert _get_document_format_from_url("https://example.com/document") == "pdf"
+
+
+def test_document_no_chat_block(tiny_pdf):
+    prompt = Prompt("{{ test }} and {{ another | document }}")
+    messages = prompt.chat_messages({"test": "hello world", "another": str(tiny_pdf)})
+    assert len(messages) == 1
+    message = messages[0]
+    assert len(message.content) == 2
+    assert message.content[0].text == "hello world and"  # type: ignore
+    assert message.content[1].type == "document"  # type:ignore