#!/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] # The value the endpoint uses when this parameter is absent; None when it has none. default: str | None = None 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", # 5. Server maintenance, not a document operation: releases finished jobs and # their stored files. Nothing an edit agent should ever call on its own. "/api/v1/general/jobs/cleanup", ) 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: parts = [f"param={json.dumps(condition['param'])}", f"matches={json.dumps(condition['matches'])}"] if "default" in condition: parts.append(f"default={json.dumps(condition['default'])}") return f"ToolIOWhen({', '.join(parts)})" 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()