Unable to stream chat using OciGenAiStreamingChatModel

[ranjkkum]

I’m working on implementing a method for streaming chat responses. The method runs without errors, but instead of streaming partial responses, it only returns the full response once the process is complete. Could you review my implementation and help me identify what might be preventing the response from streaming as expected?

pom.xml

<dependency>
            <groupId>dev.langchain4j</groupId>
            <artifactId>langchain4j-spring-boot-starter</artifactId>
            <version>1.5.0-beta11</version>
        </dependency>

        <dependency>
            <groupId>dev.langchain4j</groupId>
            <artifactId>langchain4j-community-oci-genai</artifactId>
            <version>1.5.0-beta11</version>
        </dependency>
        <dependency>
            <groupId>dev.langchain4j</groupId>
            <artifactId>langchain4j-reactor</artifactId>
            <version>1.5.0-beta11</version>
        </dependency>

StreamingAssistant.java

@SystemMessage("You are an assistant for deployment queries. For general questions, use RAG to retrieve relevant information. When invoking tools (e.g., listComputeInstances), do not retrieve or augment with RAG sources; rely solely on the tool's output for accuracy and speed.")
public interface StreamingAssistant {

	@UserMessage("{{it}}")
	TokenStream chat(String text);
}

@Bean
	public OciGenAiStreamingChatModel getStreamingChatModel() {
		log.info("Creating streaming chat model with modelName: {}, compartmentId: {}, region: {}",
				config.getModel().getName(), config.getCompartmentId(), config.getRegion());
		return OciGenAiStreamingChatModel.builder()
				.modelName(config.getModel().getName())
				.compartmentId(config.getCompartmentId())
				.authProvider(authenticationDetailsProvider)
				.region(Region.valueOf(config.getRegion()))
				.build();
	}

private StreamingAssistant getOrCreateStreamingAssistant(String sessionId) {
		return streamingAssistantCache.computeIfAbsent(sessionId, id -> {
			log.debug("Creating new assistant for sessionId: {}", id);
			MessageWindowChatMemory memory = createChatMemory(id);
			AiServices<StreamingAssistant> builder = AiServices.builder(StreamingAssistant.class);
			builder.streamingChatModel(getStreamingChatModel());
			builder.tools(assistantToolProvider.getTools());
			builder.chatMemory(memory);
			return builder.build();
		});
	}

Approach 1 -

public Flux<String> chatStream(AssistantChatRequest chatRequest) throws JsonProcessingException {

		Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();

		StreamingAssistant streamingAssistant = getOrCreateStreamingAssistant(chatRequest.getSessionId());
		streamingAssistant.chat(chatRequest.getPrompt())
				.onPartialResponse(sink::tryEmitNext)
				.onCompleteResponse(aiMessageResponse -> sink.tryEmitComplete())
				.onError(sink::tryEmitError)
				.start();
		return sink.asFlux();

	}

Approach 2 -

	@Override
	public Flux<String> chatStream(AssistantChatRequest chatRequest) throws JsonProcessingException {
		return Flux.create(sink -> {
			var model = getStreamingChatModel();
			model.chat(chatRequest.getPrompt(), new StreamingChatResponseHandler() {
				@Override
				public void onPartialResponse(String s) {
					JSONObject jsonObject = new JSONObject();
					try {
						jsonObject.put("message", s);
					} catch (JSONException e) {
						throw new RuntimeException(e);
					}
					sink.next(jsonObject.toString());
				}

				@Override
				public void onCompleteResponse(ChatResponse chatResponse) {
					log.info("onCompleteResponse: {}", chatResponse);
					sink.complete();
				}

				@Override
				public void onError(Throwable throwable) {
					log.error("onError", throwable);
					sink.error(throwable);
				}
			});
		}, FluxSink.OverflowStrategy.BUFFER);
}

AssistantChatController.java

 @PostMapping(value = RESOURCE_IDENTIFIER + "/actions/chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @Operation(summary = "Stream chat with Database Tools Chat AI", description = "Stream chat with Database Tools Chat AI")
    public Flux<String> chatStream(@RequestBody AssistantChatRequest chatRequest) throws JsonProcessingException {
        return assistantChatService.chatStream(chatRequest);
    }

[danielkec]

You don't heave to do anything, when you have dev.langchain4j:langchain4j-reactor on your classpath, just let your AiService method return Flux<String>

@SystemMessage("You are an assistant for deployment queries. For general questions, use RAG to retrieve relevant information. When invoking tools (e.g., listComputeInstances), do not retrieve or augment with RAG sources; rely solely on the tool's output for accuracy and speed.")
public interface StreamingAssistant {

	@UserMessage("{{it}}")
	Flux<String> chat(String text);
}

[ranjkkum]

@danielkec I have tried that but it does not work..

@SystemMessage("You are an assistant for deployment queries. For general questions, use RAG to retrieve relevant information. When invoking tools (e.g., listComputeInstances), do not retrieve or augment with RAG sources; rely solely on the tool's output for accuracy and speed.")
public interface StreamingAssistant {

	@UserMessage("{{it}}")
	Flux<String> chat(String text);
}

@Override
	public Flux<String> chatStream(AssistantChatRequest chatRequest) throws JsonProcessingException {
		StreamingAssistant streamingAssistant = getOrCreateStreamingAssistant(chatRequest.getSessionId());
		return streamingAssistant.chat(chatRequest.getPrompt());
}

The behaviour is still the same ..

[ranjkkum]

demo.zip

I have created this project to reproduce the issue . Please take a look at it .

Please replace **changeme**

public OciGenAiStreamingChatModel getStreamingChatModel() throws IllegalArgumentException, Exception {

		return OciGenAiStreamingChatModel.builder()
				.modelName("xai.grok-4")
				.compartmentId("ocid1.compartment.oc1..**changeme**")
				.authProvider(authenticationDetailsProvider())
				.region(Region.valueOf("us-chicago-1"))
				.build();
	}

StreamingAssistant streamingAssistant = getOrCreateStreamingAssistant(
				"ocid1.genaiagentsession.oc1.us-chicago-1.**changeme**");

[danielkec]

Yes I am able to reproduce it with Flux Token stream adapter, on the other hand it seems to be streaming just fine with StreamingChatResponseHandler:

var model = OciGenAiStreamingChatModel.builder()
        .modelName("meta.llama-3.3-70b-instruct")
        .compartmentId(System.getenv("OCI_GENAI_COMPARTMENT_ID"))
        .authProvider(new ConfigFileAuthenticationDetailsProvider("DEFAULT"))
        .build();
var request = ChatRequest.builder()
        .messages(List.of(
                dev.langchain4j.data.message.SystemMessage
                        .from("You are a pirate, who likes to brag about hist time on sea."),
                UserMessage.from("Can you tell me a really long story?")
        )).build();
CompletableFuture<Void> cf = new CompletableFuture<>();
model.chat(request, new StreamingChatResponseHandler() {
    @Override
    public void onPartialResponse(String partialResponse) {
        System.out.println(partialResponse);
    }
    @Override
    public void onCompleteResponse(ChatResponse completeResponse) {
        System.out.println(completeResponse);
        cf.complete(null);
    }
    @Override
    public void onError(Throwable error) {
        System.err.println(error);
        cf.completeExceptionally(error);
    }
});
cf.join();

[ranjkkum]

Thanks @danielkec .. i was able to make it work using Mono for now

StreamingAssistant streamingAssistant = getOrCreateStreamingAssistant(
				"ocid1.genaiagentsession.oc1.us-chicago-1.changeme");
		Sinks.Many<String> sink = Sinks.many().multicast().onBackpressureBuffer();

		Mono.fromCallable(() -> {
			streamingAssistant.chat("tell me a long story")
					.onPartialResponse(sink::tryEmitNext)
					.onCompleteResponse(aiMessageResponse -> sink.tryEmitComplete())
					.onError(sink::tryEmitError)
					.start();
			return null;
		}).subscribeOn(Schedulers.boundedElastic()).subscribe();

		// Return the Flux from the sink
		return sink.asFlux();

[danielkec]

Ah trampoline, there may be a problem after all

[danielkec]

And there is, this is a bug in streaming model, thank you for locating it @ranjkkum! Streaming models should use async OCI SDK client instead.