src/ad_creative_env/studio/electronics_agent.py

"""Strict public decision contract for the autonomous Electronics studio agent."""

from __future__ import annotations

import json
from collections.abc import Mapping
from typing import Any

ELECTRONICS_AGENT_PROMPT_VERSION = "electronics-agent-tools-v3"


def _tool_branch(tool_name: str, properties: dict[str, Any], required: list[str]) -> dict[str, Any]:
    return {
        "type": "object",
        "additionalProperties": False,
        "required": ["action", "tool_name", "arguments"],
        "properties": {
            "action": {"type": "string", "const": "call_tool"},
            "tool_name": {"type": "string", "const": tool_name},
            "arguments": {
                "type": "object",
                "additionalProperties": False,
                "required": required,
                "properties": properties,
            },
        },
    }


_STRING = {"type": "string", "minLength": 1, "maxLength": 500}
_STRING_ARRAY = {"type": "array", "minItems": 1, "items": _STRING}
_CONSTRAINTS = {
    "type": "object",
    "additionalProperties": False,
    "required": [
        "maximum_price_usd",
        "minimum_memory_gb",
        "minimum_external_4k_displays",
        "preference",
    ],
    "properties": {
        "maximum_price_usd": {"type": "number"},
        "minimum_memory_gb": {"type": "number"},
        "minimum_external_4k_displays": {"type": "number"},
        "preference": _STRING,
    },
}

ELECTRONICS_AGENT_RESPONSE_SCHEMA = {
    "name": "electronics_agent_decision",
    "strict": True,
    "schema": {
        "anyOf": [
            _tool_branch("get_customer_profile", {}, []),
            _tool_branch("search_catalog", {"constraints": _CONSTRAINTS}, ["constraints"]),
            _tool_branch("search_web", {"query": _STRING}, ["query"]),
            _tool_branch(
                "fetch_page_evidence", {"evidence_refs": _STRING_ARRAY}, ["evidence_refs"]
            ),
            _tool_branch("inspect_specs", {"candidate_refs": _STRING_ARRAY}, ["candidate_refs"]),
            _tool_branch(
                "compare_candidates", {"candidate_refs": _STRING_ARRAY}, ["candidate_refs"]
            ),
            _tool_branch("check_compatibility", {"candidate_ref": _STRING}, ["candidate_ref"]),
            _tool_branch(
                "generate_image",
                {
                    "candidate_ref": _STRING,
                    "execution_mode": {"type": "string", "const": "live"},
                },
                ["candidate_ref", "execution_mode"],
            ),
            _tool_branch(
                "compose_ad",
                {
                    "candidate_ref": _STRING,
                    "asset_ref": _STRING,
                    "action": {
                        "type": "object",
                        "additionalProperties": False,
                        "required": ["headline", "body", "cta"],
                        "properties": {
                            "headline": {"type": "string", "minLength": 1, "maxLength": 60},
                            "body": {"type": "string", "minLength": 1, "maxLength": 180},
                            "cta": {"type": "string", "const": "View details"},
                        },
                    },
                },
                ["candidate_ref", "asset_ref", "action"],
            ),
            {
                "type": "object",
                "additionalProperties": False,
                "required": ["action", "intent"],
                "properties": {
                    "action": {"type": "string", "const": "update_working_intent"},
                    "intent": _STRING,
                },
            },
            {
                "type": "object",
                "additionalProperties": False,
                "required": ["action", "candidate_ref"],
                "properties": {
                    "action": {"type": "string", "const": "select_candidate"},
                    "candidate_ref": _STRING,
                },
            },
            {
                "type": "object",
                "additionalProperties": False,
                "required": ["action", "headline", "body", "cta"],
                "properties": {
                    "action": {"type": "string", "const": "write_creative"},
                    "headline": {"type": "string", "minLength": 1, "maxLength": 60},
                    "body": {"type": "string", "minLength": 1, "maxLength": 180},
                    "cta": {"type": "string", "const": "View details"},
                },
            },
            {
                "type": "object",
                "additionalProperties": False,
                "required": ["action", "artifact_ref", "request_version"],
                "properties": {
                    "action": {"type": "string", "const": "submit"},
                    "artifact_ref": _STRING,
                    "request_version": {"type": "integer", "minimum": 1},
                },
            },
        ]
    },
}


class ElectronicsAgentDecisionError(ValueError):
    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code


def parse_electronics_decision(value: str | Mapping[str, Any]) -> dict[str, Any]:
    if isinstance(value, str):
        try:
            value = json.loads(value.strip())
        except json.JSONDecodeError as exc:
            raise ElectronicsAgentDecisionError(
                "malformed_json", "The agent response was not valid JSON."
            ) from exc
    if not isinstance(value, Mapping):
        raise ElectronicsAgentDecisionError(
            "invalid_decision", "The agent decision must be an object."
        )
    decision = dict(value)
    action = decision.get("action")
    exact_fields = {
        "call_tool": {"action", "tool_name", "arguments"},
        "update_working_intent": {"action", "intent"},
        "select_candidate": {"action", "candidate_ref"},
        "write_creative": {"action", "headline", "body", "cta"},
        "submit": {"action", "artifact_ref", "request_version"},
    }
    if action not in exact_fields:
        raise ElectronicsAgentDecisionError(
            "invalid_decision", "The agent selected an unsupported or malformed action."
        )
    if set(decision) != exact_fields[action]:
        raise ElectronicsAgentDecisionError(
            "invalid_decision",
            f"A {action} decision must contain exactly these fields: "
            + ", ".join(sorted(exact_fields[action]))
            + ".",
        )
    if action == "call_tool":
        if not isinstance(decision["tool_name"], str) or not isinstance(
            decision["arguments"], dict
        ):
            raise ElectronicsAgentDecisionError(
                "invalid_arguments", "The agent tool decision is invalid."
            )
    return decision


def build_electronics_agent_prompt(
    *,
    request: str,
    available_actions: list[dict[str, Any]],
    public_state: dict[str, Any],
    last_error: dict[str, str] | None,
) -> str:
    return (
        f"Prompt version: {ELECTRONICS_AGENT_PROMPT_VERSION}. You are the Ad Studio Agent. "
        "Choose exactly one currently available action. Valid choices are ONLY the entries in "
        "available_actions below — any action, tool, or candidate not listed there will be "
        "rejected, even if it seems like the right next step. If last_error is present, do not "
        "repeat the rejected decision; pick one of the actions it lists as currently available. "
        "You own the choice of tool/action and "
        "must use only public state and returned tool evidence. Copy opaque references exactly. "
        "Do not repeat completed work. Resolve conflicting evidence before selecting a candidate. "
        "Return one bare JSON object with no markdown, explanation, or extra fields. "
        "Reply with exactly the fields of the chosen action: "
        'call_tool {"action","tool_name","arguments"}; '
        'update_working_intent {"action","intent"} where intent is one sentence; '
        'select_candidate {"action","candidate_ref"}; '
        'write_creative {"action","headline","body","cta"}; '
        'submit {"action","artifact_ref","request_version"}. '
        "submit is only possible after the compose_ad tool has produced the final card: "
        "artifact_ref must be the artifact_ref returned by compose_ad (never a candidate_ref), "
        "and request_version must be the value shown in the available submit action.\n"
        + json.dumps(
            {
                "marketer_request": request,
                "available_actions": available_actions,
                "public_state": public_state,
                "last_error": last_error,
            },
            sort_keys=True,
            ensure_ascii=False,
        )
    )