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:
James Brunton
2026-05-01 14:59:53 +00:00
committed by GitHub
co-authored by EthanHealy01
parent 5541dd666c
commit 51f5345151
13 changed files with 418 additions and 95 deletions
+72 -11
View File
@@ -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,
+12
View File
@@ -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):
+10 -2
View File
@@ -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):
+116 -5
View File
@@ -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 == []