From 8de94ff152b4fea4ba65014401cf04ec1e86a319 Mon Sep 17 00:00:00 2001 From: Anthony Stirling <77850077+Frooodle@users.noreply.github.com> Date: Thu, 23 Jul 2026 08:29:06 +0000 Subject: [PATCH] Ai customization settings (#7069) # Description of Changes AI settings customisation in settings menu, as part of this also tested and fixed ollama and other 3rd party AI integrations - Adds an admin AI settings UI for customizing AI behaviour, including per-provider model and API-key configuration - Backend pushes AI config changes to the Python engine at runtime via a config-push bridge, so changes apply without a restart - Config-push is gated off in SaaS; engine now drains background tasks on shutdown instead of cancelling them --- ## Checklist ### General - [ ] I have read the [Contribution Guidelines](https://github.com/Stirling-Tools/Stirling-PDF/blob/main/CONTRIBUTING.md) - [ ] I have read the [Stirling-PDF Developer Guide](https://github.com/Stirling-Tools/Stirling-PDF/blob/main/DeveloperGuide.md) (if applicable) - [ ] I have read the [How to add new languages to Stirling-PDF](https://github.com/Stirling-Tools/Stirling-PDF/blob/main/devGuide/HowToAddNewLanguage.md) (if applicable) - [ ] I have performed a self-review of my own code - [ ] My changes generate no new warnings ### Documentation - [ ] I have updated relevant docs on [Stirling-PDF's doc repo](https://github.com/Stirling-Tools/Stirling-Tools.github.io/blob/main/docs/) (if functionality has heavily changed) - [ ] I have read the section [Add New Translation Tags](https://github.com/Stirling-Tools/Stirling-PDF/blob/main/devGuide/HowToAddNewLanguage.md#add-new-translation-tags) (for new translation tags only) ### Translations (if applicable) - [ ] I ran [`scripts/counter_translation.py`](https://github.com/Stirling-Tools/Stirling-PDF/blob/main/docs/counter_translation.md) ### UI Changes (if applicable) - [ ] Screenshots or videos demonstrating the UI changes are attached (e.g., as comments or direct attachments in the PR) ### Testing (if applicable) - [ ] I have run `task check` to verify linters, typechecks, and tests pass - [ ] I have tested my changes locally. Refer to the [Testing Guide](https://github.com/Stirling-Tools/Stirling-PDF/blob/main/DeveloperGuide.md#7-testing) for more details. --- .../common/model/ApplicationProperties.java | 96 ++++ .../controller/api/misc/ConfigController.java | 14 +- .../src/main/resources/settings.yml.template | 31 ++ .../controller/api/AiEngineController.java | 21 +- .../api/ClassifyLabelController.java | 5 + .../api/CreatePdfAgentController.java | 4 +- .../api/MathAuditorAgentController.java | 3 + .../api/PdfCommentAgentController.java | 3 + .../api/AdminSettingsController.java | 106 +++- .../service/AiEngineConfigSync.java | 282 ++++++++++ .../proprietary/service/AiFeatureGate.java | 56 ++ .../api/ClassifyLabelControllerTest.java | 3 + .../api/PdfCommentAgentControllerTest.java | 27 +- .../api/AdminSettingsControllerTest.java | 102 +++- .../service/AiEngineConfigSyncTest.java | 251 +++++++++ .../service/AiFeatureGateTest.java | 77 +++ .../resources/application-saas.properties | 5 + engine/.env | 4 + engine/pyproject.toml | 1 + .../src/stirling/agents/math_presentation.py | 4 + engine/src/stirling/agents/orchestrator.py | 79 ++- engine/src/stirling/agents/output_mode.py | 28 + engine/src/stirling/agents/pdf_questions.py | 10 +- engine/src/stirling/api/app.py | 182 +++++-- engine/src/stirling/api/bootstrap.py | 72 +++ engine/src/stirling/api/routes/__init__.py | 2 + engine/src/stirling/api/routes/config.py | 250 +++++++++ engine/src/stirling/config/config_cache.py | 116 ++++ engine/src/stirling/config/settings.py | 12 + engine/src/stirling/contracts/__init__.py | 12 + engine/src/stirling/contracts/config.py | 69 +++ engine/src/stirling/documents/embedder.py | 29 +- engine/src/stirling/documents/service.py | 18 + engine/src/stirling/services/runtime.py | 144 ++++- .../tests/fixtures/processor_config_push.json | 24 + engine/tests/test_config_cache.py | 111 ++++ engine/tests/test_config_contract.py | 70 +++ engine/tests/test_config_routes.py | 383 +++++++++++++ engine/uv.lock | 2 + .../public/locales/en-US/translation.toml | 167 ++++++ frontend/editor/public/og-metadata.json | 24 + frontend/editor/public/og-metadata.saas.json | 24 + .../core/components/shared/config/types.ts | 4 + .../shared/config/configNavSections.tsx | 43 ++ .../AdminAiDocumentsSection.tsx | 383 +++++++++++++ .../configSections/AdminAiGeneralSection.tsx | 511 ++++++++++++++++++ .../configSections/AdminAiLimitsSection.tsx | 198 +++++++ .../configSections/AdminAiModelsSection.tsx | 391 ++++++++++++++ .../configSections/aiEngineSettings.test.ts | 57 ++ .../config/configSections/aiEngineSettings.ts | 101 ++++ .../config/saasConfigNavSections.test.tsx | 52 ++ 51 files changed, 4568 insertions(+), 95 deletions(-) create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineConfigSync.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/service/AiFeatureGate.java create mode 100644 app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineConfigSyncTest.java create mode 100644 app/proprietary/src/test/java/stirling/software/proprietary/service/AiFeatureGateTest.java create mode 100644 engine/src/stirling/agents/output_mode.py create mode 100644 engine/src/stirling/api/bootstrap.py create mode 100644 engine/src/stirling/api/routes/config.py create mode 100644 engine/src/stirling/config/config_cache.py create mode 100644 engine/src/stirling/contracts/config.py create mode 100644 engine/tests/fixtures/processor_config_push.json create mode 100644 engine/tests/test_config_cache.py create mode 100644 engine/tests/test_config_contract.py create mode 100644 engine/tests/test_config_routes.py create mode 100644 frontend/editor/src/proprietary/components/shared/config/configSections/AdminAiDocumentsSection.tsx create mode 100644 frontend/editor/src/proprietary/components/shared/config/configSections/AdminAiGeneralSection.tsx create mode 100644 frontend/editor/src/proprietary/components/shared/config/configSections/AdminAiLimitsSection.tsx create mode 100644 frontend/editor/src/proprietary/components/shared/config/configSections/AdminAiModelsSection.tsx create mode 100644 frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.test.ts create mode 100644 frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.ts create mode 100644 frontend/editor/src/saas/components/shared/config/saasConfigNavSections.test.tsx diff --git a/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java b/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java index b55fc381bb..2aa8413cf6 100644 --- a/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java +++ b/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java @@ -305,6 +305,102 @@ public class ApplicationProperties { * explicitly requests it via {@code AiEngineClient.postWithTimeout}. */ private int longRunningTimeoutSeconds = 600; + + /** Timeout (seconds) for the SSE stream held open by long-running orchestrator runs. */ + private int streamTimeoutSeconds = 1800; + + /** + * Whether the processor pushes settings-derived AI config to the engine on startup/save. + * Pin false for env-driven deployments (SaaS) to keep the engine env-controlled. + */ + private boolean pushConfigToEngine = true; + + /** Model + provider selection, forwarded to the engine per-request. */ + private Models models = new Models(); + + /** Retrieval-augmented-generation (RAG) knobs, forwarded to the engine per-request. */ + private Rag rag = new Rag(); + + /** Request size / cost guardrails. */ + private Limits limits = new Limits(); + + /** Per-capability on/off switches so an admin can disable individual AI tools. */ + private Features features = new Features(); + + @Data + public static class Models { + /** Provider driving the model strings: 'anthropic', 'openai', 'ollama', or 'custom'. */ + private String provider = "anthropic"; + + /** High-quality tier model name (without provider prefix), e.g. 'claude-haiku-4-5'. */ + private String smartModel = "claude-haiku-4-5"; + + /** Cheap/fast tier model name (without provider prefix). */ + private String fastModel = "claude-haiku-4-5"; + + private int smartMaxTokens = 8192; + private int fastMaxTokens = 2048; + + /** + * API key for the selected provider (secret; masked). Empty means the engine uses its + * own env credential (e.g. ANTHROPIC_API_KEY). + */ + private String apiKey = ""; + + /** + * OpenAI-compatible base URL for 'ollama' / 'custom' providers (e.g. + * http://ollama:11434/v1). Ignored for anthropic/openai. SSRF-sensitive - admin only. + */ + private String baseUrl = ""; + } + + @Data + public static class Rag { + /** + * Embedding provider: 'voyageai', 'openai', 'ollama', or 'custom' (OpenAI-compatible). + */ + private String embeddingProvider = "voyageai"; + + /** Embedding model name (without provider prefix), e.g. 'voyage-4'. */ + private String embeddingModel = "voyage-4"; + + /** + * Secret API key for the embedding provider; masked + env-overridable like + * models.apiKey. + */ + private String embeddingApiKey = ""; + + /** + * OpenAI-compatible base URL for 'ollama' / 'custom' embedding providers (e.g. + * http://ollama:11434/v1). Ignored for voyageai/openai. SSRF-sensitive - admin only. + */ + private String embeddingBaseUrl = ""; + + /** How many chunks retrieval returns per search. */ + private int topK = 20; + + /** Per-run cap on knowledge-search tool calls before the agent must answer. */ + private int maxSearches = 5; + } + + @Data + public static class Limits { + private int maxPages = 200; + private int maxCharacters = 200000; + + /** Process-wide cap on concurrent model API calls (engine restart to apply). */ + private int modelMaxConcurrency = 32; + } + + @Data + public static class Features { + private boolean chat = true; + private boolean documentQuestions = true; + private boolean createPdf = true; + private boolean mathAuditor = true; + private boolean pdfComment = true; + private boolean classify = true; + } } /** diff --git a/app/core/src/main/java/stirling/software/SPDF/controller/api/misc/ConfigController.java b/app/core/src/main/java/stirling/software/SPDF/controller/api/misc/ConfigController.java index 27477f9c2a..8044050025 100644 --- a/app/core/src/main/java/stirling/software/SPDF/controller/api/misc/ConfigController.java +++ b/app/core/src/main/java/stirling/software/SPDF/controller/api/misc/ConfigController.java @@ -336,7 +336,19 @@ public class ConfigController { configData.put("premiumEnabled", applicationProperties.getPremium().isEnabled()); // AI Engine settings - configData.put("aiEngineEnabled", applicationProperties.getAiEngine().isEnabled()); + ApplicationProperties.AiEngine aiEngineConfig = applicationProperties.getAiEngine(); + configData.put("aiEngineEnabled", aiEngineConfig.isEnabled()); + // Per-capability flags let the UI hide individual AI tools an admin has turned off. + ApplicationProperties.AiEngine.Features aiFeatures = aiEngineConfig.getFeatures(); + configData.put( + "aiFeatures", + Map.ofEntries( + Map.entry("chat", aiFeatures.isChat()), + Map.entry("documentQuestions", aiFeatures.isDocumentQuestions()), + Map.entry("createPdf", aiFeatures.isCreatePdf()), + Map.entry("mathAuditor", aiFeatures.isMathAuditor()), + Map.entry("pdfComment", aiFeatures.isPdfComment()), + Map.entry("classify", aiFeatures.isClassify()))); // Timestamp TSA settings — single source of truth for presets + admin URLs ApplicationProperties.Security.Timestamp tsConfig = diff --git a/app/core/src/main/resources/settings.yml.template b/app/core/src/main/resources/settings.yml.template index 85d97158b3..2d227c34a2 100644 --- a/app/core/src/main/resources/settings.yml.template +++ b/app/core/src/main/resources/settings.yml.template @@ -366,6 +366,35 @@ aiEngine: enabled: false # Set to 'true' to enable the AI engine integration url: http://localhost:5001 # URL of the Python AI engine timeoutSeconds: 120 # Timeout in seconds for AI engine requests + longRunningTimeoutSeconds: 600 # Timeout (seconds) for heavy operations like RAG ingestion of large documents + streamTimeoutSeconds: 1800 # SSE stream timeout (seconds) for long-running orchestrator runs + pushConfigToEngine: true # Push admin AI config to the engine on startup + save; false = engine stays fully env-controlled + models: + provider: anthropic # Model provider: 'anthropic', 'openai', 'ollama', or 'custom' (OpenAI-compatible) + smartModel: claude-haiku-4-5 # High-quality tier model name (no provider prefix) + fastModel: claude-haiku-4-5 # Cheap/fast tier model name (no provider prefix) + smartMaxTokens: 8192 # Max output tokens for the smart tier + fastMaxTokens: 2048 # Max output tokens for the fast tier + apiKey: "" # API key for the selected provider (secret). Empty = engine uses its native env credentials (e.g. ANTHROPIC_API_KEY) + baseUrl: "" # OpenAI-compatible base URL for 'ollama'/'custom' providers (e.g. http://ollama:11434/v1). Ignored for anthropic/openai + rag: + embeddingProvider: voyageai # Embedding provider: 'voyageai', 'openai', 'ollama', or 'custom' (OpenAI-compatible) + embeddingModel: voyage-4 # Embedding model name (no provider prefix) + embeddingApiKey: "" # Secret API key for the embedding provider. Empty = engine uses its native env credentials (e.g. VOYAGE_API_KEY) + embeddingBaseUrl: "" # OpenAI-compatible base URL for 'ollama'/'custom' embedding providers (e.g. http://ollama:11434/v1). Ignored for voyageai/openai + topK: 20 # Number of chunks retrieval returns per search + maxSearches: 5 # Per-run cap on knowledge-search tool calls before the agent must answer + limits: + maxPages: 200 # Upper bound on PDF pages the engine will process per request + maxCharacters: 200000 # Upper bound on characters of extracted text per request + modelMaxConcurrency: 32 # Process-wide cap on concurrent model API calls (engine restart to apply) + features: # Per-capability switches; turn an individual AI tool off without disabling the whole engine + chat: true # Assistant chat + documentQuestions: true # Ask-questions-about-a-PDF + createPdf: true # Generate a PDF from a natural-language spec + mathAuditor: true # Numerical/formula contradiction auditing + pdfComment: true # AI-authored PDF comments/annotations + classify: true # Automatic document classification/labelling policies: # Folder automations can read from and write to the directories you allow here, so treat this as a @@ -385,6 +414,8 @@ policies: mcp: enabled: false # Master switch. 'false' (default) means no /mcp endpoint, no metadata, no beans wired. scopesEnabled: true # Enforce mcp.tools.read / mcp.tools.write scopes derived from operation category + maxRequestBytes: 10485760 # Max size (bytes) of an incoming MCP tool request payload (default 10 MB) + maxInlineResponseBytes: 10485760 # Max size (bytes) of an MCP tool response returned inline before it is rejected (default 10 MB) allowedOperations: [] # Tool allow-list (operation ids, e.g. ['compress-pdf']). Empty = all. When set, ONLY these are exposed over MCP. blockedOperations: [] # Tool deny-list (operation ids). Always removed from MCP even if otherwise allowed. auth: diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java index a96b3d1240..f50d53e542 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java @@ -7,7 +7,6 @@ import java.util.concurrent.Executor; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; -import org.springframework.beans.factory.annotation.Value; import org.springframework.http.HttpStatus; import org.springframework.http.MediaType; import org.springframework.http.ResponseEntity; @@ -28,6 +27,7 @@ import jakarta.validation.Valid; import lombok.extern.slf4j.Slf4j; +import stirling.software.common.model.ApplicationProperties; import stirling.software.common.model.job.ResultFile; import stirling.software.common.service.JobOwnershipService; import stirling.software.common.service.TaskManager; @@ -38,6 +38,7 @@ 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.AiFeatureGate; import stirling.software.proprietary.service.AiWorkflowService; import tools.jackson.core.JacksonException; @@ -60,15 +61,14 @@ public class AiEngineController { private final TaskManager taskManager; private final JobOwnershipService jobOwnershipService; private final AiEngineEndpointResolver endpointResolver; + private final AiFeatureGate aiFeatureGate; private final UserServiceInterface userService; /** - * SSE emitter timeout. Long enough to accommodate multi-gigabyte PDF workflows (OCR on a - * 1000-page scan, splitting a huge PDF, etc.) without the emitter completing out from under the - * executor. Configurable via {@code stirling.ai.streamTimeoutMs}. + * SSE emitter timeout (ms), long enough for multi-gigabyte PDF workflows without completing out + * from under the executor. Derived from {@code aiEngine.streamTimeoutSeconds}. */ - @Value("${stirling.ai.streamTimeoutMs:1800000}") - private long streamTimeoutMs; + private final long streamTimeoutMs; public AiEngineController( AiEngineClient aiEngineClient, @@ -78,6 +78,8 @@ public class AiEngineController { TaskManager taskManager, JobOwnershipService jobOwnershipService, AiEngineEndpointResolver endpointResolver, + AiFeatureGate aiFeatureGate, + ApplicationProperties applicationProperties, @Autowired(required = false) UserServiceInterface userService) { this.aiEngineClient = aiEngineClient; this.aiWorkflowService = aiWorkflowService; @@ -86,7 +88,10 @@ public class AiEngineController { this.taskManager = taskManager; this.jobOwnershipService = jobOwnershipService; this.endpointResolver = endpointResolver; + this.aiFeatureGate = aiFeatureGate; this.userService = userService; + this.streamTimeoutMs = + applicationProperties.getAiEngine().getStreamTimeoutSeconds() * 1000L; } private String currentUserId() { @@ -111,6 +116,7 @@ public class AiEngineController { + " system and downloadable via GET /api/v1/general/files/{fileId}.") public AiWorkflowResponse orchestrate(@Valid @ModelAttribute AiWorkflowRequest request) throws IOException { + aiFeatureGate.requireConversationalWorkflow(); AiWorkflowResponse result = aiWorkflowService.orchestrate(request); registerFileResultAsJob(result); return result; @@ -123,6 +129,7 @@ public class AiEngineController { "Accepts a PDF upload and a user message, returns SSE events with progress" + " updates followed by the final AI workflow result") public SseEmitter orchestrateStream(@Valid @ModelAttribute AiWorkflowRequest request) { + aiFeatureGate.requireConversationalWorkflow(); SseEmitter emitter = new SseEmitter(streamTimeoutMs); emitter.onTimeout( @@ -246,6 +253,8 @@ public class AiEngineController { "Sends a user message to the PDF edit agent which returns a structured plan" + " of tool operations to perform") public ResponseEntity pdfEdit(@RequestBody String requestBody) throws IOException { + // Same gate as /orchestrate: edit agent is a model call on the same conversational surface. + aiFeatureGate.requireConversationalWorkflow(); JsonNode parsed = parseJson(requestBody); if (!parsed.isObject()) { throw new ResponseStatusException( diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java index 3b7f36037b..798de71df9 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/ClassifyLabelController.java @@ -35,6 +35,7 @@ import stirling.software.proprietary.classification.ClassificationLabelProvider; import stirling.software.proprietary.classification.model.ClassificationLabel; import stirling.software.proprietary.model.api.ai.AiPageText; import stirling.software.proprietary.service.AiEngineClient; +import stirling.software.proprietary.service.AiFeatureGate; import stirling.software.proprietary.service.PdfContentExtractor; import tools.jackson.databind.JsonNode; @@ -67,6 +68,7 @@ public class ClassifyLabelController { private final PdfContentExtractor pdfContentExtractor; private final PdfMetadataService pdfMetadataService; private final AiEngineClient aiEngineClient; + private final AiFeatureGate aiFeatureGate; private final ObjectMapper objectMapper; private final UserServiceInterface userService; @@ -81,6 +83,7 @@ public class ClassifyLabelController { PdfContentExtractor pdfContentExtractor, PdfMetadataService pdfMetadataService, AiEngineClient aiEngineClient, + AiFeatureGate aiFeatureGate, ObjectMapper objectMapper, ClassificationLabelProvider labelProvider, @Autowired(required = false) UserServiceInterface userService) { @@ -89,6 +92,7 @@ public class ClassifyLabelController { this.pdfContentExtractor = pdfContentExtractor; this.pdfMetadataService = pdfMetadataService; this.aiEngineClient = aiEngineClient; + this.aiFeatureGate = aiFeatureGate; this.objectMapper = objectMapper; this.labelProvider = labelProvider; this.userService = userService; @@ -104,6 +108,7 @@ public class ClassifyLabelController { + " intended for direct client use.") public ResponseEntity classifyAndLabel( @RequestParam("fileInput") MultipartFile fileInput) throws IOException { + aiFeatureGate.requireClassify(); try (PDDocument document = pdfDocumentFactory.load(fileInput, true)) { String fileName = safeFileName(fileInput.getOriginalFilename()); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/CreatePdfAgentController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/CreatePdfAgentController.java index 5e2298937a..0a41c552aa 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/CreatePdfAgentController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/CreatePdfAgentController.java @@ -34,6 +34,7 @@ import stirling.software.common.util.TempFileManager; import stirling.software.common.util.WebResponseUtils; import stirling.software.proprietary.model.api.ai.create.AiDocument; import stirling.software.proprietary.service.AiDocumentHtmlRenderer; +import stirling.software.proprietary.service.AiFeatureGate; import tools.jackson.core.JacksonException; import tools.jackson.databind.ObjectMapper; @@ -59,6 +60,7 @@ public class CreatePdfAgentController { private final ApplicationProperties applicationProperties; private final ObjectMapper objectMapper; private final AiDocumentHtmlRenderer htmlRenderer; + private final AiFeatureGate aiFeatureGate; /** * Returns true only when WeasyPrint is definitively unavailable — either the binary could not @@ -93,10 +95,10 @@ public class CreatePdfAgentController { public ResponseEntity createPdf( @RequestParam("document") String document, @RequestParam("filename") String filename) throws Exception { - if (!applicationProperties.getAiEngine().isEnabled()) { throw new ResponseStatusException(HttpStatus.NOT_FOUND); } + aiFeatureGate.requireCreatePdf(); AiDocument model; try { diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/MathAuditorAgentController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/MathAuditorAgentController.java index 3199901d4c..8beddd5e94 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/MathAuditorAgentController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/MathAuditorAgentController.java @@ -21,6 +21,7 @@ import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import stirling.software.proprietary.model.api.ai.Verdict; +import stirling.software.proprietary.service.AiFeatureGate; import stirling.software.proprietary.service.AiToolInputValidator; import stirling.software.proprietary.service.MathAuditorOrchestrator; @@ -49,6 +50,7 @@ public class MathAuditorAgentController { private static final Pattern NEWLINE_PATTERN = Pattern.compile("[\\r\\n]"); private final MathAuditorOrchestrator orchestrator; + private final AiFeatureGate aiFeatureGate; @PostMapping(value = "/math-auditor-agent", consumes = MediaType.MULTIPART_FORM_DATA_VALUE) @Operation( @@ -79,6 +81,7 @@ public class MathAuditorAgentController { + " ignored (default: 0.01)") @RequestParam(value = "tolerance", defaultValue = "0.01") BigDecimal tolerance) { + aiFeatureGate.requireMathAuditor(); AiToolInputValidator.validatePdfUpload(fileInput); if (tolerance.compareTo(BigDecimal.ZERO) < 0) { diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/PdfCommentAgentController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/PdfCommentAgentController.java index c02a804ce9..fadfacf03a 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/PdfCommentAgentController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/PdfCommentAgentController.java @@ -21,6 +21,7 @@ import io.swagger.v3.oas.annotations.tags.Tag; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; +import stirling.software.proprietary.service.AiFeatureGate; import stirling.software.proprietary.service.AiToolResponseHeaders; import stirling.software.proprietary.service.PdfCommentAgentOrchestrator; import stirling.software.proprietary.service.PdfCommentAgentOrchestrator.AnnotatedPdf; @@ -49,6 +50,7 @@ public class PdfCommentAgentController { private static final Pattern NEWLINE_PATTERN = Pattern.compile("[\\r\\n]"); private final PdfCommentAgentOrchestrator orchestrator; private final ObjectMapper objectMapper; + private final AiFeatureGate aiFeatureGate; @PostMapping( value = "/pdf-comment-agent", @@ -79,6 +81,7 @@ public class PdfCommentAgentController { @RequestParam("prompt") String prompt) throws IOException { + aiFeatureGate.requirePdfComment(); String originalFilename = fileInput.getOriginalFilename(); String safeName = diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AdminSettingsController.java b/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AdminSettingsController.java index 4bbe7f432a..286c0f1d01 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AdminSettingsController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AdminSettingsController.java @@ -8,6 +8,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -45,6 +46,7 @@ import stirling.software.common.util.RegexPatternUtils; import stirling.software.proprietary.security.model.api.admin.SettingValueResponse; import stirling.software.proprietary.security.model.api.admin.UpdateSettingValueRequest; import stirling.software.proprietary.security.model.api.admin.UpdateSettingsRequest; +import stirling.software.proprietary.service.AiEngineConfigSync; import tools.jackson.core.type.TypeReference; import tools.jackson.databind.ObjectMapper; @@ -58,6 +60,7 @@ public class AdminSettingsController { private final ApplicationProperties applicationProperties; private final ObjectMapper objectMapper; private final ApplicationContext applicationContext; + private final AiEngineConfigSync aiEngineConfigSync; // Track settings that have been modified but not yet applied (require restart) private static final ConcurrentHashMap pendingChanges = @@ -172,6 +175,26 @@ public class AdminSettingsController { .body(Map.of("error", "No settings provided to update")); } + // Mutable copy so we can drop masked "********" values: a UI round-trip must not + // overwrite a real secret (e.g. an API key) with the placeholder from the GET. + settings = new LinkedHashMap<>(settings); + settings.entrySet() + .removeIf( + e -> { + if (!"********".equals(e.getValue())) { + return false; + } + String key = e.getKey(); + String leaf = + key.contains(".") + ? key.substring(key.lastIndexOf('.') + 1) + : key; + return isSensitiveFieldWithPath(leaf, key); + }); + if (settings.isEmpty()) { + return ResponseEntity.ok(Map.of("message", "No changed settings to update.")); + } + // Validate all settings first before applying any changes for (Map.Entry entry : settings.entrySet()) { String key = entry.getKey(); @@ -188,6 +211,9 @@ public class AdminSettingsController { // Validate pipeline path settings String validationError = validatePipelinePathSetting(key, value); + if (validationError == null) { + validationError = validateAiEngineNumericSetting(key, value); + } if (validationError != null) { return ResponseEntity.badRequest() .body(Map.of("error", HtmlUtils.htmlEscape(validationError))); @@ -202,10 +228,13 @@ public class AdminSettingsController { for (Map.Entry entry : settings.entrySet()) { String key = entry.getKey(); Object value = entry.getValue(); - log.info("Admin updating setting: {} = {}", key, value); + log.info("Admin updating setting: {} = {}", key, logSafeValue(key, value)); pendingChanges.put(key, value != null ? value : ""); } + // Push changed AI settings live so model/RAG/limit changes skip the restart. + maybePushAiEngineLive(settings); + return ResponseEntity.ok( Map.of( "message", @@ -350,7 +379,10 @@ public class AdminSettingsController { + HtmlUtils.htmlEscape(fullKey))); } - log.info("Admin updating section setting: {} = {}", fullKey, value); + log.info( + "Admin updating section setting: {} = {}", + fullKey, + logSafeValue(fullKey, value)); GeneralUtils.saveKeyToSettings(fullKey, value); // Track this as a pending change @@ -469,7 +501,7 @@ public class AdminSettingsController { } } - log.info("Admin updating single setting: {} = {}", key, value); + log.info("Admin updating single setting: {} = {}", key, logSafeValue(key, value)); GeneralUtils.saveKeyToSettings(key, value); // Track this as a pending change @@ -600,6 +632,27 @@ public class AdminSettingsController { } } + /** + * Forward pending {@code aiEngine.*} changes to the engine after a save. Sends all accumulated + * pending changes, not just this save's: the running bean doesn't reflect unrestarted values. + */ + private void maybePushAiEngineLive(Map changedSettings) { + boolean aiChangedNow = + changedSettings.keySet().stream().anyMatch(k -> k.startsWith("aiEngine.")); + if (!aiChangedNow) { + return; + } + Map aiEnginePending = new HashMap<>(); + for (Map.Entry entry : pendingChanges.entrySet()) { + if (entry.getKey().startsWith("aiEngine.")) { + aiEnginePending.put(entry.getKey(), entry.getValue()); + } + } + if (!aiEnginePending.isEmpty()) { + aiEngineConfigSync.pushLiveAfterSave(aiEnginePending); + } + } + private Object getSectionData(String sectionName) { if (sectionName == null || sectionName.trim().isEmpty()) { return null; @@ -684,6 +737,38 @@ public class AdminSettingsController { return true; } + /** + * Minimum accepted value per bounded {@code aiEngine.*} numeric. A saved out-of-range value + * would make the engine reject every later push, including the one that fixes it. + */ + private static final Map AI_ENGINE_NUMERIC_MINIMUMS = + Map.of( + "aiEngine.models.smartMaxTokens", 1, + "aiEngine.models.fastMaxTokens", 1, + "aiEngine.rag.topK", 1, + "aiEngine.rag.maxSearches", 0, + "aiEngine.limits.maxPages", 1, + "aiEngine.limits.maxCharacters", 1, + "aiEngine.limits.modelMaxConcurrency", 1); + + private String validateAiEngineNumericSetting(String key, Object value) { + Integer min = AI_ENGINE_NUMERIC_MINIMUMS.get(key); + if (min == null || value == null) { + return null; + } + long parsed; + if (value instanceof Number number) { + parsed = number.longValue(); + } else { + try { + parsed = Long.parseLong(value.toString().trim()); + } catch (NumberFormatException e) { + return key + " must be a whole number"; + } + } + return parsed < min ? key + " must be at least " + min : null; + } + private String validatePipelinePathSetting(String key, Object value) { // Validate pipeline path settings if (key.startsWith("system.customPaths.pipeline.watchedFoldersDirs") @@ -828,6 +913,15 @@ public class AdminSettingsController { return masked; } + /** + * Value to log for a settings key, with secrets redacted: API keys, client secrets and mail + * passwords travel this path and must not land in the log in cleartext. + */ + private Object logSafeValue(String key, Object value) { + String leaf = key.contains(".") ? key.substring(key.lastIndexOf('.') + 1) : key; + return isSensitiveFieldWithPath(leaf, key) ? "" : value; + } + /** Check if a field name indicates sensitive data with full path context */ private boolean isSensitiveFieldWithPath(String fieldName, String fullPath) { String lowerField = fieldName.toLowerCase(); @@ -843,10 +937,12 @@ public class AdminSettingsController { return true; } - // Check for fields containing 'password' or 'secret' or 'bottoken' + // Match secret-bearing names (apikey covers provider creds). "token" is a suffix + // match only, so it doesn't swallow numeric fields like smartMaxTokens. return lowerField.contains("password") || lowerField.contains("secret") - || lowerField.contains("bottoken"); + || lowerField.contains("apikey") + || lowerField.endsWith("token"); } /** Create a masked representation for sensitive fields */ diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineConfigSync.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineConfigSync.java new file mode 100644 index 0000000000..99f5f025ef --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineConfigSync.java @@ -0,0 +1,282 @@ +package stirling.software.proprietary.service; + +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.stream.Collectors; + +import org.springframework.boot.context.event.ApplicationReadyEvent; +import org.springframework.context.event.EventListener; +import org.springframework.stereotype.Service; + +import jakarta.annotation.PreDestroy; + +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.model.ApplicationProperties; +import stirling.software.common.model.ApplicationProperties.AiEngine; + +import tools.jackson.databind.JsonNode; +import tools.jackson.databind.ObjectMapper; +import tools.jackson.databind.node.ObjectNode; + +/** + * Pushes admin-configured AI settings to the engine on startup and after each save; non-blocking + * and best-effort. Disabled via {@code aiEngine.pushConfigToEngine} for env-driven deployments. + */ +@Slf4j +@Service +@RequiredArgsConstructor +public class AiEngineConfigSync { + + private static final int MAX_ATTEMPTS = 5; + private static final long RETRY_DELAY_MS = 3000L; + + private final ApplicationProperties applicationProperties; + private final AiEngineClient aiEngineClient; + private final ObjectMapper objectMapper; + + // Single worker keeps pushes strictly ordered; virtual (daemon) thread never blocks shutdown. + private final ExecutorService pushExecutor = + Executors.newSingleThreadExecutor( + Thread.ofVirtual().name("ai-engine-config-sync").factory()); + + @PreDestroy + void shutdown() { + pushExecutor.shutdownNow(); + } + + @EventListener(ApplicationReadyEvent.class) + public void pushConfigOnStartup() { + AiEngine cfg = applicationProperties.getAiEngine(); + if (!cfg.isEnabled()) { + return; + } + if (!cfg.isPushConfigToEngine()) { + log.debug( + "Skipping AI engine config push: aiEngine.pushConfigToEngine is disabled" + + " (the engine is configured from its own environment)"); + return; + } + // Engine may still be booting; push off-thread with retries so startup never blocks. + submit(() -> pushWithRetries(cfg)); + } + + /** + * Push AI settings to the engine after an admin save so changes apply without a restart. No-op + * unless AI is enabled and an engine-relevant {@code aiEngine.*} key changed. + */ + public void pushLiveAfterSave(Map pendingAiEngine) { + // Save already persisted; a build/dispatch failure must not fail the save. + try { + // Gate on the running bean: the client refuses calls while disabled, so a pending + // enable would always fail here; the post-restart startup push covers first enablement. + AiEngine cfg = applicationProperties.getAiEngine(); + if (pendingAiEngine == null + || pendingAiEngine.isEmpty() + || !cfg.isPushConfigToEngine() + || !cfg.isEnabled()) { + return; + } + Set engineKeys = + pendingAiEngine.keySet().stream() + .filter(AiEngineConfigSync::isEngineRelevantKey) + .collect(Collectors.toSet()); + if (engineKeys.isEmpty()) { + return; + } + ObjectNode node = buildConfigNode(cfg); + pendingAiEngine.forEach((k, v) -> overlayIfEngineRelevant(node, k, v)); + keepEnvForUnconfiguredIdentity(node, engineKeys); + String body = node.toString(); + submit(() -> pushOnce(body)); + } catch (Exception e) { + log.warn( + "Could not build the live AI engine config push: {} (settings were saved; the" + + " engine will re-sync on the next restart)", + e.getMessage(), + e); + } + } + + /** + * Run pushes on the single-threaded executor so they stay serialised: each carries the full + * config and the engine keeps whatever lands last, so overlapping pushes could leave it stale. + */ + private void submit(Runnable task) { + pushExecutor.execute(task); + } + + private void pushOnce(String body) { + try { + aiEngineClient.post("/api/v1/config", body, null); + log.info("Pushed AI engine configuration after settings change"); + } catch (Exception e) { + log.error( + "Live AI engine config push failed: {}. The engine keeps running its previous" + + " configuration; if the engine is not on localhost, set" + + " STIRLING_ENGINE_SHARED_SECRET on both the engine and the processor" + + " so it accepts the push.", + e.getMessage()); + } + } + + // Only models/rag/limits reach the engine; the rest is processor-side. + private static boolean isEngineRelevantKey(String key) { + return key.startsWith("aiEngine.models.") + || key.startsWith("aiEngine.rag.") + || key.startsWith("aiEngine.limits."); + } + + private void overlayIfEngineRelevant(ObjectNode node, String key, Object value) { + if (!isEngineRelevantKey(key)) { + return; + } + String[] parts = key.substring("aiEngine.".length()).split("\\."); + if (parts.length < 2) { + // No leaf here; writing at parts[0] would overwrite the whole section with a scalar. + return; + } + ObjectNode parent = node; + for (int i = 0; i < parts.length - 1; i++) { + JsonNode child = parent.get(parts[i]); + parent = (child instanceof ObjectNode on) ? on : parent.putObject(parts[i]); + } + parent.set(parts[parts.length - 1], objectMapper.valueToTree(value)); + } + + private void pushWithRetries(AiEngine cfg) { + ObjectNode node = buildConfigNode(cfg); + keepEnvForUnconfiguredIdentity(node, Set.of()); + String body = node.toString(); + for (int attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) { + try { + aiEngineClient.post("/api/v1/config", body, null); + log.info("Pushed AI engine configuration on startup (attempt {})", attempt); + return; + } catch (Exception e) { + log.warn( + "AI engine config push failed (attempt {}/{}): {}", + attempt, + MAX_ATTEMPTS, + e.getMessage()); + if (attempt < MAX_ATTEMPTS) { + try { + Thread.sleep(RETRY_DELAY_MS); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + return; + } + } + } + } + log.warn( + "Giving up pushing AI engine configuration after {} attempts; the engine will use" + + " its own environment configuration until the next restart.", + MAX_ATTEMPTS); + } + + private ObjectNode buildConfigNode(AiEngine cfg) { + AiEngine.Models m = cfg.getModels(); + AiEngine.Rag r = cfg.getRag(); + AiEngine.Limits l = cfg.getLimits(); + + ObjectNode root = objectMapper.createObjectNode(); + + ObjectNode models = root.putObject("models"); + models.put("provider", m.getProvider()); + models.put("smartModel", m.getSmartModel()); + models.put("fastModel", m.getFastModel()); + models.put("smartMaxTokens", m.getSmartMaxTokens()); + models.put("fastMaxTokens", m.getFastMaxTokens()); + models.put("apiKey", m.getApiKey()); + models.put("baseUrl", m.getBaseUrl()); + + ObjectNode rag = root.putObject("rag"); + rag.put("embeddingProvider", r.getEmbeddingProvider()); + rag.put("embeddingModel", r.getEmbeddingModel()); + rag.put("embeddingApiKey", r.getEmbeddingApiKey()); + rag.put("embeddingBaseUrl", r.getEmbeddingBaseUrl()); + rag.put("topK", r.getTopK()); + rag.put("maxSearches", r.getMaxSearches()); + + ObjectNode limits = root.putObject("limits"); + limits.put("maxPages", l.getMaxPages()); + limits.put("maxCharacters", l.getMaxCharacters()); + limits.put("modelMaxConcurrency", l.getModelMaxConcurrency()); + + return root; + } + + // Defaults used to detect whether a section was configured or left at built-in values. + private static final AiEngine.Models DEFAULT_MODELS = new AiEngine.Models(); + private static final AiEngine.Rag DEFAULT_RAG = new AiEngine.Rag(); + + private static boolean isBlank(String s) { + return s == null || s.isBlank(); + } + + private static String text(JsonNode section, String field) { + return section.path(field).asText(""); + } + + /** + * Blank the identity (provider/model/credentials) of unconfigured sections so the push keeps + * the engine's env values; edited sections are sent as-is so a cleared key really clears. + */ + private void keepEnvForUnconfiguredIdentity(ObjectNode root, Set touchedKeys) { + if (root.get("models") instanceof ObjectNode models) { + boolean configured = + touchedIdentity(touchedKeys, MODEL_IDENTITY_KEYS) + || !isBlank(text(models, "apiKey")) + || !isBlank(text(models, "baseUrl")) + || !DEFAULT_MODELS.getProvider().equals(text(models, "provider")) + || !DEFAULT_MODELS.getSmartModel().equals(text(models, "smartModel")) + || !DEFAULT_MODELS.getFastModel().equals(text(models, "fastModel")); + if (!configured) { + models.put("provider", ""); + models.put("smartModel", ""); + models.put("fastModel", ""); + models.put("apiKey", ""); + models.put("baseUrl", ""); + } + } + if (root.get("rag") instanceof ObjectNode rag) { + boolean configured = + touchedIdentity(touchedKeys, RAG_IDENTITY_KEYS) + || !isBlank(text(rag, "embeddingApiKey")) + || !isBlank(text(rag, "embeddingBaseUrl")) + || !DEFAULT_RAG + .getEmbeddingProvider() + .equals(text(rag, "embeddingProvider")) + || !DEFAULT_RAG.getEmbeddingModel().equals(text(rag, "embeddingModel")); + if (!configured) { + rag.put("embeddingProvider", ""); + rag.put("embeddingModel", ""); + rag.put("embeddingApiKey", ""); + rag.put("embeddingBaseUrl", ""); + } + } + } + + private static final Set MODEL_IDENTITY_KEYS = + Set.of( + "aiEngine.models.provider", + "aiEngine.models.smartModel", + "aiEngine.models.fastModel", + "aiEngine.models.apiKey", + "aiEngine.models.baseUrl"); + + private static final Set RAG_IDENTITY_KEYS = + Set.of( + "aiEngine.rag.embeddingProvider", + "aiEngine.rag.embeddingModel", + "aiEngine.rag.embeddingApiKey", + "aiEngine.rag.embeddingBaseUrl"); + + private static boolean touchedIdentity(Set touchedKeys, Set identityKeys) { + return touchedKeys.stream().anyMatch(identityKeys::contains); + } +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiFeatureGate.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiFeatureGate.java new file mode 100644 index 0000000000..ffce8a1eb3 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiFeatureGate.java @@ -0,0 +1,56 @@ +package stirling.software.proprietary.service; + +import org.springframework.http.HttpStatus; +import org.springframework.stereotype.Component; +import org.springframework.web.server.ResponseStatusException; + +import lombok.RequiredArgsConstructor; + +import stirling.software.common.model.ApplicationProperties; +import stirling.software.common.model.ApplicationProperties.AiEngine.Features; + +/** + * Central gate for the AI feature switches ({@code aiEngine.features.*}); each {@code require*} + * throws 503 when the engine is disabled or the capability is off. + */ +@Component +@RequiredArgsConstructor +public class AiFeatureGate { + + private final ApplicationProperties applicationProperties; + + private Features features() { + return applicationProperties.getAiEngine().getFeatures(); + } + + private void require(boolean featureEnabled, String feature) { + if (!applicationProperties.getAiEngine().isEnabled() || !featureEnabled) { + throw new ResponseStatusException( + HttpStatus.SERVICE_UNAVAILABLE, "AI feature '" + feature + "' is disabled"); + } + } + + /** + * Shared entry point for chat and document questions; open while either is enabled, since a + * request can't be attributed to just one. No per-capability gate exists for the same reason. + */ + public void requireConversationalWorkflow() { + require(features().isChat() || features().isDocumentQuestions(), "conversation"); + } + + public void requireCreatePdf() { + require(features().isCreatePdf(), "createPdf"); + } + + public void requireMathAuditor() { + require(features().isMathAuditor(), "mathAuditor"); + } + + public void requirePdfComment() { + require(features().isPdfComment(), "pdfComment"); + } + + public void requireClassify() { + require(features().isClassify(), "classify"); + } +} diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java index 2773ed7965..74749289a7 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/ClassifyLabelControllerTest.java @@ -29,6 +29,7 @@ import stirling.software.common.util.TempFileManager; import stirling.software.proprietary.classification.ClassificationLabelProvider; import stirling.software.proprietary.classification.model.ClassificationLabel; import stirling.software.proprietary.service.AiEngineClient; +import stirling.software.proprietary.service.AiFeatureGate; import stirling.software.proprietary.service.PdfContentExtractor; import tools.jackson.databind.JsonNode; @@ -44,6 +45,7 @@ class ClassifyLabelControllerTest { @Mock private PdfContentExtractor pdfContentExtractor; @Mock private PdfMetadataService pdfMetadataService; @Mock private AiEngineClient aiEngineClient; + @Mock private AiFeatureGate aiFeatureGate; private final ObjectMapper objectMapper = JsonMapper.builder().build(); private ClassifyLabelController controller; @@ -56,6 +58,7 @@ class ClassifyLabelControllerTest { pdfContentExtractor, pdfMetadataService, aiEngineClient, + aiFeatureGate, objectMapper, ClassificationLabelProvider.withLabels(labels), null); diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/PdfCommentAgentControllerTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/PdfCommentAgentControllerTest.java index 42edace946..6be89e99cd 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/PdfCommentAgentControllerTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/controller/api/PdfCommentAgentControllerTest.java @@ -3,6 +3,7 @@ package stirling.software.proprietary.controller.api; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -26,6 +27,7 @@ import org.springframework.web.server.ResponseStatusException; import org.springframework.web.servlet.mvc.annotation.ResponseStatusExceptionResolver; import org.springframework.web.servlet.mvc.support.DefaultHandlerExceptionResolver; +import stirling.software.proprietary.service.AiFeatureGate; import stirling.software.proprietary.service.PdfCommentAgentOrchestrator; import stirling.software.proprietary.service.PdfCommentAgentOrchestrator.AnnotatedPdf; @@ -39,13 +41,15 @@ import tools.jackson.databind.json.JsonMapper; class PdfCommentAgentControllerTest { @Mock private PdfCommentAgentOrchestrator orchestrator; + @Mock private AiFeatureGate aiFeatureGate; private MockMvc mockMvc; @BeforeEach void setUp() { PdfCommentAgentController controller = - new PdfCommentAgentController(orchestrator, JsonMapper.builder().build()); + new PdfCommentAgentController( + orchestrator, JsonMapper.builder().build(), aiFeatureGate); mockMvc = MockMvcBuilders.standaloneSetup(controller) // standaloneSetup's defaults don't handle ResponseStatusException; wire up @@ -117,6 +121,27 @@ class PdfCommentAgentControllerTest { verify(orchestrator, never()).applyComments(any(), anyString()); } + @Test + void returnsServiceUnavailableWhenPdfCommentFeatureDisabled() throws Exception { + doThrow(new ResponseStatusException(HttpStatus.SERVICE_UNAVAILABLE)) + .when(aiFeatureGate) + .requirePdfComment(); + MockMultipartFile pdfFile = + new MockMultipartFile( + "fileInput", + "input.pdf", + MediaType.APPLICATION_PDF_VALUE, + "%PDF-1.4\n%%EOF".getBytes()); + + mockMvc.perform( + multipart("/api/v1/ai/tools/pdf-comment-agent") + .file(pdfFile) + .param("prompt", "flag dates")) + .andExpect(status().isServiceUnavailable()); + + verify(orchestrator, never()).applyComments(any(), anyString()); + } + @Test void rejectsMissingPromptParameter() throws Exception { MockMultipartFile pdfFile = diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AdminSettingsControllerTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AdminSettingsControllerTest.java index 5dc0e92963..4a4e297c0b 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AdminSettingsControllerTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AdminSettingsControllerTest.java @@ -1,10 +1,13 @@ package stirling.software.proprietary.security.controller.api; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; import java.io.IOException; import java.lang.reflect.Field; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -26,6 +29,7 @@ import stirling.software.common.util.GeneralUtils; import stirling.software.proprietary.security.model.api.admin.SettingValueResponse; import stirling.software.proprietary.security.model.api.admin.UpdateSettingValueRequest; import stirling.software.proprietary.security.model.api.admin.UpdateSettingsRequest; +import stirling.software.proprietary.service.AiEngineConfigSync; import tools.jackson.databind.ObjectMapper; import tools.jackson.databind.json.JsonMapper; @@ -37,6 +41,7 @@ class AdminSettingsControllerTest { private ApplicationProperties applicationProperties; private ObjectMapper objectMapper; private ApplicationContext applicationContext; + private AiEngineConfigSync aiEngineConfigSync; private AdminSettingsController controller; @@ -45,9 +50,13 @@ class AdminSettingsControllerTest { applicationProperties = new ApplicationProperties(); objectMapper = JsonMapper.builder().build(); applicationContext = org.mockito.Mockito.mock(ApplicationContext.class); + aiEngineConfigSync = org.mockito.Mockito.mock(AiEngineConfigSync.class); controller = new AdminSettingsController( - applicationProperties, objectMapper, applicationContext); + applicationProperties, + objectMapper, + applicationContext, + aiEngineConfigSync); clearPendingChanges(); } @@ -189,6 +198,33 @@ class AdminSettingsControllerTest { assertThat(response.getBody().get("error").toString()).contains("Invalid setting key"); } + @Test + @DisplayName("rejects an out-of-range aiEngine numeric with 400") + void rejectsOutOfRangeAiEngineNumeric() { + // An out-of-range value would make the engine reject every later push, including the + // one that fixes it. + UpdateSettingsRequest request = new UpdateSettingsRequest(); + request.setSettings(Map.of("aiEngine.limits.modelMaxConcurrency", 0)); + + ResponseEntity> response = controller.updateSettings(request); + + assertThat(response.getStatusCode()).isEqualTo(HttpStatus.BAD_REQUEST); + assertThat(response.getBody().get("error").toString()).contains("at least 1"); + } + + @Test + @DisplayName("accepts zero maxSearches, which legitimately means no retrieval") + void acceptsZeroMaxSearches() { + UpdateSettingsRequest request = new UpdateSettingsRequest(); + request.setSettings(Map.of("aiEngine.rag.maxSearches", 0)); + + try (MockedStatic mocked = mockStatic(GeneralUtils.class)) { + ResponseEntity> response = controller.updateSettings(request); + + assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK); + } + } + @Test @DisplayName("rejects unknown section prefix with 400") void rejectsUnknownSection() { @@ -277,6 +313,51 @@ class AdminSettingsControllerTest { assertThat(response.getStatusCode()).isEqualTo(HttpStatus.BAD_REQUEST); } } + + @Test + @DisplayName("drops a masked ******** secret so a UI round-trip can't overwrite a real key") + void dropsMaskedSecretValue() { + UpdateSettingsRequest request = new UpdateSettingsRequest(); + Map settings = new HashMap<>(); + settings.put("aiEngine.models.apiKey", "********"); + settings.put("ui.appName", "My App"); + request.setSettings(settings); + + try (MockedStatic mocked = mockStatic(GeneralUtils.class)) { + ResponseEntity> response = controller.updateSettings(request); + + assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK); + // The masked secret is stripped; only the real change is persisted. + mocked.verify( + () -> + GeneralUtils.updateSettingsTransactional( + argThat( + (Map m) -> + !m.containsKey("aiEngine.models.apiKey") + && m.containsKey("ui.appName")))); + } + } + + @Test + @DisplayName("forwards only aiEngine.* pending keys to the engine live-push") + void forwardsOnlyAiEngineKeysToLivePush() { + UpdateSettingsRequest request = new UpdateSettingsRequest(); + Map settings = new HashMap<>(); + settings.put("aiEngine.models.provider", "ollama"); + settings.put("ui.appName", "My App"); + request.setSettings(settings); + + try (MockedStatic mocked = mockStatic(GeneralUtils.class)) { + controller.updateSettings(request); + + verify(aiEngineConfigSync) + .pushLiveAfterSave( + argThat( + (Map m) -> + m.containsKey("aiEngine.models.provider") + && !m.containsKey("ui.appName"))); + } + } } @Nested @@ -430,6 +511,25 @@ class AdminSettingsControllerTest { SettingValueResponse body = (SettingValueResponse) response.getBody(); assertThat(body.getValue()).isEqualTo("********"); } + + @Test + @DisplayName("masks aiEngine apiKey but NOT the maxTokens numeric fields") + void masksApiKeyButNotMaxTokens() { + applicationProperties.getAiEngine().getModels().setApiKey("sk-real-key"); + applicationProperties.getAiEngine().getModels().setSmartMaxTokens(8192); + + // "token" as a substring of maxTokens must not trigger masking (would break the UI + // and flip the integer to a "********" string). + ResponseEntity tokensResp = + controller.getSettingValue("aiEngine.models.smartMaxTokens"); + SettingValueResponse tokens = (SettingValueResponse) tokensResp.getBody(); + assertThat(tokens.getValue()).isEqualTo(8192); + + // The real credential is still masked. + ResponseEntity keyResp = controller.getSettingValue("aiEngine.models.apiKey"); + SettingValueResponse key = (SettingValueResponse) keyResp.getBody(); + assertThat(key.getValue()).isEqualTo("********"); + } } @Nested diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineConfigSyncTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineConfigSyncTest.java new file mode 100644 index 0000000000..41f35c9daf --- /dev/null +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineConfigSyncTest.java @@ -0,0 +1,251 @@ +package stirling.software.proprietary.service; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; + +import java.util.HashMap; +import java.util.Map; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import stirling.software.common.model.ApplicationProperties; +import stirling.software.common.model.ApplicationProperties.AiEngine; + +import tools.jackson.databind.JsonNode; +import tools.jackson.databind.ObjectMapper; +import tools.jackson.databind.json.JsonMapper; + +/** + * Config-push bridge is self-hosted-only. These lock in that the processor stays silent when {@code + * aiEngine.pushConfigToEngine} is off (env-driven/SaaS) and pushes when it is on. + */ +class AiEngineConfigSyncTest { + + private ApplicationProperties applicationProperties; + private AiEngineClient aiEngineClient; + private AiEngineConfigSync sync; + + @BeforeEach + void setUp() { + applicationProperties = new ApplicationProperties(); + applicationProperties.getAiEngine().setEnabled(true); + applicationProperties.getAiEngine().setPushConfigToEngine(true); + aiEngineClient = mock(AiEngineClient.class); + ObjectMapper objectMapper = JsonMapper.builder().build(); + sync = new AiEngineConfigSync(applicationProperties, aiEngineClient, objectMapper); + } + + @Test + void startupPushSkippedWhenPushDisabled() throws Exception { + applicationProperties.getAiEngine().setPushConfigToEngine(false); + + sync.pushConfigOnStartup(); + + // Returns synchronously before spawning the push thread, so no interaction ever happens. + verify(aiEngineClient, never()).post(anyString(), anyString(), isNull()); + } + + @Test + void startupPushSkippedWhenDisabled() throws Exception { + applicationProperties.getAiEngine().setEnabled(false); + + sync.pushConfigOnStartup(); + + verify(aiEngineClient, never()).post(anyString(), anyString(), isNull()); + } + + @Test + void startupPushSentWhenEnabledAndPushOn() throws Exception { + sync.pushConfigOnStartup(); + + // Push runs on a virtual thread; wait for the single POST to /api/v1/config. + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), anyString(), isNull()); + } + + @Test + void livePushSkippedWhenPushDisabled() throws Exception { + applicationProperties.getAiEngine().setPushConfigToEngine(false); + + sync.pushLiveAfterSave(Map.of("aiEngine.models.provider", "ollama")); + + verify(aiEngineClient, never()).post(anyString(), anyString(), isNull()); + } + + @Test + void livePushSentForEngineRelevantChangeWhenPushOn() throws Exception { + sync.pushLiveAfterSave(Map.of("aiEngine.models.provider", "ollama")); + + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), anyString(), isNull()); + } + + @Test + void startupPushSerialisesTheEngineWireContract() throws Exception { + // Distinct values so a dropped/renamed field is detectable. Keep in sync with + // engine/tests/fixtures/processor_config_push.json (the engine validates the same shape). + AiEngine ai = applicationProperties.getAiEngine(); + ai.getModels().setProvider("ollama"); + ai.getModels().setSmartModel("smart-model-x"); + ai.getModels().setFastModel("fast-model-x"); + ai.getModels().setSmartMaxTokens(1111); + ai.getModels().setFastMaxTokens(2222); + ai.getModels().setApiKey("provider-key-abc"); + ai.getModels().setBaseUrl("http://engine.example/v1"); + ai.getRag().setEmbeddingProvider("custom"); + ai.getRag().setEmbeddingModel("embed-model-x"); + ai.getRag().setEmbeddingApiKey("embed-key-abc"); + ai.getRag().setEmbeddingBaseUrl("http://embed.example/v1"); + ai.getRag().setTopK(33); + ai.getRag().setMaxSearches(7); + ai.getLimits().setMaxPages(111); + ai.getLimits().setMaxCharacters(222222); + ai.getLimits().setModelMaxConcurrency(9); + + ArgumentCaptor body = ArgumentCaptor.forClass(String.class); + sync.pushConfigOnStartup(); + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), body.capture(), isNull()); + + JsonNode root = JsonMapper.builder().build().readTree(body.getValue()); + + JsonNode models = root.get("models"); + assertEquals("ollama", models.get("provider").asText()); + assertEquals("smart-model-x", models.get("smartModel").asText()); + assertEquals("fast-model-x", models.get("fastModel").asText()); + assertEquals(1111, models.get("smartMaxTokens").asInt()); + assertEquals(2222, models.get("fastMaxTokens").asInt()); + assertEquals("provider-key-abc", models.get("apiKey").asText()); + assertEquals("http://engine.example/v1", models.get("baseUrl").asText()); + + JsonNode rag = root.get("rag"); + assertEquals("custom", rag.get("embeddingProvider").asText()); + assertEquals("embed-model-x", rag.get("embeddingModel").asText()); + assertEquals("embed-key-abc", rag.get("embeddingApiKey").asText()); + assertEquals("http://embed.example/v1", rag.get("embeddingBaseUrl").asText()); + assertEquals(33, rag.get("topK").asInt()); + assertEquals(7, rag.get("maxSearches").asInt()); + + JsonNode limits = root.get("limits"); + assertEquals(111, limits.get("maxPages").asInt()); + assertEquals(222222, limits.get("maxCharacters").asInt()); + assertEquals(9, limits.get("modelMaxConcurrency").asInt()); + } + + @Test + void livePushSkippedForNonEngineRelevantChange() throws Exception { + // features.* is processor-side only; no engine push is warranted. + sync.pushLiveAfterSave(Map.of("aiEngine.features.chat", false)); + + verify(aiEngineClient, never()).post(anyString(), anyString(), isNull()); + } + + @Test + void startupPushKeepsEnvWhenSectionsUnconfigured() throws Exception { + // All defaults, no credentials: the push must NOT override the engine's env-configured + // provider/model/embedder, so the identity fields are blanked ("keep env" on the engine). + ArgumentCaptor body = ArgumentCaptor.forClass(String.class); + sync.pushConfigOnStartup(); + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), body.capture(), isNull()); + + JsonNode root = JsonMapper.builder().build().readTree(body.getValue()); + assertEquals("", root.get("models").get("provider").asText()); + assertEquals("", root.get("models").get("smartModel").asText()); + assertEquals("", root.get("models").get("fastModel").asText()); + assertEquals("", root.get("rag").get("embeddingProvider").asText()); + assertEquals("", root.get("rag").get("embeddingModel").asText()); + } + + @Test + void startupPushSendsProviderWhenChangedFromDefault() throws Exception { + // Admin selected a non-default provider (relying on the engine's env key): send it so the + // engine actually switches provider, even though no API key was entered in the UI. + applicationProperties.getAiEngine().getModels().setProvider("openai"); + + ArgumentCaptor body = ArgumentCaptor.forClass(String.class); + sync.pushConfigOnStartup(); + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), body.capture(), isNull()); + + JsonNode models = JsonMapper.builder().build().readTree(body.getValue()).get("models"); + assertEquals("openai", models.get("provider").asText()); + } + + @Test + void startupPushSendsModelsWhenApiKeyConfigured() throws Exception { + // A configured key means the admin is driving models from the UI: send the full section. + applicationProperties.getAiEngine().getModels().setApiKey("sk-real-key"); + + ArgumentCaptor body = ArgumentCaptor.forClass(String.class); + sync.pushConfigOnStartup(); + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), body.capture(), isNull()); + + JsonNode models = JsonMapper.builder().build().readTree(body.getValue()).get("models"); + assertEquals("anthropic", models.get("provider").asText()); + assertEquals("sk-real-key", models.get("apiKey").asText()); + } + + @Test + void livePushSendsAnExplicitlyClearedApiKeyRatherThanKeepEnv() throws Exception { + // Clearing a leaked key must reach the engine as a real clear; blanking it as "keep env" + // would leave the revoked key live in the engine's cache indefinitely. + ArgumentCaptor body = ArgumentCaptor.forClass(String.class); + sync.pushLiveAfterSave(mapOf("aiEngine.models.apiKey", "")); + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), body.capture(), isNull()); + + JsonNode models = JsonMapper.builder().build().readTree(body.getValue()).get("models"); + assertEquals("", models.get("apiKey").asText()); + // Identity was NOT blanked wholesale: the provider still travels so the engine applies + // the cleared credential against the right provider. + assertEquals("anthropic", models.get("provider").asText()); + } + + @Test + void livePushKeepsEnvWhenOnlyANumericKnobChanged() throws Exception { + // The admin touched a limit, not the identity, so the engine's env-configured + // provider/model must be preserved. + ArgumentCaptor body = ArgumentCaptor.forClass(String.class); + sync.pushLiveAfterSave(Map.of("aiEngine.limits.maxPages", 42)); + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), body.capture(), isNull()); + + JsonNode root = JsonMapper.builder().build().readTree(body.getValue()); + assertEquals("", root.get("models").get("provider").asText()); + assertEquals("", root.get("models").get("apiKey").asText()); + assertEquals(42, root.get("limits").get("maxPages").asInt()); + } + + @Test + void livePushIgnoresAMalformedKeyInsteadOfClobberingTheSection() throws Exception { + // "aiEngine.models." has no leaf; writing at the section name would replace the whole + // models object with a scalar and produce an unparseable push. + ArgumentCaptor body = ArgumentCaptor.forClass(String.class); + sync.pushLiveAfterSave(mapOf("aiEngine.models.", "junk")); + verify(aiEngineClient, timeout(3000)).post(eq("/api/v1/config"), body.capture(), isNull()); + + JsonNode models = JsonMapper.builder().build().readTree(body.getValue()).get("models"); + assertTrue(models.isObject(), "models must still be an object"); + } + + @Test + void livePushNeverThrowsIntoTheCaller() throws Exception { + // The caller has already persisted settings.yml, so a push-building failure must not + // surface as a failed save. A null value inside the map is enough to break naive code. + Map pending = new HashMap<>(); + pending.put("aiEngine.models.provider", null); + + assertDoesNotThrow(() -> sync.pushLiveAfterSave(pending)); + } + + /** {@link Map#of} rejects nulls and we need entries with empty/odd values. */ + private static Map mapOf(String key, Object value) { + Map map = new HashMap<>(); + map.put(key, value); + return map; + } +} diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiFeatureGateTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiFeatureGateTest.java new file mode 100644 index 0000000000..932a56296e --- /dev/null +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiFeatureGateTest.java @@ -0,0 +1,77 @@ +package stirling.software.proprietary.service; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.http.HttpStatus; +import org.springframework.web.server.ResponseStatusException; + +import stirling.software.common.model.ApplicationProperties; + +/** + * Lock in fail-closed gating: a 503 when the engine is disabled or the feature flag is off, and a + * clean pass only when both are on. + */ +class AiFeatureGateTest { + + private ApplicationProperties props; + private AiFeatureGate gate; + + @BeforeEach + void setUp() { + props = new ApplicationProperties(); + props.getAiEngine().setEnabled(true); // features default all-on + gate = new AiFeatureGate(props); + } + + @Test + void passesWhenEngineEnabledAndFeatureOn() { + assertDoesNotThrow(() -> gate.requireClassify()); + assertDoesNotThrow(() -> gate.requireConversationalWorkflow()); + } + + @Test + void throws503WhenFeatureFlagOff() { + props.getAiEngine().getFeatures().setClassify(false); + + ResponseStatusException ex = + assertThrows(ResponseStatusException.class, () -> gate.requireClassify()); + assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); + } + + @Test + void throws503WhenEngineDisabledEvenIfFeatureOn() { + props.getAiEngine().setEnabled(false); // feature flag still true + + ResponseStatusException ex = + assertThrows( + ResponseStatusException.class, () -> gate.requireConversationalWorkflow()); + assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); + } + + @Test + void conversationalWorkflowAllowedWhileEitherChatOrDocumentQuestionsOn() { + // The orchestrate endpoint serves both, so it stays open while either is enabled. + props.getAiEngine().getFeatures().setChat(false); + props.getAiEngine().getFeatures().setDocumentQuestions(true); + assertDoesNotThrow(() -> gate.requireConversationalWorkflow()); + + props.getAiEngine().getFeatures().setChat(true); + props.getAiEngine().getFeatures().setDocumentQuestions(false); + assertDoesNotThrow(() -> gate.requireConversationalWorkflow()); + } + + @Test + void conversationalWorkflowThrows503WhenBothChatAndDocumentQuestionsOff() { + props.getAiEngine().getFeatures().setChat(false); + props.getAiEngine().getFeatures().setDocumentQuestions(false); + + ResponseStatusException ex = + assertThrows( + ResponseStatusException.class, () -> gate.requireConversationalWorkflow()); + assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); + } +} diff --git a/app/saas/src/main/resources/application-saas.properties b/app/saas/src/main/resources/application-saas.properties index 8275970929..f0630aaac1 100644 --- a/app/saas/src/main/resources/application-saas.properties +++ b/app/saas/src/main/resources/application-saas.properties @@ -1,6 +1,11 @@ # Stirling-PDF SaaS profile. Pure multi-tenant cloud. # Activated when the :saas module is on the classpath. +# ---------- AI engine ---------- +# SaaS AI backend is env-driven (reached via AiProxyController); never push settings-derived +# config to it, so pin the config-push off. +aiEngine.pushConfigToEngine=false + # ---------- Datasource ---------- system.datasource.enableCustomDatabase=true system.datasource.customDatabaseUrl=${SAAS_DB_URL:} diff --git a/engine/.env b/engine/.env index a99abf0ccf..1a90ad8c18 100644 --- a/engine/.env +++ b/engine/.env @@ -88,3 +88,7 @@ STIRLING_LOG_FILE= # Use when diagnosing worker stalls: a hung call shows a "Request" line with # no matching "Response" line. Noisy; leave off in normal use. STIRLING_HTTP_DEBUG=false + +# Let the Java processor push admin AI settings to POST /api/v1/config at startup. +# Set false in env-driven deployments so the environment is the single source of truth. +STIRLING_ALLOW_CONFIG_PUSH=true diff --git a/engine/pyproject.toml b/engine/pyproject.toml index 5113ea03ff..8cbb94daf8 100644 --- a/engine/pyproject.toml +++ b/engine/pyproject.toml @@ -4,6 +4,7 @@ version = "0.1.0" description = "AI Document Engine" requires-python = ">=3.13" dependencies = [ + "cryptography>=44.0.0", "fastapi>=0.116.0", "pgvector>=0.3.6", "psycopg[binary,pool]>=3.2", diff --git a/engine/src/stirling/agents/math_presentation.py b/engine/src/stirling/agents/math_presentation.py index bfb6fa7ef8..a8652572a6 100644 --- a/engine/src/stirling/agents/math_presentation.py +++ b/engine/src/stirling/agents/math_presentation.py @@ -15,6 +15,7 @@ from __future__ import annotations from pydantic import Field from pydantic_ai import Agent +from stirling.agents.output_mode import output_retries from stirling.contracts import ( MathAuditorToolReportArtifact, OrchestratorRequest, @@ -68,6 +69,9 @@ class MathIntentClassifier: self._agent: Agent[None, _MathIntentDecision] = Agent( model=runtime.fast_model, output_type=_MathIntentDecision, + # Local models emit valid structured output only intermittently; extra + # retries make this hot-path classifier reliable. No-op for real providers. + retries=output_retries(runtime.settings.chat_provider), system_prompt=_MATH_INTENT_SYSTEM_PROMPT, model_settings=runtime.fast_model_settings, ) diff --git a/engine/src/stirling/agents/orchestrator.py b/engine/src/stirling/agents/orchestrator.py index e07e771e34..bb7f2682df 100644 --- a/engine/src/stirling/agents/orchestrator.py +++ b/engine/src/stirling/agents/orchestrator.py @@ -2,12 +2,14 @@ from __future__ import annotations import logging from dataclasses import dataclass -from typing import assert_never +from typing import Literal, assert_never +from pydantic import ConfigDict, Field from pydantic_ai import Agent -from pydantic_ai.output import ToolOutput +from pydantic_ai.output import NativeOutput, ToolOutput from pydantic_ai.tools import RunContext +from stirling.agents.output_mode import output_retries, uses_tool_output from stirling.agents.pdf_create import PdfCreateAgent from stirling.agents.pdf_edit import PdfEditAgent from stirling.agents.pdf_questions import PdfQuestionAgent @@ -27,6 +29,7 @@ from stirling.contracts import ( format_file_names, ) from stirling.contracts.pdf_create import PdfCreateOrchestrateResponse +from stirling.models import ApiModel from stirling.services import AppRuntime logger = logging.getLogger(__name__) @@ -38,6 +41,34 @@ class OrchestratorDeps: request: OrchestratorRequest +# Enum routing for Ollama/custom local models: they pass the user message as args to the +# zero-arg tool delegates below, which reject it, so pick a capability by name and dispatch in Python. +_RouteCapability = Literal["pdf_edit", "pdf_question", "user_spec", "pdf_review", "pdf_create", "unsupported"] + + +class _RouteDecision(ApiModel): + # Local models add stray tool args and send null for optional fields; tolerate both. + model_config = ConfigDict(extra="ignore") + capability: _RouteCapability + message: str | None = Field( + default=None, + description="Only for capability='unsupported': a short, helpful message to show the user.", + ) + + +_ROUTER_SYSTEM_PROMPT = ( + "You are the top-level router. Choose exactly one capability that best handles the request:\n" + "- pdf_edit: modify or convert one or more attached PDFs.\n" + "- pdf_question: answer questions about the contents of the attached PDFs.\n" + "- user_spec: create or define an agent spec.\n" + "- pdf_review: return the PDF with review comments/annotations attached.\n" + "- pdf_create: generate a NEW document from scratch (invoice, report, letter) - no input file.\n" + "- unsupported: none of the above fit, or the user asks about the assistant itself; put a " + "helpful message in 'message'.\n" + "Respond with the capability and (only for unsupported) a message." +) + + class OrchestratorAgent: def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime @@ -86,6 +117,8 @@ class OrchestratorAgent: description="Return this when none of the delegate outputs fit the request.", ), ], + # Local models pick a delegate less reliably; extra retries. No-op for real providers. + retries=output_retries(runtime.settings.chat_provider), deps_type=OrchestratorDeps, system_prompt=( "You are the top-level orchestrator. " @@ -103,6 +136,21 @@ class OrchestratorAgent: ), model_settings=runtime.fast_model_settings, ) + # Local models can't drive the zero-arg tool delegates; route by name instead (#6163: unify these paths). + self._route_via_enum = uses_tool_output(runtime.settings.chat_provider) + # The router has no tools, so NativeOutput works on Ollama here; a lone output tool + # would tempt a local model to answer in plain text and never call it. + self._router = ( + Agent( + model=runtime.fast_model, + output_type=NativeOutput([_RouteDecision]), + retries=output_retries(runtime.settings.chat_provider), + system_prompt=_ROUTER_SYSTEM_PROMPT, + model_settings=runtime.fast_model_settings, + ) + if self._route_via_enum + else None + ) async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse: logger.info( @@ -114,6 +162,8 @@ class OrchestratorAgent: ) if request.resume_with is not None: return await self._resume(request, request.resume_with) + if self._router is not None: + return await self._route_and_dispatch(request) result = await self.agent.run( self._build_prompt(request), deps=OrchestratorDeps(runtime=self.runtime, request=request), @@ -121,6 +171,31 @@ class OrchestratorAgent: logger.info("[orchestrator] routed -> %s", type(result.output).__name__) return result.output + async def _route_and_dispatch(self, request: OrchestratorRequest) -> OrchestratorResponse: + """Local-model routing: pick a capability by name, then dispatch in Python.""" + assert self._router is not None + result = await self._router.run(self._build_prompt(request)) + decision = result.output + logger.info("[orchestrator] enum-routed -> %s", decision.capability) + match decision.capability: + case "pdf_edit": + return await self._run_pdf_edit(request) + case "pdf_question": + return await self._run_pdf_question(request) + case "user_spec": + return await self._run_agent_draft(request) + case "pdf_review": + return await self._run_pdf_review(request) + case "pdf_create": + return await self._run_pdf_create(request) + case "unsupported": + return UnsupportedCapabilityResponse( + capability="orchestrate", + message=decision.message or "I can't help with that request.", + ) + case _ as unreachable: + assert_never(unreachable) + async def _resume(self, request: OrchestratorRequest, capability: SupportedCapability) -> OrchestratorResponse: """Fast-path to get back to the correct endpoint without having to call AI. diff --git a/engine/src/stirling/agents/output_mode.py b/engine/src/stirling/agents/output_mode.py new file mode 100644 index 0000000000..8751178448 --- /dev/null +++ b/engine/src/stirling/agents/output_mode.py @@ -0,0 +1,28 @@ +"""Provider-aware output: Ollama/custom block tools under native json-schema, so use ToolOutput not NativeOutput.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from pydantic_ai.output import NativeOutput, ToolOutput + +# Providers whose OpenAI-compatible endpoint needs tool-delivered structured output. +_TOOL_OUTPUT_PROVIDERS = frozenset({"ollama", "custom"}) + + +def uses_tool_output(chat_provider: str) -> bool: + return chat_provider in _TOOL_OUTPUT_PROVIDERS + + +def structured_output(output_types: Sequence[Any], *, chat_provider: str) -> Any: + """Pick a structured-output spec compatible with the active chat provider.""" + types = list(output_types) + if uses_tool_output(chat_provider): + return [ToolOutput(t) for t in types] + return NativeOutput(types) + + +def output_retries(chat_provider: str, *, native: int = 1, tool: int = 6) -> int: + """Local models delivering via ToolOutput need more output-validation retries.""" + return tool if uses_tool_output(chat_provider) else native diff --git a/engine/src/stirling/agents/pdf_questions.py b/engine/src/stirling/agents/pdf_questions.py index 5b956c468b..f594c2fdd7 100644 --- a/engine/src/stirling/agents/pdf_questions.py +++ b/engine/src/stirling/agents/pdf_questions.py @@ -3,10 +3,10 @@ from __future__ import annotations import logging from pydantic_ai import Agent -from pydantic_ai.output import NativeOutput from stirling.agents.contradiction import ContradictionCapability, ContradictionDetector from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict +from stirling.agents.output_mode import output_retries, structured_output from stirling.agents.shared import ChunkedReasoner, WholeDocReaderCapability from stirling.contracts import ( AiFile, @@ -196,9 +196,15 @@ class PdfQuestionAgent: files=request.files, principals=principals, ) + # Ollama/custom block tool-calling under native json-schema output, so deliver the + # structured result via a tool call or the model answers ungrounded. See agents.output_mode. + provider = self.runtime.settings.chat_provider agent = Agent( model=self.runtime.smart_model, - output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]), + output_type=structured_output( + [PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse], chat_provider=provider + ), + retries=output_retries(provider), system_prompt=PDF_QUESTION_SYSTEM_PROMPT, # pydantic-ai accepts a list of (string-or-callable) instruction sources; # it resolves each at run time and concatenates them for the model. diff --git a/engine/src/stirling/api/app.py b/engine/src/stirling/api/app.py index 34b0da36fd..60de81360f 100644 --- a/engine/src/stirling/api/app.py +++ b/engine/src/stirling/api/app.py @@ -3,28 +3,20 @@ from __future__ import annotations import asyncio import logging from contextlib import asynccontextmanager -from typing import Annotated -from fastapi import Depends, FastAPI +from fastapi import Depends, FastAPI, Request from pydantic_ai import Agent +from pydantic_ai.models import Model from pydantic_ai.models.instrumented import InstrumentationSettings -from stirling.agents import ( - DocumentClassifierAgent, - ExecutionPlanningAgent, - OrchestratorAgent, - PdfEditAgent, - PdfQuestionAgent, - UserSpecAgent, -) -from stirling.agents.ledger import MathAuditorAgent -from stirling.agents.pdf_comment import PdfCommentAgent +from stirling.api.bootstrap import apply_app_state, build_app_state from stirling.api.dependencies import enforce_required_user_id from stirling.api.engine_auth import EngineSharedSecretMiddleware from stirling.api.middleware import UserIdMiddleware from stirling.api.routes import ( agent_capabilities_router, agent_draft_router, + config_router, document_classifier_router, document_router, execution_router, @@ -34,38 +26,42 @@ from stirling.api.routes import ( pdf_edit_router, pdf_question_router, ) +from stirling.api.routes.config import CONFIG_APPLY_ERRORS, apply_to_app, resolve_and_apply from stirling.config import AppSettings, load_settings +from stirling.config.config_cache import cache_stamp, load_config from stirling.contracts import HealthResponse -from stirling.documents import DocumentService -from stirling.services import build_runtime, setup_posthog_tracking +from stirling.documents import DocumentService, EmbeddingService +from stirling.services import setup_posthog_tracking logger = logging.getLogger(__name__) +# Seconds the lifespan waits for a background task to drain before cancelling it. +_BACKGROUND_TASK_DRAIN_SECONDS = 10 + + +async def _sleep_until(stop: asyncio.Event, seconds: float) -> bool: + """Wait up to ``seconds``; True if asked to stop, False if the interval elapsed.""" + try: + await asyncio.wait_for(stop.wait(), timeout=seconds) + except TimeoutError: + return False + return True + + async def _run_expired_doc_reaper( documents: DocumentService, interval_seconds: int, + stop: asyncio.Event, ) -> None: - """Periodically delete documents whose ``expires_at`` has passed. - - A reaped collection drops everything rooted at that document. Backstop - for the explicit logout purge: catches sessions that ended without a - clean logout (tab close, JWT expiry, engine restart). Persistent rows - (``expires_at`` null, the shape we use for org-shared docs) are never - touched. Runs until cancelled by the lifespan teardown. - """ + """Backstop purge of documents past ``expires_at``; persistent (null expires_at) rows are never touched.""" await _reap(documents) - while True: - await asyncio.sleep(interval_seconds) + while not await _sleep_until(stop, interval_seconds): await _reap(documents) async def _reap(documents: DocumentService) -> None: - """One reaper iteration. Logs the deleted count on success and the full - exception with traceback on failure; never re-raises non-cancel errors so - a bad iteration doesn't kill the loop. ``asyncio.CancelledError`` is - re-raised so the lifespan teardown can cancel the task cleanly. - """ + """One reaper iteration; swallows non-cancel errors so a bad iteration doesn't kill the loop.""" try: deleted = await documents.reap_expired() if deleted: @@ -76,6 +72,45 @@ async def _reap(documents: DocumentService) -> None: logger.exception("Document reaper iteration failed; will retry on next interval") +def _adopt_cached_config_if_changed(fast_api: FastAPI) -> None: + """Re-apply the persisted config when the shared cache file changed under us; never raises.""" + # Read the stamp before the payload: a write landing between the two re-applies next + # tick, whereas the reverse order would record the newer stamp and skip forever. + stamp = cache_stamp() + if stamp is None or stamp == getattr(fast_api.state, "config_cache_stamp", None): + return + # Claim the stamp up front so a cache we cannot read or apply is not retried every tick. + fast_api.state.config_cache_stamp = stamp + cached = load_config() + if cached is None: + return + try: + effective, _ = apply_to_app(fast_api, cached) + except CONFIG_APPLY_ERRORS: + logger.warning("Config pushed to another worker could not be applied here", exc_info=True) + return + logger.info( + "Adopted AI config pushed to another worker: smart_model=%s fast_model=%s", + effective.smart_model_name, + effective.fast_model_name, + ) + + +async def _run_config_cache_watcher( + fast_api: FastAPI, + interval_seconds: int, + stop: asyncio.Event, +) -> None: + """Poll the shared config cache so every worker converges on the last push.""" + while not await _sleep_until(stop, interval_seconds): + try: + _adopt_cached_config_if_changed(fast_api) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Config cache watcher iteration failed; will retry on next interval") + + def _load_startup_settings(fast_api: FastAPI) -> AppSettings: override = fast_api.dependency_overrides.get(load_settings) if override is not None: @@ -83,37 +118,84 @@ def _load_startup_settings(fast_api: FastAPI) -> AppSettings: return load_settings() +def _restore_cached_config( + settings: AppSettings, +) -> tuple[AppSettings, Model | None, Model | None, EmbeddingService | None]: + """Restore the last-applied pushed config from the encrypted cache, or env settings on any failure.""" + if not settings.allow_config_push: + return settings, None, None, None + cached = load_config() + if cached is None: + return settings, None, None, None + try: + effective, smart_model, fast_model, embedder, notes = resolve_and_apply(settings, cached) + except CONFIG_APPLY_ERRORS: + logger.warning("Cached AI config could not be applied; falling back to env settings", exc_info=True) + return settings, None, None, None + logger.info( + "Restored cached AI config: smart_model=%s fast_model=%s%s", + effective.smart_model_name, + effective.fast_model_name, + f"; {'; '.join(notes)}" if notes else "", + ) + return effective, smart_model, fast_model, embedder + + @asynccontextmanager async def lifespan(fast_api: FastAPI): # Load env vars on startup so we can immediately crash if required env vars aren't set settings = _load_startup_settings(fast_api) - runtime = build_runtime(settings) - fast_api.state.settings = settings - fast_api.state.runtime = runtime - fast_api.state.orchestrator_agent = OrchestratorAgent(runtime) - fast_api.state.pdf_edit_agent = PdfEditAgent(runtime) - fast_api.state.pdf_question_agent = PdfQuestionAgent(runtime) - fast_api.state.user_spec_agent = UserSpecAgent(runtime) - fast_api.state.execution_planning_agent = ExecutionPlanningAgent(runtime) - fast_api.state.math_auditor_agent = MathAuditorAgent(runtime) - fast_api.state.pdf_comment_agent = PdfCommentAgent(runtime) - fast_api.state.document_classifier_agent = DocumentClassifierAgent(runtime) - tracer_provider = setup_posthog_tracking(settings) + # Precedence: env < persisted cache < live push. Stamp first so a push landing mid-boot + # is re-adopted by the watcher rather than mistaken for the config we just restored. + fast_api.state.config_cache_stamp = cache_stamp() + effective, smart_model, fast_model, embedder = _restore_cached_config(settings) + app_state = build_app_state( + effective, + fast_model=fast_model, + smart_model=smart_model, + embedder=embedder, + ) + fast_api.state.settings = effective + apply_app_state(fast_api.state, app_state) + runtime = app_state.runtime + tracer_provider = setup_posthog_tracking(effective) if tracer_provider: Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider)) + stop_background = asyncio.Event() reaper_task = asyncio.create_task( _run_expired_doc_reaper( runtime.documents, interval_seconds=settings.documents_reaper_interval_seconds, + stop=stop_background, ), name="expired-document-reaper", ) + background_tasks = [reaper_task] + if effective.allow_config_push: + # A push reaches only one uvicorn worker; this watcher is how the rest of the + # pool picks it up, otherwise most requests keep running the previous models. + background_tasks.append( + asyncio.create_task( + _run_config_cache_watcher( + fast_api, + interval_seconds=settings.config_cache_poll_interval_seconds, + stop=stop_background, + ), + name="config-cache-watcher", + ) + ) yield - reaper_task.cancel() - try: - await reaper_task - except asyncio.CancelledError: - pass + # Drain the loops rather than cancel: cancelling a reaper mid `to_thread` sqlite call lets + # close() pull the connection out from under it and segfault sqlite-vec. Cancel is a backstop. + stop_background.set() + _, pending = await asyncio.wait(background_tasks, timeout=_BACKGROUND_TASK_DRAIN_SECONDS) + for task in pending: + logger.warning("Background task %s did not stop in time; cancelling", task.get_name()) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass await runtime.documents.close() if tracer_provider: tracer_provider.shutdown() @@ -135,10 +217,16 @@ app.include_router(ledger_router, dependencies=_user_gate) app.include_router(pdf_comments_router, dependencies=_user_gate) app.include_router(agent_capabilities_router, dependencies=_user_gate) app.include_router(document_classifier_router, dependencies=_user_gate) +# Config push is a system sync with no X-User-Id, so it is guarded by the shared secret +# and allow_config_push flag only, deliberately NOT the per-user identity gate. +app.include_router(config_router) @app.get("/health", response_model=HealthResponse) -async def healthcheck(settings: Annotated[AppSettings, Depends(load_settings)]) -> HealthResponse: +async def healthcheck(http_request: Request) -> HealthResponse: + # Report the LIVE config on app.state, not the boot-time env cache, so an admin + # "Test connection" shows the model actually in use after a push. Falls back to env. + settings: AppSettings = getattr(http_request.app.state, "settings", None) or load_settings() return HealthResponse( status="ok", smart_model=settings.smart_model_name, diff --git a/engine/src/stirling/api/bootstrap.py b/engine/src/stirling/api/bootstrap.py new file mode 100644 index 0000000000..241578438e --- /dev/null +++ b/engine/src/stirling/api/bootstrap.py @@ -0,0 +1,72 @@ +"""Assemble the runtime + agents into one bundle, swapped onto app.state atomically to change models at runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any + +from pydantic_ai.models import Model + +from stirling.agents import ( + DocumentClassifierAgent, + ExecutionPlanningAgent, + OrchestratorAgent, + PdfEditAgent, + PdfQuestionAgent, + UserSpecAgent, +) +from stirling.agents.ledger import MathAuditorAgent +from stirling.agents.pdf_comment import PdfCommentAgent +from stirling.config import AppSettings +from stirling.documents import DocumentService, EmbeddingService +from stirling.services import AppRuntime, build_runtime + + +@dataclass(frozen=True) +class AppState: + """Every object the lifespan assigns onto ``fast_api.state``.""" + + runtime: AppRuntime + orchestrator_agent: OrchestratorAgent + pdf_edit_agent: PdfEditAgent + pdf_question_agent: PdfQuestionAgent + user_spec_agent: UserSpecAgent + execution_planning_agent: ExecutionPlanningAgent + math_auditor_agent: MathAuditorAgent + pdf_comment_agent: PdfCommentAgent + document_classifier_agent: DocumentClassifierAgent + + +def build_app_state( + settings: AppSettings, + *, + documents: DocumentService | None = None, + fast_model: Model | None = None, + smart_model: Model | None = None, + embedder: EmbeddingService | None = None, +) -> AppState: + """Build the runtime and every agent from ``settings``.""" + runtime = build_runtime( + settings, + documents=documents, + fast_model=fast_model, + smart_model=smart_model, + embedder=embedder, + ) + return AppState( + runtime=runtime, + orchestrator_agent=OrchestratorAgent(runtime), + pdf_edit_agent=PdfEditAgent(runtime), + pdf_question_agent=PdfQuestionAgent(runtime), + user_spec_agent=UserSpecAgent(runtime), + execution_planning_agent=ExecutionPlanningAgent(runtime), + math_auditor_agent=MathAuditorAgent(runtime), + pdf_comment_agent=PdfCommentAgent(runtime), + document_classifier_agent=DocumentClassifierAgent(runtime), + ) + + +def apply_app_state(state: Any, app_state: AppState) -> None: + """Copy every field of ``app_state`` onto a Starlette ``app.state`` object.""" + for field in fields(app_state): + setattr(state, field.name, getattr(app_state, field.name)) diff --git a/engine/src/stirling/api/routes/__init__.py b/engine/src/stirling/api/routes/__init__.py index f40e9918f8..440ff18f9d 100644 --- a/engine/src/stirling/api/routes/__init__.py +++ b/engine/src/stirling/api/routes/__init__.py @@ -1,5 +1,6 @@ from .agent_capabilities import router as agent_capabilities_router from .agent_drafts import router as agent_draft_router +from .config import router as config_router from .document_classifier import router as document_classifier_router from .documents import router as document_router from .execution import router as execution_router @@ -12,6 +13,7 @@ from .pdf_questions import router as pdf_question_router __all__ = [ "agent_capabilities_router", "agent_draft_router", + "config_router", "document_classifier_router", "document_router", "execution_router", diff --git a/engine/src/stirling/api/routes/config.py b/engine/src/stirling/api/routes/config.py new file mode 100644 index 0000000000..cc06a8d87e --- /dev/null +++ b/engine/src/stirling/api/routes/config.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +import ipaddress +import logging + +from fastapi import APIRouter, FastAPI, HTTPException, Request, status +from openai import OpenAIError +from pydantic_ai.exceptions import UserError +from pydantic_ai.models import Model + +from stirling.api.bootstrap import apply_app_state, build_app_state +from stirling.config import AppSettings +from stirling.config.config_cache import cache_stamp, save_config +from stirling.contracts import ConfigApplyResponse, ConfigPushRequest +from stirling.documents import EmbeddingService +from stirling.services import AppRuntime +from stirling.services.runtime import _build_model, validate_structured_output_support + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/api/v1/config", tags=["config"]) + +# Model/provider construction + validation failures. The HTTP route maps these to a +# 400 (no swap); boot catches them to fall back to env when a cached config is bad. +CONFIG_APPLY_ERRORS = (ValueError, UserError, OpenAIError) + +_REINDEX_NOTE = ( + "Embedding model changed; existing indexed documents were embedded with the previous model and " + "must be re-indexed. If the embedding dimensionality changed, re-ingest before searching." +) + + +def _strip_provider_prefix(model_name: str) -> str: + """Drop a leading ``provider:`` from an env model string ("anthropic:x" -> "x").""" + _, sep, rest = model_name.partition(":") + return rest if sep else model_name + + +def _compose_embedding_model(provider: str, model: str) -> str: + """Compose the engine's ``provider:model`` embedding string from pushed parts.""" + provider = provider.strip() + return f"{provider}:{model}" if provider else model + + +def _split_embedding_ref(ref: str) -> tuple[str, str]: + """Split an env embedding string ("voyageai:voyage-4") into (provider, model).""" + provider, sep, model = ref.partition(":") + return (provider, model) if sep else ("", ref) + + +def _keep(pushed: int | None, current: int) -> int: + """Return the pushed value, or the current one when the push omitted it.""" + return pushed if pushed is not None else current + + +# Presence of any of these means request.client.host may be proxy-rewritten and cannot be +# trusted as the transport peer, so a spoofed X-Forwarded-For could otherwise read as loopback. +_FORWARDING_HEADERS = ("x-forwarded-for", "x-forwarded-host", "x-real-ip", "forwarded") + + +def _is_direct_loopback_client(request: Request) -> bool: + """True only for a direct local connection with no proxy; fails closed if any forwarding header is present.""" + if any(h in request.headers for h in _FORWARDING_HEADERS): + return False + client = request.client + if client is None: + return False + try: + return ipaddress.ip_address(client.host).is_loopback + except ValueError: + return client.host == "localhost" + + +def resolve_and_apply( + current: AppSettings, + request: ConfigPushRequest, +) -> tuple[AppSettings, Model, Model, EmbeddingService | None, list[str]]: + """Resolve a pushed config against the running settings; it never swaps live state (the caller does).""" + models = request.models + rag = request.rag + limits = request.limits + notes: list[str] = [] + + provider = models.provider.strip() + api_key = models.api_key + base_url = models.base_url + use_explicit_provider = bool(provider or api_key or base_url) + + if use_explicit_provider and not current.chat_provider: + # First push over an env engine: running names are still "provider:model", strip the prefix. + smart_name = models.smart_model or _strip_provider_prefix(current.smart_model_name) + fast_name = models.fast_model or _strip_provider_prefix(current.fast_model_name) + elif use_explicit_provider: + # A provider was already pushed, so the running names are bare and may legitimately + # contain a colon ("llama3.1:8b"). Stripping again would truncate them to "8b". + smart_name = models.smart_model or current.smart_model_name + fast_name = models.fast_model or current.fast_model_name + else: + # No provider/credentials pushed: keep the fully env-driven model strings. + smart_name = models.smart_model or current.smart_model_name + fast_name = models.fast_model or current.fast_model_name + + def _build(bare: str) -> Model: + if use_explicit_provider: + return _build_model(bare, provider=provider or None, api_key=api_key or None, base_url=base_url or None) + return _build_model(bare) + + smart_model = _build(smart_name) + fast_model = _build(fast_name) + validate_structured_output_support(smart_model, smart_name) + validate_structured_output_support(fast_model, fast_name) + + # Scalars: None / empty keep the current value. + smart_max_tokens = _keep(models.smart_max_tokens, current.smart_model_max_tokens) + fast_max_tokens = _keep(models.fast_max_tokens, current.fast_model_max_tokens) + top_k = _keep(rag.top_k, current.rag_default_top_k) + max_searches = _keep(rag.max_searches, current.rag_max_searches) + max_pages = _keep(limits.max_pages, current.max_pages) + max_characters = _keep(limits.max_characters, current.max_characters) + model_max_concurrency = _keep(limits.model_max_concurrency, current.model_max_concurrency) + + # Embedding: any non-empty embedding field triggers a rebuild; empty fields fall + # back to the running provider/model/creds so a partial push never clobbers env. + embedding_changed = bool( + rag.embedding_provider.strip() or rag.embedding_model.strip() or rag.embedding_api_key or rag.embedding_base_url + ) + rag_embedding_model = current.rag_embedding_model + new_embedder: EmbeddingService | None = None + if embedding_changed: + current_provider, current_model = _split_embedding_ref(current.rag_embedding_model) + embed_provider = rag.embedding_provider.strip() or current_provider + embed_model = rag.embedding_model.strip() or current_model + rag_embedding_model = _compose_embedding_model(embed_provider, embed_model) + new_embedder = EmbeddingService( + model_name=embed_model, + chunk_size=current.rag_chunk_size, + chunk_overlap=current.rag_chunk_overlap, + provider=embed_provider or None, + api_key=rag.embedding_api_key or None, + base_url=rag.embedding_base_url or None, + ) + notes.append(_REINDEX_NOTE) + + effective = current.model_copy( + update={ + "chat_provider": provider, + "smart_model_name": smart_name, + "fast_model_name": fast_name, + "smart_model_max_tokens": smart_max_tokens, + "fast_model_max_tokens": fast_max_tokens, + "rag_embedding_model": rag_embedding_model, + "rag_default_top_k": top_k, + "rag_max_searches": max_searches, + "max_pages": max_pages, + "max_characters": max_characters, + "model_max_concurrency": model_max_concurrency, + } + ) + return effective, smart_model, fast_model, new_embedder, notes + + +def apply_to_app(app: FastAPI, request: ConfigPushRequest) -> tuple[AppSettings, list[str]]: + """Resolve ``request`` and swap the bundle onto app.state; no await, so the swap is atomic wrt the event loop.""" + current: AppSettings = app.state.settings + runtime: AppRuntime = app.state.runtime + effective, smart_model, fast_model, new_embedder, notes = resolve_and_apply(current, request) + new_state = build_app_state( + effective, + documents=runtime.documents, + fast_model=fast_model, + smart_model=smart_model, + ) + app.state.settings = effective + apply_app_state(app.state, new_state) + # Retune retrieval breadth on the reused store without rebuilding it. + runtime.documents.default_top_k = effective.rag_default_top_k + if new_embedder is not None: + # Swap the embedder onto the reused DocumentService, never tearing down the live store. + runtime.documents.embedder = new_embedder + return effective, notes + + +@router.post("", response_model=ConfigApplyResponse) +async def apply_config(request: ConfigPushRequest, http_request: Request) -> ConfigApplyResponse: + """Apply admin-pushed AI settings by rebuilding the runtime + agents, persisting so it survives a restart.""" + app = http_request.app + current: AppSettings = app.state.settings + if not current.allow_config_push: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Config push is disabled on this deployment (STIRLING_ALLOW_CONFIG_PUSH is false).", + ) + # Secure-by-default: with no shared secret set, only trust a direct loopback caller, since a + # pushed base_url/model could repoint the engine to exfiltrate document content. + if not current.engine_shared_secret and not _is_direct_loopback_client(http_request): + client_host = http_request.client.host if http_request.client else "unknown" + logger.warning( + "Rejected config push from non-local/proxied caller %s with no shared secret set", + client_host, + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + "Config push from a non-local or proxied caller requires" + " STIRLING_ENGINE_SHARED_SECRET to be set on both the engine and the processor." + ), + ) + + try: + effective, notes = apply_to_app(app, request) + except CONFIG_APPLY_ERRORS as exc: + # Reject without touching the running config. + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc + + # Persist (encrypted) so the config survives a restart and sibling workers adopt it. + # Best-effort: it is already applied live, so a persist failure must never become a 500. + try: + save_config(request) + # Claim the stamp we just wrote so this worker's watcher does not rebuild for it. + app.state.config_cache_stamp = cache_stamp() + except Exception: # noqa: BLE001 - best-effort persist, never fail the applied push + logger.warning("Applied AI config but failed to persist the encrypted cache", exc_info=True) + notes.append( + "Config applied on this worker but could not be persisted; it will not survive an" + " engine restart and other workers will not pick it up." + ) + + logger.info( + "Applied pushed AI config: provider=%s smart_model=%s fast_model=%s top_k=%s", + request.models.provider.strip() or "", + effective.smart_model_name, + effective.fast_model_name, + effective.rag_default_top_k, + ) + + return ConfigApplyResponse( + status="applied", + provider=request.models.provider.strip(), + smart_model=effective.smart_model_name, + fast_model=effective.fast_model_name, + smart_max_tokens=effective.smart_model_max_tokens, + fast_max_tokens=effective.fast_model_max_tokens, + rag_embedding_model=effective.rag_embedding_model, + rag_top_k=effective.rag_default_top_k, + rag_max_searches=effective.rag_max_searches, + max_pages=effective.max_pages, + max_characters=effective.max_characters, + model_max_concurrency=effective.model_max_concurrency, + notes=notes, + ) diff --git a/engine/src/stirling/config/config_cache.py b/engine/src/stirling/config/config_cache.py new file mode 100644 index 0000000000..deb5864f97 --- /dev/null +++ b/engine/src/stirling/config/config_cache.py @@ -0,0 +1,116 @@ +"""Persistent, Fernet-encrypted cache of the last-applied config-push. Boot restores it so an +engine-only restart keeps admin config; sibling uvicorn workers poll it to adopt a push.""" + +from __future__ import annotations + +import base64 +import logging +import os +import stat +from pathlib import Path + +from cryptography.fernet import Fernet, InvalidToken +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.kdf.hkdf import HKDF + +from stirling.config.settings import ENGINE_ROOT, load_settings +from stirling.contracts import ConfigPushRequest + +logger = logging.getLogger(__name__) + +_CACHE_FILENAME = "ai_config_cache.enc" +_KEY_FILENAME = "ai_config_cache.key" +# Constant salt/info so the same shared secret always derives the same Fernet key. +_HKDF_SALT = b"stirling-ai-config-cache/v1/salt" +_HKDF_INFO = b"stirling-ai-config-cache/v1/fernet-key" + +_keyfile_warned = False + + +def _default_data_dir() -> Path: + """The engine data dir (where the sqlite store lives by default).""" + return ENGINE_ROOT / "data" + + +def _shared_secret() -> str: + return load_settings().engine_shared_secret + + +def _derive_key_from_secret(secret: str) -> bytes: + hkdf = HKDF(algorithm=hashes.SHA256(), length=32, salt=_HKDF_SALT, info=_HKDF_INFO) + return base64.urlsafe_b64encode(hkdf.derive(secret.encode("utf-8"))) + + +def _write_private_bytes(path: Path, payload: bytes) -> None: + """Write ``payload`` to ``path`` atomically (temp file + rename) and owner-only (0600) from the moment it exists.""" + tmp_path = path.with_name(f"{path.name}.tmp") + fd = os.open(tmp_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, stat.S_IRUSR | stat.S_IWUSR) + try: + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, path) + except BaseException: + tmp_path.unlink(missing_ok=True) + raise + + +def _load_or_create_keyfile(data_dir: Path) -> bytes: + global _keyfile_warned + key_path = data_dir / _KEY_FILENAME + if key_path.exists(): + return key_path.read_bytes().strip() + key = Fernet.generate_key() + data_dir.mkdir(parents=True, exist_ok=True) + _write_private_bytes(key_path, key) + if not _keyfile_warned: + logger.warning( + "STIRLING_ENGINE_SHARED_SECRET is not set; encrypting the AI config cache with a local" + " keyfile at %s (0600, best-effort). This is modest protection only - set a shared secret" + " for HKDF key derivation in any deployment where the cache must be strongly protected.", + key_path, + ) + _keyfile_warned = True + return key + + +def _fernet(data_dir: Path) -> Fernet: + secret = _shared_secret() + if secret: + return Fernet(_derive_key_from_secret(secret)) + return Fernet(_load_or_create_keyfile(data_dir)) + + +def save_config(request: ConfigPushRequest, *, data_dir: Path | None = None) -> None: + """Encrypt and persist the last-applied pushed config, overwriting any prior file.""" + data_dir = data_dir or _default_data_dir() + data_dir.mkdir(parents=True, exist_ok=True) + payload = request.model_dump_json(by_alias=True).encode("utf-8") + token = _fernet(data_dir).encrypt(payload) + _write_private_bytes(data_dir / _CACHE_FILENAME, token) + + +def cache_stamp(*, data_dir: Path | None = None) -> tuple[int, int] | None: + """Identify the current cache file as (mtime_ns, size), or None when absent; cheap enough to poll.""" + cache_path = (data_dir or _default_data_dir()) / _CACHE_FILENAME + try: + info = cache_path.stat() + except OSError: + return None + return (info.st_mtime_ns, info.st_size) + + +def load_config(*, data_dir: Path | None = None) -> ConfigPushRequest | None: + """Load + decrypt the persisted config; returns None (never raises) when absent, corrupt, or wrong key.""" + data_dir = data_dir or _default_data_dir() + cache_path = data_dir / _CACHE_FILENAME + if not cache_path.exists(): + return None + try: + token = cache_path.read_bytes() + payload = _fernet(data_dir).decrypt(token) + return ConfigPushRequest.model_validate_json(payload) + except (InvalidToken, ValueError, OSError) as exc: + logger.warning("Ignoring unreadable AI config cache at %s: %s", cache_path, exc) + return None diff --git a/engine/src/stirling/config/settings.py b/engine/src/stirling/config/settings.py index 1293cc371f..d0cf6eba2f 100644 --- a/engine/src/stirling/config/settings.py +++ b/engine/src/stirling/config/settings.py @@ -25,6 +25,9 @@ class AppSettings(BaseSettings): smart_model_name: str = Field(validation_alias="STIRLING_SMART_MODEL") fast_model_name: str = Field(validation_alias="STIRLING_FAST_MODEL") + # Provider backing the active chat models; empty for env/native, 'ollama'/'custom' by push. + # Agents read it to pick a tool-compatible output strategy since local models block tools under native json-schema. + chat_provider: str = Field(default="") smart_model_max_tokens: int = Field(validation_alias="STIRLING_SMART_MODEL_MAX_TOKENS") fast_model_max_tokens: int = Field(validation_alias="STIRLING_FAST_MODEL_MAX_TOKENS") # Process-wide ceiling on concurrent model API calls, shared by both model @@ -123,6 +126,15 @@ class AppSettings(BaseSettings): engine_shared_secret: str = Field(default="", validation_alias="STIRLING_ENGINE_SHARED_SECRET") engine_require_auth: bool = Field(default=False, validation_alias="STIRLING_ENGINE_REQUIRE_AUTH") + # When true, the Java processor may push admin AI settings to POST /api/v1/config at startup. + # Turn off in env-driven deployments so the environment is the single source of truth. + allow_config_push: bool = Field(default=True, validation_alias="STIRLING_ALLOW_CONFIG_PUSH") + # How often each worker polls the shared config cache; bounds how long the pool can disagree on the active model. + config_cache_poll_interval_seconds: int = Field( + default=15, + validation_alias="STIRLING_CONFIG_CACHE_POLL_INTERVAL_SECONDS", + ) + def _configure_logging(level_name: str, log_file: str, http_debug: bool) -> None: """Configure the ``stirling`` logger hierarchy.""" diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index 513906d610..caf41effa0 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -29,6 +29,13 @@ from .common import ( format_conversation_history, format_file_names, ) +from .config import ( + ConfigApplyResponse, + ConfigLimitsSection, + ConfigModelsSection, + ConfigPushRequest, + ConfigRagSection, +) from .contradiction import ( Claim, Contradiction, @@ -139,6 +146,11 @@ __all__ = [ "Claim", "CommentSpec", "CompletedExecutionAction", + "ConfigApplyResponse", + "ConfigLimitsSection", + "ConfigModelsSection", + "ConfigPushRequest", + "ConfigRagSection", "Contradiction", "ContradictionReport", "ContradictionSeverity", diff --git a/engine/src/stirling/contracts/config.py b/engine/src/stirling/contracts/config.py new file mode 100644 index 0000000000..fed75ac320 --- /dev/null +++ b/engine/src/stirling/contracts/config.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from pydantic import ConfigDict, Field + +from stirling.models import ApiModel + + +class TolerantApiModel(ApiModel): + """Push-contract base: unknown fields from a newer processor are ignored rather than + rejecting the whole push. Overrides only the extra policy; camelCase aliasing is inherited.""" + + model_config = ConfigDict(extra="ignore") + + +class ConfigModelsSection(TolerantApiModel): + """Model provider + credentials pushed by the Java processor; empty fields mean "keep the engine's env value".""" + + provider: str = "" + smart_model: str = "" + fast_model: str = "" + smart_max_tokens: int | None = Field(default=None, ge=1) + fast_max_tokens: int | None = Field(default=None, ge=1) + api_key: str = "" + base_url: str = "" + + +class ConfigRagSection(TolerantApiModel): + embedding_provider: str = "" + embedding_model: str = "" + embedding_api_key: str = "" + # OpenAI-compatible endpoint URL for ollama/custom embedding providers; empty keeps the env value. + embedding_base_url: str = "" + top_k: int | None = Field(default=None, ge=1) + # 0 is a legitimate "no retrieval searches" setting, so this floors at 0 not 1. + max_searches: int | None = Field(default=None, ge=0) + + +class ConfigLimitsSection(TolerantApiModel): + max_pages: int | None = Field(default=None, ge=1) + max_characters: int | None = Field(default=None, ge=1) + # Must be >= 1: it becomes an asyncio.Semaphore bound, and 0 constructs a permanently locked + # semaphore that hangs every model call; the push is persisted so a restart won't clear it. + model_max_concurrency: int | None = Field(default=None, ge=1) + + +class ConfigPushRequest(TolerantApiModel): + """Admin-configured AI settings pushed at processor startup.""" + + models: ConfigModelsSection = Field(default_factory=ConfigModelsSection) + rag: ConfigRagSection = Field(default_factory=ConfigRagSection) + limits: ConfigLimitsSection = Field(default_factory=ConfigLimitsSection) + + +class ConfigApplyResponse(ApiModel): + """Summary of the effective config after a push. Never echoes credentials.""" + + status: str + provider: str + smart_model: str + fast_model: str + smart_max_tokens: int + fast_max_tokens: int + rag_embedding_model: str + rag_top_k: int + rag_max_searches: int + max_pages: int + max_characters: int + model_max_concurrency: int + notes: list[str] = Field(default_factory=list) diff --git a/engine/src/stirling/documents/embedder.py b/engine/src/stirling/documents/embedder.py index 13e9c7018d..c92a448149 100644 --- a/engine/src/stirling/documents/embedder.py +++ b/engine/src/stirling/documents/embedder.py @@ -1,6 +1,8 @@ from __future__ import annotations from pydantic_ai import Embedder +from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel +from pydantic_ai.providers.openai import OpenAIProvider from stirling.documents.chunker import chunk_text from stirling.documents.store import Document @@ -12,6 +14,27 @@ from stirling.documents.store import Document DEFAULT_EMBED_BATCH_SIZE = 256 +def _build_embedder( + model_name: str, + *, + provider: str | None = None, + api_key: str | None = None, + base_url: str | None = None, +) -> Embedder: + """Construct an :class:`Embedder`; explicit provider/api_key/base_url is the config-push path, else env form.""" + if not provider and not api_key and not base_url: + return Embedder(model_name) + + provider_name = (provider or "").lower() + key = api_key or None + if provider_name in ("voyageai", "openai"): + return Embedder(f"{provider_name}:{model_name}") + if provider_name in ("ollama", "custom"): + openai_provider = OpenAIProvider(base_url=base_url or None, api_key=key or "ollama") + return Embedder(OpenAIEmbeddingModel(model_name, provider=openai_provider)) + raise ValueError(f"Unsupported embedding provider {provider!r}.") + + class EmbeddingService: """Wraps Pydantic AI's Embedder to provide document chunking and embedding.""" @@ -21,8 +44,12 @@ class EmbeddingService: chunk_size: int = 512, chunk_overlap: int = 64, embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE, + *, + provider: str | None = None, + api_key: str | None = None, + base_url: str | None = None, ) -> None: - self._embedder = Embedder(model_name) + self._embedder = _build_embedder(model_name, provider=provider, api_key=api_key, base_url=base_url) self._chunk_size = chunk_size self._chunk_overlap = chunk_overlap self._embed_batch_size = embed_batch_size diff --git a/engine/src/stirling/documents/service.py b/engine/src/stirling/documents/service.py index 262f5f7ff4..0415e9d140 100644 --- a/engine/src/stirling/documents/service.py +++ b/engine/src/stirling/documents/service.py @@ -43,6 +43,24 @@ class DocumentService: self._store = store self._default_top_k = default_top_k + @property + def default_top_k(self) -> int: + return self._default_top_k + + @default_top_k.setter + def default_top_k(self, value: int) -> None: + # Lets a config-push retune retrieval breadth without rebuilding the store. + self._default_top_k = value + + @property + def embedder(self) -> EmbeddingService: + return self._embedder + + @embedder.setter + def embedder(self, value: EmbeddingService) -> None: + # Lets a config-push swap the embedding model while keeping the live store; existing vectors need re-indexing. + self._embedder = value + async def ingest( self, collection: FileId, diff --git a/engine/src/stirling/services/runtime.py b/engine/src/stirling/services/runtime.py index 011eb8aaf2..06b8595afa 100644 --- a/engine/src/stirling/services/runtime.py +++ b/engine/src/stirling/services/runtime.py @@ -1,7 +1,9 @@ from __future__ import annotations import asyncio +import json import logging +import os from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass @@ -12,8 +14,10 @@ from pydantic_ai import RunContext from pydantic_ai.messages import ModelMessage, ModelResponse from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, infer_model from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.models.wrapper import WrapperModel from pydantic_ai.providers.anthropic import AnthropicProvider +from pydantic_ai.providers.openai import OpenAIProvider from pydantic_ai.settings import ModelSettings from stirling.config import ENGINE_ROOT, AppSettings, DocumentsBackend @@ -49,6 +53,63 @@ def _build_anthropic_http_client() -> httpx.AsyncClient: ) +# Placeholder when no Anthropic key is set, so a deployment using only another provider +# (Ollama etc.) can still boot. Anthropic calls then fail with a clear 401 instead. +_UNCONFIGURED_ANTHROPIC_KEY = "unconfigured" +_warned_missing_anthropic_key = False + + +def _anthropic_provider(explicit_key: str | None = None) -> AnthropicProvider: + """Build the Anthropic provider, tolerating a missing key so an Ollama/OpenAI-only deployment can still boot.""" + http_client = _build_anthropic_http_client() + key = explicit_key or os.environ.get("ANTHROPIC_API_KEY") + if key: + return AnthropicProvider(api_key=key, http_client=http_client) + global _warned_missing_anthropic_key + if not _warned_missing_anthropic_key: + _warned_missing_anthropic_key = True + logger.warning( + "ANTHROPIC_API_KEY is not set - the engine will start, but Anthropic model " + "calls will fail until a key is provided or a different provider is configured " + "(admin AI settings / config push)." + ) + return AnthropicProvider(api_key=_UNCONFIGURED_ANTHROPIC_KEY, http_client=http_client) + + +class _NullContentCoercingTransport(httpx.AsyncHTTPTransport): + """Coerce assistant ``content: null`` to ``""`` in outgoing OpenAI requests, which Ollama otherwise rejects.""" + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + if request.headers.get("content-type", "").startswith("application/json") and request.content: + try: + body = json.loads(request.content) + except ValueError: + return await super().handle_async_request(request) + messages = body.get("messages") + if isinstance(messages, list): + changed = False + for message in messages: + if isinstance(message, dict) and message.get("content", "") is None: + message["content"] = "" + changed = True + if changed: + new_body = json.dumps(body).encode("utf-8") + headers = [(k, v) for k, v in request.headers.raw if k.lower() != b"content-length"] + request = httpx.Request( + method=request.method, + url=request.url, + headers=headers, + content=new_body, + extensions=request.extensions, + ) + return await super().handle_async_request(request) + + +def _openai_compat_http_client() -> httpx.AsyncClient: + """httpx client for Ollama/custom OpenAI-compatible endpoints (null-content fix).""" + return httpx.AsyncClient(transport=_NullContentCoercingTransport()) + + class ConcurrencyLimitedModel(WrapperModel): """Caps in-flight model API calls with a semaphore shared across the process.""" @@ -131,45 +192,72 @@ def _build_document_store(settings: AppSettings) -> DocumentStore: assert_never(settings.documents_backend) -def _build_documents(settings: AppSettings) -> DocumentService: - """Build the document service used by per-request RAG capabilities.""" - logger.info("Documents: embedding_model=%s", settings.rag_embedding_model) - embedder = EmbeddingService( - model_name=settings.rag_embedding_model, - chunk_size=settings.rag_chunk_size, - chunk_overlap=settings.rag_chunk_overlap, - ) +def _build_documents(settings: AppSettings, embedder: EmbeddingService | None = None) -> DocumentService: + """Build the document service used for RAG; ``embedder`` lets a cache-restore boot inject a pre-built one.""" + if embedder is None: + logger.info("Documents: embedding_model=%s", settings.rag_embedding_model) + embedder = EmbeddingService( + model_name=settings.rag_embedding_model, + chunk_size=settings.rag_chunk_size, + chunk_overlap=settings.rag_chunk_overlap, + ) store = _build_document_store(settings) return DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k) -def build_runtime(settings: AppSettings) -> AppRuntime: - fast_model = _build_model(settings.fast_model_name) - smart_model = _build_model(settings.smart_model_name) - validate_structured_output_support(fast_model, settings.fast_model_name) - validate_structured_output_support(smart_model, settings.smart_model_name) +def build_runtime( + settings: AppSettings, + *, + documents: DocumentService | None = None, + fast_model: Model | None = None, + smart_model: Model | None = None, + embedder: EmbeddingService | None = None, +) -> AppRuntime: + """Assemble the shared runtime; the keyword args let a config-push reuse the live store and inject built models.""" + fast = fast_model if fast_model is not None else _build_model(settings.fast_model_name) + smart = smart_model if smart_model is not None else _build_model(settings.smart_model_name) + validate_structured_output_support(fast, settings.fast_model_name) + validate_structured_output_support(smart, settings.smart_model_name) # One semaphore across both tiers: the cap protects the provider account # and process resources, which the tiers share. model_semaphore = asyncio.Semaphore(settings.model_max_concurrency) return AppRuntime( settings=settings, - fast_model=ConcurrencyLimitedModel(fast_model, model_semaphore), - smart_model=ConcurrencyLimitedModel(smart_model, model_semaphore), - documents=_build_documents(settings), + fast_model=ConcurrencyLimitedModel(fast, model_semaphore), + smart_model=ConcurrencyLimitedModel(smart, model_semaphore), + documents=documents if documents is not None else _build_documents(settings, embedder), ) -def _build_model(model_name: str) -> Model: - """Construct a model, injecting our keepalive-free httpx client for - Anthropic models so workers don't pick up stale pooled connections. +def _build_model( + model_name: str, + *, + provider: str | None = None, + api_key: str | None = None, + base_url: str | None = None, +) -> Model: + """Construct a model for ``model_name``; explicit provider/api_key/base_url is the config-push path, else env.""" + if not provider and not api_key and not base_url: + if model_name.startswith("anthropic:"): + bare_name = model_name.removeprefix("anthropic:") + return AnthropicModel(bare_name, provider=_anthropic_provider()) + return infer_model(model_name) - Other providers fall back to ``infer_model`` defaults; the stale-pool - issue is specific to the Cloudflare-fronted Anthropic API in our - observations and the fix doesn't necessarily apply elsewhere. - """ - if model_name.startswith("anthropic:"): - bare_name = model_name.removeprefix("anthropic:") - provider = AnthropicProvider(http_client=_build_anthropic_http_client()) - return AnthropicModel(bare_name, provider=provider) - return infer_model(model_name) + provider_name = (provider or "").lower() + key = api_key or None + if provider_name == "anthropic": + return AnthropicModel(model_name, provider=_anthropic_provider(key)) + if provider_name == "openai": + openai_provider = OpenAIProvider(api_key=key) if key else OpenAIProvider() + return OpenAIChatModel(model_name, provider=openai_provider) + if provider_name in ("ollama", "custom"): + # OpenAI-compatible endpoint. Ollama ignores the key but the SDK needs a non-empty + # one, so default to a placeholder; the custom client coerces null assistant content. + openai_provider = OpenAIProvider( + base_url=base_url or None, + api_key=key or "ollama", + http_client=_openai_compat_http_client(), + ) + return OpenAIChatModel(model_name, provider=openai_provider) + raise ValueError(f"Unsupported model provider {provider!r}.") diff --git a/engine/tests/fixtures/processor_config_push.json b/engine/tests/fixtures/processor_config_push.json new file mode 100644 index 0000000000..b0ee37f6d5 --- /dev/null +++ b/engine/tests/fixtures/processor_config_push.json @@ -0,0 +1,24 @@ +{ + "models": { + "provider": "ollama", + "smartModel": "smart-model-x", + "fastModel": "fast-model-x", + "smartMaxTokens": 1111, + "fastMaxTokens": 2222, + "apiKey": "provider-key-abc", + "baseUrl": "http://engine.example/v1" + }, + "rag": { + "embeddingProvider": "custom", + "embeddingModel": "embed-model-x", + "embeddingApiKey": "embed-key-abc", + "embeddingBaseUrl": "http://embed.example/v1", + "topK": 33, + "maxSearches": 7 + }, + "limits": { + "maxPages": 111, + "maxCharacters": 222222, + "modelMaxConcurrency": 9 + } +} diff --git a/engine/tests/test_config_cache.py b/engine/tests/test_config_cache.py new file mode 100644 index 0000000000..dfe29d4dd1 --- /dev/null +++ b/engine/tests/test_config_cache.py @@ -0,0 +1,111 @@ +"""Tests for the encrypted config cache (stirling.config.config_cache).""" + +from __future__ import annotations + +import os +import stat +from pathlib import Path + +import pytest + +from stirling.config import config_cache +from stirling.contracts import ConfigPushRequest + + +def _sample() -> ConfigPushRequest: + return ConfigPushRequest.model_validate( + { + "models": { + "provider": "anthropic", + "smartModel": "claude-haiku-4-5", + "fastModel": "claude-haiku-4-5", + "smartMaxTokens": 8192, + "fastMaxTokens": 2048, + "apiKey": "secret-key-value", + "baseUrl": "", + }, + "rag": { + "embeddingProvider": "voyageai", + "embeddingModel": "voyage-4", + "embeddingApiKey": "embed-secret", + "embeddingBaseUrl": "", + "topK": 20, + "maxSearches": 5, + }, + "limits": {"maxPages": 200, "maxCharacters": 200000, "modelMaxConcurrency": 32}, + } + ) + + +def test_roundtrip_with_shared_secret(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "the-shared-secret") + req = _sample() + config_cache.save_config(req, data_dir=tmp_path) + # HKDF-from-secret path: no keyfile is written. + assert not (tmp_path / "ai_config_cache.key").exists() + assert config_cache.load_config(data_dir=tmp_path) == req + # Secrets are encrypted at rest. + assert b"secret-key-value" not in (tmp_path / "ai_config_cache.enc").read_bytes() + + +def test_roundtrip_with_keyfile_when_no_secret(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "") + req = _sample() + config_cache.save_config(req, data_dir=tmp_path) + # No shared secret: a random keyfile is generated and reused for decrypt. + assert (tmp_path / "ai_config_cache.key").exists() + assert config_cache.load_config(data_dir=tmp_path) == req + + +def test_load_missing_returns_none(tmp_path: Path) -> None: + assert config_cache.load_config(data_dir=tmp_path) is None + + +def test_load_corrupt_returns_none(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "a-secret") + (tmp_path / "ai_config_cache.enc").write_bytes(b"not-a-valid-fernet-token") + assert config_cache.load_config(data_dir=tmp_path) is None + + +def test_wrong_key_returns_none(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "secret-a") + config_cache.save_config(_sample(), data_dir=tmp_path) + # A different secret derives a different key -> decrypt fails, returns None. + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "secret-b") + assert config_cache.load_config(data_dir=tmp_path) is None + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX file modes are not enforced on Windows") +def test_cache_and_keyfile_are_owner_only(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Both files hold credential material, so neither may land at the umask default.""" + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "") + config_cache.save_config(_sample(), data_dir=tmp_path) + + for name in ("ai_config_cache.enc", "ai_config_cache.key"): + mode = stat.S_IMODE((tmp_path / name).stat().st_mode) + assert mode == 0o600, f"{name} is {oct(mode)}, expected 0o600" + + +def test_save_leaves_no_temp_file_behind(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """The write goes via a temp file + rename so a reader never sees a partial cache.""" + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "the-shared-secret") + config_cache.save_config(_sample(), data_dir=tmp_path) + config_cache.save_config(_sample(), data_dir=tmp_path) + + assert not list(tmp_path.glob("*.tmp")) + assert config_cache.load_config(data_dir=tmp_path) == _sample() + + +def test_cache_stamp_tracks_writes(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """cache_stamp is None before any write and changes when the file is rewritten.""" + monkeypatch.setattr(config_cache, "_shared_secret", lambda: "the-shared-secret") + assert config_cache.cache_stamp(data_dir=tmp_path) is None + + config_cache.save_config(_sample(), data_dir=tmp_path) + first = config_cache.cache_stamp(data_dir=tmp_path) + assert first is not None + + bigger = _sample() + bigger.models.smart_model = "claude-haiku-4-5-with-a-much-longer-name-so-the-size-differs" + config_cache.save_config(bigger, data_dir=tmp_path) + assert config_cache.cache_stamp(data_dir=tmp_path) != first diff --git a/engine/tests/test_config_contract.py b/engine/tests/test_config_contract.py new file mode 100644 index 0000000000..797b7bc147 --- /dev/null +++ b/engine/tests/test_config_contract.py @@ -0,0 +1,70 @@ +"""Wire-contract test pinning the camelCase processor -> engine config push; keep in sync with the Java test.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from stirling.contracts import ConfigPushRequest + +FIXTURE = Path(__file__).parent / "fixtures" / "processor_config_push.json" + + +def _load() -> dict[str, Any]: + return json.loads(FIXTURE.read_text(encoding="utf-8")) + + +def test_processor_contract_round_trips_every_field() -> None: + payload = _load() + req = ConfigPushRequest.model_validate(payload) + + m = payload["models"] + assert req.models.provider == m["provider"] + assert req.models.smart_model == m["smartModel"] + assert req.models.fast_model == m["fastModel"] + assert req.models.smart_max_tokens == m["smartMaxTokens"] + assert req.models.fast_max_tokens == m["fastMaxTokens"] + assert req.models.api_key == m["apiKey"] + assert req.models.base_url == m["baseUrl"] + + r = payload["rag"] + assert req.rag.embedding_provider == r["embeddingProvider"] + assert req.rag.embedding_model == r["embeddingModel"] + assert req.rag.embedding_api_key == r["embeddingApiKey"] + assert req.rag.embedding_base_url == r["embeddingBaseUrl"] + assert req.rag.top_k == r["topK"] + assert req.rag.max_searches == r["maxSearches"] + + limits = payload["limits"] + assert req.limits.max_pages == limits["maxPages"] + assert req.limits.max_characters == limits["maxCharacters"] + assert req.limits.model_max_concurrency == limits["modelMaxConcurrency"] + + +def test_processor_contract_has_no_unmapped_keys() -> None: + """Guard the fixture itself: every wire key must map to a model field, none silently absorbed by extra="ignore".""" + payload = _load() + expected = { + "models": { + "provider", + "smartModel", + "fastModel", + "smartMaxTokens", + "fastMaxTokens", + "apiKey", + "baseUrl", + }, + "rag": { + "embeddingProvider", + "embeddingModel", + "embeddingApiKey", + "embeddingBaseUrl", + "topK", + "maxSearches", + }, + "limits": {"maxPages", "maxCharacters", "modelMaxConcurrency"}, + } + assert set(payload) == set(expected) + for section, keys in expected.items(): + assert set(payload[section]) == keys diff --git a/engine/tests/test_config_routes.py b/engine/tests/test_config_routes.py new file mode 100644 index 0000000000..7d30cd7635 --- /dev/null +++ b/engine/tests/test_config_routes.py @@ -0,0 +1,383 @@ +"""Tests for the config-push endpoint (POST /api/v1/config), driving the real lifespan so app.state is populated.""" + +from __future__ import annotations + +import asyncio +import threading +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + +import pytest +from conftest import build_app_settings +from fastapi.testclient import TestClient + +from stirling.api import app +from stirling.api.app import _adopt_cached_config_if_changed +from stirling.config import AppSettings, config_cache, load_settings +from stirling.contracts import ConfigPushRequest +from stirling.documents import DocumentService + + +@pytest.fixture(autouse=True) +def _isolate_config_cache(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Point the encrypted config cache at a per-test tmp dir so persistence never leaks.""" + monkeypatch.setattr(config_cache, "_default_data_dir", lambda: tmp_path) + + +@contextmanager +def _client( + settings_factory: Callable[[], AppSettings], + *, + client_addr: tuple[str, int] = ("127.0.0.1", 12345), +) -> Iterator[TestClient]: + """Enter a TestClient whose lifespan builds app.state from ``settings_factory``.""" + previous = app.dependency_overrides.get(load_settings) + app.dependency_overrides[load_settings] = settings_factory + try: + with TestClient(app, client=client_addr) as client: + yield client + finally: + if previous is None: + app.dependency_overrides.pop(load_settings, None) + else: + app.dependency_overrides[load_settings] = previous + + +def _anthropic_push() -> dict[str, object]: + return { + "models": { + "provider": "anthropic", + "smartModel": "claude-haiku-4-5", + "fastModel": "claude-haiku-4-5", + "smartMaxTokens": 4096, + "fastMaxTokens": 1024, + "apiKey": "test-key-not-a-real-secret", + "baseUrl": "", + }, + "rag": { + "embeddingProvider": "", + "embeddingModel": "", + "embeddingApiKey": "", + "embeddingBaseUrl": "", + "topK": 7, + "maxSearches": 3, + }, + "limits": {"maxPages": 50, "maxCharacters": 12345, "modelMaxConcurrency": 8}, + } + + +def test_config_push_forbidden_when_disabled() -> None: + def factory() -> AppSettings: + return build_app_settings().model_copy(update={"allow_config_push": False}) + + with _client(factory) as client: + response = client.post("/api/v1/config", json=_anthropic_push()) + assert response.status_code == 403 + + +def test_config_push_from_non_local_caller_without_secret_returns_403() -> None: + """Secure-by-default: with no shared secret, a remote caller cannot push a config.""" + with _client(build_app_settings, client_addr=("203.0.113.9", 4444)) as client: + response = client.post("/api/v1/config", json=_anthropic_push()) + assert response.status_code == 403 + assert "STIRLING_ENGINE_SHARED_SECRET" in response.json()["detail"] + + +def test_config_push_from_loopback_with_forwarded_header_returns_403() -> None: + """A forwarding header means the peer address may be proxy-rewritten, so loopback isn't trusted without a secret.""" + with _client(build_app_settings) as client: # peer is loopback 127.0.0.1 + response = client.post( + "/api/v1/config", + json=_anthropic_push(), + headers={"X-Forwarded-For": "127.0.0.1"}, + ) + assert response.status_code == 403 + assert "STIRLING_ENGINE_SHARED_SECRET" in response.json()["detail"] + + +def test_config_push_persist_failure_still_applies_without_500() -> None: + """Persistence is best-effort: a persist failure must not turn an already-applied push into a 500.""" + with _client(build_app_settings) as client: + with patch( + "stirling.api.routes.config.save_config", + side_effect=ValueError("corrupt keyfile"), + ): + response = client.post("/api/v1/config", json=_anthropic_push()) + assert response.status_code == 200 + # The config WAS applied live despite the persist failure. + assert app.state.settings.smart_model_name == "claude-haiku-4-5" + assert any("could not be persisted" in note for note in response.json()["notes"]) + + +def test_config_push_applies_model_and_limits() -> None: + with _client(build_app_settings) as client: + response = client.post("/api/v1/config", json=_anthropic_push()) + assert response.status_code == 200 + body = response.json() + # Wire summary is camelCase and never echoes the api key. + assert body["smartModel"] == "claude-haiku-4-5" + assert body["fastModel"] == "claude-haiku-4-5" + assert body["smartMaxTokens"] == 4096 + assert body["ragTopK"] == 7 + assert body["maxPages"] == 50 + assert body["modelMaxConcurrency"] == 8 + assert "test-key-not-a-real-secret" not in response.text + # State was swapped: the running settings now reflect the push. + assert app.state.settings.smart_model_name == "claude-haiku-4-5" + assert app.state.settings.max_pages == 50 + assert app.state.runtime.documents.default_top_k == 7 + + +def test_config_push_unsupported_provider_returns_400_without_swap() -> None: + with _client(build_app_settings) as client: + before = app.state.runtime + payload = _anthropic_push() + payload["models"]["provider"] = "nonsense-provider" # type: ignore[index] + response = client.post("/api/v1/config", json=payload) + assert response.status_code == 400 + # Running runtime is untouched when the push is rejected. + assert app.state.runtime is before + assert app.state.settings.smart_model_name == "test" + + +def test_config_push_unsupported_model_returns_400() -> None: + """A model that fails structured-output validation is rejected with 400.""" + with _client(build_app_settings) as client: + before = app.state.runtime + with patch( + "stirling.api.routes.config.validate_structured_output_support", + side_effect=ValueError("Unsupported model foo. This model does not support structured outputs."), + ): + response = client.post("/api/v1/config", json=_anthropic_push()) + assert response.status_code == 400 + assert "does not support structured outputs" in response.json()["detail"] + assert app.state.runtime is before + + +def test_config_push_ollama_embedding_rebuilds_embedder() -> None: + """A pushed ollama/custom embedding provider swaps the embedder onto the reused store, with a re-index note.""" + with _client(build_app_settings) as client: + store_before = app.state.runtime.documents + embedder_before = app.state.runtime.documents.embedder + payload = _anthropic_push() + payload["rag"] = { + "embeddingProvider": "ollama", + "embeddingModel": "nomic-embed-text", + "embeddingApiKey": "", + "embeddingBaseUrl": "http://localhost:11434/v1", + "topK": 9, + "maxSearches": 4, + } + response = client.post("/api/v1/config", json=payload) + assert response.status_code == 200 + body = response.json() + assert body["ragEmbeddingModel"] == "ollama:nomic-embed-text" + assert any("re-index" in note for note in body["notes"]) + docs = app.state.runtime.documents + # Same store object reused (connection pool intact), embedder swapped. + assert docs is store_before + assert docs.embedder is not embedder_before + assert docs.default_top_k == 9 + + +def test_config_push_unsupported_embedding_provider_returns_400() -> None: + with _client(build_app_settings) as client: + embedder_before = app.state.runtime.documents.embedder + payload = _anthropic_push() + payload["rag"] = { + "embeddingProvider": "totally-bogus", + "embeddingModel": "x", + "embeddingApiKey": "", + "embeddingBaseUrl": "", + "topK": None, + "maxSearches": None, + } + response = client.post("/api/v1/config", json=payload) + assert response.status_code == 400 + # Embedder untouched when the push is rejected. + assert app.state.runtime.documents.embedder is embedder_before + + +def test_config_push_empty_models_keep_env_value() -> None: + """An empty models block keeps the engine's env models but still applies pushed limits.""" + with _client(build_app_settings) as client: + payload = _anthropic_push() + payload["models"] = { + "provider": "", + "smartModel": "", + "fastModel": "", + "smartMaxTokens": None, + "fastMaxTokens": None, + "apiKey": "", + "baseUrl": "", + } + response = client.post("/api/v1/config", json=payload) + assert response.status_code == 200 + # Env "test" model preserved for both tiers; limits still applied. + assert app.state.settings.smart_model_name == "test" + assert app.state.settings.fast_model_name == "test" + assert app.state.settings.max_pages == 50 + + +def test_config_push_ignores_unknown_fields() -> None: + """A newer processor pushing unknown fields must not 422; they are ignored and the rest applies.""" + with _client(build_app_settings) as client: + payload = _anthropic_push() + payload["futureTopLevelField"] = {"anything": 1} + payload["models"]["experimentalFlag"] = True # type: ignore[index] + response = client.post("/api/v1/config", json=payload) + assert response.status_code == 200 + assert app.state.settings.smart_model_name == "claude-haiku-4-5" + assert app.state.settings.max_pages == 50 + + +def test_boot_restores_cached_config() -> None: + """A persisted config is decrypted and applied on boot, overriding env.""" + config_cache.save_config(ConfigPushRequest.model_validate(_anthropic_push())) + with _client(build_app_settings): + # Env model is "test"; the cache pushed claude-haiku-4-5 + limits. + assert app.state.settings.smart_model_name == "claude-haiku-4-5" + assert app.state.settings.fast_model_name == "claude-haiku-4-5" + assert app.state.settings.max_pages == 50 + assert app.state.settings.smart_model_max_tokens == 4096 + assert app.state.runtime.documents.default_top_k == 7 + + +def test_boot_ignores_cache_when_push_disabled() -> None: + """With allow_config_push false, env wins and the cache is ignored.""" + config_cache.save_config(ConfigPushRequest.model_validate(_anthropic_push())) + + def factory() -> AppSettings: + return build_app_settings().model_copy(update={"allow_config_push": False}) + + with _client(factory): + assert app.state.settings.smart_model_name == "test" + assert app.state.settings.max_pages == 200 + + +def test_boot_proceeds_on_corrupt_cache(tmp_path: Path) -> None: + """A corrupt cache file is ignored and boot falls back to env, never crashing.""" + (tmp_path / "ai_config_cache.enc").write_bytes(b"not-a-valid-fernet-token") + with _client(build_app_settings): + assert app.state.settings.smart_model_name == "test" + assert app.state.settings.max_pages == 200 + + +def test_config_push_from_remote_caller_is_allowed_when_secret_is_set() -> None: + """The loopback gate is a fallback for the no-secret case only; a remote caller may push once a secret is set.""" + + def factory() -> AppSettings: + return build_app_settings().model_copy(update={"engine_shared_secret": "s3cret"}) + + with _client(factory, client_addr=("203.0.113.9", 4444)) as client: + response = client.post("/api/v1/config", json=_anthropic_push()) + assert response.status_code == 200 + assert app.state.settings.smart_model_name == "claude-haiku-4-5" + + +@pytest.mark.parametrize( + ("section", "field", "value"), + [ + ("limits", "modelMaxConcurrency", 0), + ("limits", "modelMaxConcurrency", -1), + ("limits", "maxPages", 0), + ("limits", "maxCharacters", 0), + ("rag", "topK", 0), + ("models", "smartMaxTokens", 0), + ], +) +def test_config_push_rejects_out_of_range_numbers(section: str, field: str, value: int) -> None: + """Out-of-range numbers are rejected by the contract before anything is applied.""" + with _client(build_app_settings) as client: + payload = _anthropic_push() + payload[section][field] = value # type: ignore[index] + response = client.post("/api/v1/config", json=payload) + assert response.status_code == 422 + # Nothing was applied: the engine is still on its env config. + assert app.state.settings.smart_model_name == "test" + assert app.state.settings.max_pages == 200 + + +def test_config_push_allows_zero_max_searches() -> None: + """0 searches is a legitimate "no retrieval" setting, not an out-of-range value.""" + with _client(build_app_settings) as client: + payload = _anthropic_push() + payload["rag"]["maxSearches"] = 0 # type: ignore[index] + response = client.post("/api/v1/config", json=payload) + assert response.status_code == 200 + assert app.state.settings.rag_max_searches == 0 + + +def test_second_push_keeps_a_colon_bearing_model_name() -> None: + """A pushed model name may contain a colon ("llama3.1:8b") and must survive a follow-up push, not be re-stripped.""" + ollama_push: dict[str, object] = { + "models": { + "provider": "ollama", + "smartModel": "llama3.1:8b", + "fastModel": "llama3.1:8b", + "baseUrl": "http://localhost:11434/v1", + }, + "rag": {}, + "limits": {}, + } + with _client(build_app_settings) as client: + assert client.post("/api/v1/config", json=ollama_push).status_code == 200 + assert app.state.settings.smart_model_name == "llama3.1:8b" + + # Second push: same provider/base URL, model left empty ("keep what you have"). + followup: dict[str, object] = { + "models": { + "provider": "ollama", + "smartModel": "", + "fastModel": "", + "baseUrl": "http://localhost:11434/v1", + }, + "rag": {}, + "limits": {"maxPages": 42}, + } + response = client.post("/api/v1/config", json=followup) + assert response.status_code == 200 + assert app.state.settings.smart_model_name == "llama3.1:8b" + assert app.state.settings.fast_model_name == "llama3.1:8b" + assert app.state.settings.max_pages == 42 + + +def test_worker_adopts_a_config_pushed_to_a_sibling_worker() -> None: + """A push reaches one uvicorn worker; the rest adopt it from the shared cache file.""" + with _client(build_app_settings): + assert app.state.settings.smart_model_name == "test" + + config_cache.save_config(ConfigPushRequest.model_validate(_anthropic_push())) + _adopt_cached_config_if_changed(app) + + assert app.state.settings.smart_model_name == "claude-haiku-4-5" + assert app.state.settings.max_pages == 50 + assert app.state.runtime.documents.default_top_k == 7 + + +def test_worker_does_not_rebuild_when_the_cache_is_unchanged() -> None: + """The watcher is a poll, so an unchanged cache must be a no-op, not a rebuild every tick.""" + config_cache.save_config(ConfigPushRequest.model_validate(_anthropic_push())) + with _client(build_app_settings): + before = app.state.orchestrator_agent + _adopt_cached_config_if_changed(app) + assert app.state.orchestrator_agent is before + + +def test_shutdown_drains_background_tasks_instead_of_cancelling_them() -> None: + """Shutdown must let the reaper iteration finish before the store closes, else close segfaults sqlite-vec.""" + reap_finished = threading.Event() + + async def slow_reap(*_args: object, **_kwargs: object) -> int: + # Long enough to still be running when the (empty) test body hands back to teardown. + await asyncio.sleep(0.2) + reap_finished.set() + return 0 + + with patch.object(DocumentService, "reap_expired", slow_reap): + with _client(build_app_settings): + pass + + assert reap_finished.is_set(), "teardown cancelled the reaper mid-iteration" diff --git a/engine/uv.lock b/engine/uv.lock index afa5f389ae..5e9d166af0 100644 --- a/engine/uv.lock +++ b/engine/uv.lock @@ -605,6 +605,7 @@ name = "engine" version = "0.1.0" source = { editable = "." } dependencies = [ + { name = "cryptography" }, { name = "fastapi" }, { name = "opentelemetry-sdk" }, { name = "pgvector" }, @@ -631,6 +632,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "cryptography", specifier = ">=44.0.0" }, { name = "fastapi", specifier = ">=0.116.0" }, { name = "opentelemetry-sdk", specifier = ">=1.39.0" }, { name = "pgvector", specifier = ">=0.3.6" }, diff --git a/frontend/editor/public/locales/en-US/translation.toml b/frontend/editor/public/locales/en-US/translation.toml index 8c2a8333b6..6cf87f4252 100644 --- a/frontend/editor/public/locales/en-US/translation.toml +++ b/frontend/editor/public/locales/en-US/translation.toml @@ -692,6 +692,166 @@ manualLinks = "Manual downloads: click the links and place the files into the te noLanguages = "No tessdata languages found in the configured directory." permissionNotice = "The tessdata path is not writable. Downloads will be opened in the browser; please save the .traineddata files manually into the tessdata folder." +# AI engine admin settings (AI nav group) +[admin.settings.ai.documents] +description = "Configure the embedding model and retrieval settings used to answer questions over documents. Applied to the AI engine when saved." +title = "Documents & RAG" + +[admin.settings.ai.documents.embeddingApiKey] +description = "Leave blank to use the engine's own environment credential. Applies to self-hosted single-engine deployments." +generic = "Embedding API key" +openai = "OpenAI API key" +setPlaceholder = "Saved - leave blank to keep the current key" +voyageai = "VoyageAI API key" + +[admin.settings.ai.documents.embeddingBaseUrl] +description = "Base URL of the OpenAI-compatible / Ollama embeddings endpoint, e.g. http://ollama:11434/v1. Must point at a trusted internal endpoint (SSRF-sensitive)." +label = "Embedding base URL" + +[admin.settings.ai.documents.embeddingModel] +description = "Embedding model name. Free text; suggestions are hints only." +label = "Embedding model" + +[admin.settings.ai.documents.embeddingProvider] +description = "Provider used to turn document text into vector embeddings." +label = "Embedding provider" + +[admin.settings.ai.documents.maxSearches] +description = "Maximum number of retrieval searches the agent may run per request." +label = "Max searches" + +[admin.settings.ai.documents.reindexNote] +body = "Changing the embedding model takes effect immediately, but documents indexed with the previous model must be re-indexed for search to return correct results." +title = "Re-index required" + +[admin.settings.ai.documents.saved] +reindexNote = "If you changed the embedding model, re-index existing documents so search uses the new model." + +[admin.settings.ai.documents.topK] +description = "Number of most-relevant chunks retrieved per search." +label = "Top K" + +[admin.settings.ai.general] +description = "Connect Stirling to the Python AI engine and choose which AI capabilities are exposed. Changes apply on restart." +title = "AI Engine" + +[admin.settings.ai.general.capabilities] +description = "Turn individual AI features on or off. Disabled features are hidden in the app." +title = "Capabilities" + +[admin.settings.ai.general.enabled] +description = "Master switch. When off, no AI tools, agents, or engine calls are available." +label = "Enable AI" + +[admin.settings.ai.general.features.chat] +description = "Conversational assistant for working with PDFs." +label = "Chat assistant" + +[admin.settings.ai.general.features.classify] +description = "Automatically categorise documents by type or content." +label = "Document classification" + +[admin.settings.ai.general.features.createPdf] +description = "Generate a new PDF (e.g. from HTML) via an AI agent." +label = "Create PDF from prompt" + +[admin.settings.ai.general.features.documentQuestions] +description = "Ask questions and get answers grounded in an uploaded document." +label = "Document questions" + +[admin.settings.ai.general.features.mathAuditor] +description = "Review documents for mathematical and numerical errors." +label = "Math auditor" + +[admin.settings.ai.general.features.pdfComment] +description = "Add AI-authored review comments and annotations to a PDF." +label = "PDF comment agent" + +[admin.settings.ai.general.longRunningTimeoutSeconds] +description = "Timeout for heavier agent operations such as document generation." +label = "Long-running timeout (seconds)" + +[admin.settings.ai.general.note] +body = "The AI engine runs as a separate service. Its shared secret" +body2 = "is set via a container environment variable. Provider API keys can be entered on these pages or supplied as engine environment variables; keys entered here are pushed to the engine when saved." +title = "About the AI engine" + +[admin.settings.ai.general.streamTimeoutSeconds] +description = "Timeout for streamed (token-by-token) chat responses." +label = "Stream timeout (seconds)" + +[admin.settings.ai.general.test] +button = "Test connection" +failBody = "The AI engine did not respond. Check the URL, that the engine container is running, and that AI is enabled (a restart is needed after enabling)." +failTitle = "AI engine unreachable" +okBody = "The AI engine responded to a health check." +okTitle = "AI engine reachable" + +[admin.settings.ai.general.timeoutSeconds] +description = "Timeout for standard AI requests to the engine." +label = "Request timeout (seconds)" + +[admin.settings.ai.general.url] +description = "Internal URL of the Python AI engine, e.g. http://stirling-pdf-engine:5001." +label = "AI engine URL" + +[admin.settings.ai.limits] +description = "Guardrails for how much work AI requests may do and how many run concurrently. Applied to the AI engine when saved." +title = "Limits & Performance" + +[admin.settings.ai.limits.maxCharacters] +description = "Guardrail: reject AI requests whose extracted text exceeds this length." +label = "Max characters per request" + +[admin.settings.ai.limits.maxPages] +description = "Guardrail: reject AI requests over this many PDF pages." +label = "Max pages per request" + +[admin.settings.ai.limits.modelMaxConcurrency] +description = "Maximum simultaneous in-flight model calls across the whole engine." +label = "Model max concurrency" + +[admin.settings.ai.models] +description = "Choose the LLM provider and the smart/fast models the AI engine uses. Applied to the AI engine when saved." +title = "Models & Providers" + +[admin.settings.ai.models.apiKey] +anthropic = "Anthropic API key" +description = "Leave blank to use the engine's own environment credential. Applies to self-hosted single-engine deployments." +generic = "API key" +openai = "OpenAI API key" +setPlaceholder = "Saved - leave blank to keep the current key" + +[admin.settings.ai.models.baseUrl] +description = "Base URL of the OpenAI-compatible / Ollama endpoint, e.g. http://ollama:11434/v1." +label = "Provider base URL" +warning = "The base URL must point at a trusted internal endpoint. The engine will make server-side requests to it, so an untrusted value is SSRF-sensitive." + +[admin.settings.ai.models.fastMaxTokens] +description = "Maximum output tokens for the fast model." +label = "Fast model max tokens" + +[admin.settings.ai.models.fastModel] +description = "Cheaper, faster model for lightweight tasks. Free text; suggestions are hints only." +label = "Fast model" + +[admin.settings.ai.models.provider] +description = "Which LLM provider the engine talks to." +label = "Provider" + +[admin.settings.ai.models.smartMaxTokens] +description = "Maximum output tokens for the smart model." +label = "Smart model max tokens" + +[admin.settings.ai.models.smartModel] +description = "High-capability model for complex reasoning. Free text; suggestions are hints only." +label = "Smart model" + +[admin.settings.ai.saved] +body = "Changes are pushed to the AI engine automatically." +bodyNoPush = "Settings saved. They will apply the next time the AI engine picks up its configuration." +title = "AI settings saved" + [admin.settings.badge] clickToUpgrade = "Click to view plan details" @@ -8816,6 +8976,13 @@ title = "Signature Order" [settings] close = "Close" +[settings.ai] +documents = "Documents & RAG" +general = "General" +limits = "Limits & Performance" +models = "Models & Providers" +title = "AI" + [settings.configuration] advanced = "Advanced" database = "Database" diff --git a/frontend/editor/public/og-metadata.json b/frontend/editor/public/og-metadata.json index 1fc552cd80..63461920ce 100644 --- a/frontend/editor/public/og-metadata.json +++ b/frontend/editor/public/og-metadata.json @@ -475,6 +475,26 @@ "title": "Admin Mcp Settings - Stirling PDF", "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" }, + "/settings/adminAiGeneral": { + "image": "/og_images/home.png", + "title": "Admin Ai General Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, + "/settings/adminAiModels": { + "image": "/og_images/home.png", + "title": "Admin Ai Models Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, + "/settings/adminAiDocuments": { + "image": "/og_images/home.png", + "title": "Admin Ai Documents Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, + "/settings/adminAiLimits": { + "image": "/og_images/home.png", + "title": "Admin Ai Limits Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, "/settings/help": { "image": "/og_images/home.png", "title": "Help Settings - Stirling PDF", @@ -654,6 +674,10 @@ "/settings/adminEndpoints": "/settings/adminEndpoints", "/settings/adminStorageSharing": "/settings/adminStorageSharing", "/settings/adminMcp": "/settings/adminMcp", + "/settings/adminAiGeneral": "/settings/adminAiGeneral", + "/settings/adminAiModels": "/settings/adminAiModels", + "/settings/adminAiDocuments": "/settings/adminAiDocuments", + "/settings/adminAiLimits": "/settings/adminAiLimits", "/settings/help": "/settings/help", "/settings/legal": "/settings/legal", "/settings/backendThirdPartyLicenses": "/settings/backendThirdPartyLicenses", diff --git a/frontend/editor/public/og-metadata.saas.json b/frontend/editor/public/og-metadata.saas.json index 2bf49f7d9f..8f581f2125 100644 --- a/frontend/editor/public/og-metadata.saas.json +++ b/frontend/editor/public/og-metadata.saas.json @@ -476,6 +476,26 @@ "title": "Admin Mcp Settings - Stirling PDF", "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" }, + "/settings/adminAiGeneral": { + "image": "/og_images/home.png", + "title": "Admin Ai General Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, + "/settings/adminAiModels": { + "image": "/og_images/home.png", + "title": "Admin Ai Models Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, + "/settings/adminAiDocuments": { + "image": "/og_images/home.png", + "title": "Admin Ai Documents Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, + "/settings/adminAiLimits": { + "image": "/og_images/home.png", + "title": "Admin Ai Limits Settings - Stirling PDF", + "description": "The Free Adobe Acrobat alternative (10M+ Downloads)" + }, "/settings/help": { "image": "/og_images/home.png", "title": "Help Settings - Stirling PDF", @@ -667,6 +687,10 @@ "/settings/adminEndpoints": "/settings/adminEndpoints", "/settings/adminStorageSharing": "/settings/adminStorageSharing", "/settings/adminMcp": "/settings/adminMcp", + "/settings/adminAiGeneral": "/settings/adminAiGeneral", + "/settings/adminAiModels": "/settings/adminAiModels", + "/settings/adminAiDocuments": "/settings/adminAiDocuments", + "/settings/adminAiLimits": "/settings/adminAiLimits", "/settings/help": "/settings/help", "/settings/legal": "/settings/legal", "/settings/backendThirdPartyLicenses": "/settings/backendThirdPartyLicenses", diff --git a/frontend/editor/src/core/components/shared/config/types.ts b/frontend/editor/src/core/components/shared/config/types.ts index 5cbdd1b169..1182933d61 100644 --- a/frontend/editor/src/core/components/shared/config/types.ts +++ b/frontend/editor/src/core/components/shared/config/types.ts @@ -32,6 +32,10 @@ export const VALID_NAV_KEYS = [ "adminEndpoints", "adminStorageSharing", "adminMcp", + "adminAiGeneral", + "adminAiModels", + "adminAiDocuments", + "adminAiLimits", "help", "legal", "backendThirdPartyLicenses", diff --git a/frontend/editor/src/proprietary/components/shared/config/configNavSections.tsx b/frontend/editor/src/proprietary/components/shared/config/configNavSections.tsx index 0081c594a3..218d4b2f57 100644 --- a/frontend/editor/src/proprietary/components/shared/config/configNavSections.tsx +++ b/frontend/editor/src/proprietary/components/shared/config/configNavSections.tsx @@ -17,6 +17,10 @@ import AdminPlanSection from "@app/components/shared/config/configSections/Admin import AdminFeaturesSection from "@app/components/shared/config/configSections/AdminFeaturesSection"; import AdminEndpointsSection from "@app/components/shared/config/configSections/AdminEndpointsSection"; import AdminMcpSection from "@app/components/shared/config/configSections/AdminMcpSection"; +import AdminAiGeneralSection from "@app/components/shared/config/configSections/AdminAiGeneralSection"; +import AdminAiModelsSection from "@app/components/shared/config/configSections/AdminAiModelsSection"; +import AdminAiDocumentsSection from "@app/components/shared/config/configSections/AdminAiDocumentsSection"; +import AdminAiLimitsSection from "@app/components/shared/config/configSections/AdminAiLimitsSection"; import AdminAuditSection from "@app/components/shared/config/configSections/AdminAuditSection"; import AdminUsageSection from "@app/components/shared/config/configSections/AdminUsageSection"; import AdminStorageSharingSection from "@app/components/shared/config/configSections/AdminStorageSharingSection"; @@ -166,6 +170,45 @@ export const useConfigNavSections = ( ], }); + // AI + sections.push({ + title: t("settings.ai.title", "AI"), + items: [ + { + key: "adminAiGeneral", + label: t("settings.ai.general", "General"), + icon: "smart-toy-rounded", + component: , + disabled: requiresLogin, + disabledTooltip: requiresLogin ? enableLoginTooltip : undefined, + }, + { + key: "adminAiModels", + label: t("settings.ai.models", "Models & Providers"), + icon: "psychology", + component: , + disabled: requiresLogin, + disabledTooltip: requiresLogin ? enableLoginTooltip : undefined, + }, + { + key: "adminAiDocuments", + label: t("settings.ai.documents", "Documents & RAG"), + icon: "description", + component: , + disabled: requiresLogin, + disabledTooltip: requiresLogin ? enableLoginTooltip : undefined, + }, + { + key: "adminAiLimits", + label: t("settings.ai.limits", "Limits & Performance"), + icon: "speed", + component: , + disabled: requiresLogin, + disabledTooltip: requiresLogin ? enableLoginTooltip : undefined, + }, + ], + }); + // Security & Authentication sections.push({ title: t("settings.securityAuth.title", "Security & Authentication"), diff --git a/frontend/editor/src/proprietary/components/shared/config/configSections/AdminAiDocumentsSection.tsx b/frontend/editor/src/proprietary/components/shared/config/configSections/AdminAiDocumentsSection.tsx new file mode 100644 index 0000000000..55c5867288 --- /dev/null +++ b/frontend/editor/src/proprietary/components/shared/config/configSections/AdminAiDocumentsSection.tsx @@ -0,0 +1,383 @@ +import { useCallback, useEffect, useState } from "react"; +import { useTranslation } from "react-i18next"; +import { + NumberInput, + PasswordInput, + Autocomplete, + Select, + TextInput, + Stack, + Paper, + Text, + Loader, + Group, + Alert, +} from "@mantine/core"; +import { alert } from "@app/components/toast"; +import LocalIcon from "@app/components/shared/LocalIcon"; +import { useAdminSettings } from "@app/hooks/useAdminSettings"; +import { useSettingsDirty } from "@app/hooks/useSettingsDirty"; +import PendingBadge from "@app/components/shared/config/PendingBadge"; +import { SettingsStickyFooter } from "@app/components/shared/config/SettingsStickyFooter"; +import apiClient from "@app/services/apiClient"; +import { useLoginRequired } from "@app/hooks/useLoginRequired"; +import { Z_INDEX_OVER_CONFIG_MODAL } from "@app/styles/zIndex"; +import { + AiEngineSettingsData, + AiEngineRag, + AiEngineApiResponse, + EMBEDDING_MODEL_SUGGESTIONS, + MASKED_SECRET, + clampMin, + savedToastBody, +} from "@app/components/shared/config/configSections/aiEngineSettings"; + +export default function AdminAiDocumentsSection() { + const { t } = useTranslation(); + const { loginEnabled } = useLoginRequired(); + + // Track edits to the masked secret so we never send "********" back. + const [embeddingApiKeyDirty, setEmbeddingApiKeyDirty] = useState(false); + + const { + settings, + setSettings, + loading, + saving, + fetchSettings, + saveSettings, + isFieldPending, + } = useAdminSettings({ + sectionName: "aiEngine", + fetchTransformer: async (): Promise< + AiEngineSettingsData & { _pending?: Partial } + > => { + const response = await apiClient.get( + "/api/v1/admin/settings/section/aiEngine", + ); + return response.data || {}; + }, + // Save ONLY this page's keys as dot-notation so sibling AI keys are preserved. + saveTransformer: (s: AiEngineSettingsData) => { + const prov = s.rag?.embeddingProvider || "voyageai"; + const usesBaseUrl = prov === "ollama" || prov === "custom"; + const usesApiKey = prov !== "ollama"; + const deltaSettings: Record = { + "aiEngine.rag.embeddingProvider": prov, + "aiEngine.rag.embeddingModel": s.rag?.embeddingModel ?? "", + // Never send a base URL for a provider that doesn't use one. + "aiEngine.rag.embeddingBaseUrl": usesBaseUrl + ? (s.rag?.embeddingBaseUrl ?? "") + : "", + "aiEngine.rag.topK": clampMin(s.rag?.topK, 1), + "aiEngine.rag.maxSearches": clampMin(s.rag?.maxSearches, 0), + }; + // Send the key when the user typed one, or a provider switch cleared it; the explicit + // "" wipes the previous provider's stored key. + if (embeddingApiKeyDirty) { + deltaSettings["aiEngine.rag.embeddingApiKey"] = usesApiKey + ? (s.rag?.embeddingApiKey ?? "") + : ""; + } + return { sectionData: {}, deltaSettings }; + }, + }); + + useEffect(() => { + fetchSettings(); + }, []); + + const { isDirty, resetToSnapshot, markSaved } = useSettingsDirty( + settings, + loading, + ); + + const handleSave = async () => { + try { + await saveSettings(); + setEmbeddingApiKeyDirty(false); + markSaved(); + // Pushed to the engine live (the embedder is hot-swapped), but documents embedded with the + // previous model must be re-indexed, so the toast flags that caveat. + alert({ + alertType: "success", + title: t("admin.settings.ai.saved.title", "AI settings saved"), + body: `${savedToastBody(settings, t)} ${t( + "admin.settings.ai.documents.saved.reindexNote", + "If you changed the embedding model, re-index existing documents so search uses the new model.", + )}`, + }); + } catch (_error) { + alert({ + alertType: "error", + title: t("admin.error", "Error"), + body: t("admin.settings.saveError", "Failed to save settings"), + }); + } + }; + + const handleDiscard = useCallback(() => { + setEmbeddingApiKeyDirty(false); + setSettings(resetToSnapshot()); + }, [resetToSnapshot, setSettings]); + + const setRag = (patch: Partial) => + setSettings({ ...settings, rag: { ...(settings.rag || {}), ...patch } }); + + if (loading) { + return ( + + + + ); + } + + const embeddingProvider = settings.rag?.embeddingProvider || "voyageai"; + const embeddingSuggestions = + EMBEDDING_MODEL_SUGGESTIONS[embeddingProvider] || []; + const showEmbeddingBaseUrl = + embeddingProvider === "ollama" || embeddingProvider === "custom"; + // Ollama's embeddings endpoint needs no API key, mirroring the Models page. + const showEmbeddingApiKey = embeddingProvider !== "ollama"; + const embeddingApiKeyLabel = + embeddingProvider === "voyageai" + ? t( + "admin.settings.ai.documents.embeddingApiKey.voyageai", + "VoyageAI API key", + ) + : embeddingProvider === "openai" + ? t( + "admin.settings.ai.documents.embeddingApiKey.openai", + "OpenAI API key", + ) + : t( + "admin.settings.ai.documents.embeddingApiKey.generic", + "Embedding API key", + ); + const embeddingApiKeyPlaceholder = + embeddingProvider === "voyageai" + ? "pa-..." + : embeddingProvider === "openai" + ? "sk-..." + : ""; + + return ( +
+ +
+ + {t("admin.settings.ai.documents.title", "Documents & RAG")} + + + {t( + "admin.settings.ai.documents.description", + "Configure the embedding model and retrieval settings used to answer questions over documents. Applied to the AI engine when saved.", + )} + +
+ + + + + + {t("admin.settings.ai.models.provider.label", "Provider")} + + + + } + description={t( + "admin.settings.ai.models.provider.description", + "Which LLM provider the engine talks to.", + )} + data={[ + { value: "anthropic", label: "Anthropic" }, + { value: "openai", label: "OpenAI" }, + { value: "ollama", label: "Ollama" }, + { value: "custom", label: "Custom (OpenAI-compatible)" }, + ]} + value={provider} + onChange={(v) => { + const next = v || "anthropic"; + const patch: Partial = { provider: next }; + // Clear fields the new provider doesn't use so a stale hidden value can't + // leak into the payload (e.g. an Ollama base URL after switching to Anthropic). + if (next !== "ollama" && next !== "custom") patch.baseUrl = ""; + if (next !== provider) { + // Only one key is stored and it belongs to its provider; carrying it across a + // switch would 401 every call while the field still read "Saved", so clear it. + patch.apiKey = ""; + setApiKeyDirty(true); + } + setModels(patch); + }} + allowDeselect={false} + comboboxProps={{ + withinPortal: true, + zIndex: Z_INDEX_OVER_CONFIG_MODAL, + }} + /> + + + + {t( + "admin.settings.ai.models.smartModel.label", + "Smart model", + )} + + + + } + description={t( + "admin.settings.ai.models.smartModel.description", + "High-capability model for complex reasoning. Free text; suggestions are hints only.", + )} + data={modelSuggestions} + value={settings.models?.smartModel || ""} + onChange={(value) => setModels({ smartModel: value })} + comboboxProps={{ + withinPortal: true, + zIndex: Z_INDEX_OVER_CONFIG_MODAL, + }} + /> + + + + {t( + "admin.settings.ai.models.fastModel.label", + "Fast model", + )} + + + + } + description={t( + "admin.settings.ai.models.fastModel.description", + "Cheaper, faster model for lightweight tasks. Free text; suggestions are hints only.", + )} + data={modelSuggestions} + value={settings.models?.fastModel || ""} + onChange={(value) => setModels({ fastModel: value })} + comboboxProps={{ + withinPortal: true, + zIndex: Z_INDEX_OVER_CONFIG_MODAL, + }} + /> + + + + {t( + "admin.settings.ai.models.smartMaxTokens.label", + "Smart model max tokens", + )} + + + + } + description={t( + "admin.settings.ai.models.smartMaxTokens.description", + "Maximum output tokens for the smart model.", + )} + value={settings.models?.smartMaxTokens ?? 0} + onChange={(value) => setModels({ smartMaxTokens: Number(value) })} + min={1} + /> + + + + {t( + "admin.settings.ai.models.fastMaxTokens.label", + "Fast model max tokens", + )} + + + + } + description={t( + "admin.settings.ai.models.fastMaxTokens.description", + "Maximum output tokens for the fast model.", + )} + value={settings.models?.fastMaxTokens ?? 0} + onChange={(value) => setModels({ fastMaxTokens: Number(value) })} + min={1} + /> + + {showApiKey && ( + + {apiKeyLabel} + + + } + description={t( + "admin.settings.ai.models.apiKey.description", + "Leave blank to use the engine's own environment credential. Applies to self-hosted single-engine deployments.", + )} + // Keep the field blank when a key is already stored (returned masked as "********"); + // a pre-filled sentinel would corrupt the key on append, so bind the real value only once typed. + value={apiKeyDirty ? (settings.models?.apiKey ?? "") : ""} + onChange={(e) => { + setApiKeyDirty(true); + setModels({ apiKey: e.target.value }); + }} + placeholder={ + !apiKeyDirty && settings.models?.apiKey === MASKED_SECRET + ? t( + "admin.settings.ai.models.apiKey.setPlaceholder", + "Saved - leave blank to keep the current key", + ) + : apiKeyPlaceholder + } + /> + )} + + {showBaseUrl && ( + + + {t( + "admin.settings.ai.models.baseUrl.label", + "Provider base URL", + )} + + + + } + description={t( + "admin.settings.ai.models.baseUrl.description", + "Base URL of the OpenAI-compatible / Ollama endpoint, e.g. http://ollama:11434/v1.", + )} + value={settings.models?.baseUrl || ""} + onChange={(e) => setModels({ baseUrl: e.target.value })} + placeholder="http://ollama:11434/v1" + /> + )} + + {showBaseUrl && ( + + } + > + + {t( + "admin.settings.ai.models.baseUrl.warning", + "The base URL must point at a trusted internal endpoint. The engine will make server-side requests to it, so an untrusted value is SSRF-sensitive.", + )} + + + )} + + + + + +
+ ); +} diff --git a/frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.test.ts b/frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.test.ts new file mode 100644 index 0000000000..5bd0da6b4f --- /dev/null +++ b/frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it } from "vitest"; +import { + clampMin, + savedToastBody, +} from "@app/components/shared/config/configSections/aiEngineSettings"; + +describe("clampMin", () => { + it("keeps a valid integer unchanged", () => { + expect(clampMin(8192, 1)).toBe(8192); + expect(clampMin(200, 1)).toBe(200); + }); + + it("floors below the minimum for empty / zero / NaN / junk input", () => { + // A cleared NumberInput yields "" -> 0; a transient "-" -> NaN; both must clamp to min. + expect(clampMin("", 1)).toBe(1); + expect(clampMin(0, 1)).toBe(1); + expect(clampMin(Number.NaN, 1)).toBe(1); + expect(clampMin(undefined, 1)).toBe(1); + expect(clampMin("-", 1)).toBe(1); + }); + + it("floors fractional values to an integer", () => { + expect(clampMin(5.7, 1)).toBe(5); + }); + + it("allows zero when the minimum is zero (e.g. maxSearches)", () => { + expect(clampMin(0, 0)).toBe(0); + expect(clampMin("", 0)).toBe(0); + expect(clampMin(4, 0)).toBe(4); + }); +}); + +describe("savedToastBody", () => { + // The helper is passed i18next's t(); echo the key back so assertions read clearly. + const t = (key: string) => key; + + it("promises a live push only when AI is on and config push is enabled", () => { + expect(savedToastBody({ enabled: true, pushConfigToEngine: true }, t)).toBe( + "admin.settings.ai.saved.body", + ); + // pushConfigToEngine defaults to true on the backend, so an absent flag still promises it. + expect(savedToastBody({ enabled: true }, t)).toBe( + "admin.settings.ai.saved.body", + ); + }); + + it("does not promise a push the processor will not make", () => { + // AI off: pushLiveAfterSave returns early, so nothing reaches the engine. + expect(savedToastBody({ enabled: false }, t)).toBe( + "admin.settings.ai.saved.bodyNoPush", + ); + // Env-driven deployment (SaaS pins this false): the engine owns its own config. + expect( + savedToastBody({ enabled: true, pushConfigToEngine: false }, t), + ).toBe("admin.settings.ai.saved.bodyNoPush"); + }); +}); diff --git a/frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.ts b/frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.ts new file mode 100644 index 0000000000..197f521d8b --- /dev/null +++ b/frontend/editor/src/proprietary/components/shared/config/configSections/aiEngineSettings.ts @@ -0,0 +1,101 @@ +// Shared types + constants for the aiEngine admin settings section. +// All four sub-pages share one GET /section/aiEngine payload; each saves only its own keys so siblings survive. + +/** Secret fields come back masked as this literal when a value is set. */ +export const MASKED_SECRET = "********"; + +/** Clamp a numeric setting to a safe integer at submit; a persisted 0 timeout/concurrency would deadlock the engine (Mantine's min only guards on blur). */ +export const clampMin = (value: unknown, min: number): number => + Math.max(min, Math.floor(Number(value) || min)); + +export interface AiEngineModels { + provider?: string; + smartModel?: string; + fastModel?: string; + smartMaxTokens?: number; + fastMaxTokens?: number; + apiKey?: string; + baseUrl?: string; +} + +export interface AiEngineRag { + embeddingProvider?: string; + embeddingModel?: string; + embeddingApiKey?: string; + embeddingBaseUrl?: string; + topK?: number; + maxSearches?: number; +} + +export interface AiEngineLimits { + maxPages?: number; + maxCharacters?: number; + modelMaxConcurrency?: number; +} + +export interface AiEngineFeatures { + chat?: boolean; + documentQuestions?: boolean; + createPdf?: boolean; + mathAuditor?: boolean; + pdfComment?: boolean; + classify?: boolean; +} + +export interface AiEngineSettingsData { + enabled?: boolean; + url?: string; + /** + * Whether the processor forwards settings to the engine; SaaS pins it false, so a save + * persists but never reaches the engine (the toast must say so, not promise a live push). + */ + pushConfigToEngine?: boolean; + timeoutSeconds?: number; + longRunningTimeoutSeconds?: number; + streamTimeoutSeconds?: number; + models?: AiEngineModels; + rag?: AiEngineRag; + limits?: AiEngineLimits; + features?: AiEngineFeatures; +} + +/** + * Post-save toast body; the processor pushes only when AI is enabled AND config push is on, + * so the message must not promise a live push unconditionally. + */ +export const savedToastBody = ( + settings: AiEngineSettingsData, + t: (key: string, fallback: string) => string, +): string => + settings.enabled && settings.pushConfigToEngine !== false + ? t( + "admin.settings.ai.saved.body", + "Changes are pushed to the AI engine automatically.", + ) + : t( + "admin.settings.ai.saved.bodyNoPush", + "Settings saved. They will apply the next time the AI engine picks up its configuration.", + ); + +export interface ApiResponseWithPending { + _pending?: Partial; +} + +export type AiEngineApiResponse = AiEngineSettingsData & + ApiResponseWithPending; + +/** Free-text model suggestions per provider (hints only). */ +export const MODEL_SUGGESTIONS: Record = { + anthropic: ["claude-haiku-4-5", "claude-sonnet-5", "claude-opus-4-8"], + openai: ["gpt-4o", "gpt-4o-mini", "o3-mini"], + ollama: ["llama3.1", "qwen2.5", "mistral"], + custom: ["llama3.1", "qwen2.5", "mistral"], +}; + +/** Free-text embedding-model suggestions per embedding provider. */ +export const EMBEDDING_MODEL_SUGGESTIONS: Record = { + voyageai: ["voyage-4", "voyage-3.5"], + openai: ["text-embedding-3-small", "text-embedding-3-large"], + ollama: ["nomic-embed-text", "mxbai-embed-large", "bge-m3"], + custom: ["nomic-embed-text", "mxbai-embed-large", "bge-m3"], +}; diff --git a/frontend/editor/src/saas/components/shared/config/saasConfigNavSections.test.tsx b/frontend/editor/src/saas/components/shared/config/saasConfigNavSections.test.tsx new file mode 100644 index 0000000000..890cb604b4 --- /dev/null +++ b/frontend/editor/src/saas/components/shared/config/saasConfigNavSections.test.tsx @@ -0,0 +1,52 @@ +import { describe, it, expect } from "vitest"; +import { type TFunction } from "i18next"; +import { allowConsole } from "@app/tests/failOnConsole"; +import { createSaasConfigNavSections } from "@app/components/shared/config/saasConfigNavSections"; + +// Passthrough i18n stub: return the provided fallback (2nd arg) or the key. +const t = ((key: string, fallback?: string) => + fallback ?? key) as unknown as TFunction<"translation", undefined>; + +const Overview = () => null; + +type Sections = ReturnType; + +function itemKeys(sections: Sections): string[] { + return sections.flatMap((s) => s.items.map((i) => i.key)); +} + +// Admin AI settings pages exist only in the self-hosted proprietary flavor; this locks in that +// the AI group can never leak into the SaaS nav (fails loudly if wired into the SaaS cascade). +describe("saasConfigNavSections", () => { + const AI_ITEM_KEYS = [ + "adminAiGeneral", + "adminAiModels", + "adminAiDocuments", + "adminAiLimits", + ]; + + it("never exposes the admin AI settings group or its pages", () => { + // The shared core nav helper warns it is deprecated; incidental to this test. + allowConsole.warn(/createConfigNavSections is deprecated/); + const sections = createSaasConfigNavSections(Overview, () => {}, { t }); + + const keys = itemKeys(sections); + for (const aiKey of AI_ITEM_KEYS) { + expect(keys).not.toContain(aiKey); + } + expect(sections.map((s) => s.title)).not.toContain("AI"); + }); + + it("also hides the AI pages for anonymous users", () => { + allowConsole.warn(/createConfigNavSections is deprecated/); + const sections = createSaasConfigNavSections(Overview, () => {}, { + t, + isAnonymous: true, + }); + + const keys = itemKeys(sections); + for (const aiKey of AI_ITEM_KEYS) { + expect(keys).not.toContain(aiKey); + } + }); +});