fix: strip code fences, add MIME/size validation, add tests
This commit is contained in:
@@ -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(
|
||||||
|
|||||||
@@ -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