mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
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:
co-authored by
EthanHealy01
parent
294b616a63
commit
575684ee4b
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user