mirror of
https://github.com/Stirling-Tools/Stirling-PDF.git
synced 2026-09-03 05:10:16 +03:00
Define tool inputs & outputs in a structured way (#7204)
# Description of Changes Change tool APIs to use structured definitions for input/output/type info because we need that info to be able to validate whether policies can actually successfully work based on whether one tool accepts the output of another. There were various bugs in the previous string definitions because of either misspellings or just incorrect definitions, so I've gone through and fixed all that I can find. <img width="729" height="271" alt="image" src="https://github.com/user-attachments/assets/08357e96-6fbb-4b9c-ba4d-8995420c7b86" /> <img width="749" height="264" alt="image" src="https://github.com/user-attachments/assets/76f46284-1866-4b64-b1ed-2480e01866e9" /> <img width="402" height="636" alt="image" src="https://github.com/user-attachments/assets/8f7a36ca-2845-4f14-a2df-ec9c772e66f6" /> <img width="393" height="317" alt="image" src="https://github.com/user-attachments/assets/46d8b891-9820-4ce3-8109-a8b782277037" /> --------- Co-authored-by: Anthony Stirling <77850077+Frooodle@users.noreply.github.com>
This commit is contained in:
co-authored by
Anthony Stirling
parent
7a5eb73c5b
commit
cd199c8659
@@ -1,8 +1,9 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate Python tool models from the Java backend's OpenAPI spec (SwaggerDoc.json).
|
||||
"""Generate the Python files derived from the Java OpenAPI spec (SwaggerDoc.json).
|
||||
|
||||
Uses datamodel-code-generator to convert OpenAPI request schemas to Pydantic models.
|
||||
Run via:
|
||||
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
|
||||
"""
|
||||
|
||||
@@ -10,6 +11,7 @@ from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import subprocess
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
@@ -24,8 +26,69 @@ 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"
|
||||
@@ -281,7 +344,7 @@ def generate_models_code(combined_schema: dict[str, Any]) -> str:
|
||||
return str(code or "")
|
||||
|
||||
|
||||
def write_output(out_path: Path, tools: list[ToolSpec], models_code: str) -> None:
|
||||
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 " "
|
||||
@@ -301,30 +364,119 @@ def write_output(out_path: Path, tools: list[ToolSpec], models_code: str) -> Non
|
||||
]
|
||||
|
||||
parts = [models_code, "\n", *union_lines, "\n", *enum_lines, "\n", *ops_lines, ""]
|
||||
out_path.write_text("\n".join(parts), encoding="utf-8")
|
||||
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 Python tool models from Java OpenAPI spec")
|
||||
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 engine:tool-models' to generate it.")
|
||||
output_path = Path(args.output)
|
||||
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()
|
||||
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}")
|
||||
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__":
|
||||
|
||||
Reference in New Issue
Block a user