Add document context for edit agent (#6152)

# Description of Changes
Adds the ability for the Edit agent to request the content of the
document before it decides which parameters it needs. This makes it able
to process requests like `Split the document after the page containing
the "My Section" section`, allowing for document context-based requests
for all[^1] tools.

I had to make a few changes elsewhere to make this work, including:
- Moving the requesting of content out of the Question Agent and into a
common location
- Added specific API docs for the Split param because the generic ones
were not specific enough for the AI to be able to reliably perform the
correct operation
- Fixed an issue in the tool models generator which caused the Redact
params to only be half-generated (causing Pydantic to crash when the AI
tried to run Redact)
- Added missing logging to a bunch of tools and hooked it up properly so
it'll print to stderr
- Made the limits for the max pages/chars to extract from PDFs
configurable via env var

[^1]: Many of the tools can't actually do anything useful with the
context at this stage, but will just need the tool API to be extended
with new features like page-specific operations to be automatically able
to do smart operations without needing to change the Edit agent itself.
This commit is contained in:
James Brunton
2026-04-23 13:19:27 +00:00
committed by GitHub
parent e087b54cf0
commit 3e94157137
23 changed files with 462 additions and 108 deletions
+39
View File
@@ -10,6 +10,7 @@ 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
@@ -85,12 +86,30 @@ class ToolDiscovery:
defs[class_name] = {"type": "object", "properties": clean_props}
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
@@ -121,6 +140,26 @@ class ToolDiscovery:
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.