Add edit text support to stirling engine (#6245)

# Description of Changes
Hooks up the (alpha) PDF Editor backend to the AI engine Edit Agent via
an intermediary API which is easier for the agent to call. It suffers
from all the same issues that the PDF Editor does in actually editing
the text, but should also benefit from any fixes to that.

It also adds protection against the underlying tools misbehaving by
hanging, and fixes a hanging bug in the PDF Editor.

---------

Co-authored-by: EthanHealy01 <[email protected]>
This commit is contained in:
James Brunton
2026-05-11 09:57:41 +00:00
committed by GitHub
co-authored by EthanHealy01
parent 294b616a63
commit 575684ee4b
17 changed files with 1937 additions and 52 deletions
+24
View File
@@ -373,6 +373,25 @@ class EditTableOfContentsParams(ApiModel):
)
class EditTextOperation(ApiModel):
find: str = Field(..., description="The literal text to find.")
replace: str = Field(..., description="The replacement text. May be empty to delete the matched text.")
class EditTextParams(ApiModel):
edits: list[EditTextOperation] | None = Field(
None,
description="Ordered list of find/replace operations. Each replaces every occurrence on the selected pages, in order; later operations see the result of earlier ones (so 'foo'->'foos' then 'foos'->'bars' turns 'foo' into 'bars').",
)
page_numbers: str | None = Field(
"all",
description="The pages to select, Supports ranges (e.g., '1,3,5-9'), or 'all' or functions in the format 'an+b' where 'a' is the multiplier of the page number 'n', and 'b' is a constant (e.g., '2n+1', '3n', '6n-5')",
)
whole_word_search: bool | None = Field(
False, description="Whether matches must be whole words (boundaries determined by non-word characters)"
)
class EmlToPdfParams(ApiModel):
download_html: bool | None = Field(
None, description="Download HTML intermediate file instead of PDF", examples=[False]
@@ -1105,6 +1124,7 @@ class Model(
| BookletImpositionParams
| CropParams
| EditTableOfContentsParams
| EditTextParams
| MergePdfsParams
| MultiPageLayoutParams
| OverlayPdfsParams
@@ -1172,6 +1192,7 @@ class Model(
| BookletImpositionParams
| CropParams
| EditTableOfContentsParams
| EditTextParams
| MergePdfsParams
| MultiPageLayoutParams
| OverlayPdfsParams
@@ -1240,6 +1261,7 @@ type ParamToolModel = (
| BookletImpositionParams
| CropParams
| EditTableOfContentsParams
| EditTextParams
| MergePdfsParams
| MultiPageLayoutParams
| OverlayPdfsParams
@@ -1309,6 +1331,7 @@ class ToolEndpoint(StrEnum):
BOOKLET_IMPOSITION = "/api/v1/general/booklet-imposition"
CROP = "/api/v1/general/crop"
EDIT_TABLE_OF_CONTENTS = "/api/v1/general/edit-table-of-contents"
EDIT_TEXT = "/api/v1/general/edit-text"
MERGE_PDFS = "/api/v1/general/merge-pdfs"
MULTI_PAGE_LAYOUT = "/api/v1/general/multi-page-layout"
OVERLAY_PDFS = "/api/v1/general/overlay-pdfs"
@@ -1376,6 +1399,7 @@ OPERATIONS: dict[ToolEndpoint, ParamToolModelType] = {
ToolEndpoint.BOOKLET_IMPOSITION: BookletImpositionParams,
ToolEndpoint.CROP: CropParams,
ToolEndpoint.EDIT_TABLE_OF_CONTENTS: EditTableOfContentsParams,
ToolEndpoint.EDIT_TEXT: EditTextParams,
ToolEndpoint.MERGE_PDFS: MergePdfsParams,
ToolEndpoint.MULTI_PAGE_LAYOUT: MultiPageLayoutParams,
ToolEndpoint.OVERLAY_PDFS: OverlayPdfsParams,
+180 -4
View File
@@ -21,8 +21,15 @@ from stirling.contracts import (
SupportedCapability,
ToolOperationStep,
)
from stirling.models import OPERATIONS, FileId
from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
from stirling.models import OPERATIONS, FileId, ParamToolModel
from stirling.models.tool_models import (
Angle,
EditTextOperation,
EditTextParams,
FlattenParams,
RotatePdfParams,
ToolEndpoint,
)
from stirling.services.runtime import AppRuntime
@@ -35,8 +42,11 @@ class ParameterSelectorCall:
class RecordingParameterSelector:
def __init__(self) -> None:
"""Test double that records calls and returns predetermined parameter objects per index."""
def __init__(self, params_by_index: list[ParamToolModel] | None = None) -> None:
self.calls: list[ParameterSelectorCall] = []
self._params_by_index = params_by_index
async def select(
self,
@@ -44,7 +54,7 @@ class RecordingParameterSelector:
operation_plan: list[ToolEndpoint],
operation_index: int,
generated_steps: list[ToolOperationStep],
) -> RotatePdfParams | FlattenParams:
) -> ParamToolModel:
self.calls.append(
ParameterSelectorCall(
request=request,
@@ -53,6 +63,8 @@ class RecordingParameterSelector:
generated_steps=list(generated_steps),
)
)
if self._params_by_index is not None:
return self._params_by_index[operation_index]
if operation_index == 0:
return RotatePdfParams(angle=Angle(90))
return FlattenParams(flatten_only_forms=False, render_dpi=None)
@@ -393,3 +405,167 @@ async def test_pdf_edit_agent_rejects_plan_referencing_unavailable_operations(
assert "not available" in response.reason
assert "COMPRESS_PDF" in response.reason
assert parameter_selector.calls == []
@pytest.mark.anyio
async def test_pdf_edit_agent_supports_literal_find_replace(runtime: AppRuntime) -> None:
params = EditTextParams(
edits=[EditTextOperation(find="2025", replace="2026")],
page_numbers="all",
whole_word_search=False,
)
parameter_selector = RecordingParameterSelector([params])
agent = StubPdfEditAgent(
runtime,
PdfEditPlanSelection(
operations=[ToolEndpoint.EDIT_TEXT],
summary="Replace 2025 with 2026 throughout the document.",
),
parameter_selector=parameter_selector,
)
response = await agent.handle(
PdfEditRequest(
user_message="Change every 2025 to 2026.",
files=[AiFile(id=FileId("contract-id"), name="contract.pdf")],
)
)
assert isinstance(response, EditPlanResponse)
assert len(response.steps) == 1
step = response.steps[0]
assert step.tool == ToolEndpoint.EDIT_TEXT
assert isinstance(step.parameters, EditTextParams)
assert step.parameters.edits == [EditTextOperation(find="2025", replace="2026")]
@pytest.mark.anyio
async def test_pdf_edit_agent_supports_copy_edit_using_page_text(runtime: AppRuntime) -> None:
page_text = [
ExtractedFileText(
file_name="memo.pdf",
pages=[
PdfTextSelection(
page_number=3,
text="The quick brown fox jumps over the lazy dog.",
)
],
)
]
params = EditTextParams(
edits=[
EditTextOperation(find="quick", replace="slow"),
EditTextOperation(find="lazy", replace="energetic"),
],
page_numbers="3",
whole_word_search=False,
)
parameter_selector = RecordingParameterSelector([params])
agent = StubPdfEditAgent(
runtime,
PdfEditPlanSelection(
operations=[ToolEndpoint.EDIT_TEXT],
summary="Fix typos on page 3.",
),
parameter_selector=parameter_selector,
)
response = await agent.handle(
PdfEditRequest(
user_message="Fix typos on page 3.",
files=[AiFile(id=FileId("memo-id"), name="memo.pdf")],
page_text=page_text,
)
)
assert isinstance(response, EditPlanResponse)
assert len(parameter_selector.calls) == 1
# The parameter selector receives the extracted page text, which is what enables free-form
# copy-editing: it can read the current text and propose specific edits.
assert parameter_selector.calls[0].request.page_text == page_text
step = response.steps[0]
assert step.tool == ToolEndpoint.EDIT_TEXT
assert isinstance(step.parameters, EditTextParams)
assert step.parameters.page_numbers == "3"
assert step.parameters.edits is not None
assert len(step.parameters.edits) == 2
@pytest.mark.anyio
async def test_pdf_edit_agent_supports_natural_language_directed_edit(runtime: AppRuntime) -> None:
page_text = [
ExtractedFileText(
file_name="agreement.pdf",
pages=[
PdfTextSelection(
page_number=1,
text="This agreement is between OldCompany Inc. and the client.",
)
],
)
]
params = EditTextParams(
edits=[EditTextOperation(find="OldCompany Inc.", replace="Acme Corp")],
page_numbers="all",
whole_word_search=False,
)
parameter_selector = RecordingParameterSelector([params])
agent = StubPdfEditAgent(
runtime,
PdfEditPlanSelection(
operations=[ToolEndpoint.EDIT_TEXT],
summary="Update the company name to Acme Corp.",
),
parameter_selector=parameter_selector,
)
response = await agent.handle(
PdfEditRequest(
user_message="Update the company name to Acme Corp.",
files=[AiFile(id=FileId("agreement-id"), name="agreement.pdf")],
page_text=page_text,
)
)
assert isinstance(response, EditPlanResponse)
step = response.steps[0]
assert step.tool == ToolEndpoint.EDIT_TEXT
assert isinstance(step.parameters, EditTextParams)
# The exact find string came from interpreting the user's intent against the extracted text.
assert step.parameters.edits is not None
assert step.parameters.edits[0].find == "OldCompany Inc."
assert step.parameters.edits[0].replace == "Acme Corp"
@pytest.mark.anyio
async def test_pdf_edit_agent_composes_edit_text_with_other_operations(runtime: AppRuntime) -> None:
"""EDIT_TEXT can appear alongside other operations in a single plan."""
edit_params = EditTextParams(
edits=[EditTextOperation(find="DRAFT", replace="")],
page_numbers="all",
whole_word_search=False,
)
parameter_selector = RecordingParameterSelector([edit_params, RotatePdfParams(angle=Angle(90))])
agent = StubPdfEditAgent(
runtime,
PdfEditPlanSelection(
operations=[ToolEndpoint.EDIT_TEXT, ToolEndpoint.ROTATE_PDF],
summary="Remove DRAFT marker, then rotate.",
),
parameter_selector=parameter_selector,
)
response = await agent.handle(
PdfEditRequest(
user_message="Remove the DRAFT watermark text and then rotate.",
files=[AiFile(id=FileId("draft-id"), name="draft.pdf")],
)
)
assert isinstance(response, EditPlanResponse)
assert [step.tool for step in response.steps] == [
ToolEndpoint.EDIT_TEXT,
ToolEndpoint.ROTATE_PDF,
]
assert isinstance(response.steps[0].parameters, EditTextParams)
assert isinstance(response.steps[1].parameters, RotatePdfParams)