Add MessageInterceptor support for streaming and introduce PartialResponseInterceptor

docs/content/services/chat-service.md

@@ -81,7 +81,8 @@ ChatService chatService = ChatService.builder()
 | `timeout` | Duration | No | Request timeout (default: 60 seconds) |
 | `parameters` | ChatParameters | No | Default parameters applied to all requests |
 | `tools` | List\<Tool\> | No | Default tools available to the model |
-| `messageInterceptor` | MessageInterceptor\<ChatRequest\> | No | Modify assistant messages before returning |
+| `messageInterceptor` | MessageInterceptor\<ChatRequest\> | No | Modify the complete assistant message before returning |
+| `partialResponseInterceptor` | PartialResponseInterceptor\<ChatRequest\> | No | Modify each streamed content token before delivery |
 | `toolInterceptor` | ToolInterceptor\<ChatRequest\> | No | Normalize/modify tool call arguments |
 | `logRequests` | Boolean | No | Enable request logging (default: false) |
 | `logResponses` | Boolean | No | Enable response logging (default: false) |
@@ -394,13 +395,13 @@ System.out.println(answer);
 
 ## Interceptors
 
-Interceptors run automatically after every non-streaming response, before the result is returned to your application. They are configured once on the service builder and apply transparently to all subsequent calls. Both are `@FunctionalInterface`, so you can pass a lambda directly.
+Interceptors are configured once on the service builder and apply transparently to all subsequent calls. Each one is a `@FunctionalInterface`, so you can pass a lambda directly.
 
 ### Message Interceptor
 
 `MessageInterceptor` lets you modify or sanitize the assistant's text content. Common uses: stripping whitespace, filtering unwanted patterns, normalizing formatting.
 
-> **Note:** `MessageInterceptor` applies to **non-streaming** requests only. For streaming, process the content directly inside `ChatHandler` callbacks.
+It is always called **once per assistant message, with the complete content**, in both modes. In **non-streaming** mode it runs before the response is returned to the caller. In **streaming** mode it runs on the aggregated message, before that message reaches `onCompleteResponse` and before the `CompletableFuture` returned by `chatStreaming` completes. Transformations that need the whole text, such as `strip()` or a pattern that spans several tokens, are therefore safe everywhere.
 
 ```java
 ChatService chatService = ChatService.builder()
@@ -409,6 +410,21 @@ ChatService chatService = ChatService.builder()
     .build();
 ```
 
+> **Note:** the tokens delivered to `onPartialResponse` are left untouched, so a caller that renders the stream sees the original text and only the final message carries the transformation. Use `PartialResponseInterceptor` when the tokens themselves have to change.
+
+### Partial Response Interceptor
+
+`PartialResponseInterceptor` intercepts each content token in **streaming** mode, before it is delivered to `onPartialResponse`. Common uses: masking or highlighting text as it appears, adapting tokens for a terminal or a UI widget.
+
+```java
+ChatService chatService = ChatService.builder()
+    // ...
+    .partialResponseInterceptor((ctx, partialResponse) -> partialResponse.toUpperCase())
+    .build();
+```
+
+Every invocation receives one token exactly as the model streamed it. Tokens are never buffered, so a transformation that has to match across token boundaries belongs in a `MessageInterceptor` instead. This interceptor has no effect on non-streaming requests and does not alter the `ChatResponse` delivered to `onCompleteResponse`, which means the two hooks can be combined freely.
+
 ### Tool Interceptor
 
 `ToolInterceptor` intercepts each completed tool call before it reaches your handler, letting you validate, normalize, or unwrap the arguments (for example, unwrapping double-encoded JSON strings that some models produce).
@@ -428,15 +444,15 @@ ChatService chatService = ChatService.builder()
 
 ### InterceptorContext
 
-Both interceptors receive an `InterceptorContext` as their first argument, which provides access to the current request, the current response, and a way to invoke the model again.
+Every interceptor receives an `InterceptorContext` as its first argument, which provides access to the current request, the current response, and a way to invoke the model again.
 
 | Method | Description |
 |--------|-------------|
 | `ctx.request()` | The original `ChatRequest` that triggered this response |
-| `ctx.response()` | An `Optional<ChatResponse>` with the current response |
+| `ctx.response()` | An `Optional<ChatResponse>` with the current response, empty in `PartialResponseInterceptor` because no response exists yet |
 | `ctx.invoke(ChatRequest)` | Sends a new request to the model and returns its response |
 
-`MessageInterceptor`, `ToolInterceptor`, and `InterceptorContext` are parameterized by the request type of the service they are registered on, so `ctx.request()` returns that concrete type with no cast: `ChatRequest` here, `DeploymentChatRequest` on [`DeploymentService`](/services/deployment-service), and `ModelGatewayChatRequest` on [`ModelGatewayChatService`](/services/model-gateway). The type argument is inferred when you pass a lambda, and only needs to be written out if you declare the interceptor separately:
+`MessageInterceptor`, `PartialResponseInterceptor`, `ToolInterceptor`, and `InterceptorContext` are parameterized by the request type of the service they are registered on, so `ctx.request()` returns that concrete type with no cast: `ChatRequest` here, `DeploymentChatRequest` on [`DeploymentService`](/services/deployment-service), and `ModelGatewayChatRequest` on [`ModelGatewayChatService`](/services/model-gateway). The type argument is inferred when you pass a lambda, and only needs to be written out if you declare the interceptor separately:
 
 ```java
 MessageInterceptor<ChatRequest> interceptor = (ctx, message) -> message == null ? "" : message.strip();

docs/content/services/deployment-service.md

@@ -69,7 +69,8 @@ All routing is done through the `deploymentId` in each request, so no `projectId
 | `verifySsl` | Boolean | No | SSL certificate verification (default: true) |
 | `parameters` | ChatParameters | No | Default chat parameters applied to every chat request |
 | `tools` | List\<Tool\> | No | Default tools available to the model |
-| `messageInterceptor` | MessageInterceptor\<DeploymentChatRequest\> | No | Post-processing hook for the assistant's text content |
+| `messageInterceptor` | MessageInterceptor\<DeploymentChatRequest\> | No | Post-processing hook for the complete assistant message |
+| `partialResponseInterceptor` | PartialResponseInterceptor\<DeploymentChatRequest\> | No | Post-processing hook for each streamed content token |
 | `toolInterceptor` | ToolInterceptor\<DeploymentChatRequest\> | No | Post-processing hook for function call arguments |
 
 > Either `apiKey` or `authenticator` must be provided. `projectId`, `spaceId`, and `modelId` are **ignored**, and a warning is logged if they are set on a request's parameters object.

docs/content/services/model-gateway/chat.md

@@ -31,7 +31,7 @@ System.out.println(response.toAssistantMessage().content());
 
 - Send synchronous and streaming chat requests to any model available through the gateway.
 - Use gateway-specific parameters such as service tier, reasoning effort, audio modalities, caching, and routing configuration.
-- Apply `MessageInterceptor` and `ToolInterceptor` for post-processing.
+- Apply `MessageInterceptor`, `PartialResponseInterceptor`, and `ToolInterceptor` for post-processing.
 - Read gateway metadata on every response: `serviceTier()`, `systemFingerprint()`, and `cached()`.
 
 ---
@@ -58,7 +58,8 @@ ModelGatewayChatService service = ModelGatewayChatService.builder()
 | `modelId` | String | Yes | Third-party model identifier (e.g., `"gpt-4o"`, `"claude-3-5-sonnet"`) |
 | `parameters` | ModelGatewayChatParameters | No | Default parameters applied to every request |
 | `tools` | List\<Tool\> | No | Default tools available to the model |
-| `messageInterceptor` | MessageInterceptor\<ModelGatewayChatRequest\> | No | Post-processing hook for the assistant's text content |
+| `messageInterceptor` | MessageInterceptor\<ModelGatewayChatRequest\> | No | Post-processing hook for the complete assistant message |
+| `partialResponseInterceptor` | PartialResponseInterceptor\<ModelGatewayChatRequest\> | No | Post-processing hook for each streamed content token |
 | `toolInterceptor` | ToolInterceptor\<ModelGatewayChatRequest\> | No | Post-processing hook for function call arguments |
 | `timeout` | Duration | No | Default request timeout (default: 60 seconds) |
 | `logRequests` | Boolean | No | Enable request logging (default: false) |
@@ -341,7 +342,7 @@ System.out.println("Total tokens: " + response.usage().totalTokens());
 
 ## Interceptors
 
-Interceptors work identically to how they work in `ChatService`. See the [Chat Service - Interceptors](../../services/chat-service#interceptors) section for the full description of `MessageInterceptor`, `ToolInterceptor`, and `InterceptorContext`.
+Interceptors work identically to how they work in `ChatService`. See the [Chat Service - Interceptors](../../services/chat-service#interceptors) section for the full description of `MessageInterceptor`, `PartialResponseInterceptor`, `ToolInterceptor`, and `InterceptorContext`.
 
 ```java
 ModelGatewayChatService service = ModelGatewayChatService.builder()

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatClientContext.java

@@ -4,6 +4,8 @@
  */
 package com.ibm.watsonx.ai.chat;
 
+import com.ibm.watsonx.ai.chat.interceptor.MessageInterceptor;
+import com.ibm.watsonx.ai.chat.interceptor.PartialResponseInterceptor;
 import com.ibm.watsonx.ai.chat.interceptor.ToolInterceptor;
 import com.ibm.watsonx.ai.chat.model.ExtractionTags;
 
@@ -15,12 +17,16 @@
 public class ChatClientContext<R extends BaseChatRequest> {
     private final ChatProvider<R, ?> chatProvider;
     private final R chatRequest;
+    private final MessageInterceptor<R> messageInterceptor;
+    private final PartialResponseInterceptor<R> partialResponseInterceptor;
     private final ToolInterceptor<R> toolInterceptor;
     private final ExtractionTags extractionTags;
 
     private ChatClientContext(Builder<R> builder) {
         chatProvider = builder.chatProvider;
         chatRequest = builder.chatRequest;
+        messageInterceptor = builder.messageInterceptor;
+        partialResponseInterceptor = builder.partialResponseInterceptor;
         toolInterceptor = builder.toolInterceptor;
         extractionTags = builder.extractionTags;
     }
@@ -43,6 +49,24 @@ public R chatRequest() {
         return chatRequest;
     }
 
+    /**
+     * Returns the message interceptor.
+     *
+     * @return the message interceptor
+     */
+    public MessageInterceptor<R> messageInterceptor() {
+        return messageInterceptor;
+    }
+
+    /**
+     * Returns the partial response interceptor.
+     *
+     * @return the partial response interceptor
+     */
+    public PartialResponseInterceptor<R> partialResponseInterceptor() {
+        return partialResponseInterceptor;
+    }
+
     /**
      * Returns the tool interceptor.
      *
@@ -79,6 +103,8 @@ public static <R extends BaseChatRequest> Builder<R> builder() {
     public static class Builder<R extends BaseChatRequest> {
         private ChatProvider<R, ?> chatProvider;
         private R chatRequest;
+        private MessageInterceptor<R> messageInterceptor;
+        private PartialResponseInterceptor<R> partialResponseInterceptor;
         private ToolInterceptor<R> toolInterceptor;
         private ExtractionTags extractionTags;
 
@@ -104,6 +130,26 @@ public Builder<R> chatRequest(R chatRequest) {
             return this;
         }
 
+        /**
+         * Sets the message interceptor.
+         *
+         * @param messageInterceptor the message interceptor
+         */
+        public Builder<R> messageInterceptor(MessageInterceptor<R> messageInterceptor) {
+            this.messageInterceptor = messageInterceptor;
+            return this;
+        }
+
+        /**
+         * Sets the partial response interceptor.
+         *
+         * @param partialResponseInterceptor the partial response interceptor
+         */
+        public Builder<R> partialResponseInterceptor(PartialResponseInterceptor<R> partialResponseInterceptor) {
+            this.partialResponseInterceptor = partialResponseInterceptor;
+            return this;
+        }
+
         /**
          * Sets the tool interceptor.
          *

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatService.java

@@ -15,6 +15,7 @@
 import com.ibm.watsonx.ai.WatsonxService.CryptoService;
 import com.ibm.watsonx.ai.chat.interceptor.InterceptorContext;
 import com.ibm.watsonx.ai.chat.interceptor.MessageInterceptor;
+import com.ibm.watsonx.ai.chat.interceptor.PartialResponseInterceptor;
 import com.ibm.watsonx.ai.chat.interceptor.ToolInterceptor;
 import com.ibm.watsonx.ai.chat.model.BaseChatParameters.ToolChoiceOption;
 import com.ibm.watsonx.ai.chat.model.ChatMessage;
@@ -51,6 +52,7 @@
 public class ChatService extends CryptoService implements ChatProvider<ChatRequest, TextChatResponse> {
     private final ChatRestClient client;
     private final MessageInterceptor<ChatRequest> messageInterceptor;
+    private final PartialResponseInterceptor<ChatRequest> partialResponseInterceptor;
     private final ToolInterceptor<ChatRequest> toolInterceptor;
     private final ChatProvider<ChatRequest, TextChatResponse> chatProvider;
     private final ChatParameters defaultParameters;
@@ -60,6 +62,7 @@ private ChatService(Builder builder) {
         super(builder);
         requireNonNull(builder.authenticator(), "authenticator cannot be null");
         messageInterceptor = builder.messageInterceptor;
+        partialResponseInterceptor = builder.partialResponseInterceptor;
         toolInterceptor = builder.toolInterceptor;
         defaultTools = isNull(builder.defaultTools) ? null : List.copyOf(builder.defaultTools);
 
@@ -85,7 +88,7 @@ private ChatService(Builder builder) {
             .authenticator(builder.authenticator())
             .build();
 
-        chatProvider = nonNull(messageInterceptor) || nonNull(toolInterceptor)
+        chatProvider = nonNull(messageInterceptor) || nonNull(partialResponseInterceptor) || nonNull(toolInterceptor)
             ? builder.copyWithoutInterceptors().parameters(defaultParameters).build()
             : null;
     }
@@ -153,6 +156,8 @@ public CompletableFuture<ChatResponse> chatStreaming(ChatRequest chatRequest, Ch
         var context = ChatClientContext.<ChatRequest>builder()
             .chatProvider(chatProvider)
             .chatRequest(chatRequest)
+            .messageInterceptor(messageInterceptor)
+            .partialResponseInterceptor(partialResponseInterceptor)
             .toolInterceptor(toolInterceptor)
             .extractionTags(extractionTags)
             .build();
@@ -391,6 +396,7 @@ public static Builder builder() {
      */
     public final static class Builder extends CryptoService.Builder<Builder> {
         private MessageInterceptor<ChatRequest> messageInterceptor;
+        private PartialResponseInterceptor<ChatRequest> partialResponseInterceptor;
         private ToolInterceptor<ChatRequest> toolInterceptor;
         private ChatParameters defaultParameters;
         private List<Tool> defaultTools;
@@ -413,14 +419,13 @@ public Builder parameters(ChatParameters parameters) {
         /**
          * Registers a {@link MessageInterceptor} used to modify or sanitize the assistant's textual content before it is returned to the caller.
          * <p>
-         * This interceptor is invoked on the final aggregated content (non-streaming responses only), allowing adjustments such as rewriting,
-         * filtering, or normalization.
+         * The interceptor is invoked once per assistant message with the complete content, both on non-streaming responses and on the aggregated
+         * message of a streaming session.
          * <p>
          * <b>Example:</b>
          *
          * <pre>{@code
-         * ChatService.builder()
-         *     .messageInterceptor((request, content) -> content.replace("error", "issue"));
+         * ChatService.builder().messageInterceptor((ctx, content) -> content.replace("error", "issue"));
          * }</pre>
          *
          * @param messageInterceptor the interceptor to apply
@@ -430,6 +435,26 @@ public Builder messageInterceptor(MessageInterceptor<ChatRequest> messageInterce
             return this;
         }
 
+        /**
+         * Registers a {@link PartialResponseInterceptor} used to modify each partial response before it is delivered to
+         * {@link ChatHandler#onPartialResponse}.
+         * <p>
+         * This interceptor has no effect on non-streaming requests and does not alter the {@link ChatResponse} delivered to
+         * {@link ChatHandler#onCompleteResponse}.
+         * <p>
+         * <b>Example:</b>
+         *
+         * <pre>{@code
+         * ChatService.builder().partialResponseInterceptor((ctx, partialResponse) -> partialResponse.toUpperCase());
+         * }</pre>
+         *
+         * @param partialResponseInterceptor the interceptor to apply
+         */
+        public Builder partialResponseInterceptor(PartialResponseInterceptor<ChatRequest> partialResponseInterceptor) {
+            this.partialResponseInterceptor = partialResponseInterceptor;
+            return this;
+        }
+
         /**
          * Registers a {@link ToolInterceptor} to modify or normalize function call arguments before tool execution.
          * <p>

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/DefaultRestClient.java

@@ -84,7 +84,8 @@ public CompletableFuture<ChatResponse> chatStreaming(
         var chatSubscriber =
             new DefaultChatSubscriber(
                 new SseEventProcessor(textChatRequest.tools(), context.extractionTags(), TextChatResponse::builder),
-                new ChatHandlerDecorator<>(handler, interceptorContext, context.toolInterceptor())
+                new ChatHandlerDecorator<>(handler, interceptorContext, context.messageInterceptor(), context.partialResponseInterceptor(),
+                    context.toolInterceptor())
             );
 
         var subscriber = chatSubscriber.asFlowSubscriber(response, !handler.failOnFirstError());