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:
James Brunton
co-authored by
EthanHealy01
EthanHealy01
parent
5541dd666c
commit
51f5345151
@@ -68,6 +68,36 @@ public class EndpointConfiguration {
|
||||
return endpoint.startsWith("/") ? endpoint.substring(1) : endpoint;
|
||||
}
|
||||
|
||||
/**
|
||||
* Translate a full request URI like {@code /api/v1/general/remove-pages} into the endpoint key
|
||||
* used by this configuration ({@code remove-pages}). Convert endpoints are a special case -
|
||||
* {@code /api/v1/convert/pdf/img} is registered as {@code pdf-to-img}. Returns {@code null} if
|
||||
* the URI is not an {@code /api/v1/<group>/<endpoint>} path.
|
||||
*/
|
||||
public static String endpointKeyForUri(String uri) {
|
||||
if (uri == null || !uri.contains("/api/v1")) {
|
||||
return null;
|
||||
}
|
||||
String[] parts = uri.split("/");
|
||||
if (parts.length <= 4) {
|
||||
return null;
|
||||
}
|
||||
if ("convert".equals(parts[3]) && parts.length > 5) {
|
||||
return parts[4] + "-to-" + parts[5];
|
||||
}
|
||||
return parts[4];
|
||||
}
|
||||
|
||||
/**
|
||||
* Convenience wrapper around {@link #isEndpointEnabled(String)} that accepts a full request URI
|
||||
* and translates it to the endpoint key. Falls back to treating the URI itself as a key for
|
||||
* non-{@code /api/v1/...} paths so callers can pass arbitrary URIs.
|
||||
*/
|
||||
public boolean isEndpointEnabledForUri(String uri) {
|
||||
String key = endpointKeyForUri(uri);
|
||||
return isEndpointEnabled(key != null ? key : uri);
|
||||
}
|
||||
|
||||
public void enableEndpoint(String endpoint) {
|
||||
String normalized = normalizeEndpoint(endpoint);
|
||||
endpointStatuses.put(normalized, true);
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
package stirling.software.SPDF.config;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertNull;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
class EndpointConfigurationTest {
|
||||
|
||||
@Test
|
||||
void endpointKeyForUriExtractsSimpleKebab() {
|
||||
assertEquals(
|
||||
"remove-pages",
|
||||
EndpointConfiguration.endpointKeyForUri("/api/v1/general/remove-pages"));
|
||||
assertEquals(
|
||||
"compress-pdf",
|
||||
EndpointConfiguration.endpointKeyForUri("/api/v1/misc/compress-pdf"));
|
||||
assertEquals(
|
||||
"add-watermark",
|
||||
EndpointConfiguration.endpointKeyForUri("/api/v1/security/add-watermark"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void endpointKeyForUriComposesConvertEndpoints() {
|
||||
assertEquals(
|
||||
"pdf-to-img", EndpointConfiguration.endpointKeyForUri("/api/v1/convert/pdf/img"));
|
||||
assertEquals(
|
||||
"pdf-to-word", EndpointConfiguration.endpointKeyForUri("/api/v1/convert/pdf/word"));
|
||||
assertEquals(
|
||||
"html-to-pdf", EndpointConfiguration.endpointKeyForUri("/api/v1/convert/html/pdf"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void endpointKeyForUriReturnsNullForNonApiPaths() {
|
||||
assertNull(EndpointConfiguration.endpointKeyForUri(null));
|
||||
assertNull(EndpointConfiguration.endpointKeyForUri("/some-page"));
|
||||
assertNull(EndpointConfiguration.endpointKeyForUri("/api/v1/general"));
|
||||
}
|
||||
}
|
||||
@@ -21,28 +21,7 @@ public class EndpointInterceptor implements HandlerInterceptor {
|
||||
HttpServletRequest request, HttpServletResponse response, Object handler)
|
||||
throws Exception {
|
||||
String requestURI = request.getRequestURI();
|
||||
boolean isEnabled;
|
||||
|
||||
// Extract the specific endpoint name (e.g: /api/v1/general/remove-pages -> remove-pages)
|
||||
if (requestURI.contains("/api/v1") && requestURI.split("/").length > 4) {
|
||||
|
||||
String[] requestURIParts = requestURI.split("/");
|
||||
String requestEndpoint;
|
||||
|
||||
// Endpoint: /api/v1/convert/pdf/img becomes pdf-to-img
|
||||
if ("convert".equals(requestURIParts[3]) && requestURIParts.length > 5) {
|
||||
requestEndpoint = requestURIParts[4] + "-to-" + requestURIParts[5];
|
||||
} else {
|
||||
requestEndpoint = requestURIParts[4];
|
||||
}
|
||||
|
||||
log.debug("Request endpoint: {}", requestEndpoint);
|
||||
isEnabled = endpointConfiguration.isEndpointEnabled(requestEndpoint);
|
||||
log.debug("Is endpoint enabled: {}", isEnabled);
|
||||
} else {
|
||||
isEnabled = endpointConfiguration.isEndpointEnabled(requestURI);
|
||||
}
|
||||
|
||||
boolean isEnabled = endpointConfiguration.isEndpointEnabledForUri(requestURI);
|
||||
if (!isEnabled) {
|
||||
response.sendError(HttpServletResponse.SC_FORBIDDEN, "This endpoint is disabled");
|
||||
return false;
|
||||
|
||||
@@ -27,60 +27,19 @@ class EndpointInterceptorTest {
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleAllowsEnabledApiEndpoint() throws Exception {
|
||||
void preHandleAllowsEnabledEndpoint() throws Exception {
|
||||
when(request.getRequestURI()).thenReturn("/api/v1/general/remove-pages");
|
||||
when(endpointConfiguration.isEndpointEnabled("remove-pages")).thenReturn(true);
|
||||
when(endpointConfiguration.isEndpointEnabledForUri("/api/v1/general/remove-pages"))
|
||||
.thenReturn(true);
|
||||
assertTrue(interceptor.preHandle(request, response, new Object()));
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleBlocksDisabledApiEndpoint() throws Exception {
|
||||
void preHandleBlocksDisabledEndpoint() throws Exception {
|
||||
when(request.getRequestURI()).thenReturn("/api/v1/general/remove-pages");
|
||||
when(endpointConfiguration.isEndpointEnabled("remove-pages")).thenReturn(false);
|
||||
when(endpointConfiguration.isEndpointEnabledForUri("/api/v1/general/remove-pages"))
|
||||
.thenReturn(false);
|
||||
assertFalse(interceptor.preHandle(request, response, new Object()));
|
||||
verify(response).sendError(HttpServletResponse.SC_FORBIDDEN, "This endpoint is disabled");
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleExtractsConvertEndpointCorrectly() throws Exception {
|
||||
when(request.getRequestURI()).thenReturn("/api/v1/convert/pdf/img");
|
||||
when(endpointConfiguration.isEndpointEnabled("pdf-to-img")).thenReturn(true);
|
||||
assertTrue(interceptor.preHandle(request, response, new Object()));
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleBlocksDisabledConvertEndpoint() throws Exception {
|
||||
when(request.getRequestURI()).thenReturn("/api/v1/convert/pdf/img");
|
||||
when(endpointConfiguration.isEndpointEnabled("pdf-to-img")).thenReturn(false);
|
||||
assertFalse(interceptor.preHandle(request, response, new Object()));
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleUsesFullUriForNonApiPaths() throws Exception {
|
||||
when(request.getRequestURI()).thenReturn("/some-page");
|
||||
when(endpointConfiguration.isEndpointEnabled("/some-page")).thenReturn(true);
|
||||
assertTrue(interceptor.preHandle(request, response, new Object()));
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleBlocksDisabledNonApiPath() throws Exception {
|
||||
when(request.getRequestURI()).thenReturn("/some-page");
|
||||
when(endpointConfiguration.isEndpointEnabled("/some-page")).thenReturn(false);
|
||||
assertFalse(interceptor.preHandle(request, response, new Object()));
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleUsesFullUriForShortApiPath() throws Exception {
|
||||
// URI with /api/v1 but not enough segments (split length <= 4)
|
||||
when(request.getRequestURI()).thenReturn("/api/v1/general");
|
||||
when(endpointConfiguration.isEndpointEnabled("/api/v1/general")).thenReturn(true);
|
||||
assertTrue(interceptor.preHandle(request, response, new Object()));
|
||||
}
|
||||
|
||||
@Test
|
||||
void preHandleExtractsNonConvertApiEndpoint() throws Exception {
|
||||
when(request.getRequestURI()).thenReturn("/api/v1/security/add-watermark");
|
||||
when(endpointConfiguration.isEndpointEnabled("add-watermark")).thenReturn(true);
|
||||
assertTrue(interceptor.preHandle(request, response, new Object()));
|
||||
}
|
||||
}
|
||||
|
||||
+30
-5
@@ -34,11 +34,14 @@ import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
|
||||
import stirling.software.proprietary.service.AiEngineClient;
|
||||
import stirling.software.proprietary.service.AiEngineEndpointResolver;
|
||||
import stirling.software.proprietary.service.AiWorkflowService;
|
||||
|
||||
import tools.jackson.core.JacksonException;
|
||||
import tools.jackson.databind.JsonNode;
|
||||
import tools.jackson.databind.ObjectMapper;
|
||||
import tools.jackson.databind.node.ArrayNode;
|
||||
import tools.jackson.databind.node.ObjectNode;
|
||||
|
||||
@Slf4j
|
||||
@RestController
|
||||
@@ -53,6 +56,7 @@ public class AiEngineController {
|
||||
private final Executor aiStreamExecutor;
|
||||
private final TaskManager taskManager;
|
||||
private final JobOwnershipService jobOwnershipService;
|
||||
private final AiEngineEndpointResolver endpointResolver;
|
||||
|
||||
/**
|
||||
* SSE emitter timeout. Long enough to accommodate multi-gigabyte PDF workflows (OCR on a
|
||||
@@ -68,13 +72,15 @@ public class AiEngineController {
|
||||
ObjectMapper objectMapper,
|
||||
@Qualifier("aiStreamExecutor") Executor aiStreamExecutor,
|
||||
TaskManager taskManager,
|
||||
JobOwnershipService jobOwnershipService) {
|
||||
JobOwnershipService jobOwnershipService,
|
||||
AiEngineEndpointResolver endpointResolver) {
|
||||
this.aiEngineClient = aiEngineClient;
|
||||
this.aiWorkflowService = aiWorkflowService;
|
||||
this.objectMapper = objectMapper;
|
||||
this.aiStreamExecutor = aiStreamExecutor;
|
||||
this.taskManager = taskManager;
|
||||
this.jobOwnershipService = jobOwnershipService;
|
||||
this.endpointResolver = endpointResolver;
|
||||
}
|
||||
|
||||
@GetMapping("/health")
|
||||
@@ -197,17 +203,36 @@ public class AiEngineController {
|
||||
"Sends a user message to the PDF edit agent which returns a structured plan"
|
||||
+ " of tool operations to perform")
|
||||
public ResponseEntity<String> pdfEdit(@RequestBody String requestBody) throws IOException {
|
||||
validateJson(requestBody);
|
||||
String response = aiEngineClient.post("/api/v1/pdf/edit", requestBody);
|
||||
JsonNode parsed = parseJson(requestBody);
|
||||
if (!parsed.isObject()) {
|
||||
throw new ResponseStatusException(
|
||||
HttpStatus.BAD_REQUEST, "Request body must be a JSON object");
|
||||
}
|
||||
String forwardedBody = withEnabledEndpoints((ObjectNode) parsed);
|
||||
String response = aiEngineClient.post("/api/v1/pdf/edit", forwardedBody);
|
||||
return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response);
|
||||
}
|
||||
|
||||
private void validateJson(String body) {
|
||||
private JsonNode parseJson(String body) {
|
||||
try {
|
||||
objectMapper.readValue(body, JsonNode.class);
|
||||
return objectMapper.readValue(body, JsonNode.class);
|
||||
} catch (JacksonException e) {
|
||||
throw new ResponseStatusException(
|
||||
HttpStatus.BAD_REQUEST, "Request body is not valid JSON");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Always overwrite {@code enabled_endpoints} with the server's view of which endpoints are
|
||||
* usable. The engine must not trust a client-supplied list - the gate is owned by the Java
|
||||
* EndpointConfiguration. Values are full URL paths (e.g. {@code /api/v1/misc/compress-pdf})
|
||||
* that the engine matches against its {@code ToolEndpoint} enum, silently dropping any it
|
||||
* doesn't recognise (which lets the two sides drift in either direction without breaking).
|
||||
*/
|
||||
private String withEnabledEndpoints(ObjectNode body) {
|
||||
ArrayNode enabled = objectMapper.createArrayNode();
|
||||
endpointResolver.getEnabledEndpointUrls().forEach(enabled::add);
|
||||
body.set("enabled_endpoints", enabled);
|
||||
return body.toString();
|
||||
}
|
||||
}
|
||||
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
package stirling.software.proprietary.service;
|
||||
|
||||
import java.lang.reflect.Method;
|
||||
import java.util.HashSet;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
import java.util.TreeSet;
|
||||
|
||||
import org.springframework.context.ApplicationContext;
|
||||
import org.springframework.context.event.ContextRefreshedEvent;
|
||||
import org.springframework.context.event.EventListener;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;
|
||||
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.SPDF.config.EndpointConfiguration;
|
||||
|
||||
/**
|
||||
* Discovers every {@code /api/v1/...} request mapping in the application and exposes the subset
|
||||
* that {@link EndpointConfiguration} reports as currently enabled. The AI engine receives this list
|
||||
* as-is and silently drops anything it doesn't recognise, so we don't try to predict what the
|
||||
* engine considers a tool - we just emit what's enabled here.
|
||||
*/
|
||||
@Slf4j
|
||||
@Service
|
||||
public class AiEngineEndpointResolver {
|
||||
|
||||
private static final String API_PREFIX = "/api/v1/";
|
||||
|
||||
private final ApplicationContext applicationContext;
|
||||
private final EndpointConfiguration endpointConfiguration;
|
||||
// Written once on the Spring startup thread during ContextRefreshedEvent, read on HTTP
|
||||
// request threads. Spring's lifecycle establishes happens-before (the servlet container
|
||||
// and its worker threads are started after refresh completes), so no volatile is needed.
|
||||
private Set<String> apiUrls = Set.of();
|
||||
|
||||
public AiEngineEndpointResolver(
|
||||
ApplicationContext applicationContext, EndpointConfiguration endpointConfiguration) {
|
||||
this.applicationContext = applicationContext;
|
||||
this.endpointConfiguration = endpointConfiguration;
|
||||
}
|
||||
|
||||
@EventListener(ContextRefreshedEvent.class)
|
||||
public void discoverApiUrls() {
|
||||
Set<String> discovered = new TreeSet<>();
|
||||
for (RequestMappingHandlerMapping mapping :
|
||||
applicationContext.getBeansOfType(RequestMappingHandlerMapping.class).values()) {
|
||||
for (RequestMappingInfo info : mapping.getHandlerMethods().keySet()) {
|
||||
for (String pattern : extractPatterns(info)) {
|
||||
if (pattern.startsWith(API_PREFIX)) {
|
||||
discovered.add(pattern);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
apiUrls = Set.copyOf(discovered);
|
||||
log.debug("Discovered {} /api/v1/ endpoint URLs for AI engine filtering", apiUrls.size());
|
||||
}
|
||||
|
||||
public List<String> getEnabledEndpointUrls() {
|
||||
return apiUrls.stream()
|
||||
.filter(endpointConfiguration::isEndpointEnabledForUri)
|
||||
.sorted()
|
||||
.toList();
|
||||
}
|
||||
|
||||
private static Set<String> extractPatterns(RequestMappingInfo info) {
|
||||
try {
|
||||
Method getDirectPaths = info.getClass().getMethod("getDirectPaths");
|
||||
Object result = getDirectPaths.invoke(info);
|
||||
if (result instanceof Set<?> set) {
|
||||
Set<String> patterns = new HashSet<>();
|
||||
for (Object value : set) {
|
||||
if (value instanceof String s) {
|
||||
patterns.add(s);
|
||||
}
|
||||
}
|
||||
return patterns;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.trace("getDirectPaths unavailable on RequestMappingInfo", e);
|
||||
}
|
||||
return Set.of();
|
||||
}
|
||||
}
|
||||
+4
@@ -71,6 +71,7 @@ public class AiWorkflowService {
|
||||
private final ToolMetadataService toolMetadataService;
|
||||
private final TempFileManager tempFileManager;
|
||||
private final FileIdStrategy fileIdStrategy;
|
||||
private final AiEngineEndpointResolver endpointResolver;
|
||||
|
||||
@FunctionalInterface
|
||||
public interface ProgressListener {
|
||||
@@ -125,6 +126,7 @@ public class AiWorkflowService {
|
||||
request.getConversationHistory() == null
|
||||
? new ArrayList<>()
|
||||
: new ArrayList<>(request.getConversationHistory()));
|
||||
initialRequest.setEnabledEndpoints(endpointResolver.getEnabledEndpointUrls());
|
||||
|
||||
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.ANALYZING));
|
||||
|
||||
@@ -219,6 +221,7 @@ public class AiWorkflowService {
|
||||
nextRequest.setConversationHistory(request.getConversationHistory());
|
||||
nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults));
|
||||
nextRequest.setResumeWith(response.getResumeWith());
|
||||
nextRequest.setEnabledEndpoints(request.getEnabledEndpoints());
|
||||
return new WorkflowState.Pending(nextRequest);
|
||||
} finally {
|
||||
for (LoadedFile lf : loadedFiles) {
|
||||
@@ -624,5 +627,6 @@ public class AiWorkflowService {
|
||||
private List<AiConversationMessage> conversationHistory = new ArrayList<>();
|
||||
private List<WorkflowArtifact> artifacts = new ArrayList<>();
|
||||
private String resumeWith;
|
||||
private List<String> enabledEndpoints = new ArrayList<>();
|
||||
}
|
||||
}
|
||||
|
||||
+4
-1
@@ -79,6 +79,7 @@ class AiWorkflowServiceTest {
|
||||
@Mock private FileStorage fileStorage;
|
||||
@Mock private ToolMetadataService toolMetadataService;
|
||||
@Mock private FileIdStrategy fileIdStrategy;
|
||||
@Mock private AiEngineEndpointResolver endpointResolver;
|
||||
|
||||
@TempDir Path tempDir;
|
||||
|
||||
@@ -110,7 +111,9 @@ class AiWorkflowServiceTest {
|
||||
fileStorage,
|
||||
toolMetadataService,
|
||||
tempFileManager,
|
||||
fileIdStrategy);
|
||||
fileIdStrategy,
|
||||
endpointResolver);
|
||||
when(endpointResolver.getEnabledEndpointUrls()).thenReturn(List.of());
|
||||
}
|
||||
|
||||
@Test
|
||||
|
||||
@@ -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