add streaming support

Author: alexcheng1982

pom.xml

@@ -6,13 +6,13 @@
   <parent>
     <groupId>org.springframework.boot</groupId>
     <artifactId>spring-boot-starter-parent</artifactId>
-    <version>3.4.1</version>
+    <version>3.4.3</version>
     <relativePath/>
   </parent>
   <groupId>io.github.javaaidev</groupId>
   <artifactId>native-simple-ai-agent</artifactId>
   <version>0.0.1-SNAPSHOT</version>
-  <name>native-simple-ai-agent</name>
+  <name>Simple AI Agent</name>
   <description>Simple AI agent</description>
   <properties>
     <java.version>21</java.version>

src/main/java/com/javaaidev/agent/AbstractChatAgentController.java

@@ -0,0 +1,41 @@
+package com.javaaidev.agent;
+
+import com.javaaidev.chatagent.model.ChatRequest;
+import com.javaaidev.chatagent.model.TextContentPart;
+import com.javaaidev.chatagent.model.ThreadAssistantMessage;
+import com.javaaidev.chatagent.model.ThreadUserMessage;
+import java.util.List;
+import java.util.stream.Stream;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.UserMessage;
+
+public abstract class AbstractChatAgentController {
+
+  public static final String SYSTEM_TEXT = """
+      You are a chef who is proficient in various cuisines. Please answer users' questions about cooking.
+      For other unrelated inputs, simply tell the user that you don't know.
+      """;
+
+
+  protected List<Message> chatRequestToMessages(ChatRequest request) {
+    return request.messages().stream().flatMap(message -> {
+      if (message instanceof ThreadUserMessage userMessage) {
+        return userMessage.content().stream().map(part -> {
+          if (part instanceof TextContentPart(String text)) {
+            return new UserMessage(text);
+          }
+          return null;
+        });
+      } else if (message instanceof ThreadAssistantMessage assistantMessage) {
+        return assistantMessage.content().stream().map(part -> {
+          if (part instanceof TextContentPart(String text)) {
+            return new AssistantMessage(text);
+          }
+          return null;
+        });
+      }
+      return Stream.<Message>of();
+    }).toList();
+  }
+}

src/main/java/com/javaaidev/agent/ChatAgentController.java

@@ -1,59 +1,37 @@
 package com.javaaidev.agent;
 
 import com.javaaidev.chatagent.model.ChatRequest;
-import com.javaaidev.chatagent.model.ChatResponse;
-import com.javaaidev.chatagent.model.TextContentPart;
-import com.javaaidev.chatagent.model.ThreadAssistantMessage;
-import com.javaaidev.chatagent.model.ThreadUserMessage;
-import java.util.List;
-import java.util.stream.Stream;
 import org.springframework.ai.chat.client.ChatClient;
-import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.messages.Message;
-import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.http.codec.ServerSentEvent;
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestBody;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
+import reactor.core.publisher.Flux;
 
 @RestController
 @RequestMapping("/chat")
-public class ChatAgentController {
+public class ChatAgentController extends AbstractChatAgentController {
 
-  private static final String SYSTEM_TEXT = "You are a chef who is proficient in various cuisines. Please answer users' questions about cooking.";
   private final ChatClient chatClient;
 
   public ChatAgentController(ChatClient.Builder builder) {
     chatClient = builder.build();
   }
 
   @PostMapping
-  public ChatResponse chat(@RequestBody ChatRequest request) {
+  public Flux<ServerSentEvent<String>> chat(@RequestBody ChatRequest request) {
     if (request == null) {
-      return new ChatResponse(List.of());
+      return Flux.empty();
     }
-    var messages = request.messages().stream().flatMap(message -> {
-      if (message instanceof ThreadUserMessage userMessage) {
-        return userMessage.content().stream().map(part -> {
-          if (part instanceof TextContentPart(String text)) {
-            return new UserMessage(text);
-          }
-          return null;
-        });
-      } else if (message instanceof ThreadAssistantMessage assistantMessage) {
-        return assistantMessage.content().stream().map(part -> {
-          if (part instanceof TextContentPart(String text)) {
-            return new AssistantMessage(text);
-          }
-          return null;
-        });
-      }
-      return Stream.of();
-    }).toList();
+    var messages = chatRequestToMessages(request);
     var output = chatClient.prompt().system(SYSTEM_TEXT)
         .messages(messages.toArray(new Message[0]))
         .call()
         .content();
-    return new ChatResponse(List.of(new TextContentPart(output)));
+    return Flux.just(ServerSentEvent.<String>builder()
+        .data("#" + output + "#")
+        .build());
   }
 }

src/main/java/com/javaaidev/agent/ChatAgentStreamingController.java

@@ -0,0 +1,32 @@
+package com.javaaidev.agent;
+
+import com.javaaidev.chatagent.model.ChatRequest;
+import org.springframework.ai.chat.client.ChatClient;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.http.codec.ServerSentEvent;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+import reactor.core.publisher.Flux;
+
+@RestController
+@RequestMapping("/chat_streaming")
+public class ChatAgentStreamingController extends AbstractChatAgentController {
+
+  private final ChatClient chatClient;
+
+  public ChatAgentStreamingController(ChatClient.Builder builder) {
+    chatClient = builder.build();
+  }
+
+  @PostMapping
+  public Flux<ServerSentEvent<String>> chatStreaming(@RequestBody ChatRequest request) {
+    var messages = chatRequestToMessages(request);
+    return chatClient.prompt().system(SYSTEM_TEXT).messages(messages.toArray(new Message[0]))
+        .stream().content()
+        .map(output -> ServerSentEvent.<String>builder()
+            .data("#" + output + "#")
+            .build());
+  }
+}