Step 03 - Streaming responses
LLM responses can be long. Imagine asking the model to generate a story. It could potentially produce hundreds of lines of text.
In the current application, the entire response is accumulated before being sent to the client. During that generation, the client is waiting for the response, and the server is waiting for the model to finish generating the response. Sure there is the “…” bubble indicating that something is happening, but it is not the best user experience.
Streaming allows us to send the response in chunks as it is generated by the model. The model sends the response in chunks (tokens) and the server sends these chunks to the client as they arrive.
The final code of this step is located in the step-03 directory. However, we recommend you to follow the instructions
below to get there, and continue extending your current application.
Asking the LLM to return chunks
The first step is to ask the LLM to return the response in chunks. Initially, our AI service looked like this:
package dev.langchain4j.workshop;
import dev.langchain4j.cdi.spi.RegisterAIService;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;
import jakarta.enterprise.context.ApplicationScoped;
@RegisterAIService(
chatModelName = "customer-support-agent",
chatMemoryProviderName = "customer-support-agent-memory",
scope = ApplicationScoped.class
)
public interface CustomerSupportAgent {
String chat(@MemoryId String sessionId, @UserMessage String userMessage);
}
Note that the return type of the chat method is String.
We will change it to TokenStream to indicate that the response will be streamed instead of returned synchronously.
package dev.langchain4j.workshop;
import dev.langchain4j.cdi.spi.RegisterAIService;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import jakarta.enterprise.context.ApplicationScoped;
@RegisterAIService(
streamingChatModelName = "customer-support-agent",
chatMemoryProviderName = "customer-support-agent-memory",
scope = ApplicationScoped.class
)
public interface CustomerSupportAgent {
TokenStream chat(@MemoryId String sessionId, @UserMessage String userMessage);
}
As it’s name suggest, TokenStream represents a stream of tokens from the model. An AI service can subscribe to the
stram and receive updates when new tokens are available. This is known as a partial response.
Serving streams from the websocket
Ok, now our AI Service returns a stream of tokens, but we need to modify our websocket endpoint to handle this stream and send it to the client.
Currently, our websocket endpoint looks like this:
package dev.langchain4j.workshop;
import java.io.IOException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import jakarta.inject.Inject;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.ServerEndpoint;
@ServerEndpoint("/customer-support-agent")
public class CustomerSupportAgentWebSocket {
// Thread-safe set to store all active sessions
private static final Set<Session> sessions = Collections.synchronizedSet(new HashSet<>());
@Inject
private CustomerSupportAgent customerSupportAgent;
@OnOpen
public void onOpen(Session session) throws IOException {
sessions.add(session);
session.getBasicRemote().sendText("Welcome to Miles of Smiles! How can I help you today?");
}
@OnMessage
public void onMessage(String message, Session session) throws IOException {
session.getBasicRemote().sendText(customerSupportAgent.chat(session.getId(), message));
}
@OnClose
public void onClose(Session session) throws IOException {
sessions.remove(session);
}
}
The TokenStream
interface defines a number of methods that allows you to process various types of response from the LLM. Our AI service
only needs to process partial, completion and error responses. In order to keep the the websocket endpoint
implementation clean, we have separated the token stream processing logic into the WebSocketTokenStreamProcessor
class.
package dev.langchain4j.workshop;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.PartialThinking;
import dev.langchain4j.model.chat.response.PartialToolCall;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.tool.BeforeToolExecution;
import dev.langchain4j.service.tool.ToolExecution;
import jakarta.websocket.Session;
import java.io.IOException;
import java.util.concurrent.CompletableFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Processes token streams from AI models and handles various callback events.
* This class encapsulates the logic for handling partial responses, tool executions,
* and other streaming events from the LangChain4j TokenStream.
*/
public class WebSocketTokenStreamProcessor {
private static final Logger logger = LoggerFactory.getLogger(WebSocketTokenStreamProcessor.class);
private final Session session;
public WebSocketTokenStreamProcessor(Session session) {
this.session = session;
}
/**
* Processes the token stream with all callback handlers configured.
*
* @param tokenStream the token stream to process
* @return a CompletableFuture that completes when the response is fully received
*/
public void process(TokenStream tokenStream) {
CompletableFuture<ChatResponse> futureResponse = new CompletableFuture<>();
try {
// Process the tokens in the stream
tokenStream
.onPartialResponse((String partialResponse) -> processPartialResponse(partialResponse, futureResponse))
.onPartialThinking((PartialThinking partialThinking) -> processPartialThinking(partialThinking, futureResponse))
.onIntermediateResponse((ChatResponse intermediateResponse) -> processIntermediateResponse(intermediateResponse, futureResponse))
.onPartialToolCall((PartialToolCall partialToolCall) -> processPartialToolCall(partialToolCall, futureResponse))
.beforeToolExecution((BeforeToolExecution beforeToolExecution) -> processBeforeToolExecution(beforeToolExecution, futureResponse))
.onToolExecuted((ToolExecution toolExecution) -> processToolExecuted(toolExecution, futureResponse))
.onCompleteResponse((ChatResponse response) -> processCompleteResponse(response, futureResponse))
.onError((Throwable error) -> processError(error, futureResponse))
.start();
// Wait for the response to complete
futureResponse.join();
} catch (Exception e) {
logger.error("Error processing message", e);
try {
session.getBasicRemote().sendText("Error processing your message: " + e.getMessage());
futureResponse.completeExceptionally(e);
} catch (IOException ioException) {
logger.error("Error sending error message to client", ioException);
}
}
}
private void processPartialResponse(String partialResponse, CompletableFuture<ChatResponse> futureResponse) {
logger.debug("Partial response: {}", partialResponse);
try {
session.getBasicRemote().sendText(partialResponse);
} catch (IOException e) {
logger.error("Error processing partial response", e);
futureResponse.completeExceptionally(e);
}
}
private void processPartialThinking(PartialThinking partialThinking, CompletableFuture<ChatResponse> futureResponse) {
logger.debug("Partial thinking: {}", partialThinking);
}
private void processIntermediateResponse(ChatResponse intermediateResponse, CompletableFuture<ChatResponse> futureResponse) {
logger.debug("Intermediate response: {}", intermediateResponse);
}
private void processPartialToolCall(PartialToolCall partialToolCall, CompletableFuture<ChatResponse> futureResponse) {
logger.debug("Partial tool call: {}", partialToolCall);
}
private void processBeforeToolExecution(BeforeToolExecution beforeToolExecution, CompletableFuture<ChatResponse> futureResponse) {
logger.debug("Before tool execution: {}", beforeToolExecution);
}
private void processToolExecuted(ToolExecution toolExecution, CompletableFuture<ChatResponse> futureResponse) {
logger.debug("Tool executed: {}", toolExecution);
}
private void processCompleteResponse(ChatResponse chatResponse, CompletableFuture<ChatResponse> futureResponse) {
logger.debug("Complete response: {}", chatResponse);
logger.info("Chat Id: {}", chatResponse.id());
logger.info("Model Name: {}", chatResponse.modelName());
logger.info("Input Token Count: {}", chatResponse.tokenUsage().inputTokenCount());
logger.info("Output Token Count: {}", chatResponse.tokenUsage().outputTokenCount());
logger.info("Total Token Count: {}", chatResponse.tokenUsage().totalTokenCount());
// { text = "Hello! How can I help you today?", thinking = null, toolExecutionRequests = [], attributes = {} }, metadata = ChatResponseMetadata{id='null', modelName='gpt-oss:20b', tokenUsage=TokenUsage { inputTokenCount = 68, outputTokenCount = 35, totalTokenCount = 103 }, finishReason=STOP} }
futureResponse.complete(chatResponse);
}
private void processError(Throwable error, CompletableFuture<ChatResponse> futureResponse) {
logger.error("Error processing message", error);
try {
session.getBasicRemote().sendText("Error processing your message: " + error.getMessage());
futureResponse.completeExceptionally(error);
} catch (IOException ioException) {
logger.error("Error sending error message to client", ioException);
futureResponse.completeExceptionally(ioException);
}
}
}
Let’s modify the onMessage method to use the WebSocketTokenStreamProcessor to process the stream of tokens.
package dev.langchain4j.workshop;
import dev.langchain4j.service.TokenStream;
import java.io.IOException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import jakarta.inject.Inject;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.ServerEndpoint;
@ServerEndpoint("/customer-support-agent")
public class CustomerSupportAgentWebSocket {
// Thread-safe set to store all active sessions
private static final Set<Session> sessions = Collections.synchronizedSet(new HashSet<>());
@Inject
private CustomerSupportAgent customerSupportAgent;
@OnOpen
public void onOpen(Session session) throws IOException {
sessions.add(session);
session.getBasicRemote().sendText("Welcome to Miles of Smiles! How can I help you today?");
}
@OnMessage
public void onMessage(String message, Session session) throws IOException {
// Get the token stream from the agent
TokenStream tokenStream = customerSupportAgent.chat(session.getId(), message);
// Use the WebSocketTokenStreamProcessor to handle the stream
WebSocketTokenStreamProcessor processor = new WebSocketTokenStreamProcessor(session);
processor.process(tokenStream);
}
@OnClose
public void onClose(Session session) throws IOException {
sessions.remove(session);
}
}
That’s it! Now the response will be streamed to the client as it arrives.
Testing the streaming
To test the streaming, you can use the same chat interface as before. The application should still be running. Go back to the browser, refresh the page, and start chatting. If you ask simple questions, you may not notice the difference.
Ask something like
and you will see the response being displayed as it arrives.
Let’s now switch to the next step!