Extract row category route policy

prompt_builder.py

@@ -35,6 +35,7 @@ try:
     from . import pov_policy
     from . import row_normalization as row_policy
     from . import row_camera as row_camera_policy
+    from . import row_category_route as row_category_route_policy
     from . import row_expression as row_expression_policy
     from . import row_generation as row_generation_policy
     from . import row_item as row_item_policy
@@ -80,6 +81,7 @@ except ImportError:  # Allows local smoke tests with `python -c`.
     import pov_policy
     import row_normalization as row_policy
     import row_camera as row_camera_policy
+    import row_category_route as row_category_route_policy
     import row_expression as row_expression_policy
     import row_generation as row_generation_policy
     import row_item as row_item_policy
@@ -772,18 +774,32 @@ def _axis_rng(seed_config: dict[str, int], axis: str, base_seed: int, row_number
 
 
 def _is_pose_content_category(category: dict[str, Any], subcategory: dict[str, Any]) -> bool:
-    haystack = " ".join(
-        str(value)
-        for value in (
-            category.get("name", ""),
-            category.get("slug", ""),
-            category.get("item_label", ""),
-            subcategory.get("name", ""),
-            subcategory.get("slug", ""),
-            subcategory.get("item_label", ""),
-        )
-    ).lower()
-    return "pose" in haystack or "sex" in haystack
+    return row_category_route_policy.is_pose_content_category(category, subcategory)
+
+
+def _select_category_item_route(
+    *,
+    category_choice: str,
+    subcategory_choice: str,
+    seed_config: dict[str, int],
+    seed: int,
+    row_number: int,
+    women_count: int,
+    men_count: int,
+    hardcore_position_config: dict[str, Any] | None = None,
+    categories: list[dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    return row_category_route_policy.select_category_item_route(
+        category_choice=category_choice,
+        subcategory_choice=subcategory_choice,
+        seed_config=seed_config,
+        seed=seed,
+        row_number=row_number,
+        women_count=women_count,
+        men_count=men_count,
+        hardcore_position_config=hardcore_position_config,
+        categories=categories,
+    )
 
 
 def _format(template: str, context: dict[str, Any]) -> str:
@@ -2004,9 +2020,6 @@ def _build_custom_row(
     location_config: str | dict[str, Any] | None = None,
     composition_config: str | dict[str, Any] | None = None,
 ) -> dict[str, Any]:
-    categories = load_category_library()
-    category_rng = _axis_rng(seed_config, "category", seed, row_number)
-    subcategory_rng = _axis_rng(seed_config, "subcategory", seed, row_number)
     person_rng = _axis_rng(seed_config, "person", seed, row_number)
     scene_rng = _axis_rng(seed_config, "scene", seed, row_number)
     pose_rng = _axis_rng(seed_config, "pose", seed, row_number)
@@ -2017,50 +2030,29 @@ def _build_custom_row(
     parsed_location_config = _parse_location_config(location_config)
     parsed_composition_config = _parse_composition_config(composition_config)
 
-    requested_women_count = women_count
-    requested_men_count = men_count
-    categories = _filter_hardcore_categories_for_position(
-        categories,
-        parsed_hardcore_position_config,
-        women_count,
-        men_count,
+    category_route = _select_category_item_route(
+        category_choice=category_choice,
+        subcategory_choice=subcategory_choice,
+        seed_config=seed_config,
+        seed=seed,
+        row_number=row_number,
+        women_count=women_count,
+        men_count=men_count,
+        hardcore_position_config=parsed_hardcore_position_config,
     )
-    category, subcategory, women_count, men_count = _find_subcategory(
-        categories,
-        category_choice,
-        subcategory_choice,
-        category_rng,
-        subcategory_rng,
-        women_count,
-        men_count,
-    )
-    count_adjustment = {}
-    if women_count != requested_women_count or men_count != requested_men_count:
-        count_adjustment = {
-            "requested_women_count": requested_women_count,
-            "requested_men_count": requested_men_count,
-            "effective_women_count": women_count,
-            "effective_men_count": men_count,
-        }
-    if _is_hardcore_sexual_category(category):
-        subcategory = _apply_hardcore_position_config_to_subcategory(subcategory, parsed_hardcore_position_config)
-    content_axis = "pose" if _is_pose_content_category(category, subcategory) else "content"
-    content_rng = _axis_rng(seed_config, content_axis, seed, row_number)
-    items = _list_from(subcategory.get("items", [subcategory["name"]]))
-    item = _weighted_choice(content_rng, items)
-    item_text, item_name, item_axis_values, item_template_metadata = _compose_item(
-        content_rng,
-        category,
-        subcategory,
-        item,
-        women_count,
-        men_count,
-    )
-    is_pose_category = _is_pose_content_category(category, subcategory)
-    if is_pose_category:
-        item_text = _sanitize_hardcore_environment_anchors(item_text)
-        item_axis_values = _sanitize_hardcore_axis_values(item_axis_values)
-    item_formatter_hints = _template_formatter_hints(item_template_metadata)
+    category = category_route["category"]
+    subcategory = category_route["subcategory"]
+    women_count = int(category_route["women_count"])
+    men_count = int(category_route["men_count"])
+    count_adjustment = dict(category_route.get("count_adjustment") or {})
+    content_axis = str(category_route.get("content_axis") or "content")
+    item = category_route["item"]
+    item_text = str(category_route.get("item_text") or "")
+    item_name = str(category_route.get("item_name") or "")
+    item_axis_values = dict(category_route.get("item_axis_values") or {})
+    item_template_metadata = dict(category_route.get("item_template_metadata") or {})
+    item_formatter_hints = dict(category_route.get("formatter_hints") or {})
+    is_pose_category = bool(category_route.get("is_pose_category"))
     subject_type = str(_merged_field(category, subcategory, item, "subject_type", "single_any"))
     context = _subject_context(person_rng, subject_type, ethnicity, figure, no_plus_women, no_black, women_count, men_count)
     character_slots = _parse_character_cast(character_cast)