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

Add document context for edit agent (#6152)

Browse Source 
# Description of Changes
Adds the ability for the Edit agent to request the content of the
document before it decides which parameters it needs. This makes it able
to process requests like `Split the document after the page containing
the "My Section" section`, allowing for document context-based requests
for all[^1] tools.

I had to make a few changes elsewhere to make this work, including:
- Moving the requesting of content out of the Question Agent and into a
common location
- Added specific API docs for the Split param because the generic ones
were not specific enough for the AI to be able to reliably perform the
correct operation
- Fixed an issue in the tool models generator which caused the Redact
params to only be half-generated (causing Pydantic to crash when the AI
tried to run Redact)
- Added missing logging to a bunch of tools and hooked it up properly so
it'll print to stderr
- Made the limits for the max pages/chars to extract from PDFs
configurable via env var

[^1]: Many of the tools can't actually do anything useful with the
context at this stage, but will just need the tool API to be extended
with new features like page-specific operations to be automatically able
to do smart operations without needing to change the Edit agent itself.
This commit is contained in:
James Brunton
2026-04-23 13:19:27 +00:00
committed by GitHub
parent e087b54cf0
commit 3e94157137
23 changed files with 462 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/.env
+4
View File
@@ -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
engine/scripts/generate_tool_models.py
+39
View File
@@ -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.
engine/src/stirling/agents/_page_text.py
+18
View File
@@ -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)
engine/src/stirling/agents/orchestrator.py
+13
View File
@@ -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 [],
)
)
engine/src/stirling/agents/pdf_edit.py
+110 -30
View File
@@ -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,
)
engine/src/stirling/agents/pdf_questions.py
+9 -18
View File
@@ -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)
engine/src/stirling/agents/user_spec.py
+4 -4
View File
@@ -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,
)
engine/src/stirling/config/settings.py
+12 -1
View File
@@ -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)
engine/src/stirling/contracts/__init__.py
+7 -3
View File
@@ -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",
engine/src/stirling/contracts/common.py
+15
View File
@@ -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
engine/src/stirling/contracts/orchestrator.py
+10 -4
View File
@@ -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"),
]
engine/src/stirling/contracts/pdf_edit.py
+5 -3
View File
@@ -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"),
]
engine/src/stirling/contracts/pdf_questions.py
+4 -19
View File
@@ -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"),
]
engine/src/stirling/logging.py
+13 -2
View File
@@ -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)
engine/src/stirling/models/tool_models.py
+10 -6
View File
@@ -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(
engine/tests/conftest.py
+2
View File
@@ -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",
engine/tests/test_pdf_edit_agent.py
+124 -2
View File
@@ -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
engine/tests/test_pdf_question_agent.py
+2 -2
View File
@@ -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
engine/tests/test_stirling_api.py
+10 -3
View File
@@ -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:
engine/tests/test_stirling_contracts.py
+2
View File
@@ -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",