Stirling-PDF/engine/scripts/generate_tool_models.py

#!/usr/bin/env python3
"""Generate the Python files derived from the Java OpenAPI spec (SwaggerDoc.json).

tool_models.py holds each tool's request model; tool_io.py holds what it accepts and produces,
from ``@ToolIO`` via the ``x-stirling-io`` extension. One pass over one spec, so the two cannot
drift apart. Run via:
    task engine:tool-models
"""

from __future__ import annotations

import argparse
import json
import subprocess
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from datamodel_code_generator import InputFileType, PythonVersion, generate
from datamodel_code_generator.enums import DataModelType
from datamodel_code_generator.format import Formatter
from referencing import Registry, Resource
from referencing.jsonschema import DRAFT202012

# Fields inherited from PDFFile base class - not tool parameters.
BASE_CLASS_FIELDS = frozenset({"fileInput", "fileId"})

IO_EXTENSION = "x-stirling-io"
IO_VOCABULARY_EXTENSION = "x-stirling-io-vocabulary"

_ENGINE_ROOT = Path(__file__).resolve().parents[1]

_IO_TEMPLATE = '''# AUTO-GENERATED FILE. DO NOT EDIT.
# Generated by scripts/generate_tool_models.py from the Java OpenAPI spec (SwaggerDoc.json).
# Regenerate with: task engine:tool-models
"""What each tool endpoint accepts and produces, so a planned chain can be checked before it
is run. Declared in Java with ``@ToolIO``; see ``stirling.services.tool_io_compat`` for the
compatibility rules that read this table."""

from enum import StrEnum

from pydantic import Field

from stirling.models.base import ApiModel
from stirling.models.tool_models import ToolEndpoint


class ToolFormat(StrEnum):
    """The kind of file a tool consumes or produces. ``ANY`` accepts or produces anything;
    ``NONE`` means no file at all, such as a report or a status."""

{formats}


class ToolArity(StrEnum):
    """How many files go in and come out (Single/Multiple In, Single/Multiple Out). A
    multi-output tool returns its results zipped, and the caller unpacks them."""

{arities}


class ToolIOWhen(ApiModel):
    """One condition on a request parameter, guarding a :class:`ToolIOCase`."""

    param: str
    matches: list[str]


class ToolIOCase(ApiModel):
    """An output that applies when every condition in ``when`` holds."""

    when: list[ToolIOWhen]
    produces: ToolFormat
    arity: ToolArity


class ToolIOSpec(ApiModel):
    """What one endpoint accepts and produces."""

    accepts: list[ToolFormat]
    produces: ToolFormat
    arity: ToolArity
    cases: list[ToolIOCase] = Field(default_factory=list)


TOOL_IO: dict[ToolEndpoint, ToolIOSpec] = {{
{declarations}
}}
'''

_FILE_HEADER = (
    "# AUTO-GENERATED FILE. DO NOT EDIT.\n"
    "# Generated by scripts/generate_tool_models.py from Java OpenAPI spec (SwaggerDoc.json).\n"
    "# ruff: noqa: E501"
)


@dataclass
class ToolSpec:
    path: str
    enum_name: str
    class_name: str


@dataclass
class DiscoveryResult:
    tools: list[ToolSpec]
    combined_schema: dict[str, Any]


class ToolDiscovery:
    """Discovers tool endpoints from an OpenAPI spec and builds a combined JSON Schema."""

    # Namespaces exposed to the LLM as callable tools. Largely matches ``InternalApiClient.java``.
    # Note: ``/api/v1/filter/`` is intentionally excluded because those APIs are for pipeline processing,
    #  not tool execution.
    ALLOWED_PATH_PREFIXES = (
        "/api/v1/general/",
        "/api/v1/misc/",
        "/api/v1/security/",
        "/api/v1/convert/",
    )

    # Endpoints under the allowed prefixes that are NOT edit-agent operations. A listed
    # path and everything nested under it is dropped. Several kinds live here:
    EXCLUDED_PATHS = (
        # 1. Cert-signing family: needs certificate/key files the agent can't supply, plus
        #    interactive session and hardware-token management. The whole subtree is dropped.
        "/api/v1/security/cert-sign",
        # 2. Interactive PDF text-editor endpoints, not one-shot operations.
        "/api/v1/convert/pdf/text-editor",
        "/api/v1/convert/text-editor/pdf",
        # 3. Introspection / query endpoints that return metadata, a listing, or a
        #    verification verdict rather than a transformed document, so they belong to
        #    the question path, not the edit agent. (decompress is a dev-only stream op.)
        "/api/v1/security/get-info-on-pdf",
        "/api/v1/security/verify-pdf",
        "/api/v1/security/validate-signature",
        "/api/v1/misc/list-attachments",
        "/api/v1/misc/show-javascript",
        "/api/v1/misc/decompress-pdf",
        "/api/v1/general/extract-bookmarks",
        # 4. Require a secondary file (image, overlay PDF, attachments) on top of the input
        #    PDF. The agent only ever supplies the input PDF(s), so these can never run.
        #    (add-stamp / add-watermark stay: their text mode needs no extra file.)
        "/api/v1/misc/add-image",
        "/api/v1/misc/add-attachments",
        "/api/v1/general/overlay-pdfs",
    )

    def _is_excluded(self, path: str) -> bool:
        return any(path == p or path.startswith(p + "/") for p in self.EXCLUDED_PATHS)

    def __init__(self, spec: dict[str, Any]):
        resource = Resource.from_contents(spec, default_specification=DRAFT202012)
        self.resolver = Registry().with_resource("", resource).resolver()
        self.spec = spec

    def discover(self) -> DiscoveryResult:
        tools: list[ToolSpec] = []
        defs: dict[str, Any] = {}
        used_enum: set[str] = set()
        used_class: set[str] = set()

        for path, path_item in sorted(self.spec.get("paths", {}).items()):
            if "{" in path or not any(path.startswith(p) for p in self.ALLOWED_PATH_PREFIXES):
                continue
            if self._is_excluded(path):
                continue
            body_schema = self._get_request_body_schema(path_item) or {}
            query_props = self._get_query_parameters(path_item)
            body_props = body_schema.get("properties") or {}
            # Body properties win on name collision — body is the canonical param source
            # for the existing tools; query params are additive.
            properties = {**query_props, **body_props}
            clean_props = self._filter_properties(properties)

            enum_name = _deduplicate(_path_to_enum_name(path), used_enum)
            class_name = _deduplicate(_path_to_class_name(path), used_class)

            entry: dict[str, Any] = {
                "type": "object",
                "properties": clean_props,
                "description": body_schema.get("description"),
            }
            # Calculate which fields are actually required (many are marked as required,
            # but have a default set, so they're not really required)
            required = [
                name
                for name in body_schema.get("required") or []
                if name in clean_props and "default" not in (clean_props[name] or {})
            ]
            if required:
                entry["required"] = required
            defs[class_name] = entry
            tools.append(ToolSpec(path, enum_name, class_name))

        self._inline_component_refs(defs)

        combined_schema: dict[str, Any] = {
            "$defs": defs,
            "anyOf": [{"$ref": f"#/$defs/{t.class_name}"} for t in tools],
        }
        return DiscoveryResult(tools=tools, combined_schema=combined_schema)

    def _inline_component_refs(self, defs: dict[str, Any]) -> None:
        """Pull every component transitively referenced from tool param schemas into ``defs``
        and rewrite the refs from ``#/components/schemas/X`` to ``#/$defs/X``.

        Without this, nested refs (e.g. ``list[RedactionArea]``) are unresolvable when the
        combined schema is handed to datamodel-code-generator, producing ``RootModel[Any]``
        shells that downstream JSON-schema strict-mode transformers reject.
        """
        schemas = self.spec.get("components", {}).get("schemas", {})
        queue: list[object] = list(defs.values())
        while queue:
            for name in _rewrite_refs(queue.pop()):
                if name not in defs and name in schemas:
                    defs[name] = schemas[name]
                    queue.append(schemas[name])

    def _resolve_ref(self, schema: dict[str, Any]) -> dict[str, Any]:
        if "$ref" in schema:
            return self.resolver.lookup(schema["$ref"]).contents
        return schema

    def _get_request_body_schema(self, path_item: dict[str, Any]) -> dict[str, Any] | None:
        post = path_item.get("post")
        if not post:
            return None
        content = post.get("requestBody", {}).get("content", {})
        for media_type in ("multipart/form-data", "application/json"):
            if media_type in content:
                schema = content[media_type].get("schema")
                if schema:
                    return self._resolve_ref(schema)
        return None

    def _get_query_parameters(self, path_item: dict[str, Any]) -> dict[str, Any]:
        """Extract query parameters as a property map — AI tools expose their main
        inputs (e.g. ``prompt``, ``tolerance``) here rather than in the request body,
        and a handful of converters use query strings alongside multipart files.
        """
        post = path_item.get("post") or {}
        props: dict[str, Any] = {}
        for param in post.get("parameters") or []:
            if param.get("in") != "query":
                continue
            name = param.get("name")
            schema = param.get("schema")
            if not name or not schema:
                continue
            resolved = dict(self._resolve_ref(schema))
            if "description" not in resolved and param.get("description"):
                resolved["description"] = param["description"]
            props[name] = resolved
        return props

    def _filter_properties(self, properties: dict[str, Any]) -> dict[str, Any]:
        """Remove base-class fields and binary upload fields, resolving any $refs."""
        clean: dict[str, Any] = {}
        for name, prop in properties.items():
            if name in BASE_CLASS_FIELDS:
                continue
            prop = self._resolve_ref(prop)
            if prop.get("type") == "string" and prop.get("format") == "binary":
                continue
            clean[name] = prop
        return clean


_COMPONENT_REF_PREFIX = "#/components/schemas/"


def _rewrite_refs(obj: object) -> Iterable[str]:
    """Rewrite ``#/components/schemas/X`` refs to ``#/$defs/X`` in place, yielding each
    component name encountered so the caller can pull referenced schemas into ``$defs``.
    """
    if isinstance(obj, dict):
        ref = obj.get("$ref")
        if isinstance(ref, str) and ref.startswith(_COMPONENT_REF_PREFIX):
            name = ref.removeprefix(_COMPONENT_REF_PREFIX)
            obj["$ref"] = "#/$defs/" + name
            yield name
        for value in obj.values():
            yield from _rewrite_refs(value)
    elif isinstance(obj, list):
        for value in obj:
            yield from _rewrite_refs(value)


def _tool_name_segments(path: str) -> str:
    """Extract a descriptive name from the endpoint path.

    Converters use two segments (e.g. /api/v1/convert/cbr/pdf → cbr-to-pdf).
    Other tools use the last segment (e.g. /api/v1/misc/compress-pdf → compress-pdf).
    """
    parts = path.rstrip("/").split("/")
    if "/api/v1/convert/" in path and len(parts) >= 6:
        return f"{parts[-2]}-to-{parts[-1]}"
    return parts[-1]


def _path_to_enum_name(path: str) -> str:
    return _tool_name_segments(path).replace("-", "_").upper()


def _path_to_class_name(path: str) -> str:
    return "".join(p.capitalize() for p in _tool_name_segments(path).split("-")) + "Params"


def _deduplicate(name: str, used: set[str]) -> str:
    """Return name, appending 2, 3, ... if already in used. Adds result to used."""
    candidate = name
    n = 2
    while candidate in used:
        candidate = f"{name}{n}"
        n += 1
    used.add(candidate)
    return candidate


def generate_models_code(combined_schema: dict[str, Any]) -> str:
    """Run datamodel-code-generator once on the combined schema."""
    code = generate(
        input_=json.dumps(combined_schema, sort_keys=True),
        input_file_type=InputFileType.JsonSchema,
        output_model_type=DataModelType.PydanticV2BaseModel,
        target_python_version=PythonVersion.PY_313,
        snake_case_field=True,
        base_class="stirling.models.base.ApiModel",
        field_constraints=True,
        no_alias=True,
        set_default_enum_member=True,
        strict_nullable=True,
        use_schema_description=True,
        additional_imports=["enum.StrEnum"],
        enable_version_header=False,
        custom_file_header=_FILE_HEADER,
        formatters=[Formatter.RUFF_FORMAT, Formatter.RUFF_CHECK],
        settings_path=_ENGINE_ROOT / "pyproject.toml",
    )
    return str(code or "")


def render_models(tools: list[ToolSpec], models_code: str) -> str:
    union_lines = ["type ParamToolModel = ("]
    for i, tool in enumerate(tools):
        prefix = "    | " if i > 0 else "    "
        union_lines.append(f"{prefix}{tool.class_name}")
    union_lines.append(")")
    union_lines.append("type ParamToolModelType = type[ParamToolModel]")

    enum_lines = [
        "class ToolEndpoint(StrEnum):",
        *(f'    {t.enum_name} = "{t.path}"' for t in tools),
    ]

    ops_lines = [
        "OPERATIONS: dict[ToolEndpoint, ParamToolModelType] = {",
        *(f"    ToolEndpoint.{t.enum_name}: {t.class_name}," for t in tools),
        "}",
    ]

    parts = [models_code, "\n", *union_lines, "\n", *enum_lines, "\n", *ops_lines, ""]
    return "\n".join(parts)


def collect_tool_io(spec: dict[str, Any]) -> dict[str, dict[str, Any]]:
    """The ``x-stirling-io`` declaration for every endpoint that carries one."""
    table: dict[str, dict[str, Any]] = {}
    for path, path_item in sorted(spec.get("paths", {}).items()):
        for operation in path_item.values():
            if isinstance(operation, dict) and operation.get(IO_EXTENSION):
                table[path] = operation[IO_EXTENSION]
    return table


def _render_when(condition: dict[str, Any]) -> str:
    return f"ToolIOWhen(param={json.dumps(condition['param'])}, matches={json.dumps(condition['matches'])})"


def _render_case(case: dict[str, Any]) -> str:
    when = ", ".join(_render_when(condition) for condition in case["when"])
    return f"ToolIOCase(when=[{when}], produces=ToolFormat.{case['produces']}, arity=ToolArity.{case['arity']})"


def _render_spec(declaration: dict[str, Any]) -> str:
    """One declaration as a constructor call, so a bad spec is a type error at import rather
    than a validation failure at first use."""
    accepts = ", ".join(f"ToolFormat.{f}" for f in declaration["accepts"])
    parts = [
        f"accepts=[{accepts}]",
        f"produces=ToolFormat.{declaration['produces']}",
        f"arity=ToolArity.{declaration['arity']}",
    ]
    cases = declaration.get("cases")
    if cases:
        parts.append("cases=[" + ", ".join(_render_case(case) for case in cases) + "]")
    return f"ToolIOSpec({', '.join(parts)})"


def _members(values: list[str]) -> str:
    return "\n".join(f'    {value} = "{value}"' for value in values)


def render_tool_io(spec: dict[str, Any], tools: list[ToolSpec]) -> str:
    """Keyed by ``ToolEndpoint`` so the endpoint strings live in one place and lookups are checked.

    Declarations outside the enum - the filters, and the introspection endpoints the agent never
    plans - are dropped.
    """
    by_path = {tool.path: tool.enum_name for tool in tools}
    table = {path: d for path, d in collect_tool_io(spec).items() if path in by_path}
    if not table:
        raise SystemExit(
            f"No {IO_EXTENSION} declarations in the spec. The backend publishes these from @ToolIO; "
            "regenerate the spec with 'task backend:swagger'."
        )
    # Published separately: deriving the enums from the declarations present would shrink them
    # whenever an endpoint is disabled in a build.
    vocabulary = spec.get(IO_VOCABULARY_EXTENSION)
    if not vocabulary:
        raise SystemExit(f"No {IO_VOCABULARY_EXTENSION} in the spec; regenerate it from a current backend.")

    rendered = _IO_TEMPLATE.format(
        formats=_members(vocabulary["formats"]),
        arities=_members(vocabulary["arities"]),
        declarations="\n".join(
            f"    ToolEndpoint.{by_path[path]}: {_render_spec(declaration)},"
            for path, declaration in sorted(table.items())
        ),
    )
    # Formatted before writing so --check compares like for like.
    return subprocess.run(
        ["ruff", "format", "--stdin-filename", "tool_io.py", "-"],
        input=rendered,
        capture_output=True,
        text=True,
        check=True,
        cwd=_ENGINE_ROOT,
    ).stdout


def write_or_check(out_path: Path, rendered: str, check: bool) -> None:
    """In check mode, fail when the committed file is out of date."""
    if check:
        current = out_path.read_text(encoding="utf-8") if out_path.exists() else ""
        if current != rendered:
            raise SystemExit(f"{out_path} is out of date. Run 'task engine:tool-models' and commit the result.")
        return
    out_path.write_text(rendered, encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate the Python files derived from the Java OpenAPI spec")
    parser.add_argument("--spec", required=True, help="Path to SwaggerDoc.json")
    parser.add_argument("--output", required=True, help="Path to output tool_models.py")
    parser.add_argument("--io-output", required=True, help="Path to output tool_io.py")
    parser.add_argument("--check", action="store_true", help="Fail if a committed file is out of date")
    args = parser.parse_args()

    spec_path = Path(args.spec)
    if not spec_path.exists():
        raise SystemExit(f"OpenAPI spec not found at {spec_path}\nRun 'task backend:swagger' to generate it.")

    with open(spec_path, encoding="utf-8") as f:
        spec = json.load(f)

    result = ToolDiscovery(spec).discover()

    io_table = render_tool_io(spec, result.tools)
    write_or_check(Path(args.io_output), io_table, args.check)
    print(f"{'Up to date' if args.check else 'Generated'}: {len(result.tools)} tool I/O declarations")

    models_code = generate_models_code(result.combined_schema)
    write_or_check(Path(args.output), render_models(result.tools, models_code), args.check)
    print(f"{'Up to date' if args.check else 'Generated'}: {len(result.tools)} tool models from {spec_path.name}")


if __name__ == "__main__":
    main()