laygen¶

class ChatMessageLike(TypedDict):
    """Chat-style message with role and content text fields."""

    role: str
    content: str

class LayoutOutputDict(TypedDict):
    """Dictionary form of the canonical layout generation schema."""

    bbox: (
        Float[np.ndarray, "batch elements 4"] | Float[torch.Tensor, "batch elements 4"]
    )
    labels: Int[np.ndarray, "batch elements"] | Int[torch.Tensor, "batch elements"]
    mask: Bool[np.ndarray, "batch elements"] | Bool[torch.Tensor, "batch elements"]
    id2label: dict[int, str]
    sequences: LayoutAuxValue | None
    scores: LayoutAuxValue | None
    trajectory: LayoutAuxValue | None
    intermediates: Mapping[str, LayoutAuxValue] | None

class ExemplarSelector(Protocol[ExampleT]):
    """Strategy that chooses in-context examples for a layout prompt."""

    def __call__(self, prompt: str, examples: Sequence[ExampleT]) -> Sequence[ExampleT]:
        """Return examples selected for ``prompt``."""
        ...

class PromptBuilder(Protocol[ExampleT]):
    """Strategy that serializes a user request and exemplars for a provider."""

    def __call__(
        self, prompt: str, exemplars: Sequence[ExampleT]
    ) -> str | Sequence[ChatMessageLike]:
        """Serialize ``prompt`` and ``exemplars`` for model execution."""
        ...

class ResponseParser(Protocol):
    """Strategy that converts provider text into a shared layout output."""

    def __call__(self, text: str, *, canvas_size: int) -> LayoutGenerationOutput:
        """Parse provider ``text`` into the shared output schema."""
        ...

@dataclass
class BaseResponseParser(Generic[ParsedOutputT], ABC):
    """Small base for parser strategies with shared repair/error hooks."""

    parser_name: str = "response parser"

    def repair_response_text(self, text: str) -> str:
        """Repair provider text before parser-specific extraction."""
        return text

    def parser_error(self, message: str) -> RuntimeError:
        """Build a consistent parser error with parser context."""
        return RuntimeError(f"{self.parser_name}: {message}")

@dataclass
class BaseExemplarSelector(Generic[ExampleT], ABC):
    """Small base for selector strategies with shared candidate validation."""

    def validate_examples(self, examples: Sequence[ExampleT]) -> None:
        """Validate the selector has at least one candidate exemplar."""
        if not examples:
            msg = "exemplar selector requires at least one candidate"
            raise ValueError(msg)

    def selection_error(self, message: str) -> ValueError:
        """Build a consistent selector error."""
        return ValueError(f"exemplar selector: {message}")

class LayoutItem2DLike(Protocol):
    """Minimal parsed 2D item required by the shared output builder."""

    @property
    def label(self) -> str:
        """Display label parsed for this item."""
        ...

    @property
    def bbox_xywh(self) -> tuple[float, float, float, float]:
        """Normalized center ``xywh`` box."""
        ...

class LayoutOutputLike(Protocol):
    """Minimal shared layout output fields used for dict serialization."""

    @property
    def bbox(
        self,
    ) -> (
        Float[np.ndarray, "batch elements 4"] | Float[torch.Tensor, "batch elements 4"]
    ):
        """Layout boxes."""
        ...

    @property
    def labels(
        self,
    ) -> Int[np.ndarray, "batch elements"] | Int[torch.Tensor, "batch elements"]:
        """Layout labels."""
        ...

    @property
    def mask(
        self,
    ) -> Bool[np.ndarray, "batch elements"] | Bool[torch.Tensor, "batch elements"]:
        """Valid element mask."""
        ...

    @property
    def id2label(self) -> dict[int, str]:
        """Dataset-local label names."""
        ...

class BaseLayoutAgent(Generic[RawResponseT], ABC):
    """Base Pydantic AI runner for text-conditioned layout agents.

    Subclasses own model-specific exemplar selection, prompt serialization, and
    response parsing. This base class centralizes provider model resolution,
    Pydantic AI ``Agent`` construction, common public request validation, and
    shared output dictionary serialization.
    """

    def __init__(
        self,
        *,
        model: ModelLike = None,
        model_env_var: str,
        raw_response_type: type[RawResponseT],
        instructions: str,
    ) -> None:
        """Initialize the provider runner.

        Args:
            model: Optional Pydantic AI model object or provider model id.
            model_env_var: Environment variable used when ``model`` is omitted.
            raw_response_type: Structured response model expected from the LLM.
            instructions: Provider instructions passed to Pydantic AI.
        """
        self.model_env_var = model_env_var
        self.raw_response_type = raw_response_type
        self.instructions = instructions
        self.agent = self.build_pydantic_agent(model=model)

    def resolve_model(self, model: ModelLike = None) -> ModelLike:
        """Resolve a per-call model override, constructor model, or env model id."""
        return model or os.getenv(self.model_env_var)

    def build_pydantic_agent(self, *, model: ModelLike = None) -> Agent[None]:
        """Build the underlying Pydantic AI agent."""
        return Agent(
            self.resolve_model(model),
            output_type=self.raw_response_type,
            instructions=self.instructions,
        )

    def run_raw_sync(
        self,
        model_prompt: str | Sequence[ChatMessageLike],
        *,
        model: ModelLike = None,
        model_settings: ModelSettings | None = None,
    ) -> RawResponseT:
        """Run the provider synchronously and return the structured raw response."""
        run_result = self.agent.run_sync(
            messages_to_text(model_prompt),
            model=model,
            model_settings=model_settings,
        )
        return cast(RawResponseT, run_result.output)

    def validate_generation_request(
        self,
        *,
        batch_size: int,
        condition_type: str | ConditionType,
        box_format: str | BoxFormat,
        canvas_size: tuple[int, int] | None,
        configured_canvas_size: int,
        supported_condition_types: tuple[
            ConditionType, ...
        ] = DEFAULT_SUPPORTED_CONDITIONS,
    ) -> tuple[ConditionType, BoxFormat]:
        """Validate shared generation arguments before provider execution."""
        normalized_condition_type = normalize_condition_type(condition_type)
        normalized_box_format = normalize_box_format(box_format)
        if batch_size != 1:
            msg = "provider-backed layout agents currently support batch_size=1."
            raise ValueError(msg)

        if normalized_condition_type not in supported_condition_types:
            msg = f"unsupported condition_type for this agent: {normalized_condition_type}"
            raise ValueError(msg)

        if canvas_size is not None and canvas_size != (
            configured_canvas_size,
            configured_canvas_size,
        ):
            msg = (
                "provider-backed layout agents use their configured square canvas_size."
            )
            raise ValueError(msg)

        return normalized_condition_type, normalized_box_format

    def output_to_dict(self, output: LayoutOutputLike) -> LayoutOutputDict:
        """Serialize the shared output with canonical layout schema keys."""
        return {
            "bbox": output.bbox,
            "labels": output.labels,
            "mask": output.mask,
            "id2label": output.id2label,
            "sequences": cast(
                LayoutAuxValue | None, getattr(output, "sequences", None)
            ),
            "scores": cast(LayoutAuxValue | None, getattr(output, "scores", None)),
            "trajectory": cast(
                LayoutAuxValue | None, getattr(output, "trajectory", None)
            ),
            "intermediates": cast(
                Mapping[str, LayoutAuxValue] | None,
                getattr(output, "intermediates", None),
            ),
        }

    def repair_response_text(self, text: str) -> str:
        """Hook for model-specific response repair before parsing."""
        return text

    def should_retry(self, exc: Exception, *, attempt: int) -> bool:
        """Hook for model-specific retry policy after provider or parse failure."""
        del exc, attempt
        return False

    def retry_delay_seconds(self, *, attempt: int) -> float:
        """Hook for retry backoff policies used by subclasses."""
        del attempt
        return 0.0

    def run_with_repair_policy(
        self,
        operation: Callable[[], RawResponseT],
        *,
        max_attempts: int = 1,
    ) -> RawResponseT:
        """Run an operation with subclass retry-policy hooks."""
        attempt = 1
        while True:
            try:
                return operation()
            except Exception as exc:
                if attempt >= max_attempts or not self.should_retry(
                    exc, attempt=attempt
                ):
                    raise

                attempt += 1