src/id/generator/originality.py

"""Originality: best-of-n premise generation + judge + divergence (Section 9.1).

Author ``n`` cheap candidate concepts (premise + culprit + twist), have the
judge tier score originality against the cliché blocklist and past worlds, and
return the winner to be expanded into a full world.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from ..llm.client import LLMClient
from ..llm.prompts import PromptRegistry


@dataclass
class Concept:
    premise: str
    setting: str
    culprit: str
    twist: str
    one_line: str
    twist_tag: str

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Concept:
        return cls(
            premise=str(d.get("premise", "")),
            setting=str(d.get("setting", "")),
            culprit=str(d.get("culprit", "")),
            twist=str(d.get("twist", "")),
            one_line=str(d.get("one_line", d.get("premise", "")))[:140],
            twist_tag=str(d.get("twist_tag", "")),
        )


def generate_concepts(
    client: LLMClient, prompts: PromptRegistry, *, seed: str, n: int,
    past: list[dict[str, str]],
) -> list[Concept]:
    concepts: list[Concept] = []
    for i in range(n):
        prompt = prompts.render(
            "author/concept.md.j2", seed=seed, past=past, variant=i + 1,
        )
        try:
            data, _ = client.complete_json(
                tier="author", task="concept_gen", user=prompt,
            )
        except Exception:
            continue
        if isinstance(data, dict):
            concepts.append(Concept.from_dict(data))
    return concepts


def judge_concepts(
    client: LLMClient, prompts: PromptRegistry, *, concepts: list[Concept],
    past: list[dict[str, str]],
) -> tuple[Concept, list[float]]:
    """Score each concept's originality; return the winner + all scores."""
    payload = [
        {"index": i, "premise": c.premise, "setting": c.setting,
         "culprit": c.culprit, "twist": c.twist}
        for i, c in enumerate(concepts)
    ]
    prompt = prompts.render("judge/originality.md.j2", concepts=payload, past=past)
    scores = [0.0] * len(concepts)
    try:
        data, _ = client.complete_json(
            tier="judge", task="originality_judge", user=prompt,
        )
        rows = data.get("scores", []) if isinstance(data, dict) else data
        for row in rows:
            idx = int(row.get("index", -1))
            if 0 <= idx < len(scores):
                scores[idx] = float(row.get("score", 0.0))
    except Exception:
        pass
    best = max(range(len(concepts)), key=lambda i: scores[i]) if concepts else 0
    return concepts[best], scores