arsvendg/Stirling-PDF
1
0
Fork 0
mirror of https://github.com/arsvendg/Stirling-PDF.git synced 2026-07-14 10:34:06 +02:00
Code Issues Actions 4 Packages Projects Releases 1 Wiki Activity

smart redaction (#6195)

Browse Source 
Co-authored-by: James Brunton <[email protected]>
This commit is contained in:
EthanHealy01
2026-06-03 16:16:33 +00:00
committed by GitHub
co-authored by James Brunton
parent 7f3ca7ea70
commit 35a712a278
26 changed files with 4397 additions and 1802 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/tests/test_pdf_edit_agent.py
+11 -16
View File
@@ -6,7 +6,7 @@ from dataclasses import dataclass
import pytest
from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
from stirling.agents.pdf_edit import PdfEditPlanOutput
from stirling.agents.pdf_edit import PdfEditNeedContentSelection, PdfEditPlanOutput
from stirling.contracts import (
AiFile,
EditCannotDoResponse,
@@ -82,12 +82,12 @@ class StubPdfEditAgent(PdfEditAgent):
if parameter_selector is not None:
self.parameter_selector = parameter_selector
def _get_supported_operations(self, request: PdfEditRequest) -> Iterable[ToolEndpoint]:
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 request.enabled_endpoints:
return request.enabled_endpoints
return OPERATIONS
if not request.enabled_endpoints:
return list(OPERATIONS), []
return super()._classify_operations(request)
async def _select_plan(
self,
@@ -194,12 +194,8 @@ async def test_pdf_edit_agent_returns_need_content_without_building_plan(runtime
parameter_selector = RecordingParameterSelector()
agent = StubPdfEditAgent(
runtime,
NeedContentResponse(
resume_with=SupportedCapability.PDF_EDIT,
PdfEditNeedContentSelection(
reason="Need page text to locate the NEW PAGE markers.",
files=[],
max_pages=0,
max_characters=0,
),
parameter_selector=parameter_selector,
)
@@ -274,8 +270,8 @@ async def test_pdf_edit_selection_agent_excludes_need_content_from_schema_when_n
can_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=True)
cannot_request = PdfEditSelectionAgent(runtime, "base", allow_need_content=False)
assert NeedContentResponse in _agent_output_types(can_request)
assert NeedContentResponse not in _agent_output_types(cannot_request)
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]:
@@ -336,7 +332,7 @@ async def test_pdf_edit_agent_supported_operations_defaults_to_empty(
runtime: AppRuntime,
) -> None:
agent = PdfEditAgent(runtime)
supported = agent._get_supported_operations(PdfEditRequest(user_message="hi"))
supported, _ = agent._classify_operations(PdfEditRequest(user_message="hi"))
assert list(supported) == []
@@ -351,7 +347,7 @@ async def test_pdf_edit_agent_supported_operations_uses_provided_list(
enabled_endpoints=[ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF],
)
supported = agent._get_supported_operations(request)
supported, _ = agent._classify_operations(request)
assert list(supported) == [ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF]
@@ -372,8 +368,7 @@ def test_pdf_edit_selection_prompt_includes_unavailable_operations(runtime: AppR
user_message="Run OCR.",
enabled_endpoints=[ToolEndpoint.FLATTEN],
)
supported = agent._get_supported_operations(request)
unavailable = agent._get_unavailable_operations(supported)
supported, unavailable = agent._classify_operations(request)
prompt = agent._build_selection_prompt(request, supported, unavailable)