mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
Add Java orchestrator to connect to the AI engine (#6003)
# Description of Changes Add Java orchestration layer which can connect and go back and forth with the AI engine to get results for the user. It's expected that the AI engine will not be publicly available and this Java layer will always be in front of it, to manage sessions and auth etc.
This commit is contained in:
James Brunton
@@ -42,9 +42,8 @@ select = [
|
||||
"W",
|
||||
"RUF100",
|
||||
"UP",
|
||||
]
|
||||
ignore = [
|
||||
"E501", # Temporarily disable line length limit until codebase conformat
|
||||
"PYI", # flake8-pyi: flags deprecated typing constructs
|
||||
"FA", # flake8-future-annotations: flags missing future annotations imports
|
||||
]
|
||||
|
||||
[tool.pyright]
|
||||
@@ -55,6 +54,7 @@ reportUnnecessaryCast = "warning"
|
||||
reportUnnecessaryTypeIgnoreComment = "warning"
|
||||
reportUnusedImport = "warning"
|
||||
reportUnknownParameterType = "warning"
|
||||
reportDeprecated = "warning"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import assert_never
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import ToolOutput
|
||||
@@ -12,12 +13,14 @@ from stirling.agents.user_spec import UserSpecAgent
|
||||
from stirling.contracts import (
|
||||
AgentDraftRequest,
|
||||
AgentDraftWorkflowResponse,
|
||||
ExtractedTextArtifact,
|
||||
OrchestratorRequest,
|
||||
OrchestratorResponse,
|
||||
PdfEditRequest,
|
||||
PdfEditResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
SupportedCapability,
|
||||
UnsupportedCapabilityResponse,
|
||||
)
|
||||
from stirling.services import AppRuntime
|
||||
@@ -61,7 +64,7 @@ class OrchestratorAgent:
|
||||
"You are the top-level orchestrator. "
|
||||
"Choose exactly one output function that best handles the request. "
|
||||
"Use delegate_pdf_edit for requested PDF modifications. "
|
||||
"Use delegate_pdf_question for questions about the contents of a PDF. "
|
||||
"Use delegate_pdf_question for questions about PDF contents. "
|
||||
"Use delegate_user_spec for requests to create or define an agent spec. "
|
||||
"Use unsupported_capability only when none of the other outputs fit."
|
||||
),
|
||||
@@ -69,27 +72,56 @@ class OrchestratorAgent:
|
||||
)
|
||||
|
||||
async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
|
||||
if request.resume_with is not None:
|
||||
return await self._resume(request, request.resume_with)
|
||||
result = await self.agent.run(
|
||||
request.user_message,
|
||||
self._build_prompt(request),
|
||||
deps=OrchestratorDeps(runtime=self.runtime, request=request),
|
||||
)
|
||||
return result.output
|
||||
|
||||
async def _resume(self, request: OrchestratorRequest, capability: SupportedCapability) -> OrchestratorResponse:
|
||||
"""Fast-path to get back to the correct endpoint without having to call AI."""
|
||||
match capability:
|
||||
case SupportedCapability.PDF_QUESTION:
|
||||
return await self._run_pdf_question(request)
|
||||
case SupportedCapability.PDF_EDIT:
|
||||
return await self._run_pdf_edit(request)
|
||||
case SupportedCapability.AGENT_DRAFT:
|
||||
return await self._run_agent_draft(request)
|
||||
case (
|
||||
SupportedCapability.ORCHESTRATE
|
||||
| SupportedCapability.AGENT_REVISE
|
||||
| SupportedCapability.AGENT_NEXT_ACTION
|
||||
):
|
||||
raise ValueError(f"Cannot resume orchestrator with capability: {capability}")
|
||||
case _ as unreachable:
|
||||
assert_never(unreachable)
|
||||
|
||||
async def delegate_pdf_edit(self, ctx: RunContext[OrchestratorDeps]) -> PdfEditResponse:
|
||||
request = ctx.deps.request
|
||||
return await PdfEditAgent(ctx.deps.runtime).handle(
|
||||
PdfEditRequest(user_message=request.user_message, conversation_id=request.conversation_id)
|
||||
)
|
||||
return await self._run_pdf_edit(ctx.deps.request)
|
||||
|
||||
async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse:
|
||||
return await PdfEditAgent(self.runtime).handle(PdfEditRequest(user_message=request.user_message))
|
||||
|
||||
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse:
|
||||
request = ctx.deps.request
|
||||
return await PdfQuestionAgent(ctx.deps.runtime).handle(
|
||||
PdfQuestionRequest(question=request.user_message, conversation_id=request.conversation_id)
|
||||
return await self._run_pdf_question(ctx.deps.request)
|
||||
|
||||
async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionResponse:
|
||||
extracted_text = self._get_extracted_text_artifact(request)
|
||||
return await PdfQuestionAgent(self.runtime).handle(
|
||||
PdfQuestionRequest(
|
||||
question=request.user_message,
|
||||
file_names=request.file_names,
|
||||
page_text=extracted_text.files if extracted_text is not None else [],
|
||||
)
|
||||
)
|
||||
|
||||
async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
|
||||
request = ctx.deps.request
|
||||
return await UserSpecAgent(ctx.deps.runtime).draft(AgentDraftRequest(user_message=request.user_message))
|
||||
return await self._run_agent_draft(ctx.deps.request)
|
||||
|
||||
async def _run_agent_draft(self, request: OrchestratorRequest) -> AgentDraftWorkflowResponse:
|
||||
return await UserSpecAgent(self.runtime).draft(AgentDraftRequest(user_message=request.user_message))
|
||||
|
||||
async def unsupported_capability(
|
||||
self,
|
||||
@@ -98,3 +130,28 @@ class OrchestratorAgent:
|
||||
message: str,
|
||||
) -> UnsupportedCapabilityResponse:
|
||||
return UnsupportedCapabilityResponse(capability=capability, message=message)
|
||||
|
||||
def _get_extracted_text_artifact(self, request: OrchestratorRequest) -> ExtractedTextArtifact | None:
|
||||
for artifact in request.artifacts:
|
||||
if isinstance(artifact, ExtractedTextArtifact):
|
||||
return artifact
|
||||
return None
|
||||
|
||||
def _build_prompt(self, request: OrchestratorRequest) -> str:
|
||||
artifact_summary = self._describe_artifacts(request)
|
||||
file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
|
||||
return f"User message: {request.user_message}\nFiles: {file_names}\nAvailable artifacts:\n{artifact_summary}"
|
||||
|
||||
def _describe_artifacts(self, request: OrchestratorRequest) -> str:
|
||||
if not request.artifacts:
|
||||
return "- none"
|
||||
|
||||
descriptions: list[str] = []
|
||||
for artifact in request.artifacts:
|
||||
if isinstance(artifact, ExtractedTextArtifact):
|
||||
total_pages = sum(len(f.pages) for f in artifact.files)
|
||||
file_names = [f.file_name for f in artifact.files]
|
||||
descriptions.append(f"- extracted_text: {total_pages} pages from {file_names}")
|
||||
continue
|
||||
descriptions.append("- unknown artifact")
|
||||
return "\n".join(descriptions)
|
||||
|
||||
@@ -4,8 +4,11 @@ from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.contracts import (
|
||||
ExtractedFileText,
|
||||
NeedContentFileRequest,
|
||||
PdfContentType,
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNeedTextResponse,
|
||||
PdfQuestionNeedContentResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
@@ -14,6 +17,9 @@ 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
|
||||
self.agent = Agent(
|
||||
@@ -25,18 +31,27 @@ class PdfQuestionAgent:
|
||||
]
|
||||
),
|
||||
system_prompt=(
|
||||
"Answer questions about a PDF using only the extracted text provided in the prompt. "
|
||||
"Answer questions about PDFs using only the extracted page text provided in the prompt. "
|
||||
"Do not guess or use outside knowledge. "
|
||||
"If the answer is not supported by the provided text, return not_found. "
|
||||
"When answering, include a short list of evidence snippets copied from the provided text."
|
||||
"When answering, include a short list of evidence snippets with their page numbers."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||
if not request.extracted_text.strip():
|
||||
return PdfQuestionNeedTextResponse(
|
||||
reason="No extracted PDF text was provided, so the question cannot be answered yet."
|
||||
if not self._has_page_text(request.page_text):
|
||||
return PdfQuestionNeedContentResponse(
|
||||
reason="No extracted PDF page text was provided, so the question cannot be answered yet.",
|
||||
files=[
|
||||
NeedContentFileRequest(
|
||||
file_name=file_name,
|
||||
content_types=[PdfContentType.PAGE_TEXT],
|
||||
)
|
||||
for file_name in request.file_names
|
||||
],
|
||||
max_pages=self.DEFAULT_MAX_PAGES,
|
||||
max_characters=self.DEFAULT_MAX_CHARACTERS,
|
||||
)
|
||||
return await self._run_answer_agent(request)
|
||||
|
||||
@@ -45,5 +60,14 @@ class PdfQuestionAgent:
|
||||
return result.output
|
||||
|
||||
def _build_prompt(self, request: PdfQuestionRequest) -> str:
|
||||
file_name = request.file_name or "Unknown file"
|
||||
return f"File: {file_name}\nQuestion: {request.question}\nExtracted text:\n{request.extracted_text}"
|
||||
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)
|
||||
return f"Files: {file_names}\nQuestion: {request.question}\nExtracted 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)
|
||||
|
||||
@@ -8,7 +8,17 @@ from .agent_drafts import (
|
||||
AgentRevisionWorkflowResponse,
|
||||
)
|
||||
from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
|
||||
from .common import ConversationMessage, PdfTextSelection, ToolOperationStep
|
||||
from .common import (
|
||||
ArtifactKind,
|
||||
ConversationMessage,
|
||||
ExtractedFileText,
|
||||
PdfContentType,
|
||||
PdfTextSelection,
|
||||
StepKind,
|
||||
SupportedCapability,
|
||||
ToolOperationStep,
|
||||
WorkflowOutcome,
|
||||
)
|
||||
from .execution import (
|
||||
AgentExecutionRequest,
|
||||
CannotContinueExecutionAction,
|
||||
@@ -19,7 +29,13 @@ from .execution import (
|
||||
ToolCallExecutionAction,
|
||||
)
|
||||
from .health import HealthResponse
|
||||
from .orchestrator import OrchestratorRequest, OrchestratorResponse, SupportedCapability, UnsupportedCapabilityResponse
|
||||
from .orchestrator import (
|
||||
ExtractedTextArtifact,
|
||||
OrchestratorRequest,
|
||||
OrchestratorResponse,
|
||||
UnsupportedCapabilityResponse,
|
||||
WorkflowArtifact,
|
||||
)
|
||||
from .pdf_edit import (
|
||||
EditCannotDoResponse,
|
||||
EditClarificationRequest,
|
||||
@@ -28,14 +44,16 @@ from .pdf_edit import (
|
||||
PdfEditResponse,
|
||||
)
|
||||
from .pdf_questions import (
|
||||
NeedContentFileRequest,
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNeedTextResponse,
|
||||
PdfQuestionNeedContentResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ArtifactKind",
|
||||
"AgentDraft",
|
||||
"AgentDraftRequest",
|
||||
"AgentDraftResponse",
|
||||
@@ -49,6 +67,7 @@ __all__ = [
|
||||
"AiToolAgentStep",
|
||||
"CannotContinueExecutionAction",
|
||||
"ConversationMessage",
|
||||
"ExtractedFileText",
|
||||
"CompletedExecutionAction",
|
||||
"EditCannotDoResponse",
|
||||
"EditClarificationRequest",
|
||||
@@ -56,19 +75,25 @@ __all__ = [
|
||||
"ExecutionContext",
|
||||
"ExecutionStepResult",
|
||||
"HealthResponse",
|
||||
"NeedContentFileRequest",
|
||||
"NextExecutionAction",
|
||||
"ExtractedTextArtifact",
|
||||
"OrchestratorRequest",
|
||||
"OrchestratorResponse",
|
||||
"PdfEditRequest",
|
||||
"PdfEditResponse",
|
||||
"PdfQuestionAnswerResponse",
|
||||
"PdfQuestionNotFoundResponse",
|
||||
"PdfQuestionNeedTextResponse",
|
||||
"PdfContentType",
|
||||
"PdfQuestionNeedContentResponse",
|
||||
"PdfQuestionRequest",
|
||||
"PdfQuestionResponse",
|
||||
"PdfTextSelection",
|
||||
"StepKind",
|
||||
"SupportedCapability",
|
||||
"ToolOperationStep",
|
||||
"ToolCallExecutionAction",
|
||||
"WorkflowOutcome",
|
||||
"UnsupportedCapabilityResponse",
|
||||
"WorkflowArtifact",
|
||||
]
|
||||
|
||||
@@ -7,12 +7,12 @@ from pydantic import Field
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .agent_specs import AgentSpecStep
|
||||
from .common import ConversationMessage
|
||||
from .common import ConversationMessage, StepKind, WorkflowOutcome
|
||||
from .pdf_edit import EditCannotDoResponse, EditClarificationRequest
|
||||
|
||||
|
||||
class AgentDraftStep(ApiModel):
|
||||
kind: Literal["tool", "ai_tool"]
|
||||
kind: Literal[StepKind.TOOL, StepKind.AI_TOOL]
|
||||
title: str
|
||||
description: str
|
||||
|
||||
@@ -30,7 +30,7 @@ class AgentDraftRequest(ApiModel):
|
||||
|
||||
|
||||
class AgentDraftResponse(ApiModel):
|
||||
outcome: Literal["draft"] = "draft"
|
||||
outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT
|
||||
draft: AgentDraft
|
||||
|
||||
|
||||
@@ -41,7 +41,7 @@ class AgentRevisionRequest(ApiModel):
|
||||
|
||||
|
||||
class AgentRevisionResponse(ApiModel):
|
||||
outcome: Literal["draft"] = "draft"
|
||||
outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT
|
||||
draft: AgentDraft
|
||||
|
||||
|
||||
|
||||
@@ -6,11 +6,11 @@ from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel, OperationId
|
||||
|
||||
from .common import ToolOperationStep
|
||||
from .common import StepKind, ToolOperationStep
|
||||
|
||||
|
||||
class AiToolAgentStep(ApiModel):
|
||||
kind: Literal["ai_tool"] = "ai_tool"
|
||||
kind: Literal[StepKind.AI_TOOL] = StepKind.AI_TOOL
|
||||
title: str
|
||||
description: str
|
||||
tool: OperationId
|
||||
|
||||
@@ -1,12 +1,89 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import model_validator
|
||||
from pydantic import Field, model_validator
|
||||
|
||||
from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
|
||||
|
||||
|
||||
class PdfContentType(StrEnum):
|
||||
"""Types of content that can be extracted from a PDF and sent to the AI.
|
||||
|
||||
Java counterpart: AiPdfContentType.java - values must stay in sync.
|
||||
"""
|
||||
|
||||
# Document-level structured data
|
||||
PAGE_LAYOUT = "page_layout"
|
||||
DOCUMENT_METADATA = "document_metadata"
|
||||
ENCRYPTION_INFO = "encryption_info"
|
||||
BOOKMARKS = "bookmarks"
|
||||
LAYERS = "layers"
|
||||
EMBEDDED_FILES = "embedded_files"
|
||||
JAVASCRIPT = "javascript"
|
||||
LINKS = "links"
|
||||
IMAGE_INFO = "image_info"
|
||||
FONTS = "fonts"
|
||||
|
||||
# Text and content
|
||||
PAGE_TEXT = "page_text"
|
||||
FULL_TEXT = "full_text"
|
||||
FORM_FIELDS = "form_fields"
|
||||
ANNOTATIONS = "annotations"
|
||||
SIGNATURES = "signatures"
|
||||
STRUCTURE_TREE = "structure_tree"
|
||||
XMP_METADATA = "xmp_metadata"
|
||||
|
||||
# Heavy content
|
||||
COMPLIANCE = "compliance"
|
||||
IMAGES = "images"
|
||||
|
||||
|
||||
class WorkflowOutcome(StrEnum):
|
||||
"""Discriminator values for all workflow response unions (outcome field).
|
||||
|
||||
Java counterpart: AiWorkflowOutcome.java - values must stay in sync.
|
||||
"""
|
||||
|
||||
ANSWER = "answer"
|
||||
NEED_CONTENT = "need_content"
|
||||
NOT_FOUND = "not_found"
|
||||
PLAN = "plan"
|
||||
NEED_CLARIFICATION = "need_clarification"
|
||||
CANNOT_DO = "cannot_do"
|
||||
DRAFT = "draft"
|
||||
TOOL_CALL = "tool_call"
|
||||
COMPLETED = "completed"
|
||||
CANNOT_CONTINUE = "cannot_continue"
|
||||
UNSUPPORTED_CAPABILITY = "unsupported_capability"
|
||||
|
||||
|
||||
class ArtifactKind(StrEnum):
|
||||
"""Discriminator values for WorkflowArtifact unions (kind field).
|
||||
|
||||
Java counterpart: PdfContentExtractor.ArtifactKind - values must stay in sync.
|
||||
"""
|
||||
|
||||
EXTRACTED_TEXT = "extracted_text"
|
||||
|
||||
|
||||
class StepKind(StrEnum):
|
||||
"""Discriminator values for AgentSpecStep unions (kind field)."""
|
||||
|
||||
TOOL = "tool"
|
||||
AI_TOOL = "ai_tool"
|
||||
|
||||
|
||||
class SupportedCapability(StrEnum):
|
||||
ORCHESTRATE = "orchestrate"
|
||||
PDF_EDIT = "pdf_edit"
|
||||
PDF_QUESTION = "pdf_question"
|
||||
AGENT_DRAFT = "agent_draft"
|
||||
AGENT_REVISE = "agent_revise"
|
||||
AGENT_NEXT_ACTION = "agent_next_action"
|
||||
|
||||
|
||||
class ConversationMessage(ApiModel):
|
||||
role: str
|
||||
content: str
|
||||
@@ -17,8 +94,13 @@ class PdfTextSelection(ApiModel):
|
||||
text: str
|
||||
|
||||
|
||||
class ExtractedFileText(ApiModel):
|
||||
file_name: str
|
||||
pages: list[PdfTextSelection] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ToolOperationStep(ApiModel):
|
||||
kind: Literal["tool"] = "tool"
|
||||
kind: Literal[StepKind.TOOL] = StepKind.TOOL
|
||||
tool: OperationId
|
||||
parameters: ParamToolModel
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ from pydantic import Field
|
||||
from stirling.models import ApiModel, OperationId, ParamToolModel
|
||||
|
||||
from .agent_specs import AgentSpec
|
||||
from .common import WorkflowOutcome
|
||||
|
||||
|
||||
class ExecutionStepResult(ApiModel):
|
||||
@@ -31,19 +32,19 @@ class AgentExecutionRequest(ApiModel):
|
||||
|
||||
|
||||
class ToolCallExecutionAction(ApiModel):
|
||||
outcome: Literal["tool_call"] = "tool_call"
|
||||
outcome: Literal[WorkflowOutcome.TOOL_CALL] = WorkflowOutcome.TOOL_CALL
|
||||
tool: OperationId
|
||||
parameters: ParamToolModel
|
||||
rationale: str | None = None
|
||||
|
||||
|
||||
class CompletedExecutionAction(ApiModel):
|
||||
outcome: Literal["completed"] = "completed"
|
||||
outcome: Literal[WorkflowOutcome.COMPLETED] = WorkflowOutcome.COMPLETED
|
||||
summary: str
|
||||
|
||||
|
||||
class CannotContinueExecutionAction(ApiModel):
|
||||
outcome: Literal["cannot_continue"] = "cannot_continue"
|
||||
outcome: Literal[WorkflowOutcome.CANNOT_CONTINUE] = WorkflowOutcome.CANNOT_CONTINUE
|
||||
reason: str
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
@@ -8,27 +7,29 @@ from pydantic import Field
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .agent_drafts import AgentDraftResponse
|
||||
from .common import ArtifactKind, ExtractedFileText, SupportedCapability, WorkflowOutcome
|
||||
from .execution import NextExecutionAction
|
||||
from .pdf_edit import PdfEditResponse
|
||||
from .pdf_questions import PdfQuestionResponse
|
||||
|
||||
|
||||
class SupportedCapability(StrEnum):
|
||||
ORCHESTRATE = "orchestrate"
|
||||
PDF_EDIT = "pdf_edit"
|
||||
PDF_QUESTION = "pdf_question"
|
||||
AGENT_DRAFT = "agent_draft"
|
||||
AGENT_REVISE = "agent_revise"
|
||||
AGENT_NEXT_ACTION = "agent_next_action"
|
||||
class ExtractedTextArtifact(ApiModel):
|
||||
kind: Literal[ArtifactKind.EXTRACTED_TEXT] = ArtifactKind.EXTRACTED_TEXT
|
||||
files: list[ExtractedFileText] = Field(default_factory=list)
|
||||
|
||||
|
||||
WorkflowArtifact = Annotated[ExtractedTextArtifact, Field(discriminator="kind")]
|
||||
|
||||
|
||||
class OrchestratorRequest(ApiModel):
|
||||
user_message: str
|
||||
conversation_id: str | None = None
|
||||
file_names: list[str]
|
||||
artifacts: list[WorkflowArtifact] = Field(default_factory=list)
|
||||
resume_with: SupportedCapability | None = None
|
||||
|
||||
|
||||
class UnsupportedCapabilityResponse(ApiModel):
|
||||
outcome: Literal["unsupported_capability"] = "unsupported_capability"
|
||||
outcome: Literal[WorkflowOutcome.UNSUPPORTED_CAPABILITY] = WorkflowOutcome.UNSUPPORTED_CAPABILITY
|
||||
capability: str
|
||||
message: str
|
||||
|
||||
|
||||
@@ -6,30 +6,29 @@ from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .common import ToolOperationStep
|
||||
from .common import ToolOperationStep, WorkflowOutcome
|
||||
|
||||
|
||||
class PdfEditRequest(ApiModel):
|
||||
user_message: str
|
||||
conversation_id: str | None = None
|
||||
file_names: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class EditPlanResponse(ApiModel):
|
||||
outcome: Literal["plan"] = "plan"
|
||||
outcome: Literal[WorkflowOutcome.PLAN] = WorkflowOutcome.PLAN
|
||||
summary: str
|
||||
rationale: str | None = None
|
||||
steps: list[ToolOperationStep]
|
||||
|
||||
|
||||
class EditClarificationRequest(ApiModel):
|
||||
outcome: Literal["need_clarification"] = "need_clarification"
|
||||
outcome: Literal[WorkflowOutcome.NEED_CLARIFICATION] = WorkflowOutcome.NEED_CLARIFICATION
|
||||
question: str
|
||||
reason: str
|
||||
|
||||
|
||||
class EditCannotDoResponse(ApiModel):
|
||||
outcome: Literal["cannot_do"] = "cannot_do"
|
||||
outcome: Literal[WorkflowOutcome.CANNOT_DO] = WorkflowOutcome.CANNOT_DO
|
||||
reason: str
|
||||
|
||||
|
||||
|
||||
@@ -6,31 +6,42 @@ from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .common import ExtractedFileText, PdfContentType, SupportedCapability, WorkflowOutcome
|
||||
|
||||
|
||||
class PdfQuestionRequest(ApiModel):
|
||||
question: str
|
||||
conversation_id: str | None = None
|
||||
extracted_text: str = ""
|
||||
file_name: str | None = None
|
||||
page_text: list[ExtractedFileText] = Field(default_factory=list)
|
||||
file_names: list[str]
|
||||
|
||||
|
||||
class PdfQuestionAnswerResponse(ApiModel):
|
||||
outcome: Literal["answer"] = "answer"
|
||||
outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER
|
||||
answer: str
|
||||
evidence: list[str] = Field(default_factory=list)
|
||||
evidence: list[ExtractedFileText] = Field(default_factory=list)
|
||||
|
||||
|
||||
class PdfQuestionNeedTextResponse(ApiModel):
|
||||
outcome: Literal["need_text"] = "need_text"
|
||||
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["not_found"] = "not_found"
|
||||
outcome: Literal[WorkflowOutcome.NOT_FOUND] = WorkflowOutcome.NOT_FOUND
|
||||
reason: str
|
||||
|
||||
|
||||
PdfQuestionResponse = Annotated[
|
||||
PdfQuestionAnswerResponse | PdfQuestionNeedTextResponse | PdfQuestionNotFoundResponse,
|
||||
PdfQuestionAnswerResponse | PdfQuestionNeedContentResponse | PdfQuestionNotFoundResponse,
|
||||
Field(discriminator="outcome"),
|
||||
]
|
||||
|
||||
@@ -5,10 +5,12 @@ import pytest
|
||||
from stirling.agents import PdfQuestionAgent
|
||||
from stirling.config import AppSettings
|
||||
from stirling.contracts import (
|
||||
ExtractedFileText,
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNeedTextResponse,
|
||||
PdfQuestionNeedContentResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfTextSelection,
|
||||
)
|
||||
from stirling.services import build_runtime
|
||||
|
||||
@@ -34,13 +36,22 @@ def build_test_settings() -> AppSettings:
|
||||
)
|
||||
|
||||
|
||||
def invoice_page() -> ExtractedFileText:
|
||||
return ExtractedFileText(
|
||||
file_name="invoice.pdf",
|
||||
pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_question_agent_requires_extracted_text() -> None:
|
||||
agent = PdfQuestionAgent(build_runtime(build_test_settings()))
|
||||
|
||||
response = await agent.handle(PdfQuestionRequest(question="What is the total?", extracted_text=""))
|
||||
response = await agent.handle(
|
||||
PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"])
|
||||
)
|
||||
|
||||
assert isinstance(response, PdfQuestionNeedTextResponse)
|
||||
assert isinstance(response, PdfQuestionNeedContentResponse)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@@ -48,15 +59,15 @@ async def test_pdf_question_agent_returns_grounded_answer() -> None:
|
||||
agent = StubPdfQuestionAgent(
|
||||
PdfQuestionAnswerResponse(
|
||||
answer="The invoice total is 120.00.",
|
||||
evidence=["Invoice total: 120.00"],
|
||||
evidence=[invoice_page()],
|
||||
)
|
||||
)
|
||||
|
||||
response = await agent.handle(
|
||||
PdfQuestionRequest(
|
||||
question="What is the total?",
|
||||
extracted_text="Invoice total: 120.00",
|
||||
file_name="invoice.pdf",
|
||||
page_text=[invoice_page()],
|
||||
file_names=["invoice.pdf"],
|
||||
)
|
||||
)
|
||||
|
||||
@@ -71,8 +82,13 @@ async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient()
|
||||
response = await agent.handle(
|
||||
PdfQuestionRequest(
|
||||
question="What is the total?",
|
||||
extracted_text="This page contains only a shipping address.",
|
||||
file_name="invoice.pdf",
|
||||
page_text=[
|
||||
ExtractedFileText(
|
||||
file_name="invoice.pdf",
|
||||
pages=[PdfTextSelection(page_number=1, text="This page contains only a shipping address.")],
|
||||
)
|
||||
],
|
||||
file_names=["invoice.pdf"],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -20,9 +20,9 @@ from stirling.contracts import (
|
||||
EditCannotDoResponse,
|
||||
OrchestratorRequest,
|
||||
PdfEditRequest,
|
||||
PdfQuestionNeedContentResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
UnsupportedCapabilityResponse,
|
||||
)
|
||||
from stirling.models.tool_models import RotateParams
|
||||
|
||||
@@ -38,8 +38,8 @@ class StubSettingsProvider:
|
||||
|
||||
|
||||
class StubOrchestratorAgent:
|
||||
async def handle(self, request: OrchestratorRequest) -> UnsupportedCapabilityResponse:
|
||||
return UnsupportedCapabilityResponse(capability="pdf_edit", message=request.user_message)
|
||||
async def handle(self, request: OrchestratorRequest) -> PdfQuestionNeedContentResponse:
|
||||
return PdfQuestionNeedContentResponse(reason=request.user_message, files=[], max_pages=1, max_characters=1000)
|
||||
|
||||
|
||||
class StubPdfEditAgent:
|
||||
@@ -115,10 +115,10 @@ def test_health_route() -> None:
|
||||
|
||||
|
||||
def test_orchestrator_route() -> None:
|
||||
response = client.post("/api/v1/orchestrator", json={"userMessage": "route this"})
|
||||
response = client.post("/api/v1/orchestrator", json={"userMessage": "route this", "fileNames": ["test.pdf"]})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "unsupported_capability"
|
||||
assert response.json()["outcome"] == "need_content"
|
||||
|
||||
|
||||
def test_pdf_edit_route() -> None:
|
||||
@@ -129,7 +129,14 @@ def test_pdf_edit_route() -> None:
|
||||
|
||||
|
||||
def test_pdf_questions_route() -> None:
|
||||
response = client.post("/api/v1/pdf/questions", json={"question": "what is this?"})
|
||||
response = client.post(
|
||||
"/api/v1/pdf/questions",
|
||||
json={
|
||||
"question": "what is this?",
|
||||
"fileNames": ["test.pdf"],
|
||||
"pageText": [{"fileName": "test.pdf", "pages": [{"pageNumber": 1, "text": "Example"}]}],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["outcome"] == "not_found"
|
||||
|
||||
@@ -9,17 +9,34 @@ from stirling.contracts import (
|
||||
AgentSpecStep,
|
||||
EditPlanResponse,
|
||||
ExecutionContext,
|
||||
ExtractedFileText,
|
||||
ExtractedTextArtifact,
|
||||
OrchestratorRequest,
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfTextSelection,
|
||||
ToolOperationStep,
|
||||
)
|
||||
from stirling.models.tool_models import OperationId, RotateParams
|
||||
|
||||
|
||||
def test_orchestrator_request_accepts_user_message() -> None:
|
||||
request = OrchestratorRequest(user_message="Rotate the PDF")
|
||||
request = OrchestratorRequest(
|
||||
user_message="Rotate the PDF",
|
||||
file_names=["test.pdf"],
|
||||
artifacts=[
|
||||
ExtractedTextArtifact(
|
||||
files=[
|
||||
ExtractedFileText(
|
||||
file_name="test.pdf",
|
||||
pages=[PdfTextSelection(page_number=1, text="Hello")],
|
||||
)
|
||||
]
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
assert request.user_message == "Rotate the PDF"
|
||||
assert len(request.artifacts) == 1
|
||||
|
||||
|
||||
def test_agent_execution_request_uses_typed_agent_spec() -> None:
|
||||
|
||||
Reference in New Issue
Block a user