Stirling-PDF/engine/tests/test_pdf_edit_agent.py

from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass

import pytest

from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
from stirling.agents.pdf_edit import PdfEditNeedContentSelection, PdfEditPlanOutput
from stirling.contracts import (
    AiFile,
    EditCannotDoResponse,
    EditClarificationRequest,
    EditPlanResponse,
    ExtractedFileText,
    NeedContentFileRequest,
    NeedContentResponse,
    PdfContentType,
    PdfEditRequest,
    PdfTextSelection,
    SupportedCapability,
    ToolOperationStep,
)
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


@dataclass(frozen=True)
class ParameterSelectorCall:
    request: PdfEditRequest
    operation_plan: list[ToolEndpoint]
    operation_index: int
    generated_steps: list[ToolOperationStep]


class RecordingParameterSelector:
    """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,
        request: PdfEditRequest,
        operation_plan: list[ToolEndpoint],
        operation_index: int,
        generated_steps: list[ToolOperationStep],
    ) -> ParamToolModel:
        self.calls.append(
            ParameterSelectorCall(
                request=request,
                operation_plan=operation_plan,
                operation_index=operation_index,
                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)


class StubPdfEditAgent(PdfEditAgent):
    def __init__(
        self,
        runtime: AppRuntime,
        selection: PdfEditPlanOutput,
        parameter_selector: RecordingParameterSelector | PdfEditParameterSelector | None = None,
    ) -> None:
        super().__init__(runtime)
        self.selection = selection
        if parameter_selector is not None:
            self.parameter_selector = parameter_selector

    def _classify_operations(self, request: PdfEditRequest) -> tuple[list[ToolEndpoint], list[ToolEndpoint]]:
        # Tests construct requests without `enabled_endpoints`; pretend everything is enabled
        # unless the test explicitly supplies an enabled set.
        if not request.enabled_endpoints:
            return list(OPERATIONS), []
        return super()._classify_operations(request)

    async def _select_plan(
        self,
        request: PdfEditRequest,
        supported_operations: Iterable[ToolEndpoint],
        unavailable_operations: Iterable[ToolEndpoint],
        *,
        allow_need_content: bool = True,
    ) -> PdfEditPlanOutput:
        return self.selection


@pytest.mark.anyio
async def test_pdf_edit_agent_builds_multi_step_plan(runtime: AppRuntime) -> None:
    parameter_selector = RecordingParameterSelector()
    agent = StubPdfEditAgent(
        runtime,
        PdfEditPlanSelection(
            operations=[ToolEndpoint.ROTATE_PDF, ToolEndpoint.FLATTEN],
            summary="Rotate the PDF, then compress it.",
            rationale="The pages need reorientation before reducing file size.",
        ),
        parameter_selector=parameter_selector,
    )

    response = await agent.handle(
        PdfEditRequest(
            user_message="Rotate the PDF clockwise and then compress it.",
            files=[AiFile(id=FileId("scan-id"), name="scan.pdf")],
        )
    )

    assert isinstance(response, EditPlanResponse)
    assert response.summary == "Rotate the PDF, then compress it."
    assert response.rationale == "The pages need reorientation before reducing file size."
    assert [step.tool for step in response.steps] == [ToolEndpoint.ROTATE_PDF, ToolEndpoint.FLATTEN]
    assert isinstance(response.steps[0].parameters, RotatePdfParams)
    assert isinstance(response.steps[1].parameters, FlattenParams)


@pytest.mark.anyio
async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector(runtime: AppRuntime) -> None:
    parameter_selector = RecordingParameterSelector()
    agent = StubPdfEditAgent(
        runtime,
        PdfEditPlanSelection(
            operations=[ToolEndpoint.ROTATE_PDF, ToolEndpoint.FLATTEN],
            summary="Rotate the PDF, then compress it.",
            rationale="test rationale",
        ),
        parameter_selector=parameter_selector,
    )

    request = PdfEditRequest(
        user_message="Rotate the PDF clockwise and then compress it.",
        files=[AiFile(id=FileId("scan-id"), name="scan.pdf")],
    )
    response = await agent.handle(request)

    assert isinstance(response, EditPlanResponse)
    assert len(parameter_selector.calls) == 2
    assert parameter_selector.calls[0].operation_index == 0
    assert parameter_selector.calls[0].generated_steps == []
    assert parameter_selector.calls[1].operation_index == 1
    assert parameter_selector.calls[1].generated_steps == [
        ToolOperationStep(
            tool=ToolEndpoint.ROTATE_PDF,
            parameters=RotatePdfParams(angle=Angle(90)),
        )
    ]


@pytest.mark.anyio
async def test_pdf_edit_agent_returns_clarification_without_partial_plan(runtime: AppRuntime) -> None:
    agent = StubPdfEditAgent(
        runtime,
        EditClarificationRequest(
            question="Which pages should be rotated?",
            reason="The request does not say which pages to change.",
        ),
    )

    response = await agent.handle(PdfEditRequest(user_message="Rotate some pages."))

    assert isinstance(response, EditClarificationRequest)


@pytest.mark.anyio
async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan(runtime: AppRuntime) -> None:
    agent = StubPdfEditAgent(
        runtime,
        EditCannotDoResponse(
            reason="This request requires OCR, which is not part of PDF edit planning.",
        ),
    )

    response = await agent.handle(PdfEditRequest(user_message="Read this scan and summarize it."))

    assert isinstance(response, EditCannotDoResponse)


@pytest.mark.anyio
async def test_pdf_edit_agent_returns_need_content_without_building_plan(runtime: AppRuntime) -> None:
    parameter_selector = RecordingParameterSelector()
    agent = StubPdfEditAgent(
        runtime,
        PdfEditNeedContentSelection(
            reason="Need page text to locate the NEW PAGE markers.",
        ),
        parameter_selector=parameter_selector,
    )

    response = await agent.handle(
        PdfEditRequest(
            user_message="Split after every page that says 'NEW PAGE'.",
            files=[AiFile(id=FileId("report-id"), name="report.pdf")],
        )
    )

    assert isinstance(response, NeedContentResponse)
    assert response.resume_with == SupportedCapability.PDF_EDIT
    assert response.files == [
        NeedContentFileRequest(
            file=AiFile(id=FileId("report-id"), name="report.pdf"),
            content_types=[PdfContentType.PAGE_TEXT],
        )
    ]
    assert response.max_pages == runtime.settings.max_pages
    assert response.max_characters == runtime.settings.max_characters
    assert parameter_selector.calls == []


@pytest.mark.anyio
async def test_pdf_edit_agent_builds_selection_agent_matching_content_availability(runtime: AppRuntime) -> None:
    from stirling.agents.pdf_edit import PdfEditSelectionAgent

    agent = PdfEditAgent(runtime)
    captured: list[bool] = []

    def record(
        supported_operations: Iterable[ToolEndpoint],
        unavailable_operations: Iterable[ToolEndpoint],
        *,
        allow_need_content: bool,
    ) -> PdfEditSelectionAgent:
        captured.append(allow_need_content)
        raise _StopSelectionError()

    agent._build_selection_agent = record

    supported = list(OPERATIONS)
    with pytest.raises(_StopSelectionError):
        await agent._select_plan(PdfEditRequest(user_message="Rotate."), supported, [])
    with pytest.raises(_StopSelectionError):
        await agent._select_plan(
            PdfEditRequest(
                user_message="Rotate.",
                page_text=[
                    ExtractedFileText(
                        file_name="report.pdf",
                        pages=[PdfTextSelection(page_number=1, text="content")],
                    )
                ],
            ),
            supported,
            [],
        )
    with pytest.raises(_StopSelectionError):
        await agent._select_plan(PdfEditRequest(user_message="Rotate."), supported, [], allow_need_content=False)

    assert captured == [True, False, False]


@pytest.mark.anyio
async def test_pdf_edit_selection_agent_excludes_need_content_from_schema_when_not_allowed(
    runtime: AppRuntime,
) -> None:
    from stirling.agents.pdf_edit import PdfEditSelectionAgent

    can_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=True)
    cannot_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=False)

    assert PdfEditNeedContentSelection in _agent_output_types(can_request)
    assert PdfEditNeedContentSelection not in _agent_output_types(cannot_request)


def _agent_output_types(agent: object) -> list[type]:
    native = getattr(getattr(agent, "agent"), "output_type")
    return list(getattr(native, "outputs", []))


class _StopSelectionError(Exception):
    pass


@pytest.mark.anyio
async def test_pdf_edit_agent_passes_page_text_to_parameter_selector(runtime: AppRuntime) -> None:
    parameter_selector = RecordingParameterSelector()
    agent = StubPdfEditAgent(
        runtime,
        PdfEditPlanSelection(
            operations=[ToolEndpoint.ROTATE_PDF],
            summary="Rotate the PDF.",
            rationale="test rationale",
        ),
        parameter_selector=parameter_selector,
    )

    page_text = [
        ExtractedFileText(
            file_name="report.pdf",
            pages=[PdfTextSelection(page_number=1, text="NEW PAGE")],
        )
    ]
    await agent.handle(
        PdfEditRequest(
            user_message="Rotate clockwise.",
            files=[AiFile(id=FileId("report-id"), name="report.pdf")],
            page_text=page_text,
        )
    )

    assert parameter_selector.calls[0].request.page_text == page_text


def test_pdf_edit_request_drops_unknown_enabled_urls() -> None:
    request = PdfEditRequest.model_validate(
        {
            "user_message": "ignore me",
            "enabled_endpoints": [
                ToolEndpoint.COMPRESS_PDF.value,
                "/api/v1/not-a-real/endpoint",
            ],
        }
    )

    assert request.enabled_endpoints == [ToolEndpoint.COMPRESS_PDF]


@pytest.mark.anyio
async def test_pdf_edit_agent_supported_operations_defaults_to_empty(
    runtime: AppRuntime,
) -> None:
    agent = PdfEditAgent(runtime)
    supported, _ = agent._classify_operations(PdfEditRequest(user_message="hi"))

    assert list(supported) == []


@pytest.mark.anyio
async def test_pdf_edit_agent_supported_operations_uses_provided_list(
    runtime: AppRuntime,
) -> None:
    agent = PdfEditAgent(runtime)
    request = PdfEditRequest(
        user_message="Compress this PDF.",
        enabled_endpoints=[ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF],
    )

    supported, _ = agent._classify_operations(request)

    assert list(supported) == [ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF]


@pytest.mark.anyio
async def test_pdf_edit_agent_returns_cannot_do_when_no_operations_enabled(
    runtime: AppRuntime,
) -> None:
    agent = PdfEditAgent(runtime)
    response = await agent.handle(PdfEditRequest(user_message="Do anything.", enabled_endpoints=[]))

    assert isinstance(response, EditCannotDoResponse)


def test_pdf_edit_selection_prompt_includes_unavailable_operations(runtime: AppRuntime) -> None:
    agent = PdfEditAgent(runtime)
    request = PdfEditRequest(
        user_message="Run OCR.",
        enabled_endpoints=[ToolEndpoint.FLATTEN],
    )
    supported, unavailable = agent._classify_operations(request)

    prompt = agent._build_selection_prompt(request, supported, unavailable)

    assert "Unavailable operations" in prompt
    assert "OCR_PDF" in prompt
    assert ToolEndpoint.OCR_PDF.value in prompt


@pytest.mark.anyio
async def test_pdf_edit_agent_rejects_plan_referencing_unavailable_operations(
    runtime: AppRuntime,
) -> None:
    parameter_selector = RecordingParameterSelector()
    agent = StubPdfEditAgent(
        runtime,
        PdfEditPlanSelection(
            operations=[ToolEndpoint.COMPRESS_PDF],
            summary="Compress.",
            rationale="test rationale",
        ),
        parameter_selector=parameter_selector,
    )

    response = await agent.handle(
        PdfEditRequest(
            user_message="Compress this PDF.",
            enabled_endpoints=[ToolEndpoint.FLATTEN],
        )
    )

    assert isinstance(response, EditCannotDoResponse)
    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.",
            rationale="test rationale",
        ),
        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.",
            rationale="test rationale",
        ),
        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.",
            rationale="test rationale",
        ),
        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.",
            rationale="test rationale",
        ),
        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)