agentic_core/schemas/database.py

"""Structured output of the Database Agent."""

from __future__ import annotations

from typing import Any, ClassVar

from pydantic import BaseModel, Field, field_validator

from .limits import CappedListModel


class DBField(BaseModel):
    name: str
    type: str
    primary_key: bool = False
    foreign_key: str | None = None
    nullable: bool = False
    unique: bool = False
    indexed: bool = False

    @field_validator("primary_key", "nullable", "unique", "indexed", mode="before")
    @classmethod
    def _normalize_bools(cls, v: Any) -> bool:
        if isinstance(v, str):
            return v.strip().lower() in ("true", "1", "yes", "t")
        return bool(v)

    @field_validator("foreign_key", mode="before")
    @classmethod
    def _normalize_fk(cls, v: Any) -> str | None:
        if v is None or v is False or v == "" or str(v).strip().lower() in ("none", "null", "false"):
            return None
        return str(v).strip()


class DBEntity(CappedListModel):
    name: str
    description: str
    fields: list[DBField] = Field(default_factory=list, max_length=10)


class DatabaseOutput(CappedListModel):
    database_technology: str = ""
    entities: list[DBEntity] = Field(..., max_length=8)
    relationships: list[str] = Field(default_factory=list, max_length=8)
    indexes: list[str] = Field(default_factory=list)
    constraints: list[str] = Field(default_factory=list)
    sql_schema: str = ""
    erd_mermaid: str = ""

    # sql_schema and erd_mermaid are derived locally (see artifacts/render.py).
    # indexes and constraints are derived from entity fields (PKs, FKs, unique flags),
    # so the model never needs to spend output tokens on them.
    # Keeping the LLM output small is the main lever for preventing database timeouts.
    llm_exclude_fields: ClassVar[frozenset[str]] = frozenset({
        "sql_schema", "erd_mermaid", "indexes", "constraints"
    })