fix: strip code fences, add MIME/size validation, add tests

This commit is contained in:
2026-05-21 10:23:18 -04:00
parent 7a3f63e377
commit c015dce51e
2 changed files with 126 additions and 2 deletions
+35 -2
View File
@@ -16,9 +16,9 @@ Usage:
import argparse import argparse
import base64 import base64
import json
import mimetypes import mimetypes
import os import os
import re
import sys import sys
import yaml import yaml
@@ -66,11 +66,30 @@ def build_file_content(filepath: str) -> list[dict]:
Returns a list with one document/image content block. Returns a list with one document/image content block.
""" """
SUPPORTED_DOCUMENT_TYPES = {"application/pdf"}
SUPPORTED_IMAGE_TYPES = {"image/jpeg", "image/png", "image/gif", "image/webp"}
mime_type, _ = mimetypes.guess_type(filepath) mime_type, _ = mimetypes.guess_type(filepath)
if mime_type is None: if mime_type is None:
# Default to octet-stream; let the API decide
mime_type = "application/octet-stream" mime_type = "application/octet-stream"
if mime_type not in SUPPORTED_DOCUMENT_TYPES | SUPPORTED_IMAGE_TYPES:
print(
f"Error: unsupported file type '{mime_type}'. "
f"Supported types: PDF, JPEG, PNG, GIF, WebP",
file=sys.stderr,
)
sys.exit(1)
MAX_FILE_SIZE = 20 * 1024 * 1024 # 20MB
file_size = os.path.getsize(filepath)
if file_size > MAX_FILE_SIZE:
print(
f"Error: file too large ({file_size} bytes). Maximum file size is 20MB.",
file=sys.stderr,
)
sys.exit(1)
with open(filepath, "rb") as f: with open(filepath, "rb") as f:
data = base64.b64encode(f.read()).decode("utf-8") data = base64.b64encode(f.read()).decode("utf-8")
@@ -114,6 +133,14 @@ def cooklang_from_anthropic(
], ],
) )
# Warn if response was truncated
if response.stop_reason == "max_tokens":
print(
"Warning: response may be truncated (reached max_tokens limit). "
"Consider splitting the recipe into smaller parts.",
file=sys.stderr,
)
# Extract text from response content blocks # Extract text from response content blocks
parts: list[str] = [] parts: list[str] = []
for block in response.content: for block in response.content:
@@ -122,6 +149,12 @@ def cooklang_from_anthropic(
cooklang = "\n".join(parts).strip() cooklang = "\n".join(parts).strip()
# Strip markdown code fences if present (the prompts may instruct the
# model to wrap output in ```cooklang blocks)
cooklang = re.sub(r"^```\w*\n?", "", cooklang)
cooklang = re.sub(r"\n```\s*$", "", cooklang)
cooklang = cooklang.strip()
# Basic validation: must start with front matter marker # Basic validation: must start with front matter marker
if not cooklang.startswith("---"): if not cooklang.startswith("---"):
print( print(
+91
View File
@@ -0,0 +1,91 @@
"""Tests for scripts/convert-to-cooklang."""
import importlib.machinery
import importlib.util
import os
import tempfile
# Load the script as a module (no .py extension, so use SourceFileLoader)
SCRIPT_PATH = os.path.join(
os.path.dirname(__file__), "..", "scripts", "convert-to-cooklang"
)
_loader = importlib.machinery.SourceFileLoader("convert_to_cooklang", SCRIPT_PATH)
_spec = importlib.util.spec_from_loader(
"convert_to_cooklang", _loader, origin=SCRIPT_PATH
)
mod = importlib.util.module_from_spec(_spec)
_loader.exec_module(mod)
def test_load_config_valid():
"""Test loading a valid config file."""
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
f.write("anthropic:\n api_key: sk-test-123\n")
config_path = f.name
try:
config = mod.load_config(config_path)
assert config["api_key"] == "sk-test-123"
finally:
os.unlink(config_path)
def test_load_config_missing_api_key():
"""Test that missing api_key exits with error."""
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
f.write("anthropic:\n model: foo\n")
config_path = f.name
try:
try:
mod.load_config(config_path)
assert False, "Expected sys.exit"
except SystemExit as e:
assert e.code == 1
finally:
os.unlink(config_path)
def test_build_text_content():
"""Test building a text content block."""
result = mod.build_text_content("Hello")
assert result == [{"type": "text", "text": "Hello"}]
def test_build_file_content_pdf():
"""Test building file content for a PDF."""
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
f.write(b"%PDF- test")
filepath = f.name
try:
result = mod.build_file_content(filepath)
assert result[0]["type"] == "document"
assert result[0]["source"]["media_type"] == "application/pdf"
finally:
os.unlink(filepath)
def test_build_file_content_jpeg():
"""Test building file content for a JPEG image."""
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
f.write(b"\xff\xd8\xff\xe0 test jpeg")
filepath = f.name
try:
result = mod.build_file_content(filepath)
assert result[0]["type"] == "image"
assert result[0]["source"]["media_type"] == "image/jpeg"
finally:
os.unlink(filepath)
def test_unsupported_file_type():
"""Test that unsupported file types are rejected."""
with tempfile.NamedTemporaryFile(suffix=".tiff", delete=False) as f:
f.write(b"test")
filepath = f.name
try:
try:
mod.build_file_content(filepath)
assert False, "Expected sys.exit for unsupported type"
except SystemExit as e:
assert e.code == 1
finally:
os.unlink(filepath)