Trying to reproduce langchain4j/langchain4j#1848

jdubois · Dec 24, 2024 · aeecf3e · aeecf3e
1 parent 8807154
commit aeecf3e
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 1 deletion.
diff --git a/src/main/java/com/example/demo/DemoController.java b/src/main/java/com/example/demo/DemoController.java
@@ -9,13 +9,16 @@
 import dev.langchain4j.data.document.splitter.DocumentSplitters;
 import dev.langchain4j.data.document.transformer.jsoup.HtmlToTextDocumentTransformer;
 import dev.langchain4j.data.embedding.Embedding;
+import dev.langchain4j.data.message.AiMessage;
 import dev.langchain4j.data.message.SystemMessage;
 import dev.langchain4j.data.message.UserMessage;
 import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.memory.ChatMemory;
 import dev.langchain4j.memory.chat.MessageWindowChatMemory;
+import dev.langchain4j.model.StreamingResponseHandler;
 import dev.langchain4j.model.azure.AzureOpenAiImageModel;
 import dev.langchain4j.model.chat.ChatLanguageModel;
+import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import dev.langchain4j.model.chat.request.ChatRequest;
 import dev.langchain4j.model.chat.request.ResponseFormat;
 import dev.langchain4j.model.chat.request.json.JsonArraySchema;
@@ -24,6 +27,7 @@
 import dev.langchain4j.model.chat.request.json.JsonStringSchema;
 import dev.langchain4j.model.embedding.EmbeddingModel;
 import dev.langchain4j.model.image.ImageModel;
+import dev.langchain4j.model.output.Response;
 import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
 import dev.langchain4j.service.AiServices;
 import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
@@ -53,14 +57,17 @@ public class DemoController {
 
     private final ChatLanguageModel chatLanguageModel;
 
+    private final StreamingChatLanguageModel streamingChatLanguageModel;
+
     private final EmbeddingModel embeddingModel;
 
     private final EmbeddingStore<TextSegment> embeddingStore;
 
     private final StockPriceService stockPriceService;
 
-    public DemoController(ImageModel imageModel, ChatLanguageModel chatLanguageModel, EmbeddingModel embeddingModel, EmbeddingStore<TextSegment> embeddingStore, StockPriceService stockPriceService) {
+    public DemoController(ImageModel imageModel, StreamingChatLanguageModel streamingChatLanguageModel, ChatLanguageModel chatLanguageModel, EmbeddingModel embeddingModel, EmbeddingStore<TextSegment> embeddingStore, StockPriceService stockPriceService) {
         this.imageModel = imageModel;
+        this.streamingChatLanguageModel = streamingChatLanguageModel;
         this.chatLanguageModel = chatLanguageModel;
         this.embeddingModel = embeddingModel;
         this.embeddingStore = embeddingStore;
@@ -90,6 +97,30 @@ String getAnswer(Model model) {
         return getView(model, "2: simple question", question, answer);
     }
 
+    @GetMapping("/2b")
+    String getStreamingAnswer(Model model) {
+        String question = "Who painted the Mona Lisa?";
+        streamingChatLanguageModel.generate(question, new StreamingResponseHandler<AiMessage>() {
+
+            @Override
+            public void onNext(String token) {
+                System.out.println("Token: " + token);
+            }
+
+            @Override
+            public void onComplete(Response<AiMessage> response) {
+                System.out.println("Response: " + response.content().text());
+                StreamingResponseHandler.super.onComplete(response);
+            }
+
+            @Override
+            public void onError(Throwable error) {
+
+            }
+        });
+        return getView(model, "2: simple question", question, "Streaming test in the logs");
+    }
+
     @GetMapping("/3")
     String reasoning(Model model) {
         String question = "Maria's father has 4 daughters: Spring, Autumn, Winter. What is the name of the fourth daughter?";

diff --git a/...main/java/com/example/demo/configuration/chatmodel/azure/AzureChatModelConfiguration.java b/...main/java/com/example/demo/configuration/chatmodel/azure/AzureChatModelConfiguration.java
@@ -1,7 +1,9 @@
 package com.example.demo.configuration.chatmodel.azure;
 
 import dev.langchain4j.model.azure.AzureOpenAiChatModel;
+import dev.langchain4j.model.azure.AzureOpenAiStreamingChatModel;
 import dev.langchain4j.model.chat.ChatLanguageModel;
+import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -32,4 +34,14 @@ ChatLanguageModel azureOpenAIChatLanguageModel() {
                 .strictJsonSchema(true)
                 .build();
     }
+
+    @Bean
+    StreamingChatLanguageModel azureOpenAIStreamingChatLanguageModel() {
+        return AzureOpenAiStreamingChatModel.builder()
+                .endpoint(azureOpenAiEndpoint)
+                .apiKey(azureOpenAiKey)
+                .deploymentName("gpt-4o")
+                .logRequestsAndResponses(true)
+                .build();
+    }
 }