Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 50 additions & 2 deletions core/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,42 @@
``SkipJsonSchema`` so they never leak into the tool schema shown to the model.
"""

import json
import os
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Literal
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field
from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic.json_schema import SkipJsonSchema


def coerce_jsonish(value: Any) -> Any:
"""If ``value`` is a JSON object/array string, parse it; otherwise return as-is.

Some chat providers stringify nested structures inside tool-call arguments
even after the top-level ``arguments`` blob has been ``json.loads``'d.
"""
if not isinstance(value, str):
return value
text = value.strip()
if not text or text[0] not in "[{":
return value
try:
return json.loads(text)
except json.JSONDecodeError:
return value


def coerce_list(value: Any) -> Any:
"""Coerce a stringified JSON list (or a lone object) into a list."""
value = coerce_jsonish(value)
if isinstance(value, dict):
return [value]
return value


_COMIC_STYLE_HINT = "manhua/comic style: clean black ink line art, soft cel shading, flat colors"

# Six resumable pipeline stages.
Expand Down Expand Up @@ -53,6 +80,12 @@ class CharacterAsset(BaseModel):
name: str
role: str = ""
appearance: Appearance = Field(default_factory=Appearance)

@field_validator("appearance", mode="before")
@classmethod
def _coerce_appearance(cls, value: Any) -> Any:
return coerce_jsonish(value)

# Hardened description inlined into every panel prompt. Prefer deriving from
# ``appearance`` via ``build_l1_from_appearance``; the model may still fill this.
l1_prompt: str = ""
Expand Down Expand Up @@ -96,6 +129,11 @@ class StoryElements(BaseModel):
),
)

@field_validator("characters", "settings", mode="before")
@classmethod
def _coerce_asset_lists(cls, value: Any) -> Any:
return coerce_list(value)


class Panel(BaseModel):
"""A single comic panel within a storyboard."""
Expand All @@ -114,6 +152,11 @@ class Panel(BaseModel):
reference_characters: list[str] = Field(default_factory=list)
size: str = "1024x1024"

@field_validator("characters_present", "reference_characters", mode="before")
@classmethod
def _coerce_name_lists(cls, value: Any) -> Any:
return coerce_list(value)


class Storyboard(BaseModel):
"""Return payload of the ``plan_storyboard`` forced function call (one chunk)."""
Expand All @@ -123,6 +166,11 @@ class Storyboard(BaseModel):
chapter_id: str
panels: list[Panel] = Field(default_factory=list)

@field_validator("panels", mode="before")
@classmethod
def _coerce_panels(cls, value: Any) -> Any:
return coerce_list(value)


class ChunkCache(BaseModel):
"""Per-chunk cache of the billable chat-API results.
Expand Down
30 changes: 30 additions & 0 deletions tests/test_schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,36 @@
to_tool_schema,
)


def test_story_elements_coerces_stringified_lists():
"""Agnes sometimes JSON-stringifies nested arrays inside tool arguments."""
elements = StoryElements.model_validate(
{
"characters": '[{"name": "张一新", "role": "protagonist", '
'"l1_prompt": "conflict and reflection."}]',
"settings": "[]",
"style_guide": "manhua style",
}
)
assert len(elements.characters) == 1
assert elements.characters[0].name == "张一新"
assert elements.settings == []


def test_storyboard_coerces_stringified_panels_and_char_lists():
board = Storyboard.model_validate(
{
"chapter_id": "c1",
"panels": (
'[{"panel_id": "c1_p01", "action": "looks up", '
'"characters_present": "[\\"张一新\\"]"}]'
),
}
)
assert board.panels[0].panel_id == "c1_p01"
assert board.panels[0].characters_present == ["张一新"]


STORY_ELEMENTS_PAYLOAD = {
"characters": [
{
Expand Down
11 changes: 11 additions & 0 deletions tests/test_screenwriter.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,17 @@ def test_extract_story_elements_parses_payload():
assert result.style_guide == "manhua style"


def test_extract_story_elements_accepts_stringified_character_list():
payload = {
"characters": '[{"name": "张一新", "l1_prompt": "conflict and reflection."}]',
"settings": "[]",
"style_guide": "manhua style",
}
result = asyncio.run(extract_story_elements("text", chat=FakeChat(payload)))
assert result.characters[0].name == "张一新"
assert result.settings == []


def test_plan_storyboard_parses_payload():
payload = {
"chapter_id": "ch01",
Expand Down