Agent/cyber_agent/application/models.py

"""Serializable application declarations and their stable identity hashes."""

from __future__ import annotations

from hashlib import sha256
import json
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator


ValidationScopeName = Literal["evidence", "hypothesis", "output", "task"]


class ApplicationModel(BaseModel):
    model_config = ConfigDict(
        extra="forbid",
        frozen=True,
        str_strip_whitespace=True,
    )


def canonical_json(value: Any) -> str:
    """Return the only JSON representation used by application hashes."""
    if isinstance(value, BaseModel):
        value = value.model_dump(mode="json")
    try:
        return json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        )
    except (TypeError, ValueError) as exc:
        raise ValueError(f"application declaration must be JSON serializable: {exc}") from exc


def stable_hash(value: Any) -> str:
    return sha256(canonical_json(value).encode("utf-8")).hexdigest()


class ApplicationRef(ApplicationModel):
    schema_version: Literal[1] = 1
    application_id: str = Field(min_length=1, max_length=200)
    application_version: str = Field(min_length=1, max_length=100)
    config_hash: str = Field(pattern=r"^[0-9a-f]{64}$")


class ProviderRef(ApplicationModel):
    provider_id: str = Field(min_length=1, max_length=200)
    provider_version: str = Field(min_length=1, max_length=100)
    public_config: dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="after")
    def validate_public_config(self) -> "ProviderRef":
        canonical_json(self.public_config)
        return self


class SkillSpec(ApplicationModel):
    name: str = Field(min_length=1, max_length=200)
    version: str = Field(min_length=1, max_length=100)
    content: str = Field(min_length=1, max_length=100_000)

    @property
    def content_hash(self) -> str:
        return stable_hash({
            "name": self.name,
            "version": self.version,
            "content": self.content,
        })


class ToolSpec(ApplicationModel):
    tool_name: str = Field(min_length=1, max_length=200)
    implementation_version: str = Field(min_length=1, max_length=100)
    side_effect_class: Literal["none", "read", "write", "external"] = "none"
    idempotent: bool = True


class ToolSet(ApplicationModel):
    tool_set_id: str = Field(min_length=1, max_length=200)
    tools: tuple[ToolSpec, ...] = ()

    @model_validator(mode="after")
    def validate_unique_tools(self) -> "ToolSet":
        names = [item.tool_name for item in self.tools]
        if len(names) != len(set(names)):
            raise ValueError(f"ToolSet {self.tool_set_id} contains duplicate tools")
        return self


class QualityRuleSpec(ApplicationModel):
    """A deterministic rule frozen into an application version."""

    rule_id: str = Field(pattern=r"^[a-z][a-z0-9_.-]{2,100}$")
    version: str = Field(min_length=1, max_length=100)
    scope: ValidationScopeName = "output"
    subject_types: tuple[Literal["trace", "candidate"], ...] = ("candidate",)
    criterion: str = Field(min_length=1, max_length=2_000)
    config: dict[str, Any] = Field(default_factory=dict)
    timeout_seconds: float = Field(default=10.0, gt=0, le=120)

    @model_validator(mode="after")
    def validate_config(self) -> "QualityRuleSpec":
        canonical_json(self.config)
        if not self.subject_types or len(self.subject_types) != len(
            set(self.subject_types)
        ):
            raise ValueError("subject_types must be non-empty and unique")
        return self


class RoleRunLimits(ApplicationModel):
    max_iterations: int | None = Field(default=None, gt=0)
    max_depth: int | None = Field(default=None, ge=0)
    max_children_per_parent: int | None = Field(default=None, ge=0)
    max_parallel_children: int | None = Field(default=None, gt=0)


class RunLimits(ApplicationModel):
    max_iterations: int = Field(default=200, gt=0)
    max_depth: int = Field(default=10, ge=0)
    max_children_per_parent: int = Field(default=20, ge=0)
    max_parallel_children: int = Field(default=2, gt=0)
    max_total_agents: int = Field(default=100, gt=0)
    max_llm_calls: int = Field(default=1_000, gt=1)
    max_total_tokens: int = Field(default=10_000_000, gt=0)
    max_total_cost_usd: float = Field(default=100.0, gt=0)
    max_duration_seconds: int = Field(default=86_400, gt=0)
    max_validation_tool_calls: int = Field(default=100, gt=0)
    max_validation_material_chars: int = Field(default=1_000_000, gt=0)

    def tighten(self, role: RoleRunLimits | None) -> "RunLimits":
        if role is None:
            return self
        updates = {
            name: min(getattr(self, name), value)
            for name, value in role.model_dump().items()
            if value is not None
        }
        return self.model_copy(update=updates)


class AgentRole(ApplicationModel):
    role_id: str = Field(min_length=1, max_length=200)
    description: str = Field(default="", max_length=2_000)
    model: str = Field(min_length=1, max_length=200)
    temperature: float = Field(default=0.3, ge=0, le=2)
    model_parameters: dict[str, Any] = Field(default_factory=dict)
    system_prompt: str = Field(min_length=1, max_length=200_000)
    skills: tuple[SkillSpec, ...] = ()
    tool_sets: tuple[str, ...] = ()
    allowed_child_roles: tuple[str, ...] = ()
    default_validation_scopes: tuple[ValidationScopeName, ...] = ()
    run_limits: RoleRunLimits | None = None

    @field_validator("tool_sets", "allowed_child_roles")
    @classmethod
    def validate_unique_refs(cls, values: tuple[str, ...]) -> tuple[str, ...]:
        if len(values) != len(set(values)):
            raise ValueError("role references must be unique")
        return values

    @model_validator(mode="after")
    def validate_model_parameters(self) -> "AgentRole":
        canonical_json(self.model_parameters)
        return self


class AgentApplication(ApplicationModel):
    schema_version: Literal[1] = 1
    application_id: str = Field(min_length=1, max_length=200)
    application_version: str = Field(min_length=1, max_length=100)
    root_role: str = Field(min_length=1, max_length=200)
    roles: tuple[AgentRole, ...] = Field(min_length=1)
    tool_sets: tuple[ToolSet, ...] = ()
    context_provider_ref: ProviderRef | None = None
    artifact_resolver_ref: ProviderRef | None = None
    candidate_repository_ref: ProviderRef | None = None
    quality_provider_ref: ProviderRef | None = None
    quality_rules: tuple[QualityRuleSpec, ...] = ()
    event_projector_ref: ProviderRef | None = None
    validation_policy: dict[str, Any] = Field(default_factory=dict)
    run_limits: RunLimits = Field(default_factory=RunLimits)
    context_inheritance_policy: Literal[
        "explicit_only",
        "allow_inheritable",
    ] = "explicit_only"

    @model_validator(mode="after")
    def validate_graph(self) -> "AgentApplication":
        role_ids = [item.role_id for item in self.roles]
        if len(role_ids) != len(set(role_ids)):
            raise ValueError("AgentApplication role_id values must be unique")
        if self.root_role not in set(role_ids):
            raise ValueError("root_role must reference an existing role")
        tool_set_ids = [item.tool_set_id for item in self.tool_sets]
        if len(tool_set_ids) != len(set(tool_set_ids)):
            raise ValueError("AgentApplication tool_set_id values must be unique")
        known_roles = set(role_ids)
        known_tool_sets = set(tool_set_ids)
        for role in self.roles:
            missing_roles = set(role.allowed_child_roles) - known_roles
            if missing_roles:
                raise ValueError(
                    f"role {role.role_id} references unknown child roles: "
                    f"{sorted(missing_roles)}"
                )
            missing_sets = set(role.tool_sets) - known_tool_sets
            if missing_sets:
                raise ValueError(
                    f"role {role.role_id} references unknown tool sets: "
                    f"{sorted(missing_sets)}"
                )
        canonical_json(self.validation_policy)
        if self.candidate_repository_ref and not self.artifact_resolver_ref:
            raise ValueError(
                "candidate_repository_ref requires artifact_resolver_ref"
            )
        rule_ids = [item.rule_id for item in self.quality_rules]
        if len(rule_ids) != len(set(rule_ids)):
            raise ValueError("quality_rules must contain unique rule_id values")
        if self.quality_rules and not self.quality_provider_ref:
            raise ValueError("quality_rules require quality_provider_ref")
        return self

    def role(self, role_id: str) -> AgentRole:
        return next(item for item in self.roles if item.role_id == role_id)

    def tool_set(self, tool_set_id: str) -> ToolSet:
        return next(item for item in self.tool_sets if item.tool_set_id == tool_set_id)


def role_manifest(
    application: AgentApplication,
    role: AgentRole,
    tool_definitions: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    selected_specs = [
        spec
        for tool_set_id in role.tool_sets
        for spec in application.tool_set(tool_set_id).tools
    ]
    return {
        "role": role.model_dump(mode="json"),
        "prompt_hash": stable_hash(role.system_prompt),
        "skill_hashes": [item.content_hash for item in role.skills],
        "tools": [
            {
                "declaration": spec.model_dump(mode="json"),
                "registry": tool_definitions[spec.tool_name],
            }
            for spec in selected_specs
        ],
    }


def application_manifest(
    application: AgentApplication,
    tool_definitions: list[dict[str, Any]],
) -> dict[str, Any]:
    by_name = {item["tool_name"]: item for item in tool_definitions}
    return {
        "application": application.model_dump(mode="json"),
        "roles": [
            role_manifest(application, role, by_name)
            for role in application.roles
        ],
    }