Add typed prompt axes route

row_prompt_axes.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any
 
 try:
@@ -18,7 +19,33 @@ except ImportError:  # Allows local smoke tests from the repository root.
     from hardcore_text_cleanup import sanitize_hardcore_environment_anchors
 
 
-def resolve_prompt_axes(
+@dataclass(frozen=True)
+class PromptAxesRoute:
+    scene_slug: str
+    scene: str
+    pose: str
+    expression: str
+    shared_expression: str
+    character_expressions: list[str]
+    character_expression_text: str
+    source_composition: str
+    composition: str
+
+    def as_dict(self) -> dict[str, Any]:
+        return {
+            "scene_slug": self.scene_slug,
+            "scene": self.scene,
+            "pose": self.pose,
+            "expression": self.expression,
+            "shared_expression": self.shared_expression,
+            "character_expressions": list(self.character_expressions),
+            "character_expression_text": self.character_expression_text,
+            "source_composition": self.source_composition,
+            "composition": self.composition,
+        }
+
+
+def resolve_prompt_axes_result(
     *,
     category: dict[str, Any],
     subcategory: dict[str, Any],
@@ -43,7 +70,7 @@ def resolve_prompt_axes(
     pov_character_labels: list[str] | None = None,
     location_config: dict[str, Any] | None = None,
     composition_config: dict[str, Any] | None = None,
-) -> dict[str, Any]:
+) -> PromptAxesRoute:
     character_slots = character_slots or []
     character_slot_map = character_slot_map or {}
     pov_character_labels = pov_character_labels or []
@@ -122,14 +149,67 @@ def resolve_prompt_axes(
         source_composition = sanitize_hardcore_environment_anchors(source_composition)
     composition = pov_policy.pov_composition_prompt(source_composition, pov_character_labels)
 
-    return {
-        "scene_slug": scene_slug,
-        "scene": scene,
-        "pose": pose,
-        "expression": expression,
-        "shared_expression": shared_expression,
-        "character_expressions": character_expressions,
-        "character_expression_text": character_expression_text,
-        "source_composition": source_composition,
-        "composition": composition,
-    }
+    return PromptAxesRoute(
+        scene_slug=scene_slug,
+        scene=scene,
+        pose=pose,
+        expression=expression,
+        shared_expression=shared_expression,
+        character_expressions=character_expressions,
+        character_expression_text=character_expression_text,
+        source_composition=source_composition,
+        composition=composition,
+    )
+
+
+def resolve_prompt_axes(
+    *,
+    category: dict[str, Any],
+    subcategory: dict[str, Any],
+    item: Any,
+    subject_type: str,
+    context: dict[str, Any],
+    poses: str,
+    women_count: int,
+    men_count: int,
+    scene_rng: Any,
+    pose_rng: Any,
+    expression_rng: Any,
+    composition_rng: Any,
+    expression_disabled: bool,
+    expression_intensity: float,
+    character_slots: list[dict[str, Any]] | None = None,
+    character_slot_map: dict[str, dict[str, Any]] | None = None,
+    expression_phase: str = "",
+    source_role_graph: Any = "",
+    item_axis_values: dict[str, Any] | None = None,
+    is_pose_category: bool = False,
+    pov_character_labels: list[str] | None = None,
+    location_config: dict[str, Any] | None = None,
+    composition_config: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    return resolve_prompt_axes_result(
+        category=category,
+        subcategory=subcategory,
+        item=item,
+        subject_type=subject_type,
+        context=context,
+        poses=poses,
+        women_count=women_count,
+        men_count=men_count,
+        scene_rng=scene_rng,
+        pose_rng=pose_rng,
+        expression_rng=expression_rng,
+        composition_rng=composition_rng,
+        expression_disabled=expression_disabled,
+        expression_intensity=expression_intensity,
+        character_slots=character_slots,
+        character_slot_map=character_slot_map,
+        expression_phase=expression_phase,
+        source_role_graph=source_role_graph,
+        item_axis_values=item_axis_values,
+        is_pose_category=is_pose_category,
+        pov_character_labels=pov_character_labels,
+        location_config=location_config,
+        composition_config=composition_config,
+    ).as_dict()