fix: strip code fences, add MIME/size validation, add tests
This commit is contained in:
@@ -16,9 +16,9 @@ Usage:
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
import yaml
|
||||
@@ -66,11 +66,30 @@ def build_file_content(filepath: str) -> list[dict]:
|
||||
|
||||
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)
|
||||
if mime_type is None:
|
||||
# Default to octet-stream; let the API decide
|
||||
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:
|
||||
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
|
||||
parts: list[str] = []
|
||||
for block in response.content:
|
||||
@@ -122,6 +149,12 @@ def cooklang_from_anthropic(
|
||||
|
||||
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
|
||||
if not cooklang.startswith("---"):
|
||||
print(
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user