diff --git a/app/core/src/main/java/stirling/software/SPDF/controller/api/SplitPDFController.java b/app/core/src/main/java/stirling/software/SPDF/controller/api/SplitPDFController.java index c0b7e23cb..ed7957a81 100644 --- a/app/core/src/main/java/stirling/software/SPDF/controller/api/SplitPDFController.java +++ b/app/core/src/main/java/stirling/software/SPDF/controller/api/SplitPDFController.java @@ -21,7 +21,7 @@ import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import stirling.software.SPDF.config.swagger.MultiFileResponse; -import stirling.software.SPDF.model.api.PDFWithPageNums; +import stirling.software.SPDF.model.api.SplitPagesRequest; import stirling.software.common.annotations.AutoJobPostMapping; import stirling.software.common.annotations.api.GeneralApi; import stirling.software.common.service.CustomPDFDocumentFactory; @@ -48,7 +48,7 @@ public class SplitPDFController { + " specified page numbers or ranges. Users can specify pages using" + " individual numbers, ranges, or 'all' for every page. Input:PDF" + " Output:PDF Type:SIMO") - public ResponseEntity splitPdf(@ModelAttribute PDFWithPageNums request) + public ResponseEntity splitPdf(@ModelAttribute SplitPagesRequest request) throws IOException { MultipartFile file = request.getFileInput(); diff --git a/app/core/src/main/java/stirling/software/SPDF/model/api/SplitPagesRequest.java b/app/core/src/main/java/stirling/software/SPDF/model/api/SplitPagesRequest.java new file mode 100644 index 000000000..0ae1f7553 --- /dev/null +++ b/app/core/src/main/java/stirling/software/SPDF/model/api/SplitPagesRequest.java @@ -0,0 +1,38 @@ +package stirling.software.SPDF.model.api; + +import java.util.List; + +import org.apache.pdfbox.pdmodel.PDDocument; + +import io.swagger.v3.oas.annotations.Hidden; +import io.swagger.v3.oas.annotations.media.Schema; +import io.swagger.v3.oas.annotations.media.Schema.RequiredMode; + +import lombok.Data; +import lombok.EqualsAndHashCode; + +import stirling.software.common.model.api.PDFFile; +import stirling.software.common.util.GeneralUtils; + +@Data +@EqualsAndHashCode(callSuper = true) +public class SplitPagesRequest extends PDFFile { + + @Schema( + description = + "Split points - page numbers after which the PDF will be cut. For example," + + " `\"2\"` produces two documents (pages 1-2 and pages 3+); `\"2,5\"`" + + " produces three (pages 1-2, 3-5, 6+). Supports ranges (e.g." + + " `\"1,3,5-9\"` splits after pages 1, 3, 5, 6, 7, 8, 9, yielding 8" + + " documents), `\"all\"` (split after every page), or functions like" + + " `\"2n+1\"`, `\"3n\"`, `\"6n-5\"`.", + defaultValue = "all", + requiredMode = RequiredMode.REQUIRED) + private String pageNumbers; + + @Hidden + public List getPageNumbersList(PDDocument doc, boolean oneBased) { + int pageCount = doc.getNumberOfPages(); + return GeneralUtils.parsePageList(pageNumbers, pageCount, oneBased); + } +} diff --git a/app/core/src/test/java/stirling/software/SPDF/controller/api/SplitPDFControllerTest.java b/app/core/src/test/java/stirling/software/SPDF/controller/api/SplitPDFControllerTest.java index 3068f17ec..01a203014 100644 --- a/app/core/src/test/java/stirling/software/SPDF/controller/api/SplitPDFControllerTest.java +++ b/app/core/src/test/java/stirling/software/SPDF/controller/api/SplitPDFControllerTest.java @@ -26,7 +26,7 @@ import org.springframework.http.MediaType; import org.springframework.mock.web.MockMultipartFile; import org.springframework.web.multipart.MultipartFile; -import stirling.software.SPDF.model.api.PDFWithPageNums; +import stirling.software.SPDF.model.api.SplitPagesRequest; import stirling.software.common.service.CustomPDFDocumentFactory; import stirling.software.common.util.TempFileManager; @@ -74,7 +74,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "input.pdf", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("3"); @@ -93,7 +93,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "input.pdf", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("1,2,3"); @@ -112,7 +112,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "input.pdf", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("1"); @@ -131,7 +131,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "input.pdf", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("3,7"); @@ -150,7 +150,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "input.pdf", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("2"); @@ -171,7 +171,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "input.pdf", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("5"); @@ -190,7 +190,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "input.pdf", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("all"); @@ -209,7 +209,7 @@ class SplitPDFControllerTest { new MockMultipartFile( "fileInput", "no_extension", MediaType.APPLICATION_PDF_VALUE, pdfBytes); - PDFWithPageNums request = new PDFWithPageNums(); + SplitPagesRequest request = new SplitPagesRequest(); request.setFileInput(file); request.setPageNumbers("1"); diff --git a/engine/.env b/engine/.env index 7d1aa1394..7c5d2b4a2 100644 --- a/engine/.env +++ b/engine/.env @@ -32,6 +32,10 @@ STIRLING_RAG_CHUNK_SIZE=512 STIRLING_RAG_CHUNK_OVERLAP=64 STIRLING_RAG_TOP_K=5 +# Upper bounds on PDF page text the engine will request per extraction round. +STIRLING_MAX_PAGES=200 +STIRLING_MAX_CHARACTERS=200000 + # PostHog analytics. Set STIRLING_POSTHOG_ENABLED=true and provide an API key to enable. STIRLING_POSTHOG_ENABLED=false STIRLING_POSTHOG_API_KEY=phc_VOdeYnlevc2T63m3myFGjeBlRcIusRgmhfx6XL5a1iz diff --git a/engine/scripts/generate_tool_models.py b/engine/scripts/generate_tool_models.py index 9c67fcae4..24d5bf180 100644 --- a/engine/scripts/generate_tool_models.py +++ b/engine/scripts/generate_tool_models.py @@ -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. diff --git a/engine/src/stirling/agents/_page_text.py b/engine/src/stirling/agents/_page_text.py new file mode 100644 index 000000000..b7b624faa --- /dev/null +++ b/engine/src/stirling/agents/_page_text.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from stirling.contracts import ExtractedFileText + + +def has_page_text(page_text: list[ExtractedFileText]) -> bool: + return any(selection.text.strip() for file_text in page_text for selection in file_text.pages) + + +def format_page_text(page_text: list[ExtractedFileText], empty: str = "None") -> str: + if not has_page_text(page_text): + return empty + sections = [ + f"[File: {file_text.file_name}, Page {selection.page_number or '?'}]\n{selection.text}" + for file_text in page_text + for selection in file_text.pages + ] + return "\n\n".join(sections) diff --git a/engine/src/stirling/agents/orchestrator.py b/engine/src/stirling/agents/orchestrator.py index c2537f78c..df920eb9a 100644 --- a/engine/src/stirling/agents/orchestrator.py +++ b/engine/src/stirling/agents/orchestrator.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from dataclasses import dataclass from typing import assert_never @@ -29,6 +30,8 @@ from stirling.contracts.pdf_edit import EditPlanResponse from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams from stirling.services import AppRuntime +logger = logging.getLogger(__name__) + @dataclass(frozen=True) class OrchestratorDeps: @@ -86,12 +89,20 @@ class OrchestratorAgent: ) async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse: + logger.info( + "[orchestrator] handle: files=%s resume_with=%s artifacts=%s msg=%r", + request.file_names, + request.resume_with, + [type(a).__name__ for a in request.artifacts], + request.user_message, + ) if request.resume_with is not None: return await self._resume(request, request.resume_with) result = await self.agent.run( self._build_prompt(request), deps=OrchestratorDeps(runtime=self.runtime, request=request), ) + logger.info("[orchestrator] routed -> %s", type(result.output).__name__) return result.output async def _resume(self, request: OrchestratorRequest, capability: SupportedCapability) -> OrchestratorResponse: @@ -117,11 +128,13 @@ class OrchestratorAgent: return await self._run_pdf_edit(ctx.deps.request) async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse: + extracted_text = self._get_extracted_text_artifact(request) return await PdfEditAgent(self.runtime).handle( PdfEditRequest( user_message=request.user_message, file_names=request.file_names, conversation_history=request.conversation_history, + page_text=extracted_text.files if extracted_text is not None else [], ) ) diff --git a/engine/src/stirling/agents/pdf_edit.py b/engine/src/stirling/agents/pdf_edit.py index 40b136578..2dc962111 100644 --- a/engine/src/stirling/agents/pdf_edit.py +++ b/engine/src/stirling/agents/pdf_edit.py @@ -1,23 +1,33 @@ from __future__ import annotations -from typing import Literal +import logging +from typing import Literal, overload from pydantic import Field from pydantic_ai import Agent from pydantic_ai.output import NativeOutput +from stirling.agents._page_text import format_page_text, has_page_text from stirling.contracts import ( EditCannotDoResponse, EditClarificationRequest, EditPlanResponse, + NeedContentFileRequest, + NeedContentResponse, + PdfContentType, PdfEditRequest, PdfEditResponse, + PdfEditTerminalResponse, + SupportedCapability, ToolOperationStep, format_conversation_history, ) +from stirling.logging import Pretty from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint from stirling.services import AppRuntime +logger = logging.getLogger(__name__) + class PdfEditPlanSelection(ApiModel): outcome: Literal["plan"] = "plan" @@ -26,6 +36,39 @@ class PdfEditPlanSelection(ApiModel): rationale: str | None = None +type PdfEditPlanOutput = PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse | NeedContentResponse + + +class PdfEditSelectionAgent: + def __init__(self, runtime: AppRuntime, base_system_prompt: str, *, allow_need_content: bool) -> None: + self.runtime = runtime + output_types: list[type[PdfEditPlanOutput]] = [ + PdfEditPlanSelection, + EditClarificationRequest, + EditCannotDoResponse, + ] + system_prompt = base_system_prompt + if allow_need_content: + output_types.append(NeedContentResponse) + system_prompt += ( + " Return need_content when planning a correct answer requires inspecting the actual PDF " + "page text (e.g. 'split after every page that says NEW PAGE', " + "'rotate pages that mention draft')." + ) + self.agent = Agent( + model=runtime.smart_model, + output_type=NativeOutput(output_types), + system_prompt=system_prompt, + model_settings=runtime.smart_model_settings, + ) + + async def select(self, prompt: str) -> PdfEditPlanOutput: + logger.debug("[pdf-edit selection] prompt:\n%s", prompt) + result = await self.agent.run(prompt) + logger.debug("[pdf-edit selection] output: %s", Pretty(result.output)) + return result.output + + class PdfEditParameterSelector: def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime @@ -34,7 +77,9 @@ class PdfEditParameterSelector: system_prompt=( "Generate only the parameter object for the selected PDF operation. " "Use reasonable defaults when the request does not specify optional details. " - "Only fill fields that belong to the selected operation's parameter model." + "Only fill fields that belong to the selected operation's parameter model. " + "When extracted page text is provided, use it to compute precise parameters " + "(e.g. exact page ranges that match a described pattern)." ), model_settings=runtime.smart_model_settings, ) @@ -48,14 +93,17 @@ class PdfEditParameterSelector: ) -> ParamToolModel: operation_id = operation_plan[operation_index] parameter_model = OPERATIONS[operation_id] + prompt = self._build_parameter_prompt(request, operation_plan, operation_index, generated_steps) + logger.debug("[pdf-edit params %s] prompt:\n%s", operation_id.name, prompt) parameter_result = await self.agent.run( - self._build_parameter_prompt(request, operation_plan, operation_index, generated_steps), + prompt, output_type=NativeOutput(parameter_model), instructions=( f"Generate only the parameters for the PDF operation `{operation_id.name}`. " "Do not include fields from any other operation." ), ) + logger.debug("[pdf-edit params %s] output: %s", operation_id.name, Pretty(parameter_result.output)) return parameter_result.output def _build_parameter_prompt( @@ -76,12 +124,14 @@ class PdfEditParameterSelector: else "None" ) return ( + f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n" f"User request: {request.user_message}\n" f"Files: {file_names}\n" f"Operation plan: {operation_list}\n" f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n" f"Selected operation: {operation_id.name}\n" f"Already generated steps:\n{generated_steps_text}\n" + f"Extracted page text:\n{format_page_text(request.page_text)}\n" "Return only the parameter object for the selected operation." ) @@ -91,32 +141,27 @@ class PdfEditAgent: self.runtime = runtime self.supported_operations = list(OPERATIONS) self.parameter_selector = PdfEditParameterSelector(runtime) - self.selection_agent = Agent( - model=runtime.smart_model, - output_type=NativeOutput( - [ - PdfEditPlanSelection, - EditClarificationRequest, - EditCannotDoResponse, - ] - ), - system_prompt=( - "Plan PDF edit requests. " - f"Supported operations are: {self._supported_operations_prompt()}. " - "Return an ordered list of one or more supported operations for the plan. " - "Do not produce operation parameters in this stage. " - "Return need_clarification when the request is genuinely ambiguous. " - "Return cannot_do when the request is outside the supported operations. " - "Return plan when a reasonable multi-step plan can be created. " - "Never return partial plans." - ), - model_settings=runtime.smart_model_settings, - ) - async def handle(self, request: PdfEditRequest) -> PdfEditResponse: - selection = await self._select_plan(request) + @overload + async def handle(self, request: PdfEditRequest, allow_need_content: Literal[False]) -> PdfEditTerminalResponse: ... + @overload + async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse: ... + async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse: + logger.info( + "[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s msg=%r", + request.file_names, + has_page_text(request.page_text), + allow_need_content, + request.user_message, + ) + selection = await self._select_plan(request, allow_need_content=allow_need_content) if isinstance(selection, EditClarificationRequest | EditCannotDoResponse): + logger.info("[pdf-edit] selection -> %s: %s", selection.outcome, Pretty(selection)) return selection + if isinstance(selection, NeedContentResponse): + logger.info("[pdf-edit] selection -> need_content: %s", selection.reason) + return self._fill_need_content_defaults(selection, request) + logger.info("[pdf-edit] plan: %s", [op.name for op in selection.operations]) steps: list[ToolOperationStep] = [] for operation_index, operation_id in enumerate(selection.operations): parameters = await self.parameter_selector.select( @@ -140,9 +185,27 @@ class PdfEditAgent: async def _select_plan( self, request: PdfEditRequest, - ) -> PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse: - selection_result = await self.selection_agent.run(self._build_selection_prompt(request)) - return selection_result.output + allow_need_content: bool = True, + ) -> PdfEditPlanOutput: + can_request_content = allow_need_content and not has_page_text(request.page_text) + agent = self._build_selection_agent(allow_need_content=can_request_content) + return await agent.select(self._build_selection_prompt(request)) + + def _build_selection_agent(self, *, allow_need_content: bool) -> PdfEditSelectionAgent: + return PdfEditSelectionAgent( + self.runtime, + base_system_prompt=( + "Plan PDF edit requests. " + f"Supported operations are: {self._supported_operations_prompt()}. " + "Return an ordered list of one or more supported operations for the plan. " + "Do not produce operation parameters in this stage. " + "Return need_clarification when the request is genuinely ambiguous. " + "Return cannot_do when the request is outside the supported operations. " + "Return plan when a reasonable multi-step plan can be created. " + "Never return partial plans." + ), + allow_need_content=allow_need_content, + ) def _build_selection_prompt(self, request: PdfEditRequest) -> str: file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided." @@ -151,8 +214,25 @@ class PdfEditAgent: f"User request: {request.user_message}\n" f"Files: {file_names}\n" f"Supported operations: {self._supported_operations_prompt()}\n" - "Plan an ordered list of supported PDF edit operations or return clarification/cannot_do." + f"Extracted page text:\n{format_page_text(request.page_text)}" ) def _supported_operations_prompt(self) -> str: return ", ".join(f"{op.name} ({op.value})" for op in self.supported_operations) + + def _fill_need_content_defaults( + self, + selection: NeedContentResponse, + request: PdfEditRequest, + ) -> NeedContentResponse: + files = selection.files or [ + NeedContentFileRequest(file_name=file_name, content_types=[PdfContentType.PAGE_TEXT]) + for file_name in request.file_names + ] + return NeedContentResponse( + resume_with=SupportedCapability.PDF_EDIT, + reason=selection.reason, + files=files, + max_pages=selection.max_pages or self.runtime.settings.max_pages, + max_characters=selection.max_characters or self.runtime.settings.max_characters, + ) diff --git a/engine/src/stirling/agents/pdf_questions.py b/engine/src/stirling/agents/pdf_questions.py index c646159da..d02376d75 100644 --- a/engine/src/stirling/agents/pdf_questions.py +++ b/engine/src/stirling/agents/pdf_questions.py @@ -3,24 +3,22 @@ from __future__ import annotations from pydantic_ai import Agent from pydantic_ai.output import NativeOutput +from stirling.agents._page_text import format_page_text, has_page_text from stirling.contracts import ( - ExtractedFileText, NeedContentFileRequest, + NeedContentResponse, PdfContentType, PdfQuestionAnswerResponse, - PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, PdfQuestionResponse, + SupportedCapability, format_conversation_history, ) from stirling.services import AppRuntime class PdfQuestionAgent: - DEFAULT_MAX_PAGES = 12 - DEFAULT_MAX_CHARACTERS = 24_000 - def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime rag = runtime.rag_capability @@ -44,8 +42,9 @@ class PdfQuestionAgent: ) async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse: - if not self._has_page_text(request.page_text): - return PdfQuestionNeedContentResponse( + if not has_page_text(request.page_text): + return NeedContentResponse( + resume_with=SupportedCapability.PDF_QUESTION, reason="No extracted PDF page text was provided, so the question cannot be answered yet.", files=[ NeedContentFileRequest( @@ -54,8 +53,8 @@ class PdfQuestionAgent: ) for file_name in request.file_names ], - max_pages=self.DEFAULT_MAX_PAGES, - max_characters=self.DEFAULT_MAX_CHARACTERS, + max_pages=self.runtime.settings.max_pages, + max_characters=self.runtime.settings.max_characters, ) return await self._run_answer_agent(request) @@ -65,12 +64,7 @@ class PdfQuestionAgent: def _build_prompt(self, request: PdfQuestionRequest) -> str: file_names = ", ".join(request.file_names) if request.file_names else "Unknown files" - sections = [ - f"[File: {file_text.file_name}, Page {selection.page_number or '?'}]\n{selection.text}" - for file_text in request.page_text - for selection in file_text.pages - ] - pages = "\n\n".join(sections) + pages = format_page_text(request.page_text, empty="") history = format_conversation_history(request.conversation_history) return ( f"Conversation history:\n{history}\n" @@ -78,6 +72,3 @@ class PdfQuestionAgent: f"Question: {request.question}\n" f"Extracted page text:\n{pages}" ) - - def _has_page_text(self, page_text: list[ExtractedFileText]) -> bool: - return any(selection.text.strip() for file_text in page_text for selection in file_text.pages) diff --git a/engine/src/stirling/agents/user_spec.py b/engine/src/stirling/agents/user_spec.py index 40a028aa1..e0cad68f2 100644 --- a/engine/src/stirling/agents/user_spec.py +++ b/engine/src/stirling/agents/user_spec.py @@ -14,10 +14,9 @@ from stirling.contracts import ( AgentRevisionWorkflowResponse, AiToolAgentStep, ConversationMessage, - EditCannotDoResponse, - EditClarificationRequest, EditPlanResponse, PdfEditRequest, + PdfEditTerminalResponse, format_conversation_history, ) from stirling.models import ApiModel @@ -102,7 +101,8 @@ class UserSpecAgent: self, user_message: str, conversation_history: list[ConversationMessage], - ) -> EditPlanResponse | EditClarificationRequest | EditCannotDoResponse: + ) -> PdfEditTerminalResponse: return await self.pdf_edit_agent.handle( - PdfEditRequest(user_message=user_message, conversation_history=conversation_history) + PdfEditRequest(user_message=user_message, conversation_history=conversation_history), + allow_need_content=False, ) diff --git a/engine/src/stirling/config/settings.py b/engine/src/stirling/config/settings.py index d4e621228..2aca8d3a0 100644 --- a/engine/src/stirling/config/settings.py +++ b/engine/src/stirling/config/settings.py @@ -37,6 +37,9 @@ class AppSettings(BaseSettings): rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP") rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K") + max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES") + max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS") + log_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL") log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE") @@ -57,6 +60,14 @@ def _configure_logging(level_name: str, log_file: str) -> None: root = logging.getLogger("stirling") root.setLevel(level) + formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s [%(funcName)s] %(message)s") + + if not any(isinstance(h, logging.StreamHandler) for h in root.handlers): + sh = logging.StreamHandler() + sh.setFormatter(formatter) + sh.setLevel(level) + root.addHandler(sh) + root.propagate = False if log_file: log_path = Path(log_file) @@ -67,7 +78,7 @@ def _configure_logging(level_name: str, log_file: str) -> None: backupCount=1, encoding="utf-8", ) - fh.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s [%(funcName)s] %(message)s")) + fh.setFormatter(formatter) fh.setLevel(level) root.addHandler(fh) diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index 6fc075b44..207773db0 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -12,6 +12,8 @@ from .common import ( ArtifactKind, ConversationMessage, ExtractedFileText, + NeedContentFileRequest, + NeedContentResponse, PdfContentType, PdfTextSelection, StepKind, @@ -54,14 +56,14 @@ from .pdf_edit import ( EditPlanResponse, PdfEditRequest, PdfEditResponse, + PdfEditTerminalResponse, ) from .pdf_questions import ( - NeedContentFileRequest, PdfQuestionAnswerResponse, - PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, PdfQuestionResponse, + PdfQuestionTerminalResponse, ) from .rag import ( MAX_INDEX_TEXT_LENGTH, @@ -108,17 +110,19 @@ __all__ = [ "format_conversation_history", "HealthResponse", "NeedContentFileRequest", + "NeedContentResponse", "NextExecutionAction", "OrchestratorRequest", "OrchestratorResponse", "PdfContentType", "PdfEditRequest", "PdfEditResponse", + "PdfEditTerminalResponse", "PdfQuestionAnswerResponse", - "PdfQuestionNeedContentResponse", "PdfQuestionNotFoundResponse", "PdfQuestionRequest", "PdfQuestionResponse", + "PdfQuestionTerminalResponse", "PdfTextSelection", "RagCollectionsResponse", "RagDeleteCollectionResponse", diff --git a/engine/src/stirling/contracts/common.py b/engine/src/stirling/contracts/common.py index 50038d150..7b4e5fd2f 100644 --- a/engine/src/stirling/contracts/common.py +++ b/engine/src/stirling/contracts/common.py @@ -107,6 +107,21 @@ class ExtractedFileText(ApiModel): pages: list[PdfTextSelection] = Field(default_factory=list) +class NeedContentFileRequest(ApiModel): + file_name: str + page_numbers: list[int] = Field(default_factory=list) + content_types: list[PdfContentType] + + +class NeedContentResponse(ApiModel): + outcome: Literal[WorkflowOutcome.NEED_CONTENT] = WorkflowOutcome.NEED_CONTENT + resume_with: SupportedCapability + reason: str + files: list[NeedContentFileRequest] = Field(default_factory=list) + max_pages: int + max_characters: int + + class ToolOperationStep(ApiModel): kind: Literal[StepKind.TOOL] = StepKind.TOOL tool: AnyToolId diff --git a/engine/src/stirling/contracts/orchestrator.py b/engine/src/stirling/contracts/orchestrator.py index 759837288..fe36719f6 100644 --- a/engine/src/stirling/contracts/orchestrator.py +++ b/engine/src/stirling/contracts/orchestrator.py @@ -11,12 +11,13 @@ from .common import ( ArtifactKind, ConversationMessage, ExtractedFileText, + NeedContentResponse, SupportedCapability, WorkflowOutcome, ) from .execution import NextExecutionAction -from .pdf_edit import PdfEditResponse -from .pdf_questions import PdfQuestionResponse +from .pdf_edit import PdfEditTerminalResponse +from .pdf_questions import PdfQuestionTerminalResponse class ExtractedTextArtifact(ApiModel): @@ -41,7 +42,12 @@ class UnsupportedCapabilityResponse(ApiModel): message: str -OrchestratorResponse = Annotated[ - PdfEditResponse | PdfQuestionResponse | AgentDraftResponse | NextExecutionAction | UnsupportedCapabilityResponse, +type OrchestratorResponse = Annotated[ + PdfEditTerminalResponse + | PdfQuestionTerminalResponse + | NeedContentResponse + | AgentDraftResponse + | NextExecutionAction + | UnsupportedCapabilityResponse, Field(discriminator="outcome"), ] diff --git a/engine/src/stirling/contracts/pdf_edit.py b/engine/src/stirling/contracts/pdf_edit.py index a7bcedfa9..9ba9b09e7 100644 --- a/engine/src/stirling/contracts/pdf_edit.py +++ b/engine/src/stirling/contracts/pdf_edit.py @@ -6,13 +6,14 @@ from pydantic import Field from stirling.models import ApiModel -from .common import ConversationMessage, ToolOperationStep, WorkflowOutcome +from .common import ConversationMessage, ExtractedFileText, NeedContentResponse, ToolOperationStep, WorkflowOutcome class PdfEditRequest(ApiModel): user_message: str file_names: list[str] = Field(default_factory=list) conversation_history: list[ConversationMessage] = Field(default_factory=list) + page_text: list[ExtractedFileText] = Field(default_factory=list) class EditPlanResponse(ApiModel): @@ -33,7 +34,8 @@ class EditCannotDoResponse(ApiModel): reason: str -PdfEditResponse = Annotated[ - EditPlanResponse | EditClarificationRequest | EditCannotDoResponse, +type PdfEditTerminalResponse = EditPlanResponse | EditClarificationRequest | EditCannotDoResponse +type PdfEditResponse = Annotated[ + PdfEditTerminalResponse | NeedContentResponse, Field(discriminator="outcome"), ] diff --git a/engine/src/stirling/contracts/pdf_questions.py b/engine/src/stirling/contracts/pdf_questions.py index 4ee0d2596..0bfece3ce 100644 --- a/engine/src/stirling/contracts/pdf_questions.py +++ b/engine/src/stirling/contracts/pdf_questions.py @@ -9,8 +9,7 @@ from stirling.models import ApiModel from .common import ( ConversationMessage, ExtractedFileText, - PdfContentType, - SupportedCapability, + NeedContentResponse, WorkflowOutcome, ) @@ -28,27 +27,13 @@ class PdfQuestionAnswerResponse(ApiModel): evidence: list[ExtractedFileText] = Field(default_factory=list) -class NeedContentFileRequest(ApiModel): - file_name: str - page_numbers: list[int] = Field(default_factory=list) - content_types: list[PdfContentType] - - -class PdfQuestionNeedContentResponse(ApiModel): - outcome: Literal[WorkflowOutcome.NEED_CONTENT] = WorkflowOutcome.NEED_CONTENT - resume_with: SupportedCapability = SupportedCapability.PDF_QUESTION - reason: str - files: list[NeedContentFileRequest] = Field(default_factory=list) - max_pages: int - max_characters: int - - class PdfQuestionNotFoundResponse(ApiModel): outcome: Literal[WorkflowOutcome.NOT_FOUND] = WorkflowOutcome.NOT_FOUND reason: str -PdfQuestionResponse = Annotated[ - PdfQuestionAnswerResponse | PdfQuestionNeedContentResponse | PdfQuestionNotFoundResponse, +type PdfQuestionTerminalResponse = PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse +type PdfQuestionResponse = Annotated[ + PdfQuestionTerminalResponse | NeedContentResponse, Field(discriminator="outcome"), ] diff --git a/engine/src/stirling/logging.py b/engine/src/stirling/logging.py index a0a6fec20..3e5fdb4d4 100644 --- a/engine/src/stirling/logging.py +++ b/engine/src/stirling/logging.py @@ -4,13 +4,16 @@ from __future__ import annotations import json +from pydantic import BaseModel + class Pretty: """Lazy JSON formatter — only serialises when ``str()`` is called. Designed for use with ``logging``'s ``%s`` formatting so that the JSON serialisation is skipped entirely when the log message is - never emitted. + never emitted. Pydantic models (at the top level or nested) are + dumped via ``model_dump``; anything else falls back to ``str``. """ __slots__ = ("_obj",) @@ -19,4 +22,12 @@ class Pretty: self._obj = obj def __str__(self) -> str: - return json.dumps(self._obj, indent=2, default=str, ensure_ascii=True) + if isinstance(self._obj, BaseModel): + return self._obj.model_dump_json(indent=2) + return json.dumps(self._obj, indent=2, default=_default, ensure_ascii=True) + + +def _default(value: object) -> object: + if isinstance(value, BaseModel): + return value.model_dump() + return str(value) diff --git a/engine/src/stirling/models/tool_models.py b/engine/src/stirling/models/tool_models.py index d67144114..5709f5a5d 100644 --- a/engine/src/stirling/models/tool_models.py +++ b/engine/src/stirling/models/tool_models.py @@ -5,7 +5,6 @@ from __future__ import annotations from enum import Enum, IntEnum, StrEnum -from typing import Any from pydantic import Field, RootModel, SecretStr @@ -756,6 +755,15 @@ class RearrangePagesParams(ApiModel): ) +class RedactionArea(ApiModel): + color: str | None = Field(None, description="The color used to redact the specified area.") + height: float | None = Field(None, description="The height of the area to be redacted.") + page: int | None = Field(None, description="The page on which the area should be redacted.") + width: float | None = Field(None, description="The width of the area to be redacted.") + x: float | None = Field(None, description="The left edge point of the area to be redacted.") + y: float | None = Field(None, description="The top edge point of the area to be redacted.") + + class RemoveBlanksParams(ApiModel): threshold: int | None = Field(10, description="The threshold value to determine blank pages", ge=0, le=255) white_percent: float | None = Field( @@ -949,7 +957,7 @@ class SplitForPosterPrintParams(ApiModel): class SplitPagesParams(ApiModel): 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')", + description='Split points - page numbers after which the PDF will be cut. For example, `"2"` produces two documents (pages 1-2 and pages 3+); `"2,5"` produces three (pages 1-2, 3-5, 6+). Supports ranges (e.g. `"1,3,5-9"` splits after pages 1, 3, 5, 6, 7, 8, 9, yielding 8 documents), `"all"` (split after every page), or functions like `"2n+1"`, `"3n"`, `"6n-5"`.', ) @@ -1041,10 +1049,6 @@ class VectorToPdfParams(ApiModel): prepress: Prepress | None = Field(Prepress.boolean_false, description="Apply Ghostscript prepress settings") -class RedactionArea(RootModel[Any]): - root: Any - - class RedactParams(ApiModel): convert_pdf_to_image: bool | None = Field(False, description="Convert the redacted PDF to an image") page_numbers: str | None = Field( diff --git a/engine/tests/conftest.py b/engine/tests/conftest.py index 3e1809f76..3bc409f61 100644 --- a/engine/tests/conftest.py +++ b/engine/tests/conftest.py @@ -30,6 +30,8 @@ def build_app_settings() -> AppSettings: rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, + max_pages=200, + max_characters=200_000, posthog_enabled=False, posthog_api_key="", posthog_host="https://eu.i.posthog.com", diff --git a/engine/tests/test_pdf_edit_agent.py b/engine/tests/test_pdf_edit_agent.py index fdbeb722e..22716a50b 100644 --- a/engine/tests/test_pdf_edit_agent.py +++ b/engine/tests/test_pdf_edit_agent.py @@ -5,11 +5,18 @@ from dataclasses import dataclass import pytest from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection +from stirling.agents.pdf_edit import PdfEditPlanOutput from stirling.contracts import ( EditCannotDoResponse, EditClarificationRequest, EditPlanResponse, + ExtractedFileText, + NeedContentFileRequest, + NeedContentResponse, + PdfContentType, PdfEditRequest, + PdfTextSelection, + SupportedCapability, ToolOperationStep, ) from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint @@ -52,7 +59,7 @@ class StubPdfEditAgent(PdfEditAgent): def __init__( self, runtime: AppRuntime, - selection: PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse, + selection: PdfEditPlanOutput, parameter_selector: RecordingParameterSelector | PdfEditParameterSelector | None = None, ) -> None: super().__init__(runtime) @@ -63,7 +70,8 @@ class StubPdfEditAgent(PdfEditAgent): async def _select_plan( self, request: PdfEditRequest, - ) -> PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse: + allow_need_content: bool = True, + ) -> PdfEditPlanOutput: return self.selection @@ -153,3 +161,117 @@ async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan(runtime: Ap 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, + NeedContentResponse( + resume_with=SupportedCapability.PDF_EDIT, + reason="Need page text to locate the NEW PAGE markers.", + files=[], + max_pages=0, + max_characters=0, + ), + parameter_selector=parameter_selector, + ) + + response = await agent.handle( + PdfEditRequest( + user_message="Split after every page that says 'NEW PAGE'.", + file_names=["report.pdf"], + ) + ) + + assert isinstance(response, NeedContentResponse) + assert response.resume_with == SupportedCapability.PDF_EDIT + assert response.files == [NeedContentFileRequest(file_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(*, allow_need_content: bool) -> PdfEditSelectionAgent: + captured.append(allow_need_content) + raise _StopSelectionError() + + agent._build_selection_agent = record + + with pytest.raises(_StopSelectionError): + await agent._select_plan(PdfEditRequest(user_message="Rotate.")) + 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")], + ) + ], + ) + ) + with pytest.raises(_StopSelectionError): + await agent._select_plan(PdfEditRequest(user_message="Rotate."), 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 NeedContentResponse in _agent_output_types(can_request) + assert NeedContentResponse 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.", + ), + 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.", + file_names=["report.pdf"], + page_text=page_text, + ) + ) + + assert parameter_selector.calls[0].request.page_text == page_text diff --git a/engine/tests/test_pdf_question_agent.py b/engine/tests/test_pdf_question_agent.py index b284870b7..b638c6bc4 100644 --- a/engine/tests/test_pdf_question_agent.py +++ b/engine/tests/test_pdf_question_agent.py @@ -5,8 +5,8 @@ import pytest from stirling.agents import PdfQuestionAgent from stirling.contracts import ( ExtractedFileText, + NeedContentResponse, PdfQuestionAnswerResponse, - PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, PdfTextSelection, @@ -41,7 +41,7 @@ async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) - PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"]) ) - assert isinstance(response, PdfQuestionNeedContentResponse) + assert isinstance(response, NeedContentResponse) @pytest.mark.anyio diff --git a/engine/tests/test_stirling_api.py b/engine/tests/test_stirling_api.py index 2de187af0..96fd98268 100644 --- a/engine/tests/test_stirling_api.py +++ b/engine/tests/test_stirling_api.py @@ -19,18 +19,25 @@ from stirling.contracts import ( AgentRevisionResponse, CannotContinueExecutionAction, EditCannotDoResponse, + NeedContentResponse, OrchestratorRequest, PdfEditRequest, - PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, + SupportedCapability, ) from stirling.models.tool_models import Angle, RotatePdfParams class StubOrchestratorAgent: - async def handle(self, request: OrchestratorRequest) -> PdfQuestionNeedContentResponse: - return PdfQuestionNeedContentResponse(reason=request.user_message, files=[], max_pages=1, max_characters=1000) + async def handle(self, request: OrchestratorRequest) -> NeedContentResponse: + return NeedContentResponse( + resume_with=SupportedCapability.PDF_QUESTION, + reason=request.user_message, + files=[], + max_pages=1, + max_characters=1000, + ) class StubPdfEditAgent: diff --git a/engine/tests/test_stirling_contracts.py b/engine/tests/test_stirling_contracts.py index 9e113015c..7fc1d9c64 100644 --- a/engine/tests/test_stirling_contracts.py +++ b/engine/tests/test_stirling_contracts.py @@ -89,6 +89,8 @@ def test_app_settings_accepts_model_configuration() -> None: rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, + max_pages=200, + max_characters=200_000, posthog_enabled=False, posthog_api_key="", posthog_host="https://eu.i.posthog.com",