mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
Inform AI engine which endpoints are disabled on the backend (#6251)
# Description of Changes Have the Java send a list of enabled endpoints to the AI engine so it can intelligently respond to the user that the tool does exist but is disabled on the server so it can't acutally run the operation, instead of the current behaviour where it sends the API call back and then 503 errors because the execution fails when the URL is disabled. <img width="380" height="208" alt="image" src="https://github.com/user-attachments/assets/5842fb2e-2e55-45a5-8205-25515636daae" /> --------- Co-authored-by: EthanHealy01 <[email protected]>
This commit is contained in:
co-authored by
EthanHealy01
parent
5541dd666c
commit
51f5345151
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Iterable
|
||||
from typing import Literal, overload
|
||||
|
||||
from pydantic import Field
|
||||
@@ -140,7 +141,6 @@ class PdfEditParameterSelector:
|
||||
class PdfEditAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.supported_operations = list(OPERATIONS)
|
||||
self.parameter_selector = PdfEditParameterSelector(runtime)
|
||||
|
||||
async def orchestrate(self, request: OrchestratorRequest) -> PdfEditResponse:
|
||||
@@ -156,6 +156,7 @@ class PdfEditAgent:
|
||||
files=request.files,
|
||||
conversation_history=request.conversation_history,
|
||||
page_text=extracted_text.files if extracted_text is not None else [],
|
||||
enabled_endpoints=request.enabled_endpoints,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -165,19 +166,37 @@ class PdfEditAgent:
|
||||
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",
|
||||
"[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s enabled=%s msg=%r",
|
||||
[file.name for file in request.files],
|
||||
has_page_text(request.page_text),
|
||||
allow_need_content,
|
||||
request.enabled_endpoints,
|
||||
request.user_message,
|
||||
)
|
||||
selection = await self._select_plan(request, allow_need_content=allow_need_content)
|
||||
supported_operations = self._get_supported_operations(request)
|
||||
unavailable_operations = self._get_unavailable_operations(supported_operations)
|
||||
if not supported_operations:
|
||||
return EditCannotDoResponse(reason="No PDF edit operations are available on this server.")
|
||||
selection = await self._select_plan(
|
||||
request, supported_operations, unavailable_operations, 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)
|
||||
enabled = set(supported_operations)
|
||||
unsupported = [op for op in selection.operations if op not in enabled]
|
||||
if unsupported:
|
||||
logger.warning("[pdf-edit] plan referenced unavailable operations: %s", [op.name for op in unsupported])
|
||||
return EditCannotDoResponse(
|
||||
reason=(
|
||||
"The following operations are not available on this server "
|
||||
"(either disabled by the administrator or not installed): "
|
||||
+ ", ".join(op.name for op in unsupported)
|
||||
)
|
||||
)
|
||||
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):
|
||||
@@ -202,18 +221,39 @@ class PdfEditAgent:
|
||||
async def _select_plan(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
supported_operations: Iterable[ToolEndpoint],
|
||||
unavailable_operations: Iterable[ToolEndpoint],
|
||||
*,
|
||||
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))
|
||||
agent = self._build_selection_agent(
|
||||
supported_operations, unavailable_operations, allow_need_content=can_request_content
|
||||
)
|
||||
return await agent.select(self._build_selection_prompt(request, supported_operations, unavailable_operations))
|
||||
|
||||
def _build_selection_agent(self, *, allow_need_content: bool) -> PdfEditSelectionAgent:
|
||||
def _build_selection_agent(
|
||||
self,
|
||||
supported_operations: Iterable[ToolEndpoint],
|
||||
unavailable_operations: Iterable[ToolEndpoint],
|
||||
*,
|
||||
allow_need_content: bool,
|
||||
) -> PdfEditSelectionAgent:
|
||||
unavailable_clause = (
|
||||
f" The following operations exist on this server but are NOT currently available "
|
||||
f"(disabled by the administrator or not installed in this build) and must NOT appear "
|
||||
f"in any plan: {self._get_operations_prompt(unavailable_operations)}. "
|
||||
"If the user asks for one of these, return cannot_do, name the operation, and explain "
|
||||
"that it exists but isn't available on this server."
|
||||
if unavailable_operations
|
||||
else ""
|
||||
)
|
||||
return PdfEditSelectionAgent(
|
||||
self.runtime,
|
||||
base_system_prompt=(
|
||||
"Plan PDF edit requests. "
|
||||
f"Supported operations are: {self._supported_operations_prompt()}. "
|
||||
f"Supported operations are: {self._get_operations_prompt(supported_operations)}."
|
||||
f"{unavailable_clause} "
|
||||
"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. "
|
||||
@@ -224,17 +264,38 @@ class PdfEditAgent:
|
||||
allow_need_content=allow_need_content,
|
||||
)
|
||||
|
||||
def _build_selection_prompt(self, request: PdfEditRequest) -> str:
|
||||
def _build_selection_prompt(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
supported_operations: Iterable[ToolEndpoint],
|
||||
unavailable_operations: Iterable[ToolEndpoint],
|
||||
) -> str:
|
||||
unavailable_line = (
|
||||
"Unavailable operations (exist but not currently usable): "
|
||||
f"{self._get_operations_prompt(unavailable_operations)}\n"
|
||||
if unavailable_operations
|
||||
else ""
|
||||
)
|
||||
return (
|
||||
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
|
||||
f"User request: {request.user_message}\n"
|
||||
f"Files: {format_file_names(request.files)}\n"
|
||||
f"Supported operations: {self._supported_operations_prompt()}\n"
|
||||
f"Supported operations: {self._get_operations_prompt(supported_operations)}\n"
|
||||
f"{unavailable_line}"
|
||||
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 _get_supported_operations(self, request: PdfEditRequest) -> Iterable[ToolEndpoint]:
|
||||
return request.enabled_endpoints
|
||||
|
||||
@staticmethod
|
||||
def _get_unavailable_operations(supported_operations: Iterable[ToolEndpoint]) -> Iterable[ToolEndpoint]:
|
||||
supported_set = set(supported_operations)
|
||||
return [op for op in OPERATIONS if op not in supported_set]
|
||||
|
||||
@staticmethod
|
||||
def _get_operations_prompt(operations: Iterable[ToolEndpoint]) -> str:
|
||||
return ", ".join(f"{op.name} ({op.value})" for op in operations)
|
||||
|
||||
def _fill_need_content_defaults(
|
||||
self,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from enum import StrEnum
|
||||
from typing import Literal, assert_never
|
||||
|
||||
@@ -197,3 +198,14 @@ class ToolOperationStep(ApiModel):
|
||||
actual_type = type(self.parameters).__name__
|
||||
raise ValueError(f"Parameters for tool {self.tool} must be {expected_type.__name__}, got {actual_type}.")
|
||||
return self
|
||||
|
||||
|
||||
def drop_unknown_tool_endpoints(value: Iterable[str | ToolEndpoint]) -> list[ToolEndpoint]:
|
||||
"""Coerce inbound endpoint identifiers into `ToolEndpoint` members, dropping unknowns.
|
||||
|
||||
Java sends the full set of endpoints it considers enabled. The engine and the Java
|
||||
backend may have version drift in either direction, so we silently drop anything we
|
||||
don't recognise rather than failing the request. Anything dropped simply doesn't
|
||||
appear as a supported tool to the planner.
|
||||
"""
|
||||
return [ToolEndpoint(item) for item in value if item in ToolEndpoint]
|
||||
|
||||
@@ -2,9 +2,9 @@ from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic import BeforeValidator, Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
from stirling.models import ApiModel, ToolEndpoint
|
||||
|
||||
from .agent_drafts import AgentDraftResponse
|
||||
from .common import (
|
||||
@@ -17,6 +17,7 @@ from .common import (
|
||||
SupportedCapability,
|
||||
ToolReportArtifact,
|
||||
WorkflowOutcome,
|
||||
drop_unknown_tool_endpoints,
|
||||
)
|
||||
from .execution import NextExecutionAction
|
||||
from .pdf_edit import PdfEditTerminalResponse
|
||||
@@ -37,6 +38,10 @@ class OrchestratorRequest(ApiModel):
|
||||
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
||||
artifacts: list[WorkflowArtifact] = Field(default_factory=list)
|
||||
resume_with: SupportedCapability | None = None
|
||||
# See `PdfEditRequest.enabled_endpoints`.
|
||||
enabled_endpoints: Annotated[list[ToolEndpoint], BeforeValidator(drop_unknown_tool_endpoints)] = Field(
|
||||
default_factory=list
|
||||
)
|
||||
|
||||
|
||||
class UnsupportedCapabilityResponse(ApiModel):
|
||||
|
||||
@@ -2,9 +2,9 @@ from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic import BeforeValidator, Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
from stirling.models import ApiModel, ToolEndpoint
|
||||
|
||||
from .common import (
|
||||
AiFile,
|
||||
@@ -14,6 +14,7 @@ from .common import (
|
||||
SupportedCapability,
|
||||
ToolOperationStep,
|
||||
WorkflowOutcome,
|
||||
drop_unknown_tool_endpoints,
|
||||
)
|
||||
|
||||
|
||||
@@ -22,6 +23,13 @@ class PdfEditRequest(ApiModel):
|
||||
files: list[AiFile] = Field(default_factory=list)
|
||||
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
||||
page_text: list[ExtractedFileText] = Field(default_factory=list)
|
||||
# The set of endpoints the Java backend considers usable. Unknown URLs are silently
|
||||
# dropped so the engine and Java can drift in either direction without breaking
|
||||
# validation. An empty list means no operations are available - the planner will
|
||||
# return `cannot_do`.
|
||||
enabled_endpoints: Annotated[list[ToolEndpoint], BeforeValidator(drop_unknown_tool_endpoints)] = Field(
|
||||
default_factory=list
|
||||
)
|
||||
|
||||
|
||||
class EditPlanResponse(ApiModel):
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
@@ -20,7 +21,7 @@ from stirling.contracts import (
|
||||
SupportedCapability,
|
||||
ToolOperationStep,
|
||||
)
|
||||
from stirling.models import FileId
|
||||
from stirling.models import OPERATIONS, FileId
|
||||
from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
|
||||
from stirling.services.runtime import AppRuntime
|
||||
|
||||
@@ -69,9 +70,19 @@ class StubPdfEditAgent(PdfEditAgent):
|
||||
if parameter_selector is not None:
|
||||
self.parameter_selector = parameter_selector
|
||||
|
||||
def _get_supported_operations(self, request: PdfEditRequest) -> Iterable[ToolEndpoint]:
|
||||
# Tests construct requests without `enabled_endpoints`; pretend everything is enabled
|
||||
# unless the test explicitly supplies an enabled set.
|
||||
if request.enabled_endpoints:
|
||||
return request.enabled_endpoints
|
||||
return OPERATIONS
|
||||
|
||||
async def _select_plan(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
supported_operations: Iterable[ToolEndpoint],
|
||||
unavailable_operations: Iterable[ToolEndpoint],
|
||||
*,
|
||||
allow_need_content: bool = True,
|
||||
) -> PdfEditPlanOutput:
|
||||
return self.selection
|
||||
@@ -207,14 +218,20 @@ async def test_pdf_edit_agent_builds_selection_agent_matching_content_availabili
|
||||
agent = PdfEditAgent(runtime)
|
||||
captured: list[bool] = []
|
||||
|
||||
def record(*, allow_need_content: bool) -> PdfEditSelectionAgent:
|
||||
def record(
|
||||
supported_operations: Iterable[ToolEndpoint],
|
||||
unavailable_operations: Iterable[ToolEndpoint],
|
||||
*,
|
||||
allow_need_content: bool,
|
||||
) -> PdfEditSelectionAgent:
|
||||
captured.append(allow_need_content)
|
||||
raise _StopSelectionError()
|
||||
|
||||
agent._build_selection_agent = record
|
||||
|
||||
supported = list(OPERATIONS)
|
||||
with pytest.raises(_StopSelectionError):
|
||||
await agent._select_plan(PdfEditRequest(user_message="Rotate."))
|
||||
await agent._select_plan(PdfEditRequest(user_message="Rotate."), supported, [])
|
||||
with pytest.raises(_StopSelectionError):
|
||||
await agent._select_plan(
|
||||
PdfEditRequest(
|
||||
@@ -225,10 +242,12 @@ async def test_pdf_edit_agent_builds_selection_agent_matching_content_availabili
|
||||
pages=[PdfTextSelection(page_number=1, text="content")],
|
||||
)
|
||||
],
|
||||
)
|
||||
),
|
||||
supported,
|
||||
[],
|
||||
)
|
||||
with pytest.raises(_StopSelectionError):
|
||||
await agent._select_plan(PdfEditRequest(user_message="Rotate."), allow_need_content=False)
|
||||
await agent._select_plan(PdfEditRequest(user_message="Rotate."), supported, [], allow_need_content=False)
|
||||
|
||||
assert captured == [True, False, False]
|
||||
|
||||
@@ -282,3 +301,95 @@ async def test_pdf_edit_agent_passes_page_text_to_parameter_selector(runtime: Ap
|
||||
)
|
||||
|
||||
assert parameter_selector.calls[0].request.page_text == page_text
|
||||
|
||||
|
||||
def test_pdf_edit_request_drops_unknown_enabled_urls() -> None:
|
||||
request = PdfEditRequest.model_validate(
|
||||
{
|
||||
"user_message": "ignore me",
|
||||
"enabled_endpoints": [
|
||||
ToolEndpoint.COMPRESS_PDF.value,
|
||||
"/api/v1/not-a-real/endpoint",
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
assert request.enabled_endpoints == [ToolEndpoint.COMPRESS_PDF]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_supported_operations_defaults_to_empty(
|
||||
runtime: AppRuntime,
|
||||
) -> None:
|
||||
agent = PdfEditAgent(runtime)
|
||||
supported = agent._get_supported_operations(PdfEditRequest(user_message="hi"))
|
||||
|
||||
assert list(supported) == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_supported_operations_uses_provided_list(
|
||||
runtime: AppRuntime,
|
||||
) -> None:
|
||||
agent = PdfEditAgent(runtime)
|
||||
request = PdfEditRequest(
|
||||
user_message="Compress this PDF.",
|
||||
enabled_endpoints=[ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF],
|
||||
)
|
||||
|
||||
supported = agent._get_supported_operations(request)
|
||||
|
||||
assert list(supported) == [ToolEndpoint.FLATTEN, ToolEndpoint.ROTATE_PDF]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_returns_cannot_do_when_no_operations_enabled(
|
||||
runtime: AppRuntime,
|
||||
) -> None:
|
||||
agent = PdfEditAgent(runtime)
|
||||
response = await agent.handle(PdfEditRequest(user_message="Do anything.", enabled_endpoints=[]))
|
||||
|
||||
assert isinstance(response, EditCannotDoResponse)
|
||||
|
||||
|
||||
def test_pdf_edit_selection_prompt_includes_unavailable_operations(runtime: AppRuntime) -> None:
|
||||
agent = PdfEditAgent(runtime)
|
||||
request = PdfEditRequest(
|
||||
user_message="Run OCR.",
|
||||
enabled_endpoints=[ToolEndpoint.FLATTEN],
|
||||
)
|
||||
supported = agent._get_supported_operations(request)
|
||||
unavailable = agent._get_unavailable_operations(supported)
|
||||
|
||||
prompt = agent._build_selection_prompt(request, supported, unavailable)
|
||||
|
||||
assert "Unavailable operations" in prompt
|
||||
assert "OCR_PDF" in prompt
|
||||
assert ToolEndpoint.OCR_PDF.value in prompt
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pdf_edit_agent_rejects_plan_referencing_unavailable_operations(
|
||||
runtime: AppRuntime,
|
||||
) -> None:
|
||||
parameter_selector = RecordingParameterSelector()
|
||||
agent = StubPdfEditAgent(
|
||||
runtime,
|
||||
PdfEditPlanSelection(
|
||||
operations=[ToolEndpoint.COMPRESS_PDF],
|
||||
summary="Compress.",
|
||||
),
|
||||
parameter_selector=parameter_selector,
|
||||
)
|
||||
|
||||
response = await agent.handle(
|
||||
PdfEditRequest(
|
||||
user_message="Compress this PDF.",
|
||||
enabled_endpoints=[ToolEndpoint.FLATTEN],
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(response, EditCannotDoResponse)
|
||||
assert "not available" in response.reason
|
||||
assert "COMPRESS_PDF" in response.reason
|
||||
assert parameter_selector.calls == []
|
||||
|
||||
Reference in New Issue
Block a user