From 64061cda5a5a31b150ffbbc094cbd916336fd6a6 Mon Sep 17 00:00:00 2001 From: William Callahan Date: Thu, 15 Oct 2026 00:10:30 -0700 Subject: [PATCH] fix(llm): attribute requests to their application purpose Require the existing typed purpose at generation and embedding entrypoints so shared gateway telemetry identifies chat, learning, enrichment, reranking, retrieval, indexing, and warmup work independently of protocol and tier. Preserve model configuration and routing while carrying purpose across async embedding waves and retries. Verify semantic headers at local HTTP boundaries. Agent-Session: 01a13e5e-838d-73a1-9bff-b9f231013f11 --- .../javachat/service/ChatService.java | 2 +- .../service/EmbeddingBatchEmbedder.java | 18 +- .../javachat/service/EmbeddingClient.java | 22 ++- .../javachat/service/EnrichmentService.java | 6 +- .../javachat/service/HybridSearchService.java | 6 +- .../javachat/service/HybridVectorService.java | 32 +++- .../javachat/service/LlmGatewayTier.java | 40 ++++- .../service/LocalEmbeddingClient.java | 7 +- .../service/OpenAIStreamingService.java | 71 ++++++-- .../OpenAiCompatibleEmbeddingClient.java | 23 ++- .../javachat/service/RerankerService.java | 7 +- .../javachat/web/ChatController.java | 2 + .../web/GuidedLearningController.java | 2 + .../service/EmbeddingBatchEmbedderTest.java | 61 +++++-- .../service/EmbeddingModelKeepAliveTest.java | 8 +- .../service/EmbeddingProviderFailureTest.java | 6 +- .../service/EnrichmentServiceCacheTest.java | 12 +- .../service/EnrichmentServiceTest.java | 25 ++- .../service/HybridSearchServiceTest.java | 93 ++++++++-- .../service/HybridVectorServiceTest.java | 10 +- .../service/LocalEmbeddingClientTest.java | 30 +++- .../service/OpenAICompletionStatusTest.java | 74 ++++++-- .../OpenAIStreamingProviderFailureTest.java | 19 +- .../service/OpenAIStreamingServiceTest.java | 70 +++++--- .../OpenAIStreamingTimeoutContractTest.java | 2 +- .../OpenAiCompatibleEmbeddingClientTest.java | 162 +++++++++++++----- ...enAiCompatibleEmbeddingClientWireTest.java | 14 +- .../javachat/service/RerankerServiceTest.java | 26 ++- .../ChatControllerStreamingFailureTest.java | 25 ++- .../ChatPreparationSseIntegrationTest.java | 4 +- ...ingControllerBackpressureOverflowTest.java | 4 +- ...earningControllerStreamingFailureTest.java | 13 +- .../web/GuidedSseCitationEventTest.java | 7 +- 33 files changed, 676 insertions(+), 227 deletions(-) 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,