mirror of
https://github.com/Stirling-Tools/Stirling-PDF.git
synced 2026-09-03 05:10:16 +03:00
James Brunton
20204f0ddc
# Description of Changes A few changes to improve things in the engine: - Changed the PDF to Markdown code to be a real tool in Java, to remove the need for the `pdf_ingest` code, which looked a bit like an agent but wasn't behaving as an agent. It's now just covered automatically by the edit agent. - Noticed that 0-parameter-endpoints were previously being ignored by the `tool_models` generator, so some tools which require no params were being mistakenly excluded. - Removed tools which currently never succeed like Add Stamp, Cert Sign, and Overlay, because they require the supporting files to be sent in a different location in the API call, which we don't currently do. Ideally, we'd add proper support for this, but we're better off now removing support for these tools rather than just have them crash. We can re-add these tools in a future PR properly.
332 lines
13 KiB
Python
332 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate Python tool models from the Java backend's OpenAPI spec (SwaggerDoc.json).
|
|
|
|
Uses datamodel-code-generator to convert OpenAPI request schemas to Pydantic models.
|
|
Run via:
|
|
task engine:tool-models
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
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"})
|
|
|
|
_ENGINE_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
_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 write_output(out_path: Path, tools: list[ToolSpec], models_code: str) -> None:
|
|
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, ""]
|
|
out_path.write_text("\n".join(parts), encoding="utf-8")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Generate Python tool models from 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")
|
|
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 engine:tool-models' to generate it.")
|
|
output_path = Path(args.output)
|
|
|
|
with open(spec_path, encoding="utf-8") as f:
|
|
spec = json.load(f)
|
|
|
|
result = ToolDiscovery(spec).discover()
|
|
models_code = generate_models_code(result.combined_schema)
|
|
write_output(output_path, result.tools, models_code)
|
|
|
|
print(f"Generated {len(result.tools)} tool models from {spec_path.name}")
|
|
for tool in result.tools:
|
|
print(f" {tool.enum_name}: {tool.path} -> {tool.class_name}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|