diff --git a/src/main/java/com/williamcallahan/javachat/service/ChatService.java b/src/main/java/com/williamcallahan/javachat/service/ChatService.java index d12fe59d..ab85eb75 100644 --- a/src/main/java/com/williamcallahan/javachat/service/ChatService.java +++ b/src/main/java/com/williamcallahan/javachat/service/ChatService.java @@ -76,7 +76,7 @@ public Flux streamAnswerWithContext( buildStructuredPromptWithContextAndGuidance(history, latestUserMessage, contextDocs, guidance); return openAIStreamingService - .streamResponse(structuredPrompt, temperature) + .streamResponse(LlmGatewayTier.Purpose.CHAT_ANSWER, structuredPrompt, temperature) .flatMapMany(streamingResult -> streamingResult.textChunks()); } diff --git a/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java b/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java index fc95452e..eb058298 100644 --- a/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java +++ b/src/main/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedder.java @@ -26,7 +26,9 @@ final class EmbeddingBatchEmbedder { private EmbeddingBatchEmbedder() {} - static List embedDocuments(EmbeddingClient embeddingClient, List documents) { + static List embedDocuments( + LlmGatewayTier.Purpose purpose, EmbeddingClient embeddingClient, List documents) { + Objects.requireNonNull(purpose, "purpose"); Objects.requireNonNull(embeddingClient, "embeddingClient"); Objects.requireNonNull(documents, "documents"); @@ -50,8 +52,8 @@ static List embedDocuments(EmbeddingClient embeddingClient, List>> embeddingWave = IntStream.range(currentWaveStartIndex, requestWaveEndIndex) - .mapToObj(requestIndex -> embeddingExecutor.submit(() -> - embedRequest(embeddingClient, documents, requestIndex, expectedEmbeddingDimensions))) + .mapToObj(requestIndex -> embeddingExecutor.submit(() -> embedRequest( + purpose, embeddingClient, documents, requestIndex, expectedEmbeddingDimensions))) .toList(); for (Future> embeddingRequest : embeddingWave) { try { @@ -79,6 +81,7 @@ static List embedDocuments(EmbeddingClient embeddingClient, List embedRequest( + LlmGatewayTier.Purpose purpose, EmbeddingClient embeddingClient, List documents, int requestIndex, @@ -86,6 +89,7 @@ private static List embedRequest( int batchStartIndex = requestIndex * EMBEDDING_REQUEST_BATCH_SIZE; int batchEndIndex = Math.min(batchStartIndex + EMBEDDING_REQUEST_BATCH_SIZE, documents.size()); return embedSingleBatch( + purpose, embeddingClient, documents.subList(batchStartIndex, batchEndIndex), batchStartIndex, @@ -100,6 +104,7 @@ private static List embedRequest( * callers can identify which documents caused the failure.

*/ private static List embedSingleBatch( + LlmGatewayTier.Purpose purpose, EmbeddingClient embeddingClient, List documentBatch, int batchStartIndex, @@ -115,7 +120,7 @@ private static List embedSingleBatch( List batchEmbeddings; try { - batchEmbeddings = embedWithRetry(embeddingClient, textBatch); + batchEmbeddings = embedWithRetry(purpose, embeddingClient, textBatch); } catch (EmbeddingServiceUnavailableException embeddingFailure) { String firstBatchUrl = extractDocumentUrl(documentBatch.getFirst(), batchStartIndex); String lastBatchUrl = extractDocumentUrl(documentBatch.getLast(), batchEndIndex - 1); @@ -151,10 +156,11 @@ private static List embedSingleBatch( return batchEmbeddings; } - private static List embedWithRetry(EmbeddingClient embeddingClient, List textBatch) { + private static List embedWithRetry( + LlmGatewayTier.Purpose purpose, EmbeddingClient embeddingClient, List textBatch) { for (int embeddingAttempt = 1; embeddingAttempt <= MAX_EMBEDDING_ATTEMPTS; embeddingAttempt++) { try { - return embeddingClient.embed(textBatch, LlmGatewayTier.BATCH); + return embeddingClient.embed(purpose, textBatch, LlmGatewayTier.BATCH); } catch (EmbeddingServiceTemporarilyUnavailableException temporaryFailure) { if (embeddingAttempt == MAX_EMBEDDING_ATTEMPTS) { throw temporaryFailure; diff --git a/src/main/java/com/williamcallahan/javachat/service/EmbeddingClient.java b/src/main/java/com/williamcallahan/javachat/service/EmbeddingClient.java index 8a3be0cc..b2c6d527 100644 --- a/src/main/java/com/williamcallahan/javachat/service/EmbeddingClient.java +++ b/src/main/java/com/williamcallahan/javachat/service/EmbeddingClient.java @@ -14,11 +14,12 @@ public interface EmbeddingClient { /** * Produces one dense embedding vector per input text, preserving input order. * + * @param purpose application job responsible for the request * @param texts input texts * @param requestTier gateway capacity tier for this embedding request * @return embedding vectors in the same order as {@code texts} */ - List embed(List texts, LlmGatewayTier requestTier); + List embed(LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier); /** * Produces one dense embedding vector per input text within a caller-owned budget. @@ -29,12 +30,14 @@ public interface EmbeddingClient { * {@code requestTimeout} and the tier's configured request timeout, keeping the * configured timeout meaningful as a per-request ceiling.

* + * @param purpose application job responsible for the request * @param texts input texts * @param requestTier gateway capacity tier for this embedding request * @param requestTimeout caller-owned budget bounding the whole embedding operation * @return embedding vectors in the same order as {@code texts} */ - List embed(List texts, LlmGatewayTier requestTier, Duration requestTimeout); + List embed( + LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier, Duration requestTimeout); /** * Returns the provider model identifier used for embedding requests. @@ -46,14 +49,16 @@ public interface EmbeddingClient { /** * Produces a dense embedding vector for a single text. * + * @param purpose application job responsible for the request * @param text input text * @param requestTier gateway capacity tier for this embedding request * @return embedding vector */ - default float[] embed(String text, LlmGatewayTier requestTier) { + default float[] embed(LlmGatewayTier.Purpose purpose, String text, LlmGatewayTier requestTier) { + Objects.requireNonNull(purpose, "purpose"); Objects.requireNonNull(requestTier, "requestTier"); String safeText = Objects.requireNonNullElse(text, ""); - List embeddingVectors = embed(List.of(safeText), requestTier); + List embeddingVectors = embed(purpose, List.of(safeText), requestTier); if (embeddingVectors.isEmpty()) { throw new EmbeddingServiceUnavailableException("Embedding response was empty"); } @@ -63,16 +68,19 @@ default float[] embed(String text, LlmGatewayTier requestTier) { /** * Produces a dense embedding vector for a single text within a caller-owned budget. * + * @param purpose application job responsible for the request * @param text input text * @param requestTier gateway capacity tier for this embedding request * @param requestTimeout caller-owned budget bounding the whole embedding operation * @return embedding vector */ - default float[] embed(String text, LlmGatewayTier requestTier, Duration requestTimeout) { + default float[] embed( + LlmGatewayTier.Purpose purpose, String text, LlmGatewayTier requestTier, Duration requestTimeout) { + Objects.requireNonNull(purpose, "purpose"); Objects.requireNonNull(requestTier, "requestTier"); Objects.requireNonNull(requestTimeout, "requestTimeout"); String safeText = Objects.requireNonNullElse(text, ""); - List embeddingVectors = embed(List.of(safeText), requestTier, requestTimeout); + List embeddingVectors = embed(purpose, List.of(safeText), requestTier, requestTimeout); if (embeddingVectors.isEmpty()) { throw new EmbeddingServiceUnavailableException("Embedding response was empty"); } @@ -90,7 +98,7 @@ default float[] embed(String text, LlmGatewayTier requestTier, Duration requestT * Issues a minimal embedding request so the provider keeps its model resident. * *

Implementations must call their provider-specific request path directly instead - * of delegating to {@link #embed(List, LlmGatewayTier)}. The RAG pipeline logging aspect advises + * of delegating to {@link #embed(LlmGatewayTier.Purpose, List, LlmGatewayTier)}. The RAG pipeline logging aspect advises * public {@code embed} executions, so routing scheduled probes around that method * keeps "STEP 1" pipeline logs scoped to real requests.

* diff --git a/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java b/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java index e33819a9..0ce302b7 100644 --- a/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java +++ b/src/main/java/com/williamcallahan/javachat/service/EnrichmentService.java @@ -73,7 +73,11 @@ public Enrichment enrich(String userQuery, String jdkVersion, List conte } String enrichmentJson = openAIStreamingService - .completeJsonObject(prompt.toString(), enrichmentTemperature, enrichmentOutputTokenBudget) + .completeJsonObject( + LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, + prompt.toString(), + enrichmentTemperature, + enrichmentOutputTokenBudget) .block(); if (enrichmentJson == null || enrichmentJson.isBlank()) { throw new IllegalStateException("LLM returned empty enrichment response"); diff --git a/src/main/java/com/williamcallahan/javachat/service/HybridSearchService.java b/src/main/java/com/williamcallahan/javachat/service/HybridSearchService.java index 7d2838d7..025f4b0b 100644 --- a/src/main/java/com/williamcallahan/javachat/service/HybridSearchService.java +++ b/src/main/java/com/williamcallahan/javachat/service/HybridSearchService.java @@ -106,7 +106,11 @@ public List> searchByConstraint( float[] denseVector = queryEncoding .embeddingClient() - .embed(query, LlmGatewayTier.LIVE, remainingStageBudget(stageDeadlineNanos)); + .embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + query, + LlmGatewayTier.LIVE, + remainingStageBudget(stageDeadlineNanos)); LexicalSparseVectorEncoder.SparseVector sparseVector = queryEncoding.sparseVectorEncoder().encode(query); Duration queryTimeout = qdrantQueryExecutor.queryTimeout(); diff --git a/src/main/java/com/williamcallahan/javachat/service/HybridVectorService.java b/src/main/java/com/williamcallahan/javachat/service/HybridVectorService.java index 42e48ad5..95cd5147 100644 --- a/src/main/java/com/williamcallahan/javachat/service/HybridVectorService.java +++ b/src/main/java/com/williamcallahan/javachat/service/HybridVectorService.java @@ -92,7 +92,7 @@ public void upsert(QdrantCollectionKind collectionKind, List documents Objects.requireNonNull(collectionKind, NULL_MESSAGE_COLLECTION_KIND); String collectionName = Objects.requireNonNull(resolveCollectionName(collectionKind), NULL_MESSAGE_COLLECTION_NAME); - doUpsert(collectionName, documents); + doUpsert(indexingPurpose(collectionKind), collectionName, documents); } /** @@ -109,7 +109,7 @@ public void upsertToCollection(String collectionName, List documents) if (collectionName.isBlank()) { throw new IllegalArgumentException("collectionName must not be blank"); } - doUpsert(collectionName, documents); + doUpsert(LlmGatewayTier.Purpose.REPOSITORY_INDEX_EMBEDDING, collectionName, documents); } /** @@ -129,7 +129,7 @@ public void replaceUrlDocuments( Objects.requireNonNull(collectionKind, NULL_MESSAGE_COLLECTION_KIND); String collectionName = Objects.requireNonNull(resolveCollectionName(collectionKind), NULL_MESSAGE_COLLECTION_NAME); - replaceUrlDocuments(collectionName, sourceUrl, replacementDocuments); + replaceUrlDocuments(indexingPurpose(collectionKind), collectionName, sourceUrl, replacementDocuments); } /** @@ -140,6 +140,15 @@ public void replaceUrlDocuments( * @param replacementDocuments complete replacement document set with stable UUID identities */ public void replaceUrlDocuments(String collectionName, String sourceUrl, List replacementDocuments) { + replaceUrlDocuments( + LlmGatewayTier.Purpose.REPOSITORY_INDEX_EMBEDDING, collectionName, sourceUrl, replacementDocuments); + } + + private void replaceUrlDocuments( + LlmGatewayTier.Purpose purpose, + String collectionName, + String sourceUrl, + List replacementDocuments) { Objects.requireNonNull(collectionName, NULL_MESSAGE_COLLECTION_NAME); if (collectionName.isBlank()) { throw new IllegalArgumentException("collectionName must not be blank"); @@ -158,7 +167,7 @@ public void replaceUrlDocuments(String collectionName, String sourceUrl, List replacementPointIds = replacementPointIds(requiredReplacementDocuments); Set priorPointIds = scrollPointIdsForUrl(collectionName, requiredSourceUrl); - doUpsert(collectionName, requiredReplacementDocuments); + doUpsert(purpose, collectionName, requiredReplacementDocuments); priorPointIds.removeAll(replacementPointIds); deletePointIds(collectionName, priorPointIds); @@ -448,7 +457,16 @@ public String resolveCollectionName(QdrantCollectionKind kind) { }; } - private void doUpsert(String collectionName, List documents) { + private static LlmGatewayTier.Purpose indexingPurpose(QdrantCollectionKind collectionKind) { + return switch (collectionKind) { + case BOOKS -> LlmGatewayTier.Purpose.BOOK_INDEX_EMBEDDING; + case DOCS -> LlmGatewayTier.Purpose.DOCUMENTATION_INDEX_EMBEDDING; + case ARTICLES -> LlmGatewayTier.Purpose.ARTICLE_INDEX_EMBEDDING; + case PDFS -> LlmGatewayTier.Purpose.PDF_INDEX_EMBEDDING; + }; + } + + private void doUpsert(LlmGatewayTier.Purpose purpose, String collectionName, List documents) { Objects.requireNonNull(documents, "documents"); if (documents.isEmpty()) { return; @@ -457,8 +475,8 @@ private void doUpsert(String collectionName, List documents) { String denseVectorName = appProperties.getQdrant().getDenseVectorName(); String sparseVectorName = appProperties.getQdrant().getSparseVectorName(); - List embeddings = - Objects.requireNonNull(EmbeddingBatchEmbedder.embedDocuments(embeddingClient, documents), "embeddings"); + List embeddings = Objects.requireNonNull( + EmbeddingBatchEmbedder.embedDocuments(purpose, embeddingClient, documents), "embeddings"); if (embeddings.size() != documents.size()) { throw new IllegalStateException( "Embedding count mismatch: expected " + documents.size() + " but received " + embeddings.size()); diff --git a/src/main/java/com/williamcallahan/javachat/service/LlmGatewayTier.java b/src/main/java/com/williamcallahan/javachat/service/LlmGatewayTier.java index 41108e16..789b9adb 100644 --- a/src/main/java/com/williamcallahan/javachat/service/LlmGatewayTier.java +++ b/src/main/java/com/williamcallahan/javachat/service/LlmGatewayTier.java @@ -34,10 +34,42 @@ public String requestHeader() { } /** Returns the stable Java Chat workload label for this request tier. */ - public String workloadHeader(String operation) { - if (operation == null || operation.isBlank()) { - throw new IllegalArgumentException("operation must be non-blank"); + public String workloadHeader(Purpose purpose) { + return "java-chat:" + java.util.Objects.requireNonNull(purpose, "purpose").label + ":" + + (this == LIVE ? "live" : "batch"); + } + + /** Identifies the application job independently of its wire protocol and capacity tier. */ + public enum Purpose { + /** Answers a chat question. */ + CHAT_ANSWER("chatAnswer"), + /** Generates a guided lesson. */ + GUIDED_LEARNING("guidedLearning"), + /** Enriches rendered markdown. */ + MARKDOWN_ENRICHMENT("markdownEnrichment"), + /** Ranks retrieved search candidates. */ + SEARCH_RERANK("searchRerank"), + /** Embeds a live retrieval query. */ + SEARCH_QUERY_EMBEDDING("searchQueryEmbedding"), + /** Embeds content supplied directly to the batch executor. */ + CONTENT_INDEX_EMBEDDING("contentIndexEmbedding"), + /** Indexes repository source files. */ + REPOSITORY_INDEX_EMBEDDING("repositoryIndexEmbedding"), + /** Indexes documentation. */ + DOCUMENTATION_INDEX_EMBEDDING("documentationIndexEmbedding"), + /** Indexes book content. */ + BOOK_INDEX_EMBEDDING("bookIndexEmbedding"), + /** Indexes articles. */ + ARTICLE_INDEX_EMBEDDING("articleIndexEmbedding"), + /** Indexes PDF content. */ + PDF_INDEX_EMBEDDING("pdfIndexEmbedding"), + /** Keeps the embedding model resident. */ + EMBEDDING_WARMUP("embeddingWarmup"); + + private final String label; + + Purpose(String label) { + this.label = label; } - return "java-chat:" + operation + ":" + (this == LIVE ? "live" : "batch"); } } diff --git a/src/main/java/com/williamcallahan/javachat/service/LocalEmbeddingClient.java b/src/main/java/com/williamcallahan/javachat/service/LocalEmbeddingClient.java index 27437629..9f4e31de 100644 --- a/src/main/java/com/williamcallahan/javachat/service/LocalEmbeddingClient.java +++ b/src/main/java/com/williamcallahan/javachat/service/LocalEmbeddingClient.java @@ -73,7 +73,8 @@ public LocalEmbeddingClient( } @Override - public List embed(List texts, LlmGatewayTier requestTier) { + public List embed(LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier) { + Objects.requireNonNull(purpose, "purpose"); Objects.requireNonNull(requestTier, "requestTier"); if (texts == null || texts.isEmpty()) { return List.of(); @@ -89,7 +90,9 @@ public List embed(List texts, LlmGatewayTier requestTier) { * template whose connect and read timeouts cannot exceed the remaining budget.

*/ @Override - public List embed(List texts, LlmGatewayTier requestTier, Duration requestTimeout) { + public List embed( + LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier, Duration requestTimeout) { + Objects.requireNonNull(purpose, "purpose"); Objects.requireNonNull(requestTier, "requestTier"); Duration requiredRequestTimeout = Objects.requireNonNull(requestTimeout, "requestTimeout"); if (texts == null || texts.isEmpty()) { diff --git a/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java b/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java index bb0db4fa..dc760946 100644 --- a/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java +++ b/src/main/java/com/williamcallahan/javachat/service/OpenAIStreamingService.java @@ -113,7 +113,9 @@ public void shutdown() { } /** Retains selected-provider and context-document provenance alongside the response stream. */ - public Mono streamResponse(StructuredPrompt structuredPrompt, double temperature) { + public Mono streamResponse( + LlmGatewayTier.Purpose purpose, StructuredPrompt structuredPrompt, double temperature) { + java.util.Objects.requireNonNull(purpose, "purpose"); log.debug("Starting OpenAI stream with structured prompt"); return Mono.defer(() -> { @@ -125,8 +127,8 @@ public Mono streamResponse(StructuredPrompt structuredPrompt, d RateLimitService.ApiProvider configuredProvider = providerRoutingService.configuredProvider(); OpenAiPreparedRequest preparedStreamingRequest = requestFactory.prepareStreamingRequest(structuredPrompt, temperature); - Flux responseTextChunks = - executeStreamingWithConfiguredProvider(preparedStreamingRequest, configuredProvider); + Flux responseTextChunks = executeStreamingWithConfiguredProvider( + preparedStreamingRequest, configuredProvider, purpose); List contextDocumentIds = preparedStreamingRequest.structuredPrompt().contextDocuments().stream() .map(ContextDocumentSegment::documentId) @@ -137,34 +139,56 @@ public Mono streamResponse(StructuredPrompt structuredPrompt, d } /** Uses the configured output budget for a non-streaming text completion. */ - public Mono complete(String prompt, double temperature) { - return executeCompletion(prompt, temperature, CompletionRequestConfiguration.defaultText()); + public Mono complete(LlmGatewayTier.Purpose purpose, String prompt, double temperature) { + return executeCompletion(purpose, prompt, temperature, CompletionRequestConfiguration.defaultText()); } /** Preserves the caller's output budget for a non-streaming text completion. */ - public Mono complete(String prompt, double temperature, int maximumOutputTokens) { + public Mono complete( + LlmGatewayTier.Purpose purpose, String prompt, double temperature, int maximumOutputTokens) { return Mono.defer(() -> executeCompletion( - prompt, temperature, CompletionRequestConfiguration.boundedText(maximumOutputTokens))); + purpose, prompt, temperature, CompletionRequestConfiguration.boundedText(maximumOutputTokens))); } /** Requires a JSON object within the caller's output budget and the default request timeout. */ - public Mono completeJsonObject(String prompt, double temperature, int maximumOutputTokens) { + public Mono completeJsonObject( + LlmGatewayTier.Purpose purpose, String prompt, double temperature, int maximumOutputTokens) { return completeJsonObject( - prompt, temperature, maximumOutputTokens, CompletionRequestConfiguration.defaultRequestTimeout()); + purpose, + prompt, + temperature, + maximumOutputTokens, + CompletionRequestConfiguration.defaultRequestTimeout()); } /** Requires a JSON object within both the caller's output budget and whole-request timeout. */ public Mono completeJsonObject( - String prompt, double temperature, int maximumOutputTokens, Duration requestTimeout) { + LlmGatewayTier.Purpose purpose, + String prompt, + double temperature, + int maximumOutputTokens, + Duration requestTimeout) { return Mono.defer(() -> executeCompletion( - prompt, temperature, CompletionRequestConfiguration.jsonObject(maximumOutputTokens, requestTimeout))); + purpose, + prompt, + temperature, + CompletionRequestConfiguration.jsonObject(maximumOutputTokens, requestTimeout))); } private Mono executeCompletion( - String prompt, double temperature, CompletionRequestConfiguration configuration) { + LlmGatewayTier.Purpose purpose, + String prompt, + double temperature, + CompletionRequestConfiguration configuration) { + java.util.Objects.requireNonNull(purpose, "purpose"); return Mono.defer(() -> { RateLimitService.ApiProvider configuredProvider = providerRoutingService.configuredProvider(); - ResponseCreateParams requestParameters = buildCompletionRequest(prompt, temperature, configuration); + ResponseCreateParams requestParameters = + buildCompletionRequest(prompt, temperature, configuration).toBuilder() + .putAdditionalHeader( + LlmGatewayTier.REQUEST_WORKLOAD_HEADER, + workloadTier().workloadHeader(purpose)) + .build(); RequestOptions requestOptions = RequestOptions.builder() .timeout(Timeout.builder() .request(configuration.requestTimeout()) @@ -219,7 +243,9 @@ public boolean isRecoverableStreamingFailure(Throwable throwable) { } private Flux executeStreamingWithConfiguredProvider( - OpenAiPreparedRequest preparedStreamingRequest, RateLimitService.ApiProvider configuredProvider) { + OpenAiPreparedRequest preparedStreamingRequest, + RateLimitService.ApiProvider configuredProvider, + LlmGatewayTier.Purpose purpose) { log.info( "[LLM] Streaming started (structured, providerId={}, model={})", configuredProvider.ordinal(), @@ -229,7 +255,13 @@ private Flux executeStreamingWithConfiguredProvider( OpenAiProviderCandidate providerAdmission = requireConfiguredProviderAdmission(); AtomicBoolean emittedVisibleText = new AtomicBoolean(false); return executeStreamingRequest( - providerAdmission.client(), preparedStreamingRequest.responseParams(), configuredProvider) + providerAdmission.client(), + preparedStreamingRequest.responseParams().toBuilder() + .putAdditionalHeader( + LlmGatewayTier.REQUEST_WORKLOAD_HEADER, + workloadTier().workloadHeader(purpose)) + .build(), + configuredProvider) .doOnNext(textChunk -> { if (UnicodeVisibleContent.hasVisibleContent(textChunk)) { emittedVisibleText.set(true); @@ -375,8 +407,6 @@ private Duration streamingRequestTimeout() { private OpenAIClient createClient(String apiKey, String baseUrl) { String requestTier = resolvedLlmGatewayTier(); - LlmGatewayTier workloadTier = - LlmGatewayTier.BATCH.requestHeader().equals(requestTier) ? LlmGatewayTier.BATCH : LlmGatewayTier.LIVE; return OpenAIOkHttpClient.builder() .apiKey(apiKey) .baseUrl(OpenAiSdkUrlNormalizer.normalize(baseUrl)) @@ -384,10 +414,15 @@ private OpenAIClient createClient(String apiKey, String baseUrl) { // SDK retry sleeps interfere with reactive cancellation. .maxRetries(0) .putHeader(LlmGatewayTier.REQUEST_TIER_HEADER, requestTier) - .putHeader(LlmGatewayTier.REQUEST_WORKLOAD_HEADER, workloadTier.workloadHeader("chat-completions")) .build(); } + private LlmGatewayTier workloadTier() { + return LlmGatewayTier.BATCH.requestHeader().equals(resolvedLlmGatewayTier()) + ? LlmGatewayTier.BATCH + : LlmGatewayTier.LIVE; + } + private String resolvedLlmGatewayTier() { return llmGatewayTier == null || llmGatewayTier.isBlank() ? LlmGatewayTier.LIVE.requestHeader() diff --git a/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java b/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java index 78c24ff7..492dddcb 100644 --- a/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java +++ b/src/main/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClient.java @@ -262,9 +262,10 @@ private OpenAiCompatibleEmbeddingClient( } @Override - public List embed(List texts, LlmGatewayTier requestTier) { + public List embed(LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier) { + Objects.requireNonNull(purpose, "purpose"); Objects.requireNonNull(requestTier, "requestTier"); - return embed(texts, requestTier, requestTimeoutFor(requestTier)); + return embed(purpose, texts, requestTier, requestTimeoutFor(requestTier)); } /** @@ -275,7 +276,9 @@ public List embed(List texts, LlmGatewayTier requestTier) { * caller observes exactly the configured timeout.

*/ @Override - public List embed(List texts, LlmGatewayTier requestTier, Duration requestTimeout) { + public List embed( + LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier, Duration requestTimeout) { + Objects.requireNonNull(purpose, "purpose"); Objects.requireNonNull(requestTier, "requestTier"); Duration requiredRequestTimeout = Objects.requireNonNull(requestTimeout, "requestTimeout"); if (texts == null || texts.isEmpty()) { @@ -283,7 +286,8 @@ public List embed(List texts, LlmGatewayTier requestTier, Durat } activeForegroundEmbeddingCount.incrementAndGet(); try { - return createEmbeddings(texts, requestTier, tighterRequestTimeout(requestTier, requiredRequestTimeout)); + return createEmbeddings( + purpose, texts, requestTier, tighterRequestTimeout(requestTier, requiredRequestTimeout)); } finally { activeForegroundEmbeddingCount.decrementAndGet(); } @@ -310,7 +314,11 @@ public void warmUp() { throw new EmbeddingProbeDeferredException("foreground embedding work is active"); } try { - createEmbeddings(List.of(EMBEDDING_WARM_UP_PROBE_TEXT), LlmGatewayTier.BATCH, batchRequestTimeout); + createEmbeddings( + LlmGatewayTier.Purpose.EMBEDDING_WARMUP, + List.of(EMBEDDING_WARM_UP_PROBE_TEXT), + LlmGatewayTier.BATCH, + batchRequestTimeout); } catch (EmbeddingProviderCooldownRejectionException cooldownRejection) { throw new EmbeddingProbeDeferredException("the batch-tier provider cooldown is active"); } @@ -321,9 +329,11 @@ public String modelName() { return modelName; } - private List createEmbeddings(List texts, LlmGatewayTier requestTier, Duration requestTimeout) { + private List createEmbeddings( + LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier, Duration requestTimeout) { EmbeddingCreateParams embeddingRequest = EmbeddingCreateParams.builder() .model(modelName) + .putAdditionalHeader(LlmGatewayTier.REQUEST_WORKLOAD_HEADER, requestTier.workloadHeader(purpose)) .inputOfArrayOfStrings(texts) .build(); Semaphore requestPermits = requestPermitsFor(requestTier); @@ -737,7 +747,6 @@ private static OpenAIClient createTieredClient(String apiKey, String baseUrl, Ll .apiKey(apiKey) .baseUrl(baseUrl) .putHeader(LlmGatewayTier.REQUEST_TIER_HEADER, requestTier.requestHeader()) - .putHeader(LlmGatewayTier.REQUEST_WORKLOAD_HEADER, requestTier.workloadHeader("embeddings")) .maxRetries(0); return clientBuilder.build(); } diff --git a/src/main/java/com/williamcallahan/javachat/service/RerankerService.java b/src/main/java/com/williamcallahan/javachat/service/RerankerService.java index be750870..d606cc13 100644 --- a/src/main/java/com/williamcallahan/javachat/service/RerankerService.java +++ b/src/main/java/com/williamcallahan/javachat/service/RerankerService.java @@ -288,7 +288,12 @@ private Duration tighterRerankTimeout(Duration remainingStageBudget) { private Optional callLlmForReranking(String prompt, Duration rerankRequestTimeout) { try { return openAIStreamingService - .completeJsonObject(prompt, rerankerTemperature, rerankerOutputTokenBudget, rerankRequestTimeout) + .completeJsonObject( + LlmGatewayTier.Purpose.SEARCH_RERANK, + prompt, + rerankerTemperature, + rerankerOutputTokenBudget, + rerankRequestTimeout) .doOnError( timeoutOrApiError -> log.debug("Reranker LLM call timed out or failed", timeoutOrApiError)) .blockOptional(); diff --git a/src/main/java/com/williamcallahan/javachat/web/ChatController.java b/src/main/java/com/williamcallahan/javachat/web/ChatController.java index 990dc82d..8b0b810f 100644 --- a/src/main/java/com/williamcallahan/javachat/web/ChatController.java +++ b/src/main/java/com/williamcallahan/javachat/web/ChatController.java @@ -16,6 +16,7 @@ import com.williamcallahan.javachat.service.ConfiguredProviderTemporarilyUnavailableException; import com.williamcallahan.javachat.service.EmbeddingServiceUnavailableException; import com.williamcallahan.javachat.service.HybridSearchPartialFailureException; +import com.williamcallahan.javachat.service.LlmGatewayTier; import com.williamcallahan.javachat.service.OpenAIStreamingService; import com.williamcallahan.javachat.service.RerankingFailureException; import com.williamcallahan.javachat.service.RetrievalService; @@ -179,6 +180,7 @@ public Flux> stream( // Stream with provider transparency - surfaces which LLM is responding return openAIStreamingService .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, promptOutcome.structuredPrompt(), appProperties.getLlm().getTemperature()) .flatMapMany(streamingResult -> { diff --git a/src/main/java/com/williamcallahan/javachat/web/GuidedLearningController.java b/src/main/java/com/williamcallahan/javachat/web/GuidedLearningController.java index 326da279..4b9b174b 100644 --- a/src/main/java/com/williamcallahan/javachat/web/GuidedLearningController.java +++ b/src/main/java/com/williamcallahan/javachat/web/GuidedLearningController.java @@ -13,6 +13,7 @@ import com.williamcallahan.javachat.service.ChatMemoryService; import com.williamcallahan.javachat.service.ConfiguredProviderTemporarilyUnavailableException; import com.williamcallahan.javachat.service.GuidedLearningService; +import com.williamcallahan.javachat.service.LlmGatewayTier; import com.williamcallahan.javachat.service.MarkdownService; import com.williamcallahan.javachat.service.OpenAIStreamingService; import com.williamcallahan.javachat.service.RetrievalService; @@ -339,6 +340,7 @@ private Flux> streamGuidedResponse( // Stream with provider transparency - surfaces which LLM is responding return openAIStreamingService .streamResponse( + LlmGatewayTier.Purpose.GUIDED_LEARNING, promptOutcome.structuredPrompt(), appProperties.getLlm().getTemperature()) .flatMapMany(streamingResult -> { diff --git a/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java b/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java index e9efdf87..45b72941 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EmbeddingBatchEmbedderTest.java @@ -34,8 +34,10 @@ void usesBoundedMultiDocumentRequests() { EMBEDDING_DIMENSIONS, (requestIndex, textBatch) -> repeatedEmbeddings(textBatch.size(), EMBEDDING_VECTOR)); - List embeddingVectors = - EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(THREE_BATCH_DOCUMENT_COUNT)); + List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + embeddingClient, + sequentialDocuments(THREE_BATCH_DOCUMENT_COUNT)); assertEquals(THREE_BATCH_DOCUMENT_COUNT, embeddingVectors.size()); assertEquals(3, embeddingClient.requestedTextBatches.size()); @@ -62,8 +64,8 @@ void preservesProviderOutputOrderingAcrossBatchBoundaries() { .map(documentText -> new float[] {documentIndexFromText(documentText)}) .toList()); - List embeddingVectors = - EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(documentCount)); + List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(documentCount)); for (int documentIndex = 0; documentIndex < documentCount; documentIndex++) { assertArrayEquals(new float[] {documentIndex}, embeddingVectors.get(documentIndex)); @@ -81,7 +83,10 @@ void rejectsResponseCountMismatch() { EmbeddingServiceUnavailableException mismatchFailure = assertThrows( EmbeddingServiceUnavailableException.class, - () -> EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(documentCount))); + () -> EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + embeddingClient, + sequentialDocuments(documentCount))); assertTrue(mismatchFailure.getMessage().contains("expected 1 but received 0")); assertTrue(mismatchFailure @@ -108,7 +113,10 @@ void propagatesBatchFailureWithoutFallback() { EmbeddingServiceUnavailableException batchFailure = assertThrows( EmbeddingServiceUnavailableException.class, - () -> EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(documentCount))); + () -> EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + embeddingClient, + sequentialDocuments(documentCount))); assertSame(providerFailure, batchFailure.getCause()); int secondBatchStartIndex = EmbeddingBatchEmbedder.EMBEDDING_REQUEST_BATCH_SIZE; @@ -130,7 +138,8 @@ void rejectsEmbeddingDimensionMismatch() { EmbeddingServiceUnavailableException dimensionFailure = assertThrows( EmbeddingServiceUnavailableException.class, - () -> EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(2))); + () -> EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(2))); assertTrue(dimensionFailure.getMessage().contains("document index 0")); assertTrue(dimensionFailure.getMessage().contains("expected 2 but received 1")); @@ -146,7 +155,8 @@ void scalesRepresentativeJavaCorpusByBatchCount() { EMBEDDING_DIMENSIONS, (requestIndex, textBatch) -> repeatedEmbeddings(textBatch.size(), EMBEDDING_VECTOR)); - List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments(embeddingClient, representativeCorpus); + List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, representativeCorpus); int expectedRequestCount = Math.ceilDiv( REPRESENTATIVE_JAVA_CORPUS_CHUNK_COUNT, EmbeddingBatchEmbedder.EMBEDDING_REQUEST_BATCH_SIZE); @@ -178,7 +188,8 @@ void executesProviderRequestsConcurrently() { return repeatedEmbeddings(textBatch.size(), EMBEDDING_VECTOR); }); - EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(documentCount)); + EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(documentCount)); assertEquals(MINIMUM_PARALLEL_REQUEST_COUNT, embeddingClient.maximumConcurrentRequests.get()); } @@ -191,7 +202,8 @@ void isolatesBlockingProviderRequestsOnVirtualThreads() { return repeatedEmbeddings(textBatch.size(), EMBEDDING_VECTOR); }); - EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(1)); + EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(1)); assertEquals(1, embeddingClient.requestedTextBatches.size()); } @@ -227,7 +239,9 @@ void interruptsBlockedSiblingWhenOneBatchFails() { assertThrows( EmbeddingServiceUnavailableException.class, () -> EmbeddingBatchEmbedder.embedDocuments( - embeddingClient, sequentialDocuments(EmbeddingBatchEmbedder.EMBEDDING_REQUEST_BATCH_SIZE * 2))); + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + embeddingClient, + sequentialDocuments(EmbeddingBatchEmbedder.EMBEDDING_REQUEST_BATCH_SIZE * 2))); assertTrue(blockedRequestInterrupted.get()); } @@ -240,7 +254,8 @@ void boundsConcurrentProviderRequests() { EMBEDDING_DIMENSIONS, (requestIndex, textBatch) -> repeatedEmbeddings(textBatch.size(), EMBEDDING_VECTOR)); - EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(documentCount)); + EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(documentCount)); assertTrue(embeddingClient.maximumConcurrentRequests.get() <= EmbeddingBatchEmbedder.MAX_CONCURRENT_EMBEDDING_REQUESTS); @@ -257,7 +272,8 @@ void retriesTemporaryProviderFailuresBeforeMutationBoundary() { return repeatedEmbeddings(textBatch.size(), EMBEDDING_VECTOR); }); - List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(1)); + List embeddingVectors = EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(1)); assertEquals(1, embeddingVectors.size()); assertEquals(EmbeddingBatchEmbedder.MAX_EMBEDDING_ATTEMPTS, providerAttempts.get()); @@ -274,7 +290,8 @@ void doesNotRetryNonTransientProviderFailures() { assertThrows( EmbeddingServiceUnavailableException.class, - () -> EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(1))); + () -> EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(1))); assertEquals(1, providerAttempts.get()); } @@ -290,7 +307,8 @@ void exhaustsBoundedRetriesForTemporaryProviderFailures() { assertThrows( EmbeddingServiceUnavailableException.class, - () -> EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(1))); + () -> EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(1))); assertEquals(EmbeddingBatchEmbedder.MAX_EMBEDDING_ATTEMPTS, providerAttempts.get()); } @@ -308,7 +326,8 @@ void preservesCallerInterruptionInFailureCauseChain() { try { EmbeddingServiceUnavailableException batchFailure = assertThrows( EmbeddingServiceUnavailableException.class, - () -> EmbeddingBatchEmbedder.embedDocuments(embeddingClient, sequentialDocuments(1))); + () -> EmbeddingBatchEmbedder.embedDocuments( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, embeddingClient, sequentialDocuments(1))); assertTrue(Thread.currentThread().isInterrupted()); assertTrue(batchFailure.getCause() instanceof InterruptedException); @@ -355,7 +374,7 @@ private RecordingEmbeddingClient( } @Override - public List embed(List texts, LlmGatewayTier requestTier) { + public List embed(LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier) { int requestIndex = requestSequence.getAndIncrement(); requestedTextBatches.add(List.copyOf(texts)); requestedTiers.add(requestTier); @@ -369,8 +388,12 @@ public List embed(List texts, LlmGatewayTier requestTier) { } @Override - public List embed(List texts, LlmGatewayTier requestTier, java.time.Duration requestTimeout) { - return embed(texts, requestTier); + public List embed( + LlmGatewayTier.Purpose purpose, + List texts, + LlmGatewayTier requestTier, + java.time.Duration requestTimeout) { + return embed(purpose, texts, requestTier); } @Override diff --git a/src/test/java/com/williamcallahan/javachat/service/EmbeddingModelKeepAliveTest.java b/src/test/java/com/williamcallahan/javachat/service/EmbeddingModelKeepAliveTest.java index 50aa52fe..49f8c420 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EmbeddingModelKeepAliveTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EmbeddingModelKeepAliveTest.java @@ -175,12 +175,16 @@ private SequencedEmbeddingClient(ProbeOutcome... probeOutcome) { } @Override - public List embed(List texts, LlmGatewayTier requestTier) { + public List embed(LlmGatewayTier.Purpose purpose, List texts, LlmGatewayTier requestTier) { throw new AssertionError("keep-alive probes must not call embed(List, LlmGatewayTier)"); } @Override - public List embed(List texts, LlmGatewayTier requestTier, java.time.Duration requestTimeout) { + public List embed( + LlmGatewayTier.Purpose purpose, + List texts, + LlmGatewayTier requestTier, + java.time.Duration requestTimeout) { throw new AssertionError("keep-alive probes must not call embed(List, LlmGatewayTier, Duration)"); } diff --git a/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java b/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java index 508c0897..721e12de 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EmbeddingProviderFailureTest.java @@ -32,7 +32,8 @@ void localEmbeddingSurfacesHttpErrors() throws IOException { new LocalEmbeddingClient(baseUrl, "local-test-model", 12, 8, new RestTemplateBuilder()); EmbeddingServiceUnavailableException thrown = assertThrows( EmbeddingServiceUnavailableException.class, - () -> localClient.embed(List.of("hello"), LlmGatewayTier.LIVE)); + () -> localClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("hello"), LlmGatewayTier.LIVE)); assertTrue(thrown.getMessage().contains("HTTP 500")); } finally { @@ -57,7 +58,8 @@ void gatewayEmbeddingHttpErrorsSurfaceStatusCodes() throws IOException { OpenAiCompatibleEmbeddingClient.create(baseUrl, "test-key", gatewaySettings)) { EmbeddingServiceUnavailableException thrown = assertThrows( EmbeddingServiceUnavailableException.class, - () -> remoteClient.embed(List.of("hello"), LlmGatewayTier.LIVE)); + () -> remoteClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("hello"), LlmGatewayTier.LIVE)); assertTrue(thrown.getMessage().contains("HTTP 401")); assertNotNull(thrown.getCause()); diff --git a/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java b/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java index b22bfa89..c555ec8d 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceCacheTest.java @@ -46,7 +46,11 @@ void shouldRecomputeEnrichmentWhenContextSnippetsChange() { EnrichmentService enrichmentService = applicationContext.getBean(EnrichmentService.class); OpenAIStreamingService streamingService = applicationContext.getBean(OpenAIStreamingService.class); given(streamingService.isAvailable()).willReturn(true); - given(streamingService.completeJsonObject(anyString(), eq(ENRICHMENT_TEMPERATURE), anyInt())) + given(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + anyString(), + eq(ENRICHMENT_TEMPERATURE), + anyInt())) .willReturn(Mono.just(ENRICHMENT_INITIAL_JSON)) .willReturn(Mono.just(ENRICHMENT_FRESH_JSON)); @@ -61,7 +65,11 @@ void shouldRecomputeEnrichmentWhenContextSnippetsChange() { assertThat(freshContextEnrichment.getHints()).containsExactly(ENRICHMENT_FRESH_HINT); assertThat(repeatedInitialContextEnrichment.getHints()).containsExactly(ENRICHMENT_INITIAL_HINT); verify(streamingService, times(ENRICHMENT_CACHE_MISS_COMPLETION_COUNT)) - .completeJsonObject(anyString(), eq(ENRICHMENT_TEMPERATURE), anyInt()); + .completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), + anyString(), + 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 6dc78721..91557be7 100644 --- a/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/EnrichmentServiceTest.java @@ -7,6 +7,7 @@ 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.verify; @@ -50,7 +51,8 @@ void setUp() { @Test void enrichUsesJsonObjectCompletionBoundary() { - when(openAIStreamingService.completeJsonObject(anyString(), anyDouble(), anyInt())) + when(openAIStreamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) .thenReturn(Mono.just("{\"hints\":[\"Use List.copyOf for an unmodifiable snapshot.\"]}")); Enrichment enrichment = @@ -58,13 +60,16 @@ void enrichUsesJsonObjectCompletionBoundary() { assertEquals(ENRICHMENT_JDK_VERSION, enrichment.getJdkVersion()); assertEquals(List.of("Use List.copyOf for an unmodifiable snapshot."), enrichment.getHints()); - verify(openAIStreamingService).completeJsonObject(anyString(), anyDouble(), anyInt()); - verify(openAIStreamingService, never()).complete(anyString(), anyDouble()); + verify(openAIStreamingService) + .completeJsonObject(eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt()); + verify(openAIStreamingService, never()) + .complete(eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble()); } @Test void enrichRejectsUnexpectedPreambleWithBraces() { - when(openAIStreamingService.completeJsonObject(anyString(), anyDouble(), anyInt())) + when(openAIStreamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) .thenReturn(Mono.just("Here is the JSON (see {docs} for details):\n{\"hints\":[\"tip\"]}")); IllegalStateException jsonContractFailure = assertThrows( @@ -78,7 +83,8 @@ void enrichRejectsUnexpectedPreambleWithBraces() { @Test void enrichRejectsMalformedJsonObject() { - when(openAIStreamingService.completeJsonObject(anyString(), anyDouble(), anyInt())) + when(openAIStreamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) .thenReturn(Mono.just("{\"hints\":[")); IllegalStateException jsonContractFailure = assertThrows( @@ -92,7 +98,8 @@ void enrichRejectsMalformedJsonObject() { @Test void enrichRejectsUnknownResponseFields() { - when(openAIStreamingService.completeJsonObject(anyString(), anyDouble(), anyInt())) + when(openAIStreamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) .thenReturn(Mono.just("{\"hint\":[\"Use the documented API.\"]}")); IllegalStateException jsonContractFailure = assertThrows( @@ -106,7 +113,8 @@ void enrichRejectsUnknownResponseFields() { @Test void enrichRejectsDuplicateResponseFields() { - when(openAIStreamingService.completeJsonObject(anyString(), anyDouble(), anyInt())) + when(openAIStreamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) .thenReturn(Mono.just("{\"hints\":[\"first\"],\"hints\":[\"second\"]}")); IllegalStateException jsonContractFailure = assertThrows( @@ -127,7 +135,8 @@ void enrichRejectsTrailingTokensAndScalarCoercion() { "{\"jdkVersion\":25}"); for (String invalidEnrichmentJson : invalidEnrichmentJsonValues) { - when(openAIStreamingService.completeJsonObject(anyString(), anyDouble(), anyInt())) + when(openAIStreamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT), anyString(), anyDouble(), anyInt())) .thenReturn(Mono.just(invalidEnrichmentJson)); IllegalStateException jsonContractFailure = assertThrows( diff --git a/src/test/java/com/williamcallahan/javachat/service/HybridSearchServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/HybridSearchServiceTest.java index bb71bbe8..5e734856 100644 --- a/src/test/java/com/williamcallahan/javachat/service/HybridSearchServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/HybridSearchServiceTest.java @@ -109,7 +109,11 @@ void appliesServerFilterToHybridPrefetchAndDispatchesEachQueryWithTheConfiguredD appProperties.getQdrant().setDensePrefetchHnswEf(TEST_DENSE_PREFETCH_HNSW_EF); appProperties.getQdrant().setQueryTimeout(DISPATCH_BUDGET_TEST_TIMEOUT); - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L, 3L), List.of(2.0f, 1.0f))); @@ -154,7 +158,11 @@ void appliesServerFilterToHybridPrefetchAndDispatchesEachQueryWithTheConfiguredD @Test void primaryHybridSearchDoesNotTreatProjectTypeSyntaxAsAnOfficialJavadocConstraint() { - when(embeddingClient.embed(eq(EXACT_JAVA_API_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(EXACT_JAVA_API_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(EXACT_JAVA_API_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -179,7 +187,11 @@ void primaryHybridSearchDoesNotTreatProjectTypeSyntaxAsAnOfficialJavadocConstrai @Test void multipleVersionScopesShareOneEncodingAndRetainIndependentResultBudgets() { - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -218,7 +230,12 @@ void multipleVersionScopesShareOneEncodingAndRetainIndependentResultBudgets() { appProperties.getQdrant().getCollections().getDocs(), appProperties.getQdrant().getCollections().getPdfs()) .contains(queryRequest.getCollectionName()))); - verify(embeddingClient, times(1)).embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class)); + verify(embeddingClient, times(1)) + .embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class)); verify(sparseEncoder, times(1)).encode(HYBRID_QUERY); verify(gitHubCollectionDiscovery, never()).getDiscoveredCollections(); } @@ -230,9 +247,17 @@ void concurrentSearchesDispatchTheirQdrantFanOutsIndependently() int collectionCount = appProperties.getQdrant().getCollections().all().size(); CountDownLatch allFanOutsDispatched = new CountDownLatch(collectionCount * 2); List>> heldQueryFutures = new CopyOnWriteArrayList<>(); - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); - when(embeddingClient.embed(eq(SECOND_HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(SECOND_HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.3f, 0.2f, 0.1f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -276,11 +301,15 @@ void nearlyExhaustedStageDeadlineTightensEmbeddingAndQdrantBudgetsBelowTheirConf appProperties.getQdrant().setQueryTimeout(GENEROUS_QUERY_TIMEOUT); List capturedEmbeddingBudgets = new ArrayList<>(); doAnswer(invocation -> { - capturedEmbeddingBudgets.add(invocation.getArgument(2)); + capturedEmbeddingBudgets.add(invocation.getArgument(3)); return new float[] {0.1f, 0.2f, 0.3f}; }) .when(embeddingClient) - .embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class)); + .embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class)); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); List capturedQueryTimeouts = new ArrayList<>(); @@ -310,7 +339,11 @@ void nearlyExhaustedStageDeadlineTightensEmbeddingAndQdrantBudgetsBelowTheirConf @Test void collectionTimeoutReportsEffectiveStageBudgetInsteadOfConfiguredFullTimeout() { appProperties.getQdrant().setQueryTimeout(GENEROUS_QUERY_TIMEOUT); - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -600,7 +633,11 @@ void citationSearchRemainsStrictWhenHybridPartialFailuresAreConfiguredAsNonFatal @Test void sharesOneDispatchDeadlineAcrossMultipleStalledHybridQueries() { appProperties.getQdrant().setQueryTimeout(SHARED_QUERY_TIMEOUT); - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -638,7 +675,11 @@ void sharesOneDispatchDeadlineAcrossMultipleStalledHybridQueries() { @Test void exhaustedDispatchBudgetFailsUndispatchedCollectionsImmediately() { appProperties.getQdrant().setQueryTimeout(EXHAUSTED_DISPATCH_QUERY_TIMEOUT); - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -664,7 +705,11 @@ void exhaustedDispatchBudgetFailsUndispatchedCollectionsImmediately() { @Test void interruptionAfterFanOutCancelsEveryPendingQuery() throws InterruptedException { - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -728,7 +773,11 @@ void throwsWhenAnyCollectionFailsInStrictMode() { @Test void unavailableGrpcFailureRemainsTypedAndRetryable() { - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -753,7 +802,11 @@ void unavailableGrpcFailureRemainsTypedAndRetryable() { @Test void permissionDeniedGrpcFailureRemainsTypedAndPermanent() { - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -778,7 +831,11 @@ void permissionDeniedGrpcFailureRemainsTypedAndPermanent() { @Test void deadlineExceededGrpcFailureBecomesTimeoutCauseAndRetainsGrpcFailure() { - when(embeddingClient.embed(eq(HYBRID_QUERY), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(HYBRID_QUERY), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.1f, 0.2f, 0.3f}); when(sparseEncoder.encode(HYBRID_QUERY)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(1L), List.of(1.0f))); @@ -822,7 +879,11 @@ private HybridSearchService buildSearchServiceWithGitHubDiscovery( } private void stubPartialFailureQueryResponses(String queryText) { - when(embeddingClient.embed(eq(queryText), eq(LlmGatewayTier.LIVE), any(Duration.class))) + when(embeddingClient.embed( + eq(LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING), + eq(queryText), + eq(LlmGatewayTier.LIVE), + any(Duration.class))) .thenReturn(new float[] {0.5f, 0.1f, 0.4f}); when(sparseEncoder.encode(queryText)) .thenReturn(new LexicalSparseVectorEncoder.SparseVector(List.of(2L), List.of(1.0f))); diff --git a/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java index 393013f1..9e5602b9 100644 --- a/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/HybridVectorServiceTest.java @@ -85,7 +85,10 @@ void setUp() { @Test void upsertStartsFreshOperationAfterTransientFailure() { - when(embeddingClient.embed(List.of(DOCUMENT_TEXT), LlmGatewayTier.BATCH)) + when(embeddingClient.embed( + org.mockito.ArgumentMatchers.any(LlmGatewayTier.Purpose.class), + eq(List.of(DOCUMENT_TEXT)), + eq(LlmGatewayTier.BATCH))) .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))); @@ -364,7 +367,10 @@ private static FacetHit facetHit(String payloadValue, long pointCount) { } private void configureReplacementEmbedding() { - when(embeddingClient.embed(List.of(DOCUMENT_TEXT), LlmGatewayTier.BATCH)) + when(embeddingClient.embed( + org.mockito.ArgumentMatchers.any(LlmGatewayTier.Purpose.class), + eq(List.of(DOCUMENT_TEXT)), + eq(LlmGatewayTier.BATCH))) .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/LocalEmbeddingClientTest.java b/src/test/java/com/williamcallahan/javachat/service/LocalEmbeddingClientTest.java index 216bf4bc..ecb60755 100644 --- a/src/test/java/com/williamcallahan/javachat/service/LocalEmbeddingClientTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/LocalEmbeddingClientTest.java @@ -67,8 +67,10 @@ void batchesRequestsAndPreservesEmbeddingOrderByIndex() throws IOException { try { LocalEmbeddingClient localEmbeddingClient = new LocalEmbeddingClient(baseUrl, "local-model", 3, 2, new RestTemplateBuilder()); - List embeddingVectors = - localEmbeddingClient.embed(List.of("alpha", "beta", "gamma"), LlmGatewayTier.LIVE); + List embeddingVectors = localEmbeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("alpha", "beta", "gamma"), + LlmGatewayTier.LIVE); assertEquals(2, requestCounter.get()); assertEquals(List.of(2, 1), observedBatchSizes); @@ -101,7 +103,8 @@ void failsWhenLocalEmbeddingDimensionsDoNotMatchConfiguredDimensions() throws IO new LocalEmbeddingClient(baseUrl, "local-model", 3, 8, new RestTemplateBuilder()); EmbeddingServiceUnavailableException thrownException = assertThrows( EmbeddingServiceUnavailableException.class, - () -> localEmbeddingClient.embed(List.of("alpha"), LlmGatewayTier.LIVE)); + () -> localEmbeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("alpha"), LlmGatewayTier.LIVE)); assertTrue(thrownException.getMessage().contains("dimension mismatch")); } finally { httpServer.stop(0); @@ -141,7 +144,10 @@ void oneCallerDeadlineBoundsEveryLocalEmbeddingBatchTogether() throws IOExceptio () -> assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, () -> localEmbeddingClient.embed( - List.of("first", "second"), LlmGatewayTier.LIVE, MULTI_BATCH_OPERATION_BUDGET))); + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("first", "second"), + LlmGatewayTier.LIVE, + MULTI_BATCH_OPERATION_BUDGET))); assertTrue(deadlineFailure.getMessage().contains("Local embedding")); TimeoutException socketDeadlineFailure = @@ -174,7 +180,11 @@ void expiredCallerDeadlineDoesNotDispatchLocalEmbeddingRequest() throws IOExcept EmbeddingServiceTemporarilyUnavailableException deadlineFailure = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> localEmbeddingClient.embed(List.of("expired"), LlmGatewayTier.LIVE, Duration.ZERO)); + () -> localEmbeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("expired"), + LlmGatewayTier.LIVE, + Duration.ZERO)); assertInstanceOf(TimeoutException.class, deadlineFailure.getCause()); assertEquals(0, requestCounter.get()); @@ -195,7 +205,10 @@ void preservesNonTimeoutLocalTransportFailuresWithoutDeadlineCause() { EmbeddingServiceTemporarilyUnavailableException transportFailure = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> localEmbeddingClient.embed(List.of("transport failure"), LlmGatewayTier.LIVE)); + () -> localEmbeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("transport failure"), + LlmGatewayTier.LIVE)); ResourceAccessException typedTransportFailure = assertInstanceOf(ResourceAccessException.class, transportFailure.getCause()); @@ -213,7 +226,10 @@ void reactorConnectTimeoutPreservesRetrievalTimeoutCause() { EmbeddingServiceTemporarilyUnavailableException transportTimeout = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> localEmbeddingClient.embed(List.of("transport timeout"), LlmGatewayTier.LIVE)); + () -> localEmbeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("transport timeout"), + LlmGatewayTier.LIVE)); TimeoutException timeoutFailure = assertInstanceOf(TimeoutException.class, transportTimeout.getCause()); ResourceAccessException typedTransportFailure = diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java index 6b615011..8ba551a5 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAICompletionStatusTest.java @@ -30,6 +30,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.StructuredPrompt; import java.io.IOException; import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; @@ -40,7 +41,7 @@ import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; +import org.junit.jupiter.params.provider.CsvSource; import org.springframework.test.util.ReflectionTestUtils; import reactor.test.StepVerifier; @@ -53,8 +54,15 @@ class OpenAICompletionStatusTest { private static final String PROVIDER_REFUSAL = "I cannot help with that request."; @ParameterizedTest - @ValueSource(strings = {"production-z", "production-a", "batch"}) - void jsonCompletionSendsMatchingWorkloadTierAndNoReasoning(String requestTier) throws IOException { + @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" + }) + void completionSendsMatchingWorkloadTier(String requestTier, LlmGatewayTier.Purpose purpose, boolean streaming) + throws IOException { AtomicReference capturedRequest = new AtomicReference<>(); HttpServer gatewayServer = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); gatewayServer.createContext("/v1/responses", exchange -> { @@ -70,7 +78,17 @@ void jsonCompletionSendsMatchingWorkloadTierAndNoReasoning(String requestTier) t "status":"completed","content":[{"type":"output_text", "text":"{\\"ok\\":true}","annotations":[]}]}]} """.getBytes(StandardCharsets.UTF_8); - exchange.getResponseHeaders().set("Content-Type", "application/json"); + if (streaming) { + String completedResponse = + new ObjectMapper().readTree(responseBody).toString(); + responseBody = ("event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\"," + + "\"sequence_number\":1,\"item_id\":\"msg_wire\",\"output_index\":0,\"content_index\":0," + + "\"delta\":\"{\\\"ok\\\":true}\"}\n\nevent: response.completed\ndata: " + + "{\"type\":\"response.completed\",\"sequence_number\":2,\"response\":" + + completedResponse + "}\n\n") + .getBytes(StandardCharsets.UTF_8); + } + exchange.getResponseHeaders().set("Content-Type", streaming ? "text/event-stream" : "application/json"); exchange.sendResponseHeaders(200, responseBody.length); try (var responseStream = exchange.getResponseBody()) { responseStream.write(responseBody); @@ -94,11 +112,22 @@ void jsonCompletionSendsMatchingWorkloadTierAndNoReasoning(String requestTier) t "openAiClient", configuredClient.withOptions(options -> options.baseUrl( "http://127.0.0.1:" + gatewayServer.getAddress().getPort() + "/v1"))); - assertEquals( - "{\"ok\":true}", - streamingService - .completeJsonObject("Return a JSON object with ok true", 0.0, 512) - .block(Duration.ofSeconds(10))); + if (streaming) { + assertEquals( + "{\"ok\":true}", + streamingService + .streamResponse(purpose, StructuredPrompt.fromRawPrompt("Return ok true", 1), 0.0) + .flatMapMany(StreamingResult::textChunks) + .collectList() + .map(chunks -> String.join("", chunks)) + .block(Duration.ofSeconds(10))); + } else { + assertEquals( + "{\"ok\":true}", + streamingService + .completeJsonObject(purpose, "Return a JSON object with ok true", 0.0, 512) + .block(Duration.ofSeconds(10))); + } } finally { streamingService.shutdown(); gatewayServer.stop(0); @@ -107,12 +136,19 @@ void jsonCompletionSendsMatchingWorkloadTierAndNoReasoning(String requestTier) t assertNotNull(gatewayRequest); assertEquals(requestTier, gatewayRequest.tierHeader()); assertEquals( - "java-chat:chat-completions:" + (requestTier.equals("batch") ? "batch" : "live"), + (requestTier.equals("batch") ? LlmGatewayTier.BATCH : LlmGatewayTier.LIVE).workloadHeader(purpose), gatewayRequest.workloadHeader()); - assertEquals(512, gatewayRequest.requestBody().path("max_output_tokens").asInt()); - assertEquals( - "none", - gatewayRequest.requestBody().path("reasoning").path("effort").asText()); + if (!streaming) { + assertEquals( + 512, gatewayRequest.requestBody().path("max_output_tokens").asInt()); + assertEquals( + "none", + gatewayRequest + .requestBody() + .path("reasoning") + .path("effort") + .asText()); + } } @Test @@ -129,7 +165,7 @@ void completionRecordsSuccessOnlyAfterCompletedStatus() { .thenReturn(providerCompletionFuture); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); - StepVerifier.create(streamingService.complete("prompt", 0.7)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7)) .then(() -> providerCompletionFuture.complete(providerCompletion)) .expectNext(PROVIDER_REFUSAL) .verifyComplete(); @@ -238,7 +274,7 @@ void cancelledCompletionRetainsSdkFutureForLateOutcomeAccounting() { .thenReturn(providerCompletionFuture); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); - StepVerifier.create(streamingService.complete("prompt", 0.7)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7)) .thenCancel() .verify(); @@ -260,7 +296,7 @@ private static void assertCompletionFailure(CompletionFailureScenario completion .thenReturn(CompletableFuture.completedFuture(completedCompletion)); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); - StepVerifier.create(streamingService.complete("prompt", 0.7)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7)) .expectErrorSatisfies(completionFailure -> { OpenAiResponseException terminalFailure = assertInstanceOf(OpenAiResponseException.class, completionFailure); @@ -276,7 +312,7 @@ private static void assertCompletionFailure(CompletionFailureScenario completion verify(completionFailureScenario.providerCompletion(), never()).output(); } if (completionFailureScenario.expectsConfiguredProviderBackoff()) { - StepVerifier.create(streamingService.complete("prompt", 0.7)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7)) .expectError(ConfiguredProviderTemporarilyUnavailableException.class) .verify(); verify(responseService).create(any(ResponseCreateParams.class), any(RequestOptions.class)); @@ -284,7 +320,7 @@ private static void assertCompletionFailure(CompletionFailureScenario completion return; } - StepVerifier.create(streamingService.complete("prompt", 0.7)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7)) .expectNext(VISIBLE_COMPLETION) .verifyComplete(); verify(responseService, times(EXPECTED_PROVIDER_REQUEST_COUNT_AFTER_READMISSION)) diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java index da65ec6e..291deaaa 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingProviderFailureTest.java @@ -90,7 +90,7 @@ void completionTransportFailureIsTerminal() { .thenReturn(CompletableFuture.failedFuture(upstreamFailure)); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); - StepVerifier.create(streamingService.complete("test", TEST_TEMPERATURE)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "test", TEST_TEMPERATURE)) .expectErrorSatisfies(completionFailure -> assertSame(upstreamFailure, completionFailure)) .verify(); @@ -132,7 +132,7 @@ void completionPreservesUpstreamAndAttachesUnexpectedRateLimitRecordingFailure() RateLimitService.ApiProvider.OPENAI, upstreamRateLimitFailure); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); - StepVerifier.create(streamingService.complete("test", TEST_TEMPERATURE)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "test", TEST_TEMPERATURE)) .expectErrorSatisfies(completionFailure -> { assertSame(upstreamRateLimitFailure, completionFailure); assertEquals(1, completionFailure.getSuppressed().length); @@ -187,7 +187,10 @@ void subscribedTerminalStreamFailureLogsOneBoundedAlert() { ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), TEST_TEMPERATURE) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, + StructuredPrompt.fromRawPrompt("test", 1), + TEST_TEMPERATURE) .flatMapMany(StreamingResult::textChunks)) .expectErrorSatisfies(streamingFailure -> { OpenAiStreamingFailureException terminalFailure = @@ -212,7 +215,8 @@ void unavailableStreamDefersErrorSeverityToRequestBoundary() { OpenAIStreamingService streamingService = configuredStreamingService(mock(RateLimitService.class)); IllegalStateException unavailableFailure = assertThrows(IllegalStateException.class, () -> streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), TEST_TEMPERATURE) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), TEST_TEMPERATURE) .block()); assertFalse(streamingService.isRecoverableStreamingFailure(unavailableFailure)); @@ -271,7 +275,7 @@ private static RateLimitException headerlessRateLimitFailure() { private static void assertCompletionPreservesUpstreamFailure( OpenAIStreamingService streamingService, RateLimitException expectedUpstreamFailure) { - StepVerifier.create(streamingService.complete("test", TEST_TEMPERATURE)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "test", TEST_TEMPERATURE)) .expectErrorSatisfies(completionFailure -> { assertSame(expectedUpstreamFailure, completionFailure); assertSuppressedDecisionFailure(expectedUpstreamFailure, RATE_LIMIT_HEADERS_MISSING_MESSAGE); @@ -282,7 +286,10 @@ private static void assertCompletionPreservesUpstreamFailure( private static void assertStreamingPreservesUpstreamFailure( OpenAIStreamingService streamingService, RateLimitException expectedUpstreamFailure) { StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), TEST_TEMPERATURE) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, + StructuredPrompt.fromRawPrompt("test", 1), + TEST_TEMPERATURE) .flatMapMany(StreamingResult::textChunks)) .expectErrorSatisfies(streamingFailure -> { OpenAiStreamingFailureException terminalFailure = diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java index 95335a1d..ea7ee994 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingServiceTest.java @@ -152,7 +152,8 @@ void emptyTextDeltaBeforeStreamingFailureIsTerminal() { ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectNext("") .expectErrorSatisfies(failure -> { @@ -181,7 +182,8 @@ void providerStreamWithoutVisibleTextIsTerminal() { streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectNext(invisibleText) .expectErrorSatisfies(failure -> @@ -274,7 +276,8 @@ void terminalProviderEventsFailTheStreamBeforeSuccess() { streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectErrorSatisfies(failure -> assertInstanceOf( OpenAIException.class, @@ -295,7 +298,8 @@ void visibleProviderTextWithoutCompletedEventIsTerminal() { OpenAIStreamingService streamingService = streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectNext("truncated response") .expectErrorSatisfies(failure -> { @@ -323,7 +327,8 @@ void visibleProviderTextCompletesSuccessfully() { OpenAIStreamingService streamingService = streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectNext("visible response") .verifyComplete(); @@ -344,7 +349,8 @@ void refusalTextCompletesAsVisibleAssistantOutput() { OpenAIStreamingService streamingService = streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectNext("I cannot help with that request.") .verifyComplete(); @@ -368,7 +374,8 @@ void failedServerResponseStartsBackoffAndRemainsRetryable() { OpenAIStreamingService streamingService = streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectErrorSatisfies(failure -> { OpenAiResponseException upstreamFailure = assertInstanceOf( @@ -381,7 +388,8 @@ void failedServerResponseStartsBackoffAndRemainsRetryable() { .verify(); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectError(ConfiguredProviderTemporarilyUnavailableException.class) .verify(); @@ -404,7 +412,8 @@ void failedRateLimitResponseStartsBackoffWithoutImmediateRetry() { OpenAIStreamingService streamingService = streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectErrorSatisfies(failure -> { OpenAiResponseException upstreamFailure = assertInstanceOf( @@ -419,7 +428,8 @@ void failedRateLimitResponseStartsBackoffWithoutImmediateRetry() { .verify(); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectError(ConfiguredProviderTemporarilyUnavailableException.class) .verify(); @@ -446,7 +456,8 @@ void incompleteResponseReasonsRemainDistinctAndNonRetryable() { streamingServiceForProviderStream(rateLimitService, providerStream); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectErrorSatisfies(failure -> { OpenAiResponseException upstreamFailure = assertInstanceOf( @@ -479,14 +490,15 @@ void deniedConfiguredProviderReservationTerminatesStreamingAndCompletionAsRetrya ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectErrorSatisfies(failure -> { assertInstanceOf(ConfiguredProviderTemporarilyUnavailableException.class, failure); assertTrue(streamingService.isRecoverableStreamingFailure(failure)); }) .verify(); - StepVerifier.create(streamingService.complete("test", 0.7)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "test", 0.7)) .expectErrorSatisfies(failure -> { assertInstanceOf(ConfiguredProviderTemporarilyUnavailableException.class, failure); assertTrue(streamingService.isRecoverableStreamingFailure(failure)); @@ -510,7 +522,7 @@ void cooldownRecordedAfterStreamCreationPreventsDispatchAtTextSubscription() { ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); StreamingResult streamingResult = Objects.requireNonNull(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse(LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .block()); InternalServerException gatewayTimeout = InternalServerException.builder() .statusCode(504) @@ -551,8 +563,9 @@ void exposesOnlyDocumentIdentitiesRetainedByGpt54Truncation() { List.of(), new CurrentQuerySegment("Explain the retained reference", 50)); - StreamingResult streamingResult = Objects.requireNonNull( - streamingService.streamResponse(oversizedPrompt, 0.7).block()); + StreamingResult streamingResult = Objects.requireNonNull(streamingService + .streamResponse(LlmGatewayTier.Purpose.CHAT_ANSWER, oversizedPrompt, 0.7) + .block()); assertEquals(List.of("retained-context-document"), streamingResult.contextDocumentIds()); verifyNoInteractions(openAiClient); @@ -578,7 +591,7 @@ void preTextStreamingTransportFailureIsTerminal() { ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); StreamingResult streamingResult = Objects.requireNonNull(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse(LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .block()); assertEquals(RateLimitService.ApiProvider.OPENAI, streamingResult.provider()); @@ -621,7 +634,8 @@ void streamingFailureAfterTextIsTerminal() { ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); StepVerifier.create(streamingService - .streamResponse(StructuredPrompt.fromRawPrompt("test", 1), 0.7) + .streamResponse( + LlmGatewayTier.Purpose.CHAT_ANSWER, StructuredPrompt.fromRawPrompt("test", 1), 0.7) .flatMapMany(StreamingResult::textChunks)) .expectNext("first token") .expectErrorSatisfies(failure -> { @@ -654,7 +668,7 @@ void completionPreservesAsyncTransportFailure() { .thenReturn(providerCompletionFuture); ReflectionTestUtils.setField(streamingService, "openAiClient", openAiClient); - StepVerifier.create(streamingService.complete("prompt", 0.7)) + StepVerifier.create(streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7)) .then(() -> providerCompletionFuture.completeExceptionally(upstreamFailure)) .expectErrorSatisfies(completionFailure -> assertSame(upstreamFailure, completionFailure)) .verify(); @@ -666,15 +680,23 @@ void completionPreservesAsyncTransportFailure() { void invalidCompletionConfigurationFailsOnSubscription() { OpenAIStreamingService streamingService = createStreamingService(); - assertCompletionFailure(streamingService.complete("prompt", 0.7, 0), "maximumOutputTokens must be positive"); - assertCompletionFailure(streamingService.complete("prompt", 0.7, -1), "maximumOutputTokens must be positive"); assertCompletionFailure( - streamingService.completeJsonObject("prompt", 0.7, 128, null), "requestTimeout must be positive"); + streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7, 0), + "maximumOutputTokens must be positive"); + assertCompletionFailure( + streamingService.complete(LlmGatewayTier.Purpose.CHAT_ANSWER, "prompt", 0.7, -1), + "maximumOutputTokens must be positive"); + assertCompletionFailure( + streamingService.completeJsonObject( + LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, "prompt", 0.7, 128, null), + "requestTimeout must be positive"); assertCompletionFailure( - streamingService.completeJsonObject("prompt", 0.7, 128, Duration.ZERO), + streamingService.completeJsonObject( + LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, "prompt", 0.7, 128, Duration.ZERO), "requestTimeout must be positive"); assertCompletionFailure( - streamingService.completeJsonObject("prompt", 0.7, 128, Duration.ofSeconds(-1)), + streamingService.completeJsonObject( + LlmGatewayTier.Purpose.MARKDOWN_ENRICHMENT, "prompt", 0.7, 128, Duration.ofSeconds(-1)), "requestTimeout must be positive"); } diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java index f8ab56df..11b0991a 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAIStreamingTimeoutContractTest.java @@ -65,7 +65,7 @@ void streamingReadDeadlineInheritsWholeRequestDeadline() { ReflectionTestUtils.setField(streamingService, "streamingRequestTimeoutSeconds", 600L); assertThrows(OpenAiStreamingFailureException.class, () -> streamingService - .streamResponse(structuredPrompt, 0.4) + .streamResponse(LlmGatewayTier.Purpose.CHAT_ANSWER, structuredPrompt, 0.4) .flatMapMany(StreamingResult::textChunks) .blockLast(Duration.ofSeconds(10))); diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java index 5db73804..50be1e8b 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientTest.java @@ -133,7 +133,8 @@ void callUsesSdkAndPreservesIndexOrdering() { try (OpenAiCompatibleEmbeddingClient clientAdapter = OpenAiCompatibleEmbeddingClient.create(client, gatewaySettings())) { - List vectors = clientAdapter.embed(List.of("a", "b"), LlmGatewayTier.LIVE); + List vectors = clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("a", "b"), LlmGatewayTier.LIVE); assertEquals(2, vectors.size()); assertEquals(0.25f, vectors.get(0)[0]); @@ -169,7 +170,8 @@ void throwsWhenEmbeddingDimensionDoesNotMatchConfiguration() { OpenAiCompatibleEmbeddingClient.create(client, gatewaySettings())) { EmbeddingServiceUnavailableException thrownException = assertThrows( EmbeddingServiceUnavailableException.class, - () -> clientAdapter.embed(List.of("a"), LlmGatewayTier.LIVE)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("a"), LlmGatewayTier.LIVE)); assertTrue(thrownException.getMessage().contains("dimension mismatch")); verify(embeddingService, times(1)).create(any(), any(RequestOptions.class)); } @@ -199,7 +201,8 @@ void doesNotRetryResponseValidationFailuresOutsideTheSdk() { OpenAiCompatibleEmbeddingClient.create(client, gatewaySettings())) { assertThrows( EmbeddingServiceUnavailableException.class, - () -> clientAdapter.embed(List.of("single"), LlmGatewayTier.LIVE)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("single"), LlmGatewayTier.LIVE)); verify(embeddingService).create(any(), any(RequestOptions.class)); } } @@ -216,7 +219,8 @@ void sdkTransportTimeoutPreservesRetrievalTimeoutCause() { OpenAiCompatibleEmbeddingClient.create(client, gatewaySettings())) { EmbeddingServiceTemporarilyUnavailableException embeddingTimeout = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> clientAdapter.embed(List.of("single"), LlmGatewayTier.LIVE)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("single"), LlmGatewayTier.LIVE)); TimeoutException timeoutFailure = assertInstanceOf(TimeoutException.class, embeddingTimeout.getCause()); assertSame(sdkTimeout, timeoutFailure.getCause()); @@ -246,7 +250,10 @@ void throwsWhenEmbeddingResponseOmitsIndex() { OpenAiCompatibleEmbeddingClient.create(client, gatewaySettings())) { EmbeddingServiceUnavailableException thrownException = assertThrows( EmbeddingServiceUnavailableException.class, - () -> clientAdapter.embed(List.of("missing index"), LlmGatewayTier.LIVE)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("missing index"), + LlmGatewayTier.LIVE)); assertTrue(thrownException.getMessage().contains("omitted index")); } @@ -273,7 +280,8 @@ void embed_omitsDimensionsForNonTextEmbedding3Models() { try (OpenAiCompatibleEmbeddingClient clientAdapter = OpenAiCompatibleEmbeddingClient.create(client, gatewaySettings())) { - clientAdapter.embed(List.of("dimension check"), LlmGatewayTier.LIVE); + clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("dimension check"), LlmGatewayTier.LIVE); ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(EmbeddingCreateParams.class); verify(embeddingService).create(requestCaptor.capture(), any(RequestOptions.class)); @@ -307,7 +315,8 @@ void routesEmbeddingRequestsToTierSpecificSdkClients() { try (OpenAiCompatibleEmbeddingClient clientAdapter = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { ArgumentCaptor liveRequestOptionsCaptor = ArgumentCaptor.forClass(RequestOptions.class); - clientAdapter.embed(List.of("live query"), LlmGatewayTier.LIVE); + clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("live query"), LlmGatewayTier.LIVE); verify(liveEmbeddingService).create(any(), liveRequestOptionsCaptor.capture()); verifyNoInteractions(batchEmbeddingService); assertRemainingTransportBudget( @@ -321,7 +330,8 @@ void routesEmbeddingRequestsToTierSpecificSdkClients() { liveRequestOptionsCaptor.getValue().getTimeout().read()); ArgumentCaptor batchRequestOptionsCaptor = ArgumentCaptor.forClass(RequestOptions.class); - clientAdapter.embed(List.of("batch document"), LlmGatewayTier.BATCH); + clientAdapter.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("batch document"), LlmGatewayTier.BATCH); verify(batchEmbeddingService).create(any(), batchRequestOptionsCaptor.capture()); assertEquals( EXPECTED_BATCH_EMBEDDING_CONNECT_TIMEOUT, @@ -358,7 +368,11 @@ void callerOwnedBudgetTightensTheLiveRequestTimeoutBelowTheConfiguredCap() { try (OpenAiCompatibleEmbeddingClient clientAdapter = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { ArgumentCaptor liveRequestOptionsCaptor = ArgumentCaptor.forClass(RequestOptions.class); - clientAdapter.embed(List.of("live query"), LlmGatewayTier.LIVE, CALLER_OWNED_LIVE_REQUEST_BUDGET); + clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("live query"), + LlmGatewayTier.LIVE, + CALLER_OWNED_LIVE_REQUEST_BUDGET); verify(liveEmbeddingService).create(any(), liveRequestOptionsCaptor.capture()); assertRemainingTransportBudget( CALLER_OWNED_LIVE_REQUEST_BUDGET, @@ -381,7 +395,11 @@ void expiredCallerBudgetFailsAsTimeoutBeforeGatewayDispatch() { new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { EmbeddingServiceTemporarilyUnavailableException deadlineFailure = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> clientAdapter.embed(List.of("expired query"), LlmGatewayTier.LIVE, Duration.ZERO)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("expired query"), + LlmGatewayTier.LIVE, + Duration.ZERO)); assertInstanceOf(TimeoutException.class, deadlineFailure.getCause()); verifyNoInteractions(liveClient); @@ -406,8 +424,8 @@ void defersProbeWhenForegroundEmbeddingIsAlreadyActive() try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor(); OpenAiCompatibleEmbeddingClient clientAdapter = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { - Future> foregroundEmbedding = - executor.submit(() -> clientAdapter.embed(List.of("live query"), LlmGatewayTier.LIVE)); + Future> foregroundEmbedding = executor.submit(() -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("live query"), LlmGatewayTier.LIVE)); try { assertTrue(foregroundStarted.await(5, TimeUnit.SECONDS)); @@ -448,7 +466,10 @@ void foregroundDoesNotWaitForProbeAdmittedBeforeItArrives() assertEquals( 1, clientAdapter - .embed(List.of("live query"), LlmGatewayTier.LIVE) + .embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("live query"), + LlmGatewayTier.LIVE) .size()); verify(liveEmbeddingService).create(any(), any(RequestOptions.class)); } finally { @@ -486,17 +507,20 @@ void limitsBatchRequestConcurrencyWithoutConsumingLiveRequestCapacity() try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor(); OpenAiCompatibleEmbeddingClient embeddingClient = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { - Future> firstBatch = - executor.submit(() -> embeddingClient.embed(List.of("first batch"), LlmGatewayTier.BATCH)); + Future> firstBatch = executor.submit(() -> embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("first batch"), LlmGatewayTier.BATCH)); assertTrue(firstBatchStarted.await(5, TimeUnit.SECONDS)); - Future> secondBatch = - executor.submit(() -> embeddingClient.embed(List.of("second batch"), LlmGatewayTier.BATCH)); + Future> secondBatch = executor.submit(() -> embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("second batch"), LlmGatewayTier.BATCH)); assertFalse(secondBatchStarted.await(200, TimeUnit.MILLISECONDS)); assertEquals( 1, embeddingClient - .embed(List.of("live query"), LlmGatewayTier.LIVE) + .embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("live query"), + LlmGatewayTier.LIVE) .size()); firstBatchResponseFuture.complete(successfulResponse()); @@ -532,11 +556,17 @@ void releasesConcurrencyWithoutConsumingPacingWhenDispatchFailsBeforeLaunch() { new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, pacedSettings)) { assertThrows( EmbeddingServiceUnavailableException.class, - () -> embeddingClient.embed(List.of("failed batch"), LlmGatewayTier.BATCH)); + () -> embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("failed batch"), + LlmGatewayTier.BATCH)); assertEquals( 1, embeddingClient - .embed(List.of("successful batch"), LlmGatewayTier.BATCH) + .embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("successful batch"), + LlmGatewayTier.BATCH) .size()); assertEquals(2, batchRequestCount.get()); } @@ -573,9 +603,12 @@ void preventsSlowProviderCallsFromAccumulatingBurstCredit() throws InterruptedEx try (OpenAiCompatibleEmbeddingClient embeddingClient = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, pacedSettings)) { - embeddingClient.embed(List.of("slow batch"), LlmGatewayTier.BATCH); - embeddingClient.embed(List.of("second batch"), LlmGatewayTier.BATCH); - embeddingClient.embed(List.of("third batch"), LlmGatewayTier.BATCH); + embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("slow batch"), LlmGatewayTier.BATCH); + embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("second batch"), LlmGatewayTier.BATCH); + embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("third batch"), LlmGatewayTier.BATCH); } Duration firstToSecondLaunchInterval = @@ -605,12 +638,16 @@ void pacingDeadlineExpiresWithoutDispatchingAnotherRequest() { try (OpenAiCompatibleEmbeddingClient embeddingClient = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, deadlineSettings)) { - embeddingClient.embed(List.of("admitted batch"), LlmGatewayTier.BATCH); + embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("admitted batch"), LlmGatewayTier.BATCH); long rejectedRequestStartNanos = System.nanoTime(); EmbeddingServiceTemporarilyUnavailableException pacingDeadlineFailure = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> embeddingClient.embed(List.of("expired batch"), LlmGatewayTier.BATCH)); + () -> embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("expired batch"), + LlmGatewayTier.BATCH)); Duration rejectionDelay = Duration.ofNanos(System.nanoTime() - rejectedRequestStartNanos); assertInstanceOf(TimeoutException.class, pacingDeadlineFailure.getCause()); @@ -640,10 +677,14 @@ void interruptionDuringPacingPreventsDispatchAndReleasesAdmissionPermits() throw try (OpenAiCompatibleEmbeddingClient embeddingClient = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, interruptibleSettings)) { - embeddingClient.embed(List.of("admitted batch"), LlmGatewayTier.BATCH); + embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, List.of("admitted batch"), LlmGatewayTier.BATCH); Thread pacingRequestThread = Thread.ofVirtual().start(() -> { try { - embeddingClient.embed(List.of("interrupted batch"), LlmGatewayTier.BATCH); + embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("interrupted batch"), + LlmGatewayTier.BATCH); } catch (EmbeddingServiceUnavailableException expectedException) { interruptedStatusRestored.set(Thread.currentThread().isInterrupted()); } finally { @@ -660,7 +701,10 @@ void interruptionDuringPacingPreventsDispatchAndReleasesAdmissionPermits() throw assertEquals( 1, embeddingClient - .embed(List.of("recovered batch"), LlmGatewayTier.BATCH) + .embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("recovered batch"), + LlmGatewayTier.BATCH) .size()); verify(batchEmbeddingService, times(2)).create(any(), any(RequestOptions.class)); } @@ -690,14 +734,17 @@ void boundsLiveAdmissionWaitingByTheWholeRequestDeadline() try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor(); OpenAiCompatibleEmbeddingClient embeddingClient = new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, saturatedSettings)) { - Future> admittedRequest = - executor.submit(() -> embeddingClient.embed(List.of("admitted query"), LlmGatewayTier.LIVE)); + Future> admittedRequest = executor.submit(() -> embeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("admitted query"), LlmGatewayTier.LIVE)); assertTrue(firstLiveRequestStarted.await(5, TimeUnit.SECONDS)); long saturatedRequestStartNanos = System.nanoTime(); EmbeddingServiceTemporarilyUnavailableException concurrencyDeadlineFailure = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> embeddingClient.embed(List.of("saturated query"), LlmGatewayTier.LIVE)); + () -> embeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("saturated query"), + LlmGatewayTier.LIVE)); Duration saturationDelay = Duration.ofNanos(System.nanoTime() - saturatedRequestStartNanos); assertInstanceOf(TimeoutException.class, concurrencyDeadlineFailure.getCause()); @@ -712,7 +759,10 @@ void boundsLiveAdmissionWaitingByTheWholeRequestDeadline() assertEquals( 1, embeddingClient - .embed(List.of("recovered query"), LlmGatewayTier.LIVE) + .embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("recovered query"), + LlmGatewayTier.LIVE) .size()); verify(liveEmbeddingService, times(2)).create(any(), any(RequestOptions.class)); } finally { @@ -732,7 +782,10 @@ void defersProbeDuringBatchCooldownWithoutContactingProvider() { new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { assertThrows( EmbeddingServiceUnavailableException.class, - () -> clientAdapter.embed(List.of("batch document"), LlmGatewayTier.BATCH)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("batch document"), + LlmGatewayTier.BATCH)); assertThrows(OpenAiCompatibleEmbeddingClient.EmbeddingProbeDeferredException.class, clientAdapter::warmUp); verify(batchEmbeddingService, times(1)).create(any(), any(RequestOptions.class)); @@ -769,16 +822,25 @@ void batchRetryAfterCooldownDoesNotBlockLiveRequests() { new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { assertThrows( EmbeddingServiceUnavailableException.class, - () -> clientAdapter.embed(List.of("batch document"), LlmGatewayTier.BATCH)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("batch document"), + LlmGatewayTier.BATCH)); assertEquals( 1, clientAdapter - .embed(List.of("live query"), LlmGatewayTier.LIVE) + .embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("live query"), + LlmGatewayTier.LIVE) .size()); EmbeddingServiceUnavailableException batchCooldownFailure = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> clientAdapter.embed(List.of("another batch document"), LlmGatewayTier.BATCH)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("another batch document"), + LlmGatewayTier.BATCH)); assertTrue(batchCooldownFailure.getMessage().contains("rate limited")); verify(liveEmbeddingService).create(any(), any(RequestOptions.class)); @@ -801,16 +863,23 @@ void liveRetryAfterCooldownDoesNotBlockBatchRequests() { new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, gatewaySettings())) { assertThrows( EmbeddingServiceUnavailableException.class, - () -> clientAdapter.embed(List.of("live query"), LlmGatewayTier.LIVE)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("live query"), LlmGatewayTier.LIVE)); assertEquals( 1, clientAdapter - .embed(List.of("batch document"), LlmGatewayTier.BATCH) + .embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, + List.of("batch document"), + LlmGatewayTier.BATCH) .size()); EmbeddingServiceUnavailableException liveCooldownFailure = assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> clientAdapter.embed(List.of("another live query"), LlmGatewayTier.LIVE)); + () -> clientAdapter.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("another live query"), + LlmGatewayTier.LIVE)); assertTrue(liveCooldownFailure.getMessage().contains("rate limited")); verify(batchEmbeddingService).create(any(), any(RequestOptions.class)); @@ -846,7 +915,10 @@ void lateRateLimitAfterCallerInterruptionPublishesCooldownAndReleasesPermit() th new OpenAiCompatibleEmbeddingClient(liveClient, batchClient, singlePermitSettings)) { Thread interruptedCaller = Thread.ofVirtual().start(() -> { try { - embeddingClient.embed(List.of("interrupted query"), LlmGatewayTier.LIVE); + embeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("interrupted query"), + LlmGatewayTier.LIVE); } catch (EmbeddingServiceTemporarilyUnavailableException expectedInterruption) { callerInterruptRestored.set(Thread.currentThread().isInterrupted()); } finally { @@ -866,7 +938,10 @@ void lateRateLimitAfterCallerInterruptionPublishesCooldownAndReleasesPermit() th long cooldownRejectionStartNanos = System.nanoTime(); assertThrows( EmbeddingServiceTemporarilyUnavailableException.class, - () -> embeddingClient.embed(List.of("cooldown query"), LlmGatewayTier.LIVE)); + () -> embeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("cooldown query"), + LlmGatewayTier.LIVE)); Duration cooldownRejectionDelay = Duration.ofNanos(System.nanoTime() - cooldownRejectionStartNanos); assertTrue(cooldownRejectionDelay.compareTo(Duration.ofMillis(250)) < 0); @@ -876,7 +951,10 @@ void lateRateLimitAfterCallerInterruptionPublishesCooldownAndReleasesPermit() th assertEquals( 1, embeddingClient - .embed(List.of("recovered query"), LlmGatewayTier.LIVE) + .embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, + List.of("recovered query"), + LlmGatewayTier.LIVE) .size()); verify(liveEmbeddingService, times(2)).create(any(), any(RequestOptions.class)); } diff --git a/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientWireTest.java b/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientWireTest.java index 644a4baf..825d9522 100644 --- a/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientWireTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/OpenAiCompatibleEmbeddingClientWireTest.java @@ -49,13 +49,16 @@ void sendsVersionedGatewayRequestsWithIntentSpecificTiersAndNativeDimensions() t TEST_LIVE_MAX_CONCURRENT_REQUESTS, TEST_UNTHROTTLED_REQUESTS_PER_SECOND), OpenAiCompatibleEmbeddingClient.RequestLimits.batch( TEST_BATCH_MAX_CONCURRENT_REQUESTS, TEST_UNTHROTTLED_REQUESTS_PER_SECOND)))) { - List liveEmbeddings = embeddingClient.embed(List.of("user retrieval"), LlmGatewayTier.LIVE); + List liveEmbeddings = embeddingClient.embed( + LlmGatewayTier.Purpose.SEARCH_QUERY_EMBEDDING, List.of("user retrieval"), LlmGatewayTier.LIVE); List batchInputs = new ArrayList<>(); for (int inputIndex = 0; inputIndex < BATCH_INPUT_COUNT; inputIndex++) { batchInputs.add("ingestion input " + inputIndex); } - List batchEmbeddings = embeddingClient.embed(batchInputs, LlmGatewayTier.BATCH); + List batchEmbeddings = embeddingClient.embed( + LlmGatewayTier.Purpose.CONTENT_INDEX_EMBEDDING, batchInputs, LlmGatewayTier.BATCH); + embeddingClient.warmUp(); assertEquals(1, liveEmbeddings.size()); assertEquals(EMBEDDING_DIMENSIONS, liveEmbeddings.getFirst().length); assertEquals(BATCH_INPUT_COUNT, batchEmbeddings.size()); @@ -64,7 +67,8 @@ void sendsVersionedGatewayRequestsWithIntentSpecificTiersAndNativeDimensions() t embeddingServer.stop(0); } - assertEquals(2, capturedRequests.size()); + assertEquals(3, capturedRequests.size()); + assertEquals("java-chat:embeddingWarmup:batch", capturedRequests.get(2).workloadHeader()); assertGatewayRequest(capturedRequests.get(0), "production-z", 1); assertGatewayRequest(capturedRequests.get(1), "batch", BATCH_INPUT_COUNT); } @@ -117,7 +121,9 @@ private static void assertGatewayRequest( assertEquals("Bearer " + TEST_API_KEY, capturedRequest.authorizationHeader()); assertEquals(expectedTier, capturedRequest.tierHeader()); assertEquals( - "java-chat:embeddings:" + (expectedTier.equals("batch") ? "batch" : "live"), + (expectedTier.equals("batch") + ? "java-chat:contentIndexEmbedding:batch" + : "java-chat:searchQueryEmbedding:live"), capturedRequest.workloadHeader()); assertEquals( EMBEDDING_MODEL, diff --git a/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java b/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java index 2cf41a93..11401359 100644 --- a/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java +++ b/src/test/java/com/williamcallahan/javachat/service/RerankerServiceTest.java @@ -7,6 +7,7 @@ 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.eq; import static org.mockito.Mockito.anyString; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; @@ -60,6 +61,7 @@ void rerankPreservesAtomicProviderAdmissionFailure() { ConfiguredProviderTemporarilyUnavailableException admissionFailure = new ConfiguredProviderTemporarilyUnavailableException(RateLimitService.ApiProvider.OPENAI); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -81,6 +83,7 @@ void rerankPreservesAtomicProviderAdmissionFailure() { void rerankUsesConfiguredCompletionBudgetAndTimeout() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -97,11 +100,13 @@ void rerankUsesConfiguredCompletionBudgetAndTimeout() { ArgumentCaptor outputBudgetCaptor = ArgumentCaptor.forClass(Integer.class); verify(streamingService) .completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), promptCaptor.capture(), eq(TEST_RERANKER_TEMPERATURE), outputBudgetCaptor.capture(), eq(TEST_RERANKER_TIMEOUT)); - verify(streamingService, never()).complete(anyString(), eq(TEST_RERANKER_TEMPERATURE)); + 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")); @@ -113,6 +118,7 @@ void rerankUsesConfiguredCompletionBudgetAndTimeout() { void rerankSelectsRelevantSubsetAndDropsUnrelatedDocuments() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -132,6 +138,7 @@ void rerankSelectsRelevantSubsetAndDropsUnrelatedDocuments() { void rerankReturnsEmptySelectionWhenNoDocumentIsRelevant() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -151,6 +158,7 @@ void rerankReturnsEmptySelectionWhenNoDocumentIsRelevant() { void rerankRetainsOneRelevantDocument() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -166,6 +174,7 @@ void rerankRetainsOneRelevantDocument() { assertEquals(sourceDocuments, selectedDocuments); verify(streamingService) .completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -192,6 +201,7 @@ void rerankRejectsDuplicateInvalidAndMalformedOrderings() { "{\"order\":[\"0\",\"1\",\"2\",\"3\"]}"); for (String invalidOrderingJson : invalidOrderingJsonValues) { when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -209,6 +219,7 @@ void rerankRejectsDuplicateInvalidAndMalformedOrderings() { void emptyCompletionIsTerminalWithoutCallingAnotherCompletionPath() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -223,13 +234,15 @@ void emptyCompletionIsTerminalWithoutCallingAnotherCompletionPath() { "query", sourceDocuments, 2, stageDeadlineNanos(REMAINING_STAGE_BUDGET_TEST_TIMEOUT))); assertEquals("Reranking response was empty", rerankingFailure.getMessage()); - verify(streamingService, never()).complete(anyString(), eq(TEST_RERANKER_TEMPERATURE)); + verify(streamingService, never()) + .complete(eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE)); } @Test void rerankTightensItsTimeoutToTheRemainingStageBudget() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -245,6 +258,7 @@ void rerankTightensItsTimeoutToTheRemainingStageBudget() { ArgumentCaptor remainingBudgetCaptor = ArgumentCaptor.forClass(Duration.class); verify(streamingService) .completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -272,6 +286,7 @@ void expiredStageBudgetFailsAsTimeoutBeforeRerankerDispatch() { void expiredDeadlineRejectsCachedHit() { OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -288,6 +303,7 @@ void expiredDeadlineRejectsCachedHit() { assertInstanceOf(TimeoutException.class, deadlineFailure.getCause()); verify(streamingService, times(1)) .completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -302,6 +318,7 @@ void sameKeyCacheWaitCannotExtendCallerDeadline() CountDownLatch waitingCallerStarted = new CountDownLatch(1); CountDownLatch releaseFirstDispatch = new CountDownLatch(1); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -342,6 +359,7 @@ void sameKeyCacheWaitCannotExtendCallerDeadline() firstRerank.get(CACHE_TEST_COMPLETION_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)); verify(streamingService, times(1)) .completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -358,6 +376,7 @@ void waiterRetriesWhenCoalescedOwnerDiesOnItsStageDeadline() OpenAIStreamingService streamingService = mock(OpenAIStreamingService.class); AtomicInteger dispatchCount = new AtomicInteger(); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -425,6 +444,7 @@ void waiterInheritsProviderTimeoutOwnerFailureWithoutRetrying() CountDownLatch releaseOwnerTimeout = new CountDownLatch(1); AtomicInteger dispatchCount = new AtomicInteger(); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -487,6 +507,7 @@ void waiterInheritsNonTimeoutOwnerFailureWithoutRetrying() CountDownLatch releaseOwnerFailure = new CountDownLatch(1); AtomicInteger dispatchCount = new AtomicInteger(); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), @@ -540,6 +561,7 @@ void wholeCallTimeoutPreservesRetrievalTimeoutCause() { InterruptedIOException okHttpCallTimeout = new InterruptedIOException("timeout"); OpenAIIoException providerTimeout = new OpenAIIoException("Request failed", okHttpCallTimeout); when(streamingService.completeJsonObject( + eq(LlmGatewayTier.Purpose.SEARCH_RERANK), anyString(), eq(TEST_RERANKER_TEMPERATURE), eq(TEST_RERANKER_OUTPUT_TOKEN_BUDGET), diff --git a/src/test/java/com/williamcallahan/javachat/web/ChatControllerStreamingFailureTest.java b/src/test/java/com/williamcallahan/javachat/web/ChatControllerStreamingFailureTest.java index 5681f9cc..45091986 100644 --- a/src/test/java/com/williamcallahan/javachat/web/ChatControllerStreamingFailureTest.java +++ b/src/test/java/com/williamcallahan/javachat/web/ChatControllerStreamingFailureTest.java @@ -42,6 +42,7 @@ import com.williamcallahan.javachat.service.ChatService; import com.williamcallahan.javachat.service.ConfiguredProviderTemporarilyUnavailableException; import com.williamcallahan.javachat.service.HybridSearchPartialFailureException; +import com.williamcallahan.javachat.service.LlmGatewayTier; import com.williamcallahan.javachat.service.OpenAIStreamingService; import com.williamcallahan.javachat.service.RateLimitService; import com.williamcallahan.javachat.service.RerankingFailureException; @@ -321,7 +322,8 @@ void wrappedQdrantDeadlineUsesTheRetrievalTimeoutContract() throws JsonProcessin assertEquals(Boolean.TRUE, timeoutError.retryable()); assertEquals(STATUS_STAGE_RETRIEVAL, timeoutError.stage()); assertFalse(serializedTimeoutError.contains(UPSTREAM_SECRET_MESSAGE)); - verify(streamingService, never()).streamResponse(any(StructuredPrompt.class), anyDouble()); + verify(streamingService, never()) + .streamResponse(eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble()); List retrievalTimeoutWarnings = pipelineLogEvents.events().stream() .filter(logEvent -> logEvent.getLevel() == Level.WARN && "Response preparation timeout".equals(logEvent.getFormattedMessage())) @@ -423,7 +425,8 @@ void streamSubscriptionReturnsAfterPreparationStatusWhileRetrievalIsBlocked() when(chatService.buildStructuredPromptWithContextOutcome(anyList(), eq(USER_QUERY), any(), anyLong())) .thenReturn( new ChatService.StructuredPromptOutcome(StructuredPrompt.fromRawPrompt("test", 1), List.of())); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.never()); ExecutorService subscriptionExecutor = Executors.newSingleThreadExecutor(); @@ -528,7 +531,8 @@ void disposingStreamInterruptsBlockedPreparationAndPreventsDownstreamWork() thro "blocked retrieval should terminate after cancellation"); assertNull(streamFailure.get()); verifyNoInteractions(chatService, retrievalService); - verify(streamingService, never()).streamResponse(any(StructuredPrompt.class), anyDouble()); + verify(streamingService, never()) + .streamResponse(eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble()); verify(chatMemoryService, never()).addExchange(eq(SESSION_ID), eq(USER_QUERY), any()); } finally { releaseRetrieval.countDown(); @@ -608,7 +612,8 @@ void chatStreamCitesTheExactDocumentsUsedForPromptContextWithoutSecondDiscovery( .thenReturn(new RetrievalService.CitationOutcome(List.of(officialCitation), 0)); when(streamingService.canAttemptRequest()).thenReturn(true); when(streamingService.isAvailable()).thenReturn(true); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( Flux.just("Hello"), RateLimitService.ApiProvider.OPENAI, List.of("official-prompt-document")))); @@ -666,7 +671,8 @@ void chatStreamEmitsExactOverloadSelectionFromPromptContext() throws JsonProcess .thenReturn(new RetrievalService.CitationOutcome(List.of(exactOverloadCitation), 1)); when(streamingService.canAttemptRequest()).thenReturn(true); when(streamingService.isAvailable()).thenReturn(true); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( Flux.just("Hello"), RateLimitService.ApiProvider.OPENAI, List.of("broad-prompt-document")))); @@ -733,7 +739,8 @@ void chatStreamSurfacesNonzeroCitationFailuresBeforeCitationEvent() throws JsonP eq(USER_QUERY), any(ChatService.StructuredPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome( List.of(new Citation("https://example.com", "Example", "", "")), 2)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just( new StreamingResult(Flux.just("Hello"), RateLimitService.ApiProvider.OPENAI, List.of()))); @@ -795,7 +802,8 @@ void streamBufferOverflowDoesNotPersistPartialChatAnswer() throws JsonProcessing when(chatService.citationOutcomeForRetainedContext( eq(USER_QUERY), any(ChatService.StructuredPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome(List.of(), 0)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( partialAnswerThenOverflow, RateLimitService.ApiProvider.OPENAI, List.of()))); when(streamingService.isRecoverableStreamingFailure(streamBufferOverflowFailure)) @@ -847,7 +855,8 @@ private List> streamFailure( when(chatService.citationOutcomeForRetainedContext( eq(USER_QUERY), any(ChatService.StructuredPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome(List.of(), 0)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( Flux.error(streamingFailure), RateLimitService.ApiProvider.OPENAI, List.of()))); when(streamingService.isRecoverableStreamingFailure(streamingFailure)).thenReturn(retryable); diff --git a/src/test/java/com/williamcallahan/javachat/web/ChatPreparationSseIntegrationTest.java b/src/test/java/com/williamcallahan/javachat/web/ChatPreparationSseIntegrationTest.java index 298254bb..7a3709a6 100644 --- a/src/test/java/com/williamcallahan/javachat/web/ChatPreparationSseIntegrationTest.java +++ b/src/test/java/com/williamcallahan/javachat/web/ChatPreparationSseIntegrationTest.java @@ -24,6 +24,7 @@ import com.williamcallahan.javachat.service.EmbeddingClient; import com.williamcallahan.javachat.service.EmbeddingModelKeepAlive; import com.williamcallahan.javachat.service.ExternalServiceHealth; +import com.williamcallahan.javachat.service.LlmGatewayTier; import com.williamcallahan.javachat.service.OpenAIStreamingService; import com.williamcallahan.javachat.service.RateLimitService; import com.williamcallahan.javachat.service.RetrievalService; @@ -130,7 +131,8 @@ void flushesPreparationEventWhileRetrievalRemainsBlocked() throws JsonProcessing when(chatService.citationOutcomeForRetainedContext( eq(USER_QUERY), any(ChatService.StructuredPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome(List.of(), 0)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.CHAT_ANSWER), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just( new StreamingResult(Flux.just("Complete"), RateLimitService.ApiProvider.OPENAI, List.of()))); doAnswer(ignoredInvocation -> { diff --git a/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerBackpressureOverflowTest.java b/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerBackpressureOverflowTest.java index 2a4f9662..0fb3b8df 100644 --- a/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerBackpressureOverflowTest.java +++ b/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerBackpressureOverflowTest.java @@ -28,6 +28,7 @@ import com.williamcallahan.javachat.model.GuidedLesson; import com.williamcallahan.javachat.service.ChatMemoryService; import com.williamcallahan.javachat.service.GuidedLearningService; +import com.williamcallahan.javachat.service.LlmGatewayTier; import com.williamcallahan.javachat.service.MarkdownService; import com.williamcallahan.javachat.service.OpenAIStreamingService; import com.williamcallahan.javachat.service.RateLimitService; @@ -106,7 +107,8 @@ void streamBufferOverflowDoesNotPersistPartialGuidedAnswer() throws JsonProcessi when(guidedLearningService.citationOutcomeForRetainedContext( any(GuidedLearningService.GuidedChatPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome(List.of(), 0)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.GUIDED_LEARNING), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( partialAnswerThenOverflow, RateLimitService.ApiProvider.OPENAI, List.of()))); when(streamingService.isRecoverableStreamingFailure(streamBufferOverflowFailure)) diff --git a/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerStreamingFailureTest.java b/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerStreamingFailureTest.java index 8f66e033..f51ddded 100644 --- a/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerStreamingFailureTest.java +++ b/src/test/java/com/williamcallahan/javachat/web/GuidedLearningControllerStreamingFailureTest.java @@ -36,6 +36,7 @@ import com.williamcallahan.javachat.service.ConfiguredProviderTemporarilyUnavailableException; import com.williamcallahan.javachat.service.EmbeddingServiceTemporarilyUnavailableException; import com.williamcallahan.javachat.service.GuidedLearningService; +import com.williamcallahan.javachat.service.LlmGatewayTier; import com.williamcallahan.javachat.service.MarkdownService; import com.williamcallahan.javachat.service.OpenAIStreamingService; import com.williamcallahan.javachat.service.RateLimitService; @@ -115,7 +116,8 @@ void guidedChatKeepsTerminalExceptionTypeOutOfRetryableClientError() throws Json when(guidedLearningService.citationOutcomeForRetainedContext( any(GuidedLearningService.GuidedChatPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome(List.of(), 0)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.GUIDED_LEARNING), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( Flux.error(terminalFailure), RateLimitService.ApiProvider.OPENAI, List.of()))); when(streamingService.isRecoverableStreamingFailure(terminalFailure)).thenReturn(true); @@ -156,7 +158,8 @@ void configuredProviderCooldownUsesTheSharedRetryableClientError() throws JsonPr when(guidedLearningService.citationOutcomeForRetainedContext( any(GuidedLearningService.GuidedChatPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome(List.of(), 0)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.GUIDED_LEARNING), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( Flux.error(terminalFailure), RateLimitService.ApiProvider.OPENAI, List.of()))); @@ -296,7 +299,8 @@ void wrappedLocalEmbeddingSocketDeadlineUsesTheRetrievalTimeoutContract() throws .counter() .count()); assertEquals(0, controllerErrorCount()); - verify(streamingService, never()).streamResponse(any(StructuredPrompt.class), anyDouble()); + verify(streamingService, never()) + .streamResponse(eq(LlmGatewayTier.Purpose.GUIDED_LEARNING), any(StructuredPrompt.class), anyDouble()); } @Test @@ -312,7 +316,8 @@ void guidedChatKeepsNonTerminalExceptionTypeInStructuredLogOnly() throws JsonPro when(guidedLearningService.citationOutcomeForRetainedContext( any(GuidedLearningService.GuidedChatPromptOutcome.class), anyList())) .thenReturn(new RetrievalService.CitationOutcome(List.of(), 0)); - when(streamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + when(streamingService.streamResponse( + eq(LlmGatewayTier.Purpose.GUIDED_LEARNING), any(StructuredPrompt.class), anyDouble())) .thenReturn(Mono.just(new StreamingResult( Flux.error(upstreamFailure), RateLimitService.ApiProvider.OPENAI, List.of()))); when(streamingService.isRecoverableStreamingFailure(upstreamFailure)).thenReturn(false); diff --git a/src/test/java/com/williamcallahan/javachat/web/GuidedSseCitationEventTest.java b/src/test/java/com/williamcallahan/javachat/web/GuidedSseCitationEventTest.java index a51d1e5b..0ca89744 100644 --- a/src/test/java/com/williamcallahan/javachat/web/GuidedSseCitationEventTest.java +++ b/src/test/java/com/williamcallahan/javachat/web/GuidedSseCitationEventTest.java @@ -31,6 +31,7 @@ import com.williamcallahan.javachat.model.GuidedLesson; import com.williamcallahan.javachat.service.ChatMemoryService; import com.williamcallahan.javachat.service.GuidedLearningService; +import com.williamcallahan.javachat.service.LlmGatewayTier; import com.williamcallahan.javachat.service.MarkdownService; import com.williamcallahan.javachat.service.OpenAIStreamingService; import com.williamcallahan.javachat.service.QdrantPayloadFieldSchema; @@ -113,7 +114,8 @@ void guidedStreamEmitsCitationEvent() throws Exception { GuidedLearningService.GuidedChatPromptOutcome promptOutcome = new GuidedLearningService.GuidedChatPromptOutcome( StructuredPrompt.fromRawPrompt("test", 1), List.of(lessonContextDocument, truncatedLessonContextDocument)); - given(openAIStreamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + given(openAIStreamingService.streamResponse( + eq(LlmGatewayTier.Purpose.GUIDED_LEARNING), any(StructuredPrompt.class), anyDouble())) .willReturn(Mono.just(new StreamingResult( Flux.just("Hello"), RateLimitService.ApiProvider.OPENAI, @@ -162,7 +164,8 @@ void guidedStreamSurfacesPartialCitationConversionWithTheCanonicalStatus() throw given(chatMemoryService.getHistory(anyString())).willReturn(List.of()); GuidedLearningService.GuidedChatPromptOutcome promptOutcome = new GuidedLearningService.GuidedChatPromptOutcome( StructuredPrompt.fromRawPrompt("test", 1), List.of(lessonContextDocument)); - given(openAIStreamingService.streamResponse(any(StructuredPrompt.class), anyDouble())) + given(openAIStreamingService.streamResponse( + eq(LlmGatewayTier.Purpose.GUIDED_LEARNING), any(StructuredPrompt.class), anyDouble())) .willReturn(Mono.just(new StreamingResult( Flux.just("Hello"), RateLimitService.ApiProvider.OPENAI,