diff --git a/Dockerfile b/Dockerfile index 41343fa8..85c9c2f7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,13 +5,13 @@ # ================================ # Multi-stage build for Frontend and Backend # Switched to Debian/Ubuntu-based images to resolve Alpine network issues -# Node is sourced from public.ecr.aws; Liberica Java images come from BellSoft's Docker Hub repositories +# Base images pull through the fleet's dockerhub.haiku.host mirror; anonymous public.ecr.aws pulls hit 429 data limits in CI. # Note: Requires DOCKER_BUILDKIT=1 for cache mount support # ================================ # FRONTEND BUILD STAGE # ================================ -FROM public.ecr.aws/docker/library/node:24.18.0-bookworm-slim AS frontend-builder +FROM dockerhub.haiku.host/library/node:24.18.0-bookworm-slim AS frontend-builder WORKDIR /app/frontend # Copy dependency definitions first for cache layer @@ -39,7 +39,7 @@ RUN npm run validate && npm run test && npm run build # ================================ # BACKEND BUILD STAGE # ================================ -FROM bellsoft/liberica-openjdk-debian:25.0.3-11 AS builder +FROM dockerhub.haiku.host/bellsoft/liberica-openjdk-debian:25.0.3-11 AS builder ARG SOURCE_COMMIT=unknown WORKDIR /app @@ -124,7 +124,7 @@ ENTRYPOINT ["/app/scripts/process_all_to_qdrant.sh", "--app-jar=/app/build/app.j # ================================ # RUNTIME STAGE # ================================ -FROM bellsoft/liberica-openjre-debian:25.0.3-11 AS runtime +FROM dockerhub.haiku.host/bellsoft/liberica-openjre-debian:25.0.3-11 AS runtime LABEL io.iocloudhost.logs.owner=split # 1. System packages (never changes) - FIRST for maximum cache reuse diff --git a/frontend/scripts/verify-node-version.mjs b/frontend/scripts/verify-node-version.mjs index cdbd5359..90a033d2 100644 --- a/frontend/scripts/verify-node-version.mjs +++ b/frontend/scripts/verify-node-version.mjs @@ -24,7 +24,7 @@ const gettingStartedGuide = await readFile( const expectedNodeVersionFile = `v${canonicalNodeVersion}`; const projectedLockfileVersion = frontendLockfile.packages?.[""]?.engines?.node; -const expectedDockerImage = `public.ecr.aws/docker/library/node:${canonicalNodeVersion}-bookworm-slim`; +const expectedDockerImage = `dockerhub.haiku.host/library/node:${canonicalNodeVersion}-bookworm-slim`; const expectedDocumentationVersion = `Node.js ${canonicalNodeVersion}`; if (nodeVersionFile !== expectedNodeVersionFile) { diff --git a/src/main/java/com/williamcallahan/javachat/application/prompt/PromptTruncator.java b/src/main/java/com/williamcallahan/javachat/application/prompt/PromptTruncator.java index 91d50aa7..f29bbf25 100644 --- a/src/main/java/com/williamcallahan/javachat/application/prompt/PromptTruncator.java +++ b/src/main/java/com/williamcallahan/javachat/application/prompt/PromptTruncator.java @@ -45,8 +45,9 @@ public class PromptTruncator { * @throws AuthoritativeContextDoesNotFitException when no HIGH-priority context segment fits */ public TruncatedPrompt truncate(StructuredPrompt prompt, int maxTokens) { - int reservedTokens = - prompt.system().estimatedTokens() + prompt.currentQuery().estimatedTokens(); + int reservedTokens = prompt.system().estimatedTokens() + + prompt.currentQuery().estimatedTokens() + + prompt.turnGuidance().estimatedTokens(); List authoritativeContextDocuments = prompt.contextDocuments().stream() .filter(contextDocument -> contextDocument.priority() == PromptSegmentPriority.HIGH) .toList(); @@ -58,8 +59,8 @@ public TruncatedPrompt truncate(StructuredPrompt prompt, int maxTokens) { prompt.currentQuery().estimatedTokens(), maxTokens); // Return prompt with only system and query - no room for context or history - StructuredPrompt minimalPrompt = - new StructuredPrompt(prompt.system(), List.of(), List.of(), prompt.currentQuery()); + StructuredPrompt minimalPrompt = new StructuredPrompt( + prompt.system(), List.of(), List.of(), prompt.currentQuery(), prompt.turnGuidance()); return new TruncatedPrompt(minimalPrompt, true); } @@ -96,8 +97,8 @@ public TruncatedPrompt truncate(StructuredPrompt prompt, int maxTokens) { log.debug("Truncated context documents from {} to {}", originalDocCount, fittingDocs.size()); } - StructuredPrompt truncated = - new StructuredPrompt(prompt.system(), reindexedDocuments, fittingTurns, prompt.currentQuery()); + StructuredPrompt truncated = new StructuredPrompt( + prompt.system(), reindexedDocuments, fittingTurns, prompt.currentQuery(), prompt.turnGuidance()); if (wasTruncated) { log.info( diff --git a/src/main/java/com/williamcallahan/javachat/domain/prompt/CompletionPrompt.java b/src/main/java/com/williamcallahan/javachat/domain/prompt/CompletionPrompt.java new file mode 100644 index 00000000..068d0764 --- /dev/null +++ b/src/main/java/com/williamcallahan/javachat/domain/prompt/CompletionPrompt.java @@ -0,0 +1,12 @@ +package com.williamcallahan.javachat.domain.prompt; + +import java.util.Objects; + +/** Separates reusable task instructions from the current completion's untrusted input. */ +public record CompletionPrompt(String instructions, String input) { + /** Requires both boundary fields so instruction omission is explicit. */ + public CompletionPrompt { + Objects.requireNonNull(instructions, "instructions"); + Objects.requireNonNull(input, "input"); + } +} diff --git a/src/main/java/com/williamcallahan/javachat/domain/prompt/StructuredPrompt.java b/src/main/java/com/williamcallahan/javachat/domain/prompt/StructuredPrompt.java index 65d187a9..5e611e63 100644 --- a/src/main/java/com/williamcallahan/javachat/domain/prompt/StructuredPrompt.java +++ b/src/main/java/com/williamcallahan/javachat/domain/prompt/StructuredPrompt.java @@ -14,16 +14,32 @@ * @param contextDocuments authoritative HIGH or ordinary LOW context documents * @param conversationHistory past conversation turns (MEDIUM priority) * @param currentQuery the user's current question (HIGH priority) + * @param turnGuidance application-owned guidance specific to the current turn */ public record StructuredPrompt( SystemSegment system, List contextDocuments, List conversationHistory, - CurrentQuerySegment currentQuery) { + CurrentQuerySegment currentQuery, + SystemSegment turnGuidance) { /** Separator between prompt segments. */ public static final String SEGMENT_SEPARATOR = "\n\n"; + /** Creates a prompt without turn-specific guidance. */ + public StructuredPrompt( + SystemSegment system, + List contextDocuments, + List conversationHistory, + CurrentQuerySegment currentQuery) { + this(system, contextDocuments, conversationHistory, currentQuery, new SystemSegment("", 0)); + } + + /** Keeps turn-specific instructions outside the reusable system prefix. */ + public StructuredPrompt withTurnGuidance(SystemSegment guidance) { + return new StructuredPrompt(system, contextDocuments, conversationHistory, currentQuery, guidance); + } + /** * Creates a structured prompt with defensive copies of lists. * @@ -36,6 +52,7 @@ public record StructuredPrompt( if (currentQuery == null) { throw new IllegalArgumentException("Current query segment cannot be null"); } + java.util.Objects.requireNonNull(turnGuidance, "turnGuidance"); contextDocuments = contextDocuments == null ? List.of() : List.copyOf(contextDocuments); conversationHistory = conversationHistory == null ? List.of() : List.copyOf(conversationHistory); } @@ -46,7 +63,7 @@ public record StructuredPrompt( * @return sum of all segment token estimates */ public int totalEstimatedTokens() { - int total = system.estimatedTokens() + currentQuery.estimatedTokens(); + int total = system.estimatedTokens() + currentQuery.estimatedTokens() + turnGuidance.estimatedTokens(); for (ContextDocumentSegment contextDocument : contextDocuments) { total += contextDocument.estimatedTokens(); } @@ -60,7 +77,7 @@ public int totalEstimatedTokens() { * Renders the complete prompt to a string for consumers that require one text value. * *

Segments are joined with paragraph separators in order: system instructions, - * context documents, conversation history, and the current query.

+ * conversation history, turn guidance, context documents, and the current query.

* * @return the fully assembled prompt string */ @@ -79,13 +96,15 @@ public String render() { public String renderInput() { StringJoiner promptInput = new StringJoiner(SEGMENT_SEPARATOR); - for (ContextDocumentSegment contextDocument : contextDocuments) { - promptInput.add(contextDocument.content()); - } - for (ConversationTurnSegment conversationTurn : conversationHistory) { promptInput.add(conversationTurn.content()); } + if (!turnGuidance.content().isBlank()) { + promptInput.add(turnGuidance.content()); + } + for (ContextDocumentSegment contextDocument : contextDocuments) { + promptInput.add(contextDocument.content()); + } promptInput.add(currentQuery.content()); @@ -99,7 +118,7 @@ public String renderInput() { * @return new prompt with updated documents */ public StructuredPrompt withContextDocuments(List newContextDocuments) { - return new StructuredPrompt(system, newContextDocuments, conversationHistory, currentQuery); + return new StructuredPrompt(system, newContextDocuments, conversationHistory, currentQuery, turnGuidance); } /** @@ -109,7 +128,7 @@ public StructuredPrompt withContextDocuments(List newCon * @return new prompt with updated history */ public StructuredPrompt withConversationHistory(List newConversationHistory) { - return new StructuredPrompt(system, contextDocuments, newConversationHistory, currentQuery); + return new StructuredPrompt(system, contextDocuments, newConversationHistory, currentQuery, turnGuidance); } /** diff --git a/src/main/java/com/williamcallahan/javachat/service/ChatService.java b/src/main/java/com/williamcallahan/javachat/service/ChatService.java index ab85eb75..93c2f928 100644 --- a/src/main/java/com/williamcallahan/javachat/service/ChatService.java +++ b/src/main/java/com/williamcallahan/javachat/service/ChatService.java @@ -118,13 +118,9 @@ public RetrievalService.CitationOutcome citationOutcomeForRetainedContext( public StructuredPrompt buildStructuredPromptWithContextAndGuidance( List history, String latestUserMessage, List contextDocs, String guidance) { - // Build system prompt with guidance String basePrompt = systemPromptConfig.getCoreSystemPrompt(); - String completePrompt = guidance != null && !guidance.isBlank() - ? systemPromptConfig.buildFullPrompt(basePrompt, guidance) - : basePrompt; - - SystemSegment systemSegment = new SystemSegment(completePrompt, estimateTokens(completePrompt)); + SystemSegment systemSegment = new SystemSegment(basePrompt, estimateTokens(basePrompt)); + String turnGuidance = guidance == null ? "" : guidance; List contextSegments = buildContextSegments(contextDocs != null ? contextDocs : List.of(), latestUserMessage); @@ -134,7 +130,8 @@ public StructuredPrompt buildStructuredPromptWithContextAndGuidance( CurrentQuerySegment querySegment = new CurrentQuerySegment(latestUserMessage, estimateTokens(latestUserMessage)); - return new StructuredPrompt(systemSegment, contextSegments, conversationSegments, querySegment); + return new StructuredPrompt(systemSegment, contextSegments, conversationSegments, querySegment) + .withTurnGuidance(new SystemSegment(turnGuidance, estimateTokens(turnGuidance))); } /** @@ -195,18 +192,16 @@ public StructuredPromptOutcome buildStructuredPromptWithContextOutcome( String searchQualityNote = determineSearchQuality(contextDocs); - // Build system prompt with search quality context - StringBuilder systemPromptBuilder = new StringBuilder(systemPromptConfig.getCoreSystemPrompt()); + String coreSystemPrompt = systemPromptConfig.getCoreSystemPrompt(); + StringBuilder searchGuidance = new StringBuilder(); if (!searchQualityNote.isEmpty()) { - systemPromptBuilder.append("\n\nSEARCH CONTEXT: ").append(searchQualityNote); + searchGuidance.append("SEARCH CONTEXT: ").append(searchQualityNote); if (searchQualityNote.contains("less relevant") || searchQualityNote.contains("keyword search")) { - systemPromptBuilder.append("\n").append(systemPromptConfig.getLowQualitySearchPrompt()); + searchGuidance.append("\n").append(systemPromptConfig.getLowQualitySearchPrompt()); } } - String systemPromptText = systemPromptBuilder.toString(); - - // Build structured segments - SystemSegment systemSegment = new SystemSegment(systemPromptText, estimateTokens(systemPromptText)); + SystemSegment systemSegment = new SystemSegment(coreSystemPrompt, estimateTokens(coreSystemPrompt)); + String turnGuidance = searchGuidance.toString(); List contextSegments = buildContextSegments(contextDocs, latestUserMessage); List conversationSegments = buildConversationSegments(history); @@ -214,8 +209,9 @@ public StructuredPromptOutcome buildStructuredPromptWithContextOutcome( CurrentQuerySegment querySegment = new CurrentQuerySegment(latestUserMessage, estimateTokens(latestUserMessage)); - StructuredPrompt structuredPrompt = - new StructuredPrompt(systemSegment, contextSegments, conversationSegments, querySegment); + StructuredPrompt structuredPrompt = new StructuredPrompt( + systemSegment, contextSegments, conversationSegments, querySegment) + .withTurnGuidance(new SystemSegment(turnGuidance, estimateTokens(turnGuidance))); return new StructuredPromptOutcome(structuredPrompt, contextDocs); } diff --git a/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java b/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java index eb058298..58553b15 100644 --- a/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java +++ b/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java @@ -23,6 +23,7 @@ final class EmbeddingBatchEmbedder { static final int MAX_CONCURRENT_EMBEDDING_REQUESTS = 8; static final int MAX_EMBEDDING_ATTEMPTS = 3; private static final Duration EMBEDDING_RETRY_BASE_DELAY = Duration.ofSeconds(5); + private static final Duration EMBEDDING_BATCH_RETRY_TIMEOUT = Duration.ofMinutes(10); private EmbeddingBatchEmbedder() {} @@ -158,23 +159,42 @@ private static List embedSingleBatch( private static List embedWithRetry( LlmGatewayTier.Purpose purpose, EmbeddingClient embeddingClient, List textBatch) { - for (int embeddingAttempt = 1; embeddingAttempt <= MAX_EMBEDDING_ATTEMPTS; embeddingAttempt++) { + long retryDeadlineNanos = System.nanoTime() + EMBEDDING_BATCH_RETRY_TIMEOUT.toNanos(); + int providerAttemptCount = 0; + while (providerAttemptCount < MAX_EMBEDDING_ATTEMPTS) { + Duration remainingBudget = Duration.ofNanos(Math.max(0, retryDeadlineNanos - System.nanoTime())); + if (remainingBudget.isZero()) { + throw new EmbeddingServiceTemporarilyUnavailableException("Embedding batch retry deadline elapsed"); + } try { - return embeddingClient.embed(purpose, textBatch, LlmGatewayTier.BATCH); + return embeddingClient.embed(purpose, textBatch, LlmGatewayTier.BATCH, remainingBudget); } catch (EmbeddingServiceTemporarilyUnavailableException temporaryFailure) { - if (embeddingAttempt == MAX_EMBEDDING_ATTEMPTS) { + if (temporaryFailure.providerContacted()) { + providerAttemptCount++; + } + if (providerAttemptCount == MAX_EMBEDDING_ATTEMPTS) { + throw temporaryFailure; + } + Duration retryDelay = temporaryFailure.retryDelay(); + Duration backoff = temporaryFailure.providerContacted() + ? EMBEDDING_RETRY_BASE_DELAY.multipliedBy(providerAttemptCount) + : Duration.ZERO; + if (backoff.compareTo(retryDelay) > 0) { + retryDelay = backoff; + } + if (retryDelay.toNanos() >= retryDeadlineNanos - System.nanoTime()) { throw temporaryFailure; } - awaitRetryDelay(embeddingAttempt, temporaryFailure); + awaitRetryDelay(retryDelay, temporaryFailure); } } throw new IllegalStateException("Embedding retry loop exhausted without a terminal outcome"); } private static void awaitRetryDelay( - int completedAttemptCount, EmbeddingServiceTemporarilyUnavailableException temporaryFailure) { + Duration retryDelay, EmbeddingServiceTemporarilyUnavailableException temporaryFailure) { try { - Thread.sleep(EMBEDDING_RETRY_BASE_DELAY.multipliedBy(completedAttemptCount)); + Thread.sleep(retryDelay); } catch (InterruptedException interruptedException) { Thread.currentThread().interrupt(); EmbeddingServiceTemporarilyUnavailableException interruptionFailure = diff --git a/src/main/java/com/williamcallahan/javachat/service/EmbeddingServiceTemporarilyUnavailableException.java b/src/main/java/com/williamcallahan/javachat/service/EmbeddingServiceTemporarilyUnavailableException.java index 8b224775..5491cbb8 100644 --- a/src/main/java/com/williamcallahan/javachat/service/EmbeddingServiceTemporarilyUnavailableException.java +++ b/src/main/java/com/williamcallahan/javachat/service/EmbeddingServiceTemporarilyUnavailableException.java @@ -1,5 +1,7 @@ package com.williamcallahan.javachat.service; +import java.time.Duration; + /** * Signals that embedding work can be retried after a transient admission or provider failure. * @@ -7,6 +9,7 @@ * (no provider contact made) from a genuine transient provider failure through subtyping.

*/ public class EmbeddingServiceTemporarilyUnavailableException extends EmbeddingServiceUnavailableException { + private final long retryNotBeforeNanos; /** * Creates a retryable embedding failure with a human-readable explanation. @@ -14,7 +17,7 @@ public class EmbeddingServiceTemporarilyUnavailableException extends EmbeddingSe * @param message explanation of the transient failure */ public EmbeddingServiceTemporarilyUnavailableException(String message) { - super(message); + this(message, null, Duration.ZERO); } /** @@ -24,6 +27,22 @@ public EmbeddingServiceTemporarilyUnavailableException(String message) { * @param cause underlying provider or transport failure */ public EmbeddingServiceTemporarilyUnavailableException(String message, Throwable cause) { + this(message, cause, Duration.ZERO); + } + + /** Preserves authoritative provider retry timing across the batch boundary. */ + public EmbeddingServiceTemporarilyUnavailableException(String message, Throwable cause, Duration retryDelay) { super(message, cause); + retryNotBeforeNanos = System.nanoTime() + retryDelay.toNanos(); + } + + /** Returns the remaining authoritative cooldown without parsing diagnostic text. */ + public Duration retryDelay() { + return Duration.ofNanos(Math.max(0, retryNotBeforeNanos - System.nanoTime())); + } + + /** Distinguishes provider failures from local admission rejection for retry budgeting. */ + public boolean providerContacted() { + return true; } } diff --git a/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java b/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java index 0ce302b7..631f48a1 100644 --- a/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java +++ b/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java @@ -8,6 +8,7 @@ import com.fasterxml.jackson.databind.cfg.CoercionInputShape; import com.fasterxml.jackson.databind.type.LogicalType; import com.williamcallahan.javachat.config.AppProperties; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.model.Enrichment; import java.util.List; import org.slf4j.Logger; @@ -48,25 +49,29 @@ public EnrichmentService( @Cacheable("enrichment-cache") public Enrichment enrich(String userQuery, String jdkVersion, List contextSnippets) { logger.debug("EnrichmentService.enrich called"); - StringBuilder prompt = new StringBuilder(); - // Add base system context for consistency - prompt.append("You are operating as part of the Java learning assistant system.\n"); - prompt.append("Extract enrichment metadata to enhance learning from the provided context.\n\n"); - prompt.append("User query: ").append(userQuery).append("\n\n"); - prompt.append("Context snippets from Java documentation:\n"); - for (int snippetIndex = 0; snippetIndex < contextSnippets.size(); snippetIndex++) { - prompt.append("- ").append(contextSnippets.get(snippetIndex)).append("\n"); + String instructions = """ + You are operating as part of the Java learning assistant system. + Extract enrichment metadata to enhance learning from the provided context. + Respond as JSON with these fields: + - hints[]: Practical tips and best practices (2-3 items max) + - reminders[]: Critical things to remember (2-3 items max) + - background[]: Conceptual explanations for deeper understanding (2-3 items max) + - packageName: The main Java package being discussed + - jdkVersion: The requested JDK version unless docs indicate otherwise + - resource: Primary documentation source + - resourceVersion: Version of the resource + Make each item concise but informative. Focus on learning value. + Use your knowledge of Java best practices and common pitfalls. + """; + StringBuilder promptInput = new StringBuilder("Requested JDK version: ") + .append(jdkVersion) + .append("\nUser query: ") + .append(userQuery) + .append("\n\nContext snippets from Java documentation:\n"); + for (String contextSnippet : contextSnippets) { + promptInput.append("- ").append(contextSnippet).append("\n"); } - prompt.append("\nRespond as JSON with these fields:\n"); - prompt.append("- hints[]: Practical tips and best practices (2-3 items max)\n"); - prompt.append("- reminders[]: Critical things to remember (2-3 items max)\n"); - prompt.append("- background[]: Conceptual explanations for deeper understanding (2-3 items max)\n"); - prompt.append("- packageName: The main Java package being discussed\n"); - prompt.append("- jdkVersion: '").append(jdkVersion).append("' unless docs indicate otherwise\n"); - prompt.append("- resource: Primary documentation source\n"); - prompt.append("- resourceVersion: Version of the resource\n"); - prompt.append("\nMake each item concise but informative. Focus on learning value.\n"); - prompt.append("Use your knowledge of Java best practices and common pitfalls.\n"); + CompletionPrompt prompt = new CompletionPrompt(instructions, promptInput.toString()); if (openAIStreamingService == null || !openAIStreamingService.isAvailable()) { throw new IllegalStateException("OpenAIStreamingService unavailable for enrichment"); @@ -75,7 +80,7 @@ public Enrichment enrich(String userQuery, String jdkVersion, List conte String enrichmentJson = openAIStreamingService .completeJsonObject( LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, - prompt.toString(), + prompt, enrichmentTemperature, enrichmentOutputTokenBudget) .block(); diff --git a/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java b/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java index dc760946..5cc0fd91 100644 --- a/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java +++ b/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java @@ -6,8 +6,11 @@ import com.openai.core.Timeout; import com.openai.core.http.AsyncStreamResponse; import com.openai.models.responses.Response; +import com.openai.models.responses.ResponseCompletedEvent; import com.openai.models.responses.ResponseCreateParams; import com.openai.models.responses.ResponseError; +import com.openai.models.responses.ResponseFailedEvent; +import com.openai.models.responses.ResponseIncompleteEvent; import com.openai.models.responses.ResponseOutputItem; import com.openai.models.responses.ResponseOutputMessage; import com.openai.models.responses.ResponseOutputText; @@ -17,10 +20,12 @@ import com.williamcallahan.javachat.application.completion.CompletionRequestConfiguration; import com.williamcallahan.javachat.application.streaming.ReportedStreamingFailure; import com.williamcallahan.javachat.application.streaming.StreamingFailureReporter; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.domain.prompt.ContextDocumentSegment; import com.williamcallahan.javachat.domain.prompt.StructuredPrompt; import com.williamcallahan.javachat.domain.text.UnicodeVisibleContent; import com.williamcallahan.javachat.support.OpenAiSdkUrlNormalizer; +import io.micrometer.core.instrument.MeterRegistry; import jakarta.annotation.PostConstruct; import jakarta.annotation.PreDestroy; import java.time.Duration; @@ -50,7 +55,7 @@ public class OpenAIStreamingService { private OpenAIClient openAiClient; private volatile boolean isAvailable; - private final RateLimitService rateLimitService; + private final MeterRegistry meterRegistry; private final OpenAiRequestFactory requestFactory; private final OpenAiProviderRoutingService providerRoutingService; private final StreamingFailureReporter streamingFailureReporter; @@ -70,11 +75,11 @@ public class OpenAIStreamingService { /** Keeps provider admission, request construction, and failure reporting independently owned. */ public OpenAIStreamingService( - RateLimitService rateLimitService, + MeterRegistry meterRegistry, OpenAiRequestFactory requestFactory, OpenAiProviderRoutingService providerRoutingService, StreamingFailureReporter streamingFailureReporter) { - this.rateLimitService = rateLimitService; + this.meterRegistry = meterRegistry; this.requestFactory = requestFactory; this.providerRoutingService = providerRoutingService; this.streamingFailureReporter = streamingFailureReporter; @@ -126,7 +131,7 @@ public Mono streamResponse( RateLimitService.ApiProvider configuredProvider = providerRoutingService.configuredProvider(); OpenAiPreparedRequest preparedStreamingRequest = - requestFactory.prepareStreamingRequest(structuredPrompt, temperature); + requestFactory.prepareStreamingRequest(purpose, structuredPrompt, temperature); Flux responseTextChunks = executeStreamingWithConfiguredProvider( preparedStreamingRequest, configuredProvider, purpose); List contextDocumentIds = @@ -140,19 +145,23 @@ public Mono streamResponse( /** Uses the configured output budget for a non-streaming text completion. */ public Mono complete(LlmGatewayTier.Purpose purpose, String prompt, double temperature) { - return executeCompletion(purpose, prompt, temperature, CompletionRequestConfiguration.defaultText()); + return executeCompletion( + purpose, new CompletionPrompt("", prompt), temperature, CompletionRequestConfiguration.defaultText()); } /** Preserves the caller's output budget for a non-streaming text completion. */ public Mono complete( LlmGatewayTier.Purpose purpose, String prompt, double temperature, int maximumOutputTokens) { return Mono.defer(() -> executeCompletion( - purpose, prompt, temperature, CompletionRequestConfiguration.boundedText(maximumOutputTokens))); + purpose, + new CompletionPrompt("", prompt), + temperature, + CompletionRequestConfiguration.boundedText(maximumOutputTokens))); } /** Requires a JSON object within the caller's output budget and the default request timeout. */ public Mono completeJsonObject( - LlmGatewayTier.Purpose purpose, String prompt, double temperature, int maximumOutputTokens) { + LlmGatewayTier.Purpose purpose, CompletionPrompt prompt, double temperature, int maximumOutputTokens) { return completeJsonObject( purpose, prompt, @@ -164,7 +173,7 @@ public Mono completeJsonObject( /** Requires a JSON object within both the caller's output budget and whole-request timeout. */ public Mono completeJsonObject( LlmGatewayTier.Purpose purpose, - String prompt, + CompletionPrompt prompt, double temperature, int maximumOutputTokens, Duration requestTimeout) { @@ -177,7 +186,7 @@ public Mono completeJsonObject( private Mono executeCompletion( LlmGatewayTier.Purpose purpose, - String prompt, + CompletionPrompt prompt, double temperature, CompletionRequestConfiguration configuration) { java.util.Objects.requireNonNull(purpose, "purpose"); @@ -185,6 +194,7 @@ private Mono executeCompletion( RateLimitService.ApiProvider configuredProvider = providerRoutingService.configuredProvider(); ResponseCreateParams requestParameters = buildCompletionRequest(prompt, temperature, configuration).toBuilder() + .promptCacheKey(requestFactory.promptCacheKey(purpose, prompt.instructions())) .putAdditionalHeader( LlmGatewayTier.REQUEST_WORKLOAD_HEADER, workloadTier().workloadHeader(purpose)) @@ -202,7 +212,10 @@ private Mono executeCompletion( .async() .responses() .create(requestParameters, requestOptions) - .thenApply(this::extractTextFromCompletedResponse); + .thenApply(providerResponse -> { + recordUsage(providerResponse, purpose); + return extractTextFromCompletedResponse(providerResponse); + }); CompletableFuture accountedCompletionFuture = completionFuture.whenComplete((completionText, completionFailure) -> { if (completionFailure == null) { @@ -222,16 +235,16 @@ private Mono executeCompletion( } private ResponseCreateParams buildCompletionRequest( - String prompt, double temperature, CompletionRequestConfiguration configuration) { + CompletionPrompt prompt, double temperature, CompletionRequestConfiguration configuration) { if (configuration.requireJsonObject()) { return requestFactory.buildJsonCompletionRequest( prompt, temperature, configuration.maximumOutputTokens().orElseThrow()); } if (configuration.maximumOutputTokens().isEmpty()) { - return requestFactory.buildCompletionRequest(prompt, temperature); + return requestFactory.buildCompletionRequest(prompt.input(), temperature); } return requestFactory.buildCompletionRequest( - prompt, temperature, configuration.maximumOutputTokens().orElseThrow()); + prompt.input(), temperature, configuration.maximumOutputTokens().orElseThrow()); } /** Classifies the upstream failure even when the reporting boundary has wrapped it. */ @@ -261,7 +274,8 @@ private Flux executeStreamingWithConfiguredProvider( LlmGatewayTier.REQUEST_WORKLOAD_HEADER, workloadTier().workloadHeader(purpose)) .build(), - configuredProvider) + configuredProvider, + purpose) .doOnNext(textChunk -> { if (UnicodeVisibleContent.hasVisibleContent(textChunk)) { emittedVisibleText.set(true); @@ -281,7 +295,10 @@ private Flux executeStreamingWithConfiguredProvider( } private Flux executeStreamingRequest( - OpenAIClient client, ResponseCreateParams requestParameters, RateLimitService.ApiProvider activeProvider) { + OpenAIClient client, + ResponseCreateParams requestParameters, + RateLimitService.ApiProvider activeProvider, + LlmGatewayTier.Purpose purpose) { RequestOptions requestOptions = RequestOptions.builder() .timeout(Timeout.builder().request(streamingRequestTimeout()).build()) .build(); @@ -289,11 +306,22 @@ private Flux executeStreamingRequest( return Flux.defer(() -> { AtomicBoolean emittedVisibleText = new AtomicBoolean(false); AtomicBoolean observedCompletedEvent = new AtomicBoolean(false); + AtomicBoolean recordedTerminalUsage = new AtomicBoolean(false); Flux responseTextChunks = asyncResponseEvents(client, requestParameters, requestOptions) .concatMap(responseStreamEvent -> { if (responseStreamEvent.completed().isPresent()) { observedCompletedEvent.set(true); } + responseStreamEvent + .completed() + .map(ResponseCompletedEvent::response) + .or(() -> responseStreamEvent.failed().map(ResponseFailedEvent::response)) + .or(() -> responseStreamEvent.incomplete().map(ResponseIncompleteEvent::response)) + .ifPresent(terminalResponse -> { + if (recordedTerminalUsage.compareAndSet(false, true)) { + recordUsage(terminalResponse, purpose); + } + }); return extractTextOrTerminalFailure(responseStreamEvent); }); return enforceVisibleOutputDeadline(responseTextChunks) @@ -469,6 +497,29 @@ private Mono extractTextOrTerminalFailure(ResponseStreamEvent responseSt .or(() -> responseStreamEvent.refusalDelta().map(refusalDeltaEvent -> refusalDeltaEvent.delta()))); } + private void recordUsage(Response providerResponse, LlmGatewayTier.Purpose purpose) { + if (providerResponse == null) { + return; + } + // Gateways may omit usage detail fields; the SDK's required accessors would throw and fail + // a completed call, so only reported counts are recorded. + providerResponse.usage().ifPresent(usage -> { + String workload = purpose.name(); + usage._inputTokens().asKnown().ifPresent(count -> countTokens(workload, "input", count)); + usage._inputTokensDetails().asKnown().ifPresent(details -> { + details._cachedTokens().asKnown().ifPresent(count -> countTokens(workload, "cached_input", count)); + details._cacheWriteTokens().asKnown().ifPresent(count -> countTokens(workload, "cache_write", count)); + }); + usage._outputTokens().asKnown().ifPresent(count -> countTokens(workload, "output", count)); + }); + } + + private void countTokens(String workload, String tokenType, long count) { + meterRegistry + .counter("llm.tokens", "purpose", workload, "type", tokenType) + .increment(count); + } + private String extractTextFromCompletedResponse(Response providerResponse) { if (providerResponse == null) { throw OpenAiResponseException.missingCompletion(); diff --git a/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java b/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java index 492dddcb..082ef238 100644 --- a/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java +++ b/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java @@ -187,7 +187,8 @@ private synchronized void rejectActive() { if (remainingCooldownNanos > 0) { long retryAfterSeconds = Math.max(1L, TimeUnit.NANOSECONDS.toSeconds(remainingCooldownNanos)); throw new EmbeddingProviderCooldownRejectionException( - "Remote embedding provider is rate limited; retry after " + retryAfterSeconds + " seconds"); + "Remote embedding provider is rate limited; retry after " + retryAfterSeconds + " seconds", + Duration.ofNanos(remainingCooldownNanos)); } } @@ -424,7 +425,12 @@ private EmbeddingServiceUnavailableException wrapServiceError(OpenAIServiceExcep int statusCode = exception.statusCode(); if (statusCode == HttpStatus.TOO_MANY_REQUESTS.value()) { try { - rateLimitDecisionResolver.resolveFromOpenAiRetryAfterHeaders(exception.headers()); + RateLimitDecision retryDecision = + rateLimitDecisionResolver.resolveFromOpenAiRetryAfterHeaders(exception.headers()); + return new EmbeddingServiceTemporarilyUnavailableException( + "Remote embedding provider returned HTTP 429", + exception, + Duration.ofSeconds(retryDecision.retryAfterSeconds())); } catch (RateLimitDecisionException rateLimitDecisionFailure) { rateLimitDecisionFailure.addSuppressed(exception); return new EmbeddingServiceTemporarilyUnavailableException( @@ -810,8 +816,13 @@ static final class EmbeddingProviderCooldownRejectionException extends EmbeddingServiceTemporarilyUnavailableException { private static final long serialVersionUID = 1L; - EmbeddingProviderCooldownRejectionException(String message) { - super(message); + EmbeddingProviderCooldownRejectionException(String message, Duration retryDelay) { + super(message, null, retryDelay); + } + + @Override + public boolean providerContacted() { + return false; } } } diff --git a/src/main/java/com/williamcallahan/javachat/service/OpenAiRequestFactory.java b/src/main/java/com/williamcallahan/javachat/service/OpenAiRequestFactory.java index b00a3ffa..4c03ac9d 100644 --- a/src/main/java/com/williamcallahan/javachat/service/OpenAiRequestFactory.java +++ b/src/main/java/com/williamcallahan/javachat/service/OpenAiRequestFactory.java @@ -12,14 +12,17 @@ import com.williamcallahan.javachat.application.prompt.PromptTruncator; import com.williamcallahan.javachat.config.AppProperties; import com.williamcallahan.javachat.config.ModelConfiguration; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.domain.prompt.ContextDocumentSegment; import com.williamcallahan.javachat.domain.prompt.ConversationTurnSegment; import com.williamcallahan.javachat.domain.prompt.StructuredPrompt; import com.williamcallahan.javachat.support.AsciiTextNormalizer; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.UUID; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.factory.annotation.Value; @@ -80,11 +83,13 @@ public OpenAiRequestFactory( /** * Builds a streaming request payload and returns the resolved model identifier. * + * @param purpose application operation whose stable prefix owns cache affinity * @param structuredPrompt typed prompt segments to stream * @param temperature response temperature * @return request parameters with resolved model id */ - public OpenAiPreparedRequest prepareStreamingRequest(StructuredPrompt structuredPrompt, double temperature) { + public OpenAiPreparedRequest prepareStreamingRequest( + LlmGatewayTier.Purpose purpose, StructuredPrompt structuredPrompt, double temperature) { String modelId = configuredModelId(); PromptTruncator.TruncatedPrompt truncatedPrompt = @@ -100,6 +105,8 @@ public OpenAiPreparedRequest prepareStreamingRequest(StructuredPrompt structured List responseInputItems = buildResponseInputItems(truncatedPrompt); ResponseCreateParams responseParams = buildResponseParams(responseInputItems, temperature, modelId).toBuilder() .instructions(truncatedPrompt.prompt().system().content()) + .promptCacheKey(promptCacheKey( + purpose, truncatedPrompt.prompt().system().content())) .build(); return new OpenAiPreparedRequest(responseParams, modelId, truncatedPrompt.prompt()); } @@ -114,17 +121,21 @@ public OpenAiPreparedRequest prepareStreamingRequest(StructuredPrompt structured private static List buildResponseInputItems(PromptTruncator.TruncatedPrompt truncatedPrompt) { StructuredPrompt prompt = truncatedPrompt.prompt(); List responseInputItems = new ArrayList<>(); + for (ConversationTurnSegment conversationTurn : prompt.conversationHistory()) { + EasyInputMessage.Role role = + conversationTurn.isAssistantTurn() ? EasyInputMessage.Role.ASSISTANT : EasyInputMessage.Role.USER; + responseInputItems.add(textInputItem(role, conversationTurn.messageText())); + } + if (!prompt.turnGuidance().content().isBlank()) { + responseInputItems.add(textInputItem( + EasyInputMessage.Role.DEVELOPER, prompt.turnGuidance().content())); + } if (truncatedPrompt.wasTruncated()) { responseInputItems.add(textInputItem(EasyInputMessage.Role.DEVELOPER, TRUNCATION_NOTICE_GENERIC.strip())); } for (ContextDocumentSegment contextDocument : prompt.contextDocuments()) { responseInputItems.add(textInputItem(EasyInputMessage.Role.USER, contextDocument.content())); } - for (ConversationTurnSegment conversationTurn : prompt.conversationHistory()) { - EasyInputMessage.Role role = - conversationTurn.isAssistantTurn() ? EasyInputMessage.Role.ASSISTANT : EasyInputMessage.Role.USER; - responseInputItems.add(textInputItem(role, conversationTurn.messageText())); - } responseInputItems.add( textInputItem(EasyInputMessage.Role.USER, prompt.currentQuery().queryText())); return List.copyOf(responseInputItems); @@ -170,26 +181,32 @@ public ResponseCreateParams buildCompletionRequest(String prompt, double tempera * @param maximumOutputTokens maximum output tokens needed by this caller * @return request payload with a declared JSON-object output contract */ - public ResponseCreateParams buildJsonCompletionRequest(String prompt, double temperature, int maximumOutputTokens) { + public ResponseCreateParams buildJsonCompletionRequest( + CompletionPrompt prompt, double temperature, int maximumOutputTokens) { if (maximumOutputTokens <= 0) { throw new IllegalArgumentException("maximumOutputTokens must be positive"); } - return buildCompletionRequest(prompt, temperature, Integer.valueOf(maximumOutputTokens), true); + int availableInputTokens = Math.max(0, promptContentTokenBudget - chunker.countTokens(prompt.instructions())); + String truncatedInput = truncatePromptForCompletion(prompt.input(), availableInputTokens); + return buildResponseParams(truncatedInput, temperature, configuredModelId(), maximumOutputTokens, true) + .toBuilder() + .instructions(prompt.instructions()) + .build(); } private ResponseCreateParams buildCompletionRequest( String prompt, double temperature, Integer maximumOutputTokens, boolean requireJsonObject) { String modelId = configuredModelId(); - String truncatedPrompt = truncatePromptForCompletion(prompt); + String truncatedPrompt = truncatePromptForCompletion(prompt, promptContentTokenBudget); return buildResponseParams(truncatedPrompt, temperature, modelId, maximumOutputTokens, requireJsonObject); } - private String truncatePromptForCompletion(String prompt) { + private String truncatePromptForCompletion(String prompt, int inputTokenBudget) { if (prompt == null || prompt.isEmpty()) { return prompt; } - String truncatedPrompt = chunker.keepLastTokens(prompt, promptContentTokenBudget); + String truncatedPrompt = chunker.keepLastTokens(prompt, inputTokenBudget); if (truncatedPrompt.length() < prompt.length()) { return TRUNCATION_NOTICE_GENERIC + truncatedPrompt; @@ -248,6 +265,12 @@ private ResponseCreateParams configureResponseParams( return builder.build(); } + /** Derives routing affinity from the operation and reusable prefix, excluding current input. */ + String promptCacheKey(LlmGatewayTier.Purpose purpose, String stableInstructions) { + String prefixIdentity = purpose.name() + "\n" + openaiModel + "\n" + stableInstructions; + return "java-chat:" + UUID.nameUUIDFromBytes(prefixIdentity.getBytes(StandardCharsets.UTF_8)); + } + String configuredModelId() { return openaiModel; } diff --git a/src/main/java/com/williamcallahan/javachat/service/RerankerService.java b/src/main/java/com/williamcallahan/javachat/service/RerankerService.java index d606cc13..f270e53b 100644 --- a/src/main/java/com/williamcallahan/javachat/service/RerankerService.java +++ b/src/main/java/com/williamcallahan/javachat/service/RerankerService.java @@ -9,6 +9,7 @@ import com.fasterxml.jackson.databind.cfg.CoercionInputShape; import com.fasterxml.jackson.databind.type.LogicalType; import com.williamcallahan.javachat.config.AppProperties; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import java.io.InterruptedIOException; import java.net.SocketTimeoutException; import java.nio.charset.StandardCharsets; @@ -182,7 +183,7 @@ private List rerankUncached( requireRemainingStageBudget(stageDeadlineNanos); log.debug("Reranking {} documents", documents.size()); - String prompt = buildRerankPrompt(query, documents); + CompletionPrompt prompt = buildRerankPrompt(query, documents); Duration rerankRequestTimeout = tighterRerankTimeout(requireRemainingStageBudget(stageDeadlineNanos)); Optional llmOutputOptional = callLlmForReranking(prompt, rerankRequestTimeout); if (llmOutputOptional.isEmpty() || llmOutputOptional.get().isBlank()) { @@ -285,7 +286,7 @@ private Duration tighterRerankTimeout(Duration remainingStageBudget) { * * @return reranking response when the configured provider completes within its timeout */ - private Optional callLlmForReranking(String prompt, Duration rerankRequestTimeout) { + private Optional callLlmForReranking(CompletionPrompt prompt, Duration rerankRequestTimeout) { try { return openAIStreamingService .completeJsonObject( @@ -336,7 +337,7 @@ private static boolean containsProviderTransportTimeout(Throwable failure) { /** * Build the prompt for the reranking LLM call. */ - private String buildRerankPrompt(String query, List documents) { + private CompletionPrompt buildRerankPrompt(String query, List documents) { StringBuilder prompt = new StringBuilder(); prompt.append("You are a document re-ranker for the Java learning assistant system.\n"); prompt.append("Select the documents that are relevant to the query and order them by relevance.\n"); @@ -345,15 +346,17 @@ private String buildRerankPrompt(String query, List documents) { prompt.append("Prefer stable release documentation over early-access or preview content.\n"); prompt.append("Exclude every document that is not actually about the query's subject; "); prompt.append("a document that merely shares an isolated word with the query is not relevant.\n"); + prompt.append("Include each relevant index at most once and do not return any other values.\n"); + prompt.append("Return only JSON: {\"order\":[indices...]} with 0-based indices, most relevant first.\n"); + prompt.append("Return {\"order\":[]} when no document is relevant to the query.\n"); + prompt.append("Do not include markdown, prose, or explanations.\n\n"); + String instructions = prompt.toString(); + prompt.setLength(0); prompt.append("There are exactly ") .append(documents.size()) .append(" documents. Valid indices are 0 through ") .append(documents.size() - 1) .append(".\n"); - prompt.append("Include each relevant index at most once and do not return any other values.\n"); - prompt.append("Return only JSON: {\"order\":[indices...]} with 0-based indices, most relevant first.\n"); - prompt.append("Return {\"order\":[]} when no document is relevant to the query.\n"); - prompt.append("Do not include markdown, prose, or explanations.\n\n"); prompt.append("Query: ").append(query).append("\n\n"); for (int docIndex = 0; docIndex < documents.size(); docIndex++) { @@ -373,7 +376,7 @@ private String buildRerankPrompt(String query, List documents) { .append("\n\n"); } - return prompt.toString(); + return new CompletionPrompt(instructions, prompt.toString()); } /** diff --git a/src/test/java/com/williamcallahan/javachat/application/prompt/PromptTruncatorTest.java b/src/test/java/com/williamcallahan/javachat/application/prompt/PromptTruncatorTest.java index a9504850..d3bca9a3 100644 --- a/src/test/java/com/williamcallahan/javachat/application/prompt/PromptTruncatorTest.java +++ b/src/test/java/com/williamcallahan/javachat/application/prompt/PromptTruncatorTest.java @@ -49,6 +49,27 @@ void noTruncationWhenWithinLimit() { assertEquals(1, result.conversationTurnCount()); } + @Test + void reservesAndRetainsCurrentTurnGuidanceWhenTruncatingContext() { + StructuredPrompt prompt = new StructuredPrompt( + new SystemSegment("Stable teaching policy", 10), + List.of(new ContextDocumentSegment(1, "reference", "https://docs.test", "Retrieved text", 50)), + List.of(new ConversationTurnSegment("user", "Earlier question", 20)), + new CurrentQuerySegment("Current question", 10)) + .withTurnGuidance(new SystemSegment("Use the requested version", 30)); + + PromptTruncator.TruncatedPrompt truncatedPrompt = truncator.truncate(prompt, 70); + + assertTrue(truncatedPrompt.wasTruncated()); + assertEquals(0, truncatedPrompt.contextDocumentCount()); + assertEquals(1, truncatedPrompt.conversationTurnCount()); + assertEquals( + "Use the requested version", + truncatedPrompt.prompt().turnGuidance().content()); + assertEquals(70, truncatedPrompt.prompt().totalEstimatedTokens()); + assertTrue(truncatedPrompt.renderInput().contains("Use the requested version")); + } + @Test void truncatesContextDocumentsFirst() { StructuredPrompt prompt = new StructuredPrompt( diff --git a/src/test/java/com/williamcallahan/javachat/service/ChatServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/ChatServiceTest.java index b09684b1..1708db03 100644 --- a/src/test/java/com/williamcallahan/javachat/service/ChatServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/ChatServiceTest.java @@ -193,6 +193,27 @@ void structuredPromptDoesNotTreatQuantityAsRequestedJavaRelease() { assertFalse(sourceRecord.contains("evidenceRelation")); } + @Test + void guidedTurnGuidanceDoesNotChangeCoreSystemInstructions() { + SystemPromptConfig systemPromptConfig = mock(SystemPromptConfig.class); + when(systemPromptConfig.getCoreSystemPrompt()).thenReturn("Stable Java tutor policy"); + ChatService chatService = new ChatService( + mock(OpenAIStreamingService.class), + mock(RetrievalService.class), + systemPromptConfig, + new AppProperties()); + + var firstPrompt = chatService.buildStructuredPromptWithContextAndGuidance( + List.of(), "Explain records", List.of(), "Use Java 21"); + var nextPrompt = chatService.buildStructuredPromptWithContextAndGuidance( + List.of(), "Explain records", List.of(), "Use Java 25"); + + assertEquals("Stable Java tutor policy", firstPrompt.system().content()); + assertEquals(firstPrompt.system(), nextPrompt.system()); + assertEquals("Use Java 21", firstPrompt.turnGuidance().content()); + assertEquals("Use Java 25", nextPrompt.turnGuidance().content()); + } + private static String sourceRecordFor(String query, Document documentation) { RetrievalService retrievalService = mock(RetrievalService.class); SystemPromptConfig systemPromptConfig = mock(SystemPromptConfig.class); diff --git a/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java b/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java index 45b72941..6be45852 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java @@ -6,6 +6,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.time.Duration; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -279,6 +280,41 @@ void retriesTemporaryProviderFailuresBeforeMutationBoundary() { assertEquals(EmbeddingBatchEmbedder.MAX_EMBEDDING_ATTEMPTS, providerAttempts.get()); } + @Test + void localCooldownRejectionsDoNotExhaustProviderAttempts() { + AtomicInteger admissionCount = new AtomicInteger(); + RecordingEmbeddingClient embeddingClient = + new RecordingEmbeddingClient(EMBEDDING_DIMENSIONS, (requestIndex, textBatch) -> { + if (admissionCount.incrementAndGet() <= EmbeddingBatchEmbedder.MAX_EMBEDDING_ATTEMPTS) { + throw new OpenAiCompatibleEmbeddingClient.EmbeddingProviderCooldownRejectionException( + "A sibling request established the provider cooldown", Duration.ofMillis(10)); + } + return repeatedEmbeddings(textBatch.size(), EMBEDDING_VECTOR); + }); + + List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(1)); + + assertArrayEquals(EMBEDDING_VECTOR, embeddingVectors.getFirst()); + assertEquals(EmbeddingBatchEmbedder.MAX_EMBEDDING_ATTEMPTS + 1, admissionCount.get()); + } + + @Test + void providerCooldownBeyondOverallDeadlineFailsWithoutAnotherAttempt() { + RecordingEmbeddingClient embeddingClient = + new RecordingEmbeddingClient(EMBEDDING_DIMENSIONS, (requestIndex, textBatch) -> { + throw new EmbeddingServiceTemporarilyUnavailableException( + "Retry-After exceeds this batch's budget", null, Duration.ofHours(1)); + }); + + assertThrows( + EmbeddingServiceUnavailableException.class, + () -> EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(1))); + + assertEquals(1, embeddingClient.requestedTextBatches.size()); + } + @Test void doesNotRetryNonTransientProviderFailures() { AtomicInteger providerAttempts = new AtomicInteger(); diff --git a/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java b/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java index 721e12de..4ab28f4e 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java @@ -61,9 +61,9 @@ void gatewayEmbeddingHttpErrorsSurfaceStatusCodes() throws IOException { () -> remoteClient.embed( LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("hello"), LlmGatewayTier.LIVE)); - assertTrue(thrown.getMessage().contains("HTTP 401")); + assertTrue(thrown.getMessage().contains("HTTP 401"), thrown::getMessage); assertNotNull(thrown.getCause()); - assertTrue(thrown.getCause().getMessage().contains("401")); + assertTrue(thrown.getCause().getMessage().contains("401"), thrown.getCause()::getMessage); } finally { httpServer.stop(0); serverExecutor.shutdownNow(); diff --git a/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java b/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java index c555ec8d..3c7c79a2 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java @@ -1,8 +1,8 @@ package com.williamcallahan.javachat.service; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.BDDMockito.given; import static org.mockito.Mockito.mock; @@ -12,6 +12,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.williamcallahan.javachat.config.AppProperties; import com.williamcallahan.javachat.config.CacheConfig; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.model.Enrichment; import java.util.List; import org.junit.jupiter.api.Test; @@ -48,7 +49,7 @@ void shouldRecomputeEnrichmentWhenContextSnippetsChange() { given(streamingService.isAvailable()).willReturn(true); given(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), - anyString(), + any(CompletionPrompt.class), eq(ENRICHMENT_TEMPERATURE), anyInt())) .willReturn(Mono.just(ENRICHMENT_INITIAL_JSON)) @@ -67,7 +68,7 @@ void shouldRecomputeEnrichmentWhenContextSnippetsChange() { verify(streamingService, times(ENRICHMENT_CACHE_MISS_COMPLETION_COUNT)) .completeJsonObject( eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), - anyString(), + any(CompletionPrompt.class), eq(ENRICHMENT_TEMPERATURE), anyInt()); } diff --git a/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceTest.java index 91557be7..0941461f 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceTest.java @@ -4,22 +4,26 @@ import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyDouble; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.williamcallahan.javachat.config.AppProperties; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.model.Enrichment; import java.util.List; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.autoconfigure.jackson.JacksonAutoConfiguration; import org.springframework.test.context.junit.jupiter.SpringJUnitConfig; @@ -52,7 +56,10 @@ void setUp() { @Test void enrichUsesJsonObjectCompletionBoundary() { when(openAIStreamingService.completeJsonObject( - eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt())) .thenReturn(Mono.just("{\"hints\":[\"Use List.copyOf for an unmodifiable snapshot.\"]}")); Enrichment enrichment = @@ -61,15 +68,46 @@ void enrichUsesJsonObjectCompletionBoundary() { assertEquals(ENRICHMENT_JDK_VERSION, enrichment.getJdkVersion()); assertEquals(List.of("Use List.copyOf for an unmodifiable snapshot."), enrichment.getHints()); verify(openAIStreamingService) - .completeJsonObject(eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt()); + .completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt()); verify(openAIStreamingService, never()) .complete(eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble()); } + @Test + void enrichmentKeepsReusableInstructionsIndependentOfCurrentQueryVersionAndSnippets() { + when(openAIStreamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt())) + .thenReturn(Mono.just("{}")); + + enrichmentService.enrich("Explain records", "21", List.of("Record reference")); + enrichmentService.enrich("Explain sealed classes", "25", List.of("Sealed class reference")); + + ArgumentCaptor prompts = ArgumentCaptor.forClass(CompletionPrompt.class); + verify(openAIStreamingService, times(2)) + .completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), prompts.capture(), anyDouble(), anyInt()); + assertEquals( + prompts.getAllValues().get(0).instructions(), + prompts.getAllValues().get(1).instructions()); + assertTrue(prompts.getAllValues().get(0).input().contains("Requested JDK version: 21")); + assertTrue(prompts.getAllValues().get(1).input().contains("Requested JDK version: 25")); + assertTrue(prompts.getAllValues().get(1).input().contains("Sealed class reference")); + } + @Test void enrichRejectsUnexpectedPreambleWithBraces() { when(openAIStreamingService.completeJsonObject( - eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt())) .thenReturn(Mono.just("Here is the JSON (see {docs} for details):\n{\"hints\":[\"tip\"]}")); IllegalStateException jsonContractFailure = assertThrows( @@ -84,7 +122,10 @@ void enrichRejectsUnexpectedPreambleWithBraces() { @Test void enrichRejectsMalformedJsonObject() { when(openAIStreamingService.completeJsonObject( - eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt())) .thenReturn(Mono.just("{\"hints\":[")); IllegalStateException jsonContractFailure = assertThrows( @@ -99,7 +140,10 @@ void enrichRejectsMalformedJsonObject() { @Test void enrichRejectsUnknownResponseFields() { when(openAIStreamingService.completeJsonObject( - eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt())) .thenReturn(Mono.just("{\"hint\":[\"Use the documented API.\"]}")); IllegalStateException jsonContractFailure = assertThrows( @@ -114,7 +158,10 @@ void enrichRejectsUnknownResponseFields() { @Test void enrichRejectsDuplicateResponseFields() { when(openAIStreamingService.completeJsonObject( - eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt())) .thenReturn(Mono.just("{\"hints\":[\"first\"],\"hints\":[\"second\"]}")); IllegalStateException jsonContractFailure = assertThrows( @@ -136,7 +183,10 @@ void enrichRejectsTrailingTokensAndScalarCoercion() { for (String invalidEnrichmentJson : invalidEnrichmentJsonValues) { when(openAIStreamingService.completeJsonObject( - eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + any(CompletionPrompt.class), + anyDouble(), + anyInt())) .thenReturn(Mono.just(invalidEnrichmentJson)); IllegalStateException jsonContractFailure = assertThrows( diff --git a/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java index 9e5602b9..e62025fb 100644 --- a/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java @@ -88,7 +88,8 @@ void upsertStartsFreshOperationAfterTransientFailure() { when(embeddingClient.embed( org.mockito.ArgumentMatchers.any(LlmGatewayTier.Purpose.class), eq(List.of(DOCUMENT_TEXT)), - eq(LlmGatewayTier.BATCH))) + eq(LlmGatewayTier.BATCH), + org.mockito.ArgumentMatchers.any(java.time.Duration.class))) .thenReturn(List.of(new float[] {0.1f, 0.2f})); when(sparseVectorEncoder.encode(DOCUMENT_TEXT)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -370,7 +371,8 @@ private void configureReplacementEmbedding() { when(embeddingClient.embed( org.mockito.ArgumentMatchers.any(LlmGatewayTier.Purpose.class), eq(List.of(DOCUMENT_TEXT)), - eq(LlmGatewayTier.BATCH))) + eq(LlmGatewayTier.BATCH), + org.mockito.ArgumentMatchers.any(java.time.Duration.class))) .thenReturn(List.of(new float[] {0.1f, 0.2f})); when(sparseVectorEncoder.encode(DOCUMENT_TEXT)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java index 8ba551a5..5279d6fb 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java @@ -4,6 +4,8 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -13,6 +15,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; import com.openai.client.OpenAIClient; import com.openai.client.OpenAIClientAsync; import com.openai.core.RequestOptions; @@ -30,7 +33,9 @@ import com.williamcallahan.javachat.application.prompt.PromptTruncator; import com.williamcallahan.javachat.config.AppProperties; import com.williamcallahan.javachat.config.ModelConfiguration; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.domain.prompt.StructuredPrompt; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import java.io.IOException; import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; @@ -55,13 +60,22 @@ class OpenAICompletionStatusTest { @ParameterizedTest @CsvSource({ - "production-z,MARKDOWN_ENRICHMENT,false", - "production-a,SEARCH_RERANK,false", - "batch,MARKDOWN_ENRICHMENT,false", - "production-z,CHAT_ANSWER,true", - "production-z,GUIDED_LEARNING,true" + "production-z,MARKDOWN_ENRICHMENT,false,true,true", + "production-a,SEARCH_RERANK,false,true,true", + "batch,MARKDOWN_ENRICHMENT,false,true,true", + "production-z,CHAT_ANSWER,true,true,true", + "production-z,GUIDED_LEARNING,true,true,true", + "production-z,CHAT_ANSWER,true,true,false", + "production-z,MARKDOWN_ENRICHMENT,false,true,false", + "production-z,CHAT_ANSWER,true,false,false", + "production-z,MARKDOWN_ENRICHMENT,false,false,false" }) - void completionSendsMatchingWorkloadTier(String requestTier, LlmGatewayTier.Purpose purpose, boolean streaming) + void completionSendsMatchingWorkloadTier( + String requestTier, + LlmGatewayTier.Purpose purpose, + boolean streaming, + boolean includeUsage, + boolean includeTokenDetails) throws IOException { AtomicReference capturedRequest = new AtomicReference<>(); HttpServer gatewayServer = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); @@ -74,10 +88,23 @@ void completionSendsMatchingWorkloadTier(String requestTier, LlmGatewayTier.Purp byte[] responseBody = """ {"id":"resp_wire","object":"response","created_at":0, "model":"wire-test-model","status":"completed", + "usage":{"input_tokens":100,"input_tokens_details":{"cached_tokens":40,"cache_write_tokens":20}, + "output_tokens":7,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":107}, "output":[{"id":"msg_wire","type":"message","role":"assistant", "status":"completed","content":[{"type":"output_text", "text":"{\\"ok\\":true}","annotations":[]}]}]} """.getBytes(StandardCharsets.UTF_8); + ObjectNode providerCompletion = new ObjectMapper().readValue(responseBody, ObjectNode.class); + if (!includeUsage) { + providerCompletion.remove("usage"); + } else if (!includeTokenDetails) { + // A gateway may omit input_tokens_details; the call must still complete. + ObjectNode tokenUsage = + new ObjectMapper().convertValue(providerCompletion.path("usage"), ObjectNode.class); + tokenUsage.remove("input_tokens_details"); + providerCompletion.set("usage", tokenUsage); + } + responseBody = new ObjectMapper().writeValueAsBytes(providerCompletion); if (streaming) { String completedResponse = new ObjectMapper().readTree(responseBody).toString(); @@ -97,7 +124,8 @@ void completionSendsMatchingWorkloadTier(String requestTier, LlmGatewayTier.Purp RateLimitService rateLimitService = mock(RateLimitService.class); when(rateLimitService.tryReserveRequest(RateLimitService.ApiProvider.OPENAI)) .thenReturn(true); - OpenAIStreamingService streamingService = configuredStreamingService(rateLimitService); + SimpleMeterRegistry meterRegistry = new SimpleMeterRegistry(); + OpenAIStreamingService streamingService = configuredStreamingService(rateLimitService, meterRegistry); ReflectionTestUtils.setField(streamingService, "openaiApiKey", "wire-test-api-key"); ReflectionTestUtils.setField(streamingService, "openaiBaseUrl", "https://api.llm-gateway.iocloudhost.net/v1"); ReflectionTestUtils.setField(streamingService, "llmGatewayTier", requestTier); @@ -125,7 +153,12 @@ void completionSendsMatchingWorkloadTier(String requestTier, LlmGatewayTier.Purp assertEquals( "{\"ok\":true}", streamingService - .completeJsonObject(purpose, "Return a JSON object with ok true", 0.0, 512) + .completeJsonObject( + purpose, + new CompletionPrompt( + "Stable task instructions", "Return a JSON object with ok true"), + 0.0, + 512) .block(Duration.ofSeconds(10))); } } finally { @@ -138,6 +171,51 @@ void completionSendsMatchingWorkloadTier(String requestTier, LlmGatewayTier.Purp assertEquals( (requestTier.equals("batch") ? LlmGatewayTier.BATCH : LlmGatewayTier.LIVE).workloadHeader(purpose), gatewayRequest.workloadHeader()); + assertFalse( + gatewayRequest.requestBody().path("prompt_cache_key").asText().isBlank()); + if (includeUsage) { + assertEquals( + 100, + meterRegistry + .get("llm.tokens") + .tags("purpose", purpose.name(), "type", "input") + .counter() + .count()); + if (includeTokenDetails) { + assertEquals( + 40, + meterRegistry + .get("llm.tokens") + .tags("purpose", purpose.name(), "type", "cached_input") + .counter() + .count()); + assertEquals( + 20, + meterRegistry + .get("llm.tokens") + .tags("purpose", purpose.name(), "type", "cache_write") + .counter() + .count()); + } else { + assertNull(meterRegistry + .find("llm.tokens") + .tags("purpose", purpose.name(), "type", "cached_input") + .counter()); + assertNull(meterRegistry + .find("llm.tokens") + .tags("purpose", purpose.name(), "type", "cache_write") + .counter()); + } + assertEquals( + 7, + meterRegistry + .get("llm.tokens") + .tags("purpose", purpose.name(), "type", "output") + .counter() + .count()); + } else { + assertTrue(meterRegistry.getMeters().isEmpty()); + } if (!streaming) { assertEquals( 512, gatewayRequest.requestBody().path("max_output_tokens").asInt()); @@ -382,8 +460,13 @@ private static Response incompleteCompletion(Response.IncompleteDetails.Reason i } private static OpenAIStreamingService configuredStreamingService(RateLimitService rateLimitService) { + return configuredStreamingService(rateLimitService, new SimpleMeterRegistry()); + } + + private static OpenAIStreamingService configuredStreamingService( + RateLimitService rateLimitService, SimpleMeterRegistry meterRegistry) { return new OpenAIStreamingService( - rateLimitService, + meterRegistry, testRequestFactory(), new OpenAiProviderRoutingService(rateLimitService, configuredLlmProperties()), new OpenAiStreamingFailureReporter()); diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java index 291deaaa..d1aa9db6 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java @@ -32,6 +32,7 @@ import com.williamcallahan.javachat.config.AppProperties; import com.williamcallahan.javachat.domain.prompt.StructuredPrompt; import com.williamcallahan.javachat.support.logging.ExpectedLogEvents; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import java.util.List; import java.util.Optional; import java.util.concurrent.CompletableFuture; @@ -79,7 +80,10 @@ void completionTransportFailureIsTerminal() { .thenReturn(true); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, testRequestFactory(), providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + testRequestFactory(), + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ResponseServiceAsync responseService = mockAsyncResponseService(openAiClient); InternalServerException upstreamFailure = InternalServerException.builder() @@ -226,7 +230,7 @@ void unavailableStreamDefersErrorSeverityToRequestBoundary() { private static OpenAIStreamingService configuredStreamingService(RateLimitService rateLimitService) { return new OpenAIStreamingService( - rateLimitService, + new SimpleMeterRegistry(), testRequestFactory(), configuredProviderRoutingService(rateLimitService), new OpenAiStreamingFailureReporter()); diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java index ea7ee994..4fbde71d 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java @@ -38,10 +38,12 @@ import com.williamcallahan.javachat.application.prompt.PromptTruncator; import com.williamcallahan.javachat.application.streaming.StreamingFailureReporter; import com.williamcallahan.javachat.config.AppProperties; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.domain.prompt.ContextDocumentSegment; import com.williamcallahan.javachat.domain.prompt.CurrentQuerySegment; import com.williamcallahan.javachat.domain.prompt.StructuredPrompt; import com.williamcallahan.javachat.domain.prompt.SystemSegment; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import java.time.Duration; import java.util.List; import java.util.Objects; @@ -81,7 +83,10 @@ private OpenAIStreamingService createStreamingService() { OpenAiRequestFactory requestFactory = testRequestFactory(); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); return new OpenAIStreamingService( - rateLimitService, requestFactory, providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + requestFactory, + providerRoutingService, + new OpenAiStreamingFailureReporter()); } @Test @@ -131,7 +136,10 @@ void emptyTextDeltaBeforeStreamingFailureIsTerminal() { OpenAiRequestFactory requestFactory = testRequestFactory(); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, requestFactory, providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + requestFactory, + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ResponseServiceAsync responseService = mockAsyncResponseService(openAiClient); StreamResponse providerStream = mock(); @@ -485,7 +493,10 @@ void deniedConfiguredProviderReservationTerminatesStreamingAndCompletionAsRetrya OpenAiRequestFactory requestFactory = testRequestFactory(); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, requestFactory, providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + requestFactory, + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); @@ -517,7 +528,10 @@ void cooldownRecordedAfterStreamCreationPreventsDispatchAtTextSubscription() { OpenAiRequestFactory requestFactory = testRequestFactory(); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, requestFactory, providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + requestFactory, + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); @@ -542,7 +556,10 @@ void exposesOnlyDocumentIdentitiesRetainedByGpt54Truncation() { RateLimitService rateLimitService = mock(RateLimitService.class); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, testRequestFactory(), providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + testRequestFactory(), + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); StructuredPrompt oversizedPrompt = new StructuredPrompt( @@ -579,7 +596,10 @@ void preTextStreamingTransportFailureIsTerminal() { OpenAiRequestFactory requestFactory = testRequestFactory(); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, requestFactory, providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + requestFactory, + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ResponseServiceAsync responseService = mockAsyncResponseService(openAiClient); InternalServerException upstreamFailure = InternalServerException.builder() @@ -614,7 +634,10 @@ void streamingFailureAfterTextIsTerminal() { OpenAiRequestFactory requestFactory = testRequestFactory(); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, requestFactory, providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + requestFactory, + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ResponseServiceAsync responseService = mockAsyncResponseService(openAiClient); StreamResponse providerStream = mock(); @@ -656,7 +679,10 @@ void completionPreservesAsyncTransportFailure() { .thenReturn(true); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, testRequestFactory(), providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + testRequestFactory(), + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ResponseServiceAsync responseService = mockAsyncResponseService(openAiClient); CompletableFuture providerCompletionFuture = new CompletableFuture<>(); @@ -688,15 +714,27 @@ void invalidCompletionConfigurationFailsOnSubscription() { "maximumOutputTokens must be positive"); assertCompletionFailure( streamingService.completeJsonObject( - LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, "prompt", 0.7, 128, null), + LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, + new CompletionPrompt("Stable task instructions", "prompt"), + 0.7, + 128, + null), "requestTimeout must be positive"); assertCompletionFailure( streamingService.completeJsonObject( - LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, "prompt", 0.7, 128, Duration.ZERO), + LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, + new CompletionPrompt("Stable task instructions", "prompt"), + 0.7, + 128, + Duration.ZERO), "requestTimeout must be positive"); assertCompletionFailure( streamingService.completeJsonObject( - LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, "prompt", 0.7, 128, Duration.ofSeconds(-1)), + LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, + new CompletionPrompt("Stable task instructions", "prompt"), + 0.7, + 128, + Duration.ofSeconds(-1)), "requestTimeout must be positive"); } @@ -721,7 +759,10 @@ private static OpenAIStreamingService streamingServiceForProviderStream( .thenReturn(true); OpenAiProviderRoutingService providerRoutingService = configuredProviderRoutingService(rateLimitService); OpenAIStreamingService streamingService = new OpenAIStreamingService( - rateLimitService, testRequestFactory(), providerRoutingService, new OpenAiStreamingFailureReporter()); + new SimpleMeterRegistry(), + testRequestFactory(), + providerRoutingService, + new OpenAiStreamingFailureReporter()); OpenAIClient openAiClient = mock(OpenAIClient.class); ResponseServiceAsync responseService = mockAsyncResponseService(openAiClient); when(responseService.createStreaming(any(ResponseCreateParams.class), any(RequestOptions.class))) diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java index 11b0991a..67a4c216 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java @@ -17,6 +17,7 @@ import com.williamcallahan.javachat.adapters.out.llm.openai.OpenAiStreamingFailureReporter; import com.williamcallahan.javachat.config.ModelConfiguration; import com.williamcallahan.javachat.domain.prompt.StructuredPrompt; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import java.time.Duration; import java.util.Objects; import java.util.Optional; @@ -32,7 +33,7 @@ void streamingReadDeadlineInheritsWholeRequestDeadline() { OpenAIClient gatewayClient = mock(OpenAIClient.class); StructuredPrompt structuredPrompt = StructuredPrompt.fromRawPrompt("Confirm the timeout budget", 6); OpenAiRequestFactory requestFactory = mock(OpenAiRequestFactory.class); - when(requestFactory.prepareStreamingRequest(structuredPrompt, 0.4)) + when(requestFactory.prepareStreamingRequest(LlmGatewayTier.Purpose.CHAT_ANSWER, structuredPrompt, 0.4)) .thenReturn(new OpenAiPreparedRequest( ResponseCreateParams.builder() .model(ModelConfiguration.DEFAULT_MODEL) @@ -47,7 +48,7 @@ void streamingReadDeadlineInheritsWholeRequestDeadline() { .thenReturn( Optional.of(new OpenAiProviderCandidate(gatewayClient, RateLimitService.ApiProvider.OPENAI))); OpenAIStreamingService streamingService = new OpenAIStreamingService( - mock(RateLimitService.class), + new SimpleMeterRegistry(), requestFactory, providerRoutingService, new OpenAiStreamingFailureReporter()); diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java index 50be1e8b..187bc7d3 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java @@ -1,5 +1,6 @@ package com.williamcallahan.javachat.service; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -807,6 +808,66 @@ void probeFailureFromProviderIsNotDeferred() { } } + @Test + void batchRetriesAfterAuthoritativeCooldownInsteadOfExhaustingLocalAttempts() { + OpenAIClient liveClient = mock(OpenAIClient.class); + OpenAIClient batchClient = mock(OpenAIClient.class); + EmbeddingServiceAsync batchEmbeddingService = mockAsyncEmbeddingService(batchClient); + Duration providerCooldown = Duration.ofSeconds(16); + when(batchEmbeddingService.create(any(), any(RequestOptions.class))) + .thenReturn(CompletableFuture.failedFuture(RateLimitException.builder() + .headers(Headers.builder().put("Retry-After", "16").build()) + .build())) + .thenReturn(CompletableFuture.completedFuture(successfulResponse())); + + try (OpenAiCompatibleEmbeddingClient embeddingClient = + new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { + long batchStartedNanos = System.nanoTime(); + List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + embeddingClient, + List.of(new org.springframework.ai.document.Document("batch document"))); + + assertEquals(1, embeddingVectors.size()); + assertArrayEquals(new float[] {0.4f, 0.6f}, embeddingVectors.getFirst()); + assertTrue(Duration.ofNanos(System.nanoTime() - batchStartedNanos).compareTo(providerCooldown) >= 0); + } + } + + @Test + void preservesAuthoritativeRetryTimingForProviderAndLocalCooldownFailures() { + OpenAIClient liveClient = mock(OpenAIClient.class); + OpenAIClient batchClient = mock(OpenAIClient.class); + EmbeddingServiceAsync batchEmbeddingService = mockAsyncEmbeddingService(batchClient); + when(batchEmbeddingService.create(any(), any(RequestOptions.class))) + .thenReturn(CompletableFuture.failedFuture(RateLimitException.builder() + .headers(Headers.builder().put("Retry-After", "60").build()) + .build())); + + try (OpenAiCompatibleEmbeddingClient embeddingClient = + new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { + EmbeddingServiceTemporarilyUnavailableException providerFailure = assertThrows( + EmbeddingServiceTemporarilyUnavailableException.class, + () -> embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("batch document"), + LlmGatewayTier.BATCH)); + assertTrue(providerFailure.providerContacted()); + assertTrue(providerFailure.retryDelay().compareTo(Duration.ofSeconds(50)) > 0); + assertTrue(providerFailure.retryDelay().compareTo(Duration.ofSeconds(60)) <= 0); + + EmbeddingServiceTemporarilyUnavailableException admissionFailure = assertThrows( + EmbeddingServiceTemporarilyUnavailableException.class, + () -> embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("next document"), + LlmGatewayTier.BATCH)); + assertFalse(admissionFailure.providerContacted()); + assertTrue(admissionFailure.retryDelay().compareTo(Duration.ofSeconds(50)) > 0); + assertTrue(admissionFailure.retryDelay().compareTo(Duration.ofSeconds(60)) <= 0); + } + } + @Test void batchRetryAfterCooldownDoesNotBlockLiveRequests() { OpenAIClient liveClient = mock(OpenAIClient.class); diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAiRequestFactoryTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAiRequestFactoryTest.java index f753e1a2..c41dc2a4 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAiRequestFactoryTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAiRequestFactoryTest.java @@ -13,6 +13,7 @@ import com.williamcallahan.javachat.application.prompt.PromptTruncator; import com.williamcallahan.javachat.config.AppProperties; import com.williamcallahan.javachat.config.ModelConfiguration; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import com.williamcallahan.javachat.domain.prompt.ContextDocumentSegment; import com.williamcallahan.javachat.domain.prompt.ConversationTurnSegment; import com.williamcallahan.javachat.domain.prompt.CurrentQuerySegment; @@ -132,8 +133,9 @@ void prepareStreamingRequestAppliesTheConfiguredReasoningEffort() { List.of(), new CurrentQuerySegment("Explain records", 2)); - ResponseCreateParams reasoningRequestParams = - requestFactory.prepareStreamingRequest(structuredPrompt, 0.4).responseParams(); + ResponseCreateParams reasoningRequestParams = requestFactory + .prepareStreamingRequest(LlmGatewayTier.Purpose.CHAT_ANSWER, structuredPrompt, 0.4) + .responseParams(); assertEquals( ReasoningEffort.HIGH, @@ -144,12 +146,17 @@ void prepareStreamingRequestAppliesTheConfiguredReasoningEffort() { void jsonCompletionReservesItsOutputBudgetForVisibleJson() { OpenAiRequestFactory requestFactory = createRequestFactory("max"); - ResponseCreateParams jsonCompletionParams = - requestFactory.buildJsonCompletionRequest("Rank these documents", 0.0, 128); + ResponseCreateParams jsonCompletionParams = requestFactory.buildJsonCompletionRequest( + new CompletionPrompt("Stable task instructions", "Rank these documents"), 0.0, 128); assertTrue( jsonCompletionParams.text().orElseThrow().format().orElseThrow().isJsonObject()); assertEquals(128L, jsonCompletionParams.maxOutputTokens().orElseThrow()); + assertEquals( + "Stable task instructions", jsonCompletionParams.instructions().orElseThrow()); + assertEquals( + "Rank these documents", + jsonCompletionParams.input().orElseThrow().asText()); assertEquals( ReasoningEffort.NONE, jsonCompletionParams.reasoning().orElseThrow().effort().orElseThrow()); @@ -183,8 +190,9 @@ void prepareStreamingRequestKeepsRetrievedReferenceTextOutOfDeveloperAuthority() new ConversationTurnSegment(ConversationTurnSegment.ROLE_ASSISTANT, priorAssistantMessage, 8)), new CurrentQuerySegment("Explain records", 2)); - ResponseCreateParams responseCreateParams = - requestFactory.prepareStreamingRequest(structuredPrompt, 0.4).responseParams(); + ResponseCreateParams responseCreateParams = requestFactory + .prepareStreamingRequest(LlmGatewayTier.Purpose.CHAT_ANSWER, structuredPrompt, 0.4) + .responseParams(); assertEquals( "Follow the Java teaching policy", @@ -193,12 +201,12 @@ void prepareStreamingRequestKeepsRetrievedReferenceTextOutOfDeveloperAuthority() assertTrue(responseInput.isResponse()); List responseInputItems = responseInput.asResponse(); assertEquals(4, responseInputItems.size()); + assertInputMessage(responseInputItems.get(0), EasyInputMessage.Role.USER, "Earlier question"); + assertInputMessage(responseInputItems.get(1), EasyInputMessage.Role.ASSISTANT, priorAssistantMessage); assertInputMessage( - responseInputItems.get(0), + responseInputItems.get(2), EasyInputMessage.Role.USER, "[CTX 1] https://docs.oracle.com/java\n" + retrievedReferenceText); - assertInputMessage(responseInputItems.get(1), EasyInputMessage.Role.USER, "Earlier question"); - assertInputMessage(responseInputItems.get(2), EasyInputMessage.Role.ASSISTANT, priorAssistantMessage); assertInputMessage(responseInputItems.get(3), EasyInputMessage.Role.USER, "Explain records"); assertFalse(responseInputItems.stream() .map(ResponseInputItem::asEasyInputMessage) @@ -207,6 +215,64 @@ void prepareStreamingRequestKeepsRetrievedReferenceTextOutOfDeveloperAuthority() .anyMatch(messageText -> messageText.contains("Follow the Java teaching policy"))); } + @Test + void streamingCacheKeyAndHistoryPrefixSurviveChangingRetrievalAndGuidance() { + OpenAiRequestFactory requestFactory = createRequestFactory(""); + StructuredPrompt firstPrompt = new StructuredPrompt( + new SystemSegment("Stable Java teaching instructions", 8), + List.of(new ContextDocumentSegment( + 1, "first", "https://docs.test/first", "First references", 5)), + List.of(new ConversationTurnSegment(ConversationTurnSegment.ROLE_USER, "Earlier question", 3)), + new CurrentQuerySegment("Explain records", 2)) + .withTurnGuidance(new SystemSegment("Use Java 21 documentation", 5)); + StructuredPrompt secondPrompt = new StructuredPrompt( + firstPrompt.system(), + List.of(new ContextDocumentSegment( + 1, "second", "https://docs.test/second", "Different references", 5)), + firstPrompt.conversationHistory(), + new CurrentQuerySegment("Explain sealed classes", 4)) + .withTurnGuidance(new SystemSegment("Use Java 25 documentation", 5)); + + ResponseCreateParams firstRequest = requestFactory + .prepareStreamingRequest(LlmGatewayTier.Purpose.CHAT_ANSWER, firstPrompt, 0.4) + .responseParams(); + ResponseCreateParams secondRequest = requestFactory + .prepareStreamingRequest(LlmGatewayTier.Purpose.CHAT_ANSWER, secondPrompt, 0.4) + .responseParams(); + + assertEquals(firstRequest.instructions(), secondRequest.instructions()); + assertFalse(firstRequest.promptCacheKey().orElseThrow().isBlank()); + assertEquals(firstRequest.promptCacheKey(), secondRequest.promptCacheKey()); + List firstInput = firstRequest.input().orElseThrow().asResponse(); + List secondInput = + secondRequest.input().orElseThrow().asResponse(); + assertEquals(firstInput.getFirst(), secondInput.getFirst()); + assertInputMessage(firstInput.get(1), EasyInputMessage.Role.DEVELOPER, "Use Java 21 documentation"); + assertInputMessage(secondInput.get(1), EasyInputMessage.Role.DEVELOPER, "Use Java 25 documentation"); + assertFalse(firstRequest + .promptCacheKey() + .equals(requestFactory + .prepareStreamingRequest(LlmGatewayTier.Purpose.GUIDED_LEARNING, firstPrompt, 0.4) + .responseParams() + .promptCacheKey())); + OpenAiRequestFactory differentModel = new OpenAiRequestFactory( + new Chunker(), new PromptTruncator(), "another-gateway-route", appProperties("")); + assertFalse(firstRequest + .promptCacheKey() + .equals(differentModel + .prepareStreamingRequest(LlmGatewayTier.Purpose.CHAT_ANSWER, firstPrompt, 0.4) + .responseParams() + .promptCacheKey())); + StructuredPrompt changedSystem = new StructuredPrompt( + new SystemSegment("Revised teaching policy", 5), List.of(), List.of(), firstPrompt.currentQuery()); + assertFalse(firstRequest + .promptCacheKey() + .equals(requestFactory + .prepareStreamingRequest(LlmGatewayTier.Purpose.CHAT_ANSWER, changedSystem, 0.4) + .responseParams() + .promptCacheKey())); + } + @Test void blankOpenAiModelFailsConstruction() { IllegalArgumentException configurationFailure = assertThrows( @@ -251,8 +317,8 @@ void buildJsonCompletionRequestDeclaresJsonObjectOutput() { OpenAiRequestFactory requestFactory = new OpenAiRequestFactory( new Chunker(), new PromptTruncator(), ModelConfiguration.DEFAULT_MODEL, appProperties("")); - ResponseCreateParams responseCreateParams = - requestFactory.buildJsonCompletionRequest("Rank these documents", 0.0, 128); + ResponseCreateParams responseCreateParams = requestFactory.buildJsonCompletionRequest( + new CompletionPrompt("Stable task instructions", "Rank these documents"), 0.0, 128); assertTrue( responseCreateParams.text().orElseThrow().format().orElseThrow().isJsonObject()); diff --git a/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java index 11401359..2f7ce1df 100644 --- a/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.openai.errors.OpenAIIoException; import com.williamcallahan.javachat.config.AppProperties; +import com.williamcallahan.javachat.domain.prompt.CompletionPrompt; import java.io.InterruptedIOException; import java.time.Duration; import java.util.Collection; @@ -62,7 +63,7 @@ void rerankPreservesAtomicProviderAdmissionFailure() { new ConfiguredProviderTemporarilyUnavailableException(RateLimitService.ApiProvider.OPENAI); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -84,7 +85,7 @@ void rerankUsesConfiguredCompletionBudgetAndTimeout() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -96,7 +97,7 @@ void rerankUsesConfiguredCompletionBudgetAndTimeout() { List rankedDocuments = rerankerService.rerank( "query", sourceDocuments, 2, stageDeadlineNanos(REMAINING_STAGE_BUDGET_TEST_TIMEOUT)); - ArgumentCaptor promptCaptor = ArgumentCaptor.forClass(String.class); + ArgumentCaptor promptCaptor = ArgumentCaptor.forClass(CompletionPrompt.class); ArgumentCaptor outputBudgetCaptor = ArgumentCaptor.forClass(Integer.class); verify(streamingService) .completeJsonObject( @@ -107,9 +108,10 @@ void rerankUsesConfiguredCompletionBudgetAndTimeout() { eq(TEST_RERANKER_TIMEOUT)); verify(streamingService, never()) .complete(eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE)); - assertTrue(promptCaptor.getValue().contains("Valid indices are 0 through 1.")); - assertTrue(promptCaptor.getValue().contains("Include each relevant index at most once")); - assertTrue(promptCaptor.getValue().contains("Return {\"order\":[]} when no document is relevant")); + assertTrue(promptCaptor.getValue().input().contains("Valid indices are 0 through 1.")); + assertTrue(promptCaptor.getValue().instructions().contains("Include each relevant index at most once")); + assertTrue( + promptCaptor.getValue().instructions().contains("Return {\"order\":[]} when no document is relevant")); assertEquals(TEST_RERANKER_OUTPUT_TOKEN_BUDGET, outputBudgetCaptor.getValue()); assertEquals(List.of(sourceDocuments.get(1), sourceDocuments.get(0)), rankedDocuments); } @@ -119,7 +121,7 @@ void rerankSelectsRelevantSubsetAndDropsUnrelatedDocuments() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -139,7 +141,7 @@ void rerankReturnsEmptySelectionWhenNoDocumentIsRelevant() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -159,7 +161,7 @@ void rerankRetainsOneRelevantDocument() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -175,7 +177,7 @@ void rerankRetainsOneRelevantDocument() { verify(streamingService) .completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT)); @@ -202,7 +204,7 @@ void rerankRejectsDuplicateInvalidAndMalformedOrderings() { for (String invalidOrderingJson : invalidOrderingJsonValues) { when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -220,7 +222,7 @@ void emptyCompletionIsTerminalWithoutCallingAnotherCompletionPath() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -243,7 +245,7 @@ void rerankTightensItsTimeoutToTheRemainingStageBudget() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), any(Duration.class))) @@ -259,7 +261,7 @@ void rerankTightensItsTimeoutToTheRemainingStageBudget() { verify(streamingService) .completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), remainingBudgetCaptor.capture()); @@ -287,7 +289,7 @@ void expiredDeadlineRejectsCachedHit() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT))) @@ -304,7 +306,7 @@ void expiredDeadlineRejectsCachedHit() { verify(streamingService, times(1)) .completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT)); @@ -319,7 +321,7 @@ void sameKeyCacheWaitCannotExtendCallerDeadline() CountDownLatch releaseFirstDispatch = new CountDownLatch(1); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), any(Duration.class))) @@ -360,7 +362,7 @@ void sameKeyCacheWaitCannotExtendCallerDeadline() verify(streamingService, times(1)) .completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), any(Duration.class)); @@ -377,7 +379,7 @@ void waiterRetriesWhenCoalescedOwnerDiesOnItsStageDeadline() AtomicInteger dispatchCount = new AtomicInteger(); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), any(Duration.class))) @@ -445,7 +447,7 @@ void waiterInheritsProviderTimeoutOwnerFailureWithoutRetrying() AtomicInteger dispatchCount = new AtomicInteger(); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), any(Duration.class))) @@ -508,7 +510,7 @@ void waiterInheritsNonTimeoutOwnerFailureWithoutRetrying() AtomicInteger dispatchCount = new AtomicInteger(); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), any(Duration.class))) @@ -562,7 +564,7 @@ void wholeCallTimeoutPreservesRetrievalTimeoutCause() { OpenAIIoException providerTimeout = new OpenAIIoException("Request failed", okHttpCallTimeout); when(streamingService.completeJsonObject( eq(LlmGatewayTier.Purpose.SEARCH_RERANK), - anyString(), + any(CompletionPrompt.class), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), eq(TEST_RERANKER_TIMEOUT)))