Skip to content

Commit 2c1e274

Browse files
ashakirinashakirin
committed
Fixed conversation id, streaming, added testcontainers
1 parent 318cd78 commit 2c1e274

5 files changed

Lines changed: 43 additions & 32 deletions

File tree

‎samples/spring-ai-agent/pom.xml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,24 @@
3838
<groupId>org.springframework.boot</groupId>
3939
<artifactId>spring-boot-starter-web</artifactId>
4040
</dependency>
41+
<dependency>
42+
<groupId>org.springframework.ai</groupId>
43+
<artifactId>spring-ai-starter-model-bedrock</artifactId>
44+
</dependency>
4145
<dependency>
4246
<groupId>org.springframework.ai</groupId>
4347
<artifactId>spring-ai-starter-model-bedrock-converse</artifactId>
4448
</dependency>
4549

4650
<!-- RAG Dependencies -->
51+
<dependency>
52+
<groupId>org.springframework.ai</groupId>
53+
<artifactId>spring-ai-vector-store</artifactId>
54+
</dependency>
55+
<dependency>
56+
<groupId>org.springframework.ai</groupId>
57+
<artifactId>spring-ai-advisors-vector-store</artifactId>
58+
</dependency>
4759
<dependency>
4860
<groupId>org.springframework.ai</groupId>
4961
<artifactId>spring-ai-starter-vector-store-pgvector</artifactId>

‎samples/spring-ai-agent/src/main/java/com/aws/workshop/ai/agent/controller/ExternalizedMemoryAgentController.java‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,11 @@
33
import com.aws.workshop.ai.agent.memory.ExternalChatMemoryRepository;
44
import org.springframework.ai.chat.client.ChatClient;
55
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
6+
import org.springframework.ai.chat.memory.ChatMemory;
67
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
78
import org.springframework.ai.chat.prompt.ChatOptions;
89
import org.springframework.web.bind.annotation.*;
9-
import software.amazon.awssdk.utils.StringInputStream;
10-
11-
import java.io.InputStream;
12-
import java.util.Objects;
10+
import reactor.core.publisher.Flux;
1311

1412
@RestController
1513
@RequestMapping("/ext-memory")
@@ -32,19 +30,21 @@ public String chat(@RequestBody String prompt) {
3230
return chatClient
3331
.prompt()
3432
.advisors(promptChatMemoryAdvisor)
33+
.advisors(advisor -> advisor.param(ChatMemory.CONVERSATION_ID, "logged-user-account"))
3534
.user(prompt)
3635
.call()
3736
.content();
3837
}
3938

4039
@PostMapping("/chat/stream")
41-
public InputStream chatStream(@RequestBody String prompt) {
42-
return new StringInputStream(Objects.requireNonNull(chatClient
40+
public Flux<String> chatStream(@RequestBody String prompt) {
41+
return chatClient
4342
.prompt()
4443
.advisors(promptChatMemoryAdvisor)
4544
.user(prompt)
46-
.call()
47-
.content()));
45+
.advisors(advisor -> advisor.param(ChatMemory.CONVERSATION_ID, "logged-user-account"))
46+
.stream()
47+
.content();
4848
}
4949

5050
@PostMapping("/model")

‎samples/spring-ai-agent/src/main/java/com/aws/workshop/ai/agent/controller/MemoryAgentController.java‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,12 @@
22

33
import org.springframework.ai.chat.client.ChatClient;
44
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
5+
import org.springframework.ai.chat.memory.ChatMemory;
56
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
67
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
78
import org.springframework.ai.chat.prompt.ChatOptions;
89
import org.springframework.web.bind.annotation.*;
9-
import software.amazon.awssdk.utils.StringInputStream;
10-
11-
import java.io.InputStream;
12-
import java.util.Objects;
10+
import reactor.core.publisher.Flux;
1311

1412
@RestController
1513
@RequestMapping("/memory")
@@ -33,19 +31,22 @@ public String chat(@RequestBody String prompt) {
3331
return chatClient
3432
.prompt()
3533
.advisors(promptChatMemoryAdvisor)
34+
.advisors(advisor -> advisor.param(ChatMemory.CONVERSATION_ID, "logged-user-account"))
3635
.user(prompt)
3736
.call()
3837
.content();
3938
}
4039

4140
@PostMapping("/chat/stream")
42-
public InputStream chatStream(@RequestBody String prompt) {
43-
return new StringInputStream(Objects.requireNonNull(chatClient
41+
public Flux<String> chatStream(@RequestBody String prompt) {
42+
return chatClient
4443
.prompt()
4544
.advisors(promptChatMemoryAdvisor)
45+
.advisors(advisor -> advisor.param(ChatMemory.CONVERSATION_ID, "logged-user-account"))
4646
.user(prompt)
47-
.call()
48-
.content()));
47+
.advisors(advisor -> advisor.param(ChatMemory.CONVERSATION_ID, "logged-user-account"))
48+
.stream()
49+
.content();
4950
}
5051

5152
@PostMapping("/model")

‎samples/spring-ai-agent/src/main/java/com/aws/workshop/ai/agent/controller/RagPgVectorAgentController.java‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,9 @@
66
import org.springframework.ai.document.Document;
77
import org.springframework.ai.vectorstore.VectorStore;
88
import org.springframework.web.bind.annotation.*;
9-
import software.amazon.awssdk.utils.StringInputStream;
9+
import reactor.core.publisher.Flux;
1010

11-
import java.io.InputStream;
1211
import java.util.List;
13-
import java.util.Objects;
1412

1513
@RestController
1614
@RequestMapping("/rag-pgvector")
@@ -43,13 +41,13 @@ public void loadDataToVectorStore(@RequestBody String content) {
4341
}
4442

4543
@PostMapping("/chat/stream")
46-
public InputStream chatStream(@RequestBody String prompt) {
47-
return new StringInputStream(Objects.requireNonNull(chatClient
44+
public Flux<String> chatStream(@RequestBody String prompt) {
45+
return chatClient
4846
.prompt()
4947
.advisors(ragAdvisor)
5048
.user(prompt)
51-
.call()
52-
.content()));
49+
.stream()
50+
.content();
5351
}
5452

5553
@PostMapping("/model")

‎samples/spring-ai-agent/src/main/java/com/aws/workshop/ai/agent/controller/StatelessAgentController.java‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22

33
import org.springframework.ai.chat.client.ChatClient;
44
import org.springframework.ai.chat.prompt.ChatOptions;
5-
import org.springframework.web.bind.annotation.*;
6-
import software.amazon.awssdk.utils.StringInputStream;
7-
8-
import java.io.InputStream;
9-
import java.util.Objects;
5+
import org.springframework.web.bind.annotation.PostMapping;
6+
import org.springframework.web.bind.annotation.RequestBody;
7+
import org.springframework.web.bind.annotation.RequestParam;
8+
import org.springframework.web.bind.annotation.RestController;
9+
import reactor.core.publisher.Flux;
1010

1111
@RestController
1212
public class StatelessAgentController {
@@ -28,12 +28,12 @@ public String chat(@RequestBody String prompt) {
2828
}
2929

3030
@PostMapping("/chat/stream")
31-
public InputStream chatStream(@RequestBody String prompt) {
32-
return new StringInputStream(Objects.requireNonNull(chatClient
31+
public Flux<String> chatStream(@RequestBody String prompt) {
32+
return chatClient
3333
.prompt()
3434
.user(prompt)
35-
.call()
36-
.content()));
35+
.stream()
36+
.content();
3737
}
3838

3939
@PostMapping("/model")

0 commit comments

Comments
 (0)