Skip to content

Commit 9924753

Browse files
jasperbluesjasperblues
committed
Fix slow tests.
1 parent 3ecf4ca commit 9924753

File tree

9 files changed

+145
-45
lines changed

9 files changed

+145
-45
lines changed

‎src/main/kotlin/com/embabel/guide/chat/service/FakeRagServiceAdapter.kt‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,24 +21,26 @@ class FakeRagServiceAdapter : RagServiceAdapter {
2121
private val logger = LoggerFactory.getLogger(FakeRagServiceAdapter::class.java)
2222

2323
override suspend fun sendMessage(
24+
threadId: String,
2425
message: String,
2526
fromUserId: String,
27+
priorMessages: List<PriorMessage>,
2628
onEvent: (String) -> Unit
2729
): String {
28-
logger.info("Processing fake RAG request from user: {}", fromUserId)
30+
logger.info("Processing fake RAG request from user: {} in thread: {}", fromUserId, threadId)
2931

30-
// Simulate processing stages with events and delays
32+
// Simulate processing stages with events and minimal delays for testing
3133
onEvent("Analyzing your question...")
32-
delay(600)
34+
delay(10)
3335

3436
onEvent("Searching knowledge base...")
35-
delay(950)
37+
delay(10)
3638

3739
onEvent("Retrieving relevant documents...")
38-
delay(1400)
40+
delay(10)
3941

4042
onEvent("Generating response...")
41-
delay(2000)
43+
delay(10)
4244

4345
// Generate response based on message content (moved from JesseService)
4446
val response = when {

‎src/main/kotlin/com/embabel/guide/chat/service/GuideRagServiceAdapter.kt‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ import java.util.concurrent.ConcurrentHashMap
1919
*
2020
* This adapter bridges the WebSocket chat system with the Guide's RAG-powered chatbot,
2121
* enabling real-time AI responses through the web interface.
22+
*
23+
* Sessions are cached per thread to maintain separate conversation contexts.
2224
*/
2325
@Service
2426
@ConditionalOnProperty(
@@ -33,8 +35,8 @@ class GuideRagServiceAdapter(
3335

3436
private val logger = LoggerFactory.getLogger(GuideRagServiceAdapter::class.java)
3537

36-
// Session cache to maintain conversation continuity per user
37-
private val userSessions = ConcurrentHashMap<String, SessionContext>()
38+
// Session cache to maintain conversation continuity per thread
39+
private val threadSessions = ConcurrentHashMap<String, SessionContext>()
3840

3941
companion object {
4042
private const val RESPONSE_TIMEOUT_MS = 120000 // 2 minutes
@@ -65,11 +67,13 @@ class GuideRagServiceAdapter(
6567
}
6668

6769
override suspend fun sendMessage(
70+
threadId: String,
6871
message: String,
6972
fromUserId: String,
73+
priorMessages: List<PriorMessage>,
7074
onEvent: (String) -> Unit
7175
): String = withContext(Dispatchers.IO) {
72-
logger.info("Processing Guide RAG request from user: {}", fromUserId)
76+
logger.info("Processing Guide RAG request from user: {} in thread: {}", fromUserId, threadId)
7377

7478
val responseBuilder = StringBuilder()
7579
var isComplete = false
@@ -81,16 +85,29 @@ class GuideRagServiceAdapter(
8185
val guideUser = guideUserRepository.findById(fromUserId)
8286
.orElseThrow { RuntimeException("No user found with id: $fromUserId") }
8387

84-
// Get or create session context for this user to maintain conversation continuity
85-
val sessionContext = userSessions.computeIfAbsent(fromUserId) {
86-
logger.info("Creating new chat session for user: {}", fromUserId)
88+
// Get or create session context for this thread to maintain conversation continuity
89+
var isNewSession = false
90+
val sessionContext = threadSessions.computeIfAbsent(threadId) {
91+
logger.info("Creating new chat session for thread: {} (user: {})", threadId, fromUserId)
92+
isNewSession = true
8793
val dynamicChannel = DynamicOutputChannel()
88-
// Set the delegate before creating the session to avoid initialization errors
8994
dynamicChannel.currentDelegate = messageOutputChannel
9095
val session = chatbot.createSession(guideUser, dynamicChannel, null)
9196
SessionContext(session, dynamicChannel)
9297
}
9398

99+
// Load prior messages into the conversation if this is a new session
100+
// Messages are added directly to the conversation without being processed by AI
101+
if (isNewSession && priorMessages.isNotEmpty()) {
102+
logger.info("Loading {} prior messages into conversation for thread: {}", priorMessages.size, threadId)
103+
for (prior in priorMessages) {
104+
when (prior.role) {
105+
"user" -> sessionContext.session.conversation.addMessage(UserMessage(prior.content))
106+
"assistant" -> sessionContext.session.conversation.addMessage(AssistantMessage(prior.content))
107+
}
108+
}
109+
}
110+
94111
// Update the dynamic channel to point to this message's output channel
95112
// (for existing sessions, or if this is a new session this ensures it's set)
96113
sessionContext.dynamicChannel.currentDelegate = messageOutputChannel
@@ -102,7 +119,7 @@ class GuideRagServiceAdapter(
102119

103120
responseBuilder.toString().ifBlank { DEFAULT_ERROR_MESSAGE }
104121
} catch (e: Exception) {
105-
logger.error("Error processing message from user {}: {}", fromUserId, e.message, e)
122+
logger.error("Error processing message from user {} in thread {}: {}", fromUserId, threadId, e.message, e)
106123
throw e
107124
}
108125
}

‎src/main/kotlin/com/embabel/guide/chat/service/JesseService.kt‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,21 @@ class JesseService(
8383
authorId = guideUserId
8484
)
8585

86+
// Load existing thread messages for context
87+
val timeline = threadService.findByThreadId(threadId).orElse(null)
88+
val priorMessages = timeline?.messages
89+
?.dropLast(1) // Exclude the message we just added
90+
?.map { PriorMessage(it.message.role, it.current.text) }
91+
?: emptyList()
92+
8693
// Send status updates to the user while processing
8794
// Use WebUser ID for WebSocket delivery (that's the principal in the session)
88-
val response = ragAdapter.sendMessage(message, guideUserId) { event ->
95+
val response = ragAdapter.sendMessage(
96+
threadId = threadId,
97+
message = message,
98+
fromUserId = guideUserId,
99+
priorMessages = priorMessages
100+
) { event ->
89101
logger.debug("RAG event for user {}: {}", fromWebUserId, event)
90102
sendStatusToUser(fromWebUserId, event)
91103
}

‎src/main/kotlin/com/embabel/guide/chat/service/RagServiceAdapter.kt‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,13 @@
11
package com.embabel.guide.chat.service
22

3+
/**
4+
* Represents a prior message in a conversation thread for context loading.
5+
*/
6+
data class PriorMessage(
7+
val role: String, // "user" or "assistant"
8+
val content: String
9+
)
10+
311
/**
412
* Interface for integrating with RAG (Retrieval-Augmented Generation) systems.
513
*
@@ -15,31 +23,34 @@ interface RagServiceAdapter {
1523

1624
/**
1725
* Sends a message to the RAG system and returns the response.
26+
* Each thread maintains its own conversation context.
1827
*
1928
* This is a suspending function to avoid blocking threads during potentially
2029
* long-running RAG operations (document retrieval, LLM inference, etc.).
2130
*
31+
* @param threadId The thread ID for maintaining separate conversation contexts
2232
* @param message The user's message to process
2333
* @param fromUserId The ID of the user sending the message (for context/logging)
34+
* @param priorMessages Prior messages to load into context if this is a new session for this thread
2435
* @param onEvent Callback function to receive real-time status updates during processing
2536
* (e.g., "Planning response", "Querying database", "Generating answer")
2637
* @return The RAG system's response message
2738
*/
2839
suspend fun sendMessage(
40+
threadId: String,
2941
message: String,
3042
fromUserId: String,
43+
priorMessages: List<PriorMessage> = emptyList(),
3144
onEvent: (String) -> Unit = {}
3245
): String
3346

3447
/**
3548
* Generates a short title from message content.
49+
* Uses a one-shot context (no thread association).
3650
*
3751
* @param content The message content to generate a title from
3852
* @param fromUserId The ID of the user (for session context)
3953
* @return A short title (typically 3-6 words)
4054
*/
41-
suspend fun generateTitle(content: String, fromUserId: String): String {
42-
val response = sendMessage(TITLE_PROMPT + content, fromUserId)
43-
return response.trim().take(100) // Safety limit
44-
}
55+
suspend fun generateTitle(content: String, fromUserId: String): String
4556
}

‎src/main/kotlin/com/embabel/guide/chat/service/ThreadService.kt‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,14 +90,19 @@ class ThreadService(
9090
ownerId: String,
9191
displayName: String
9292
): ThreadTimeline = withContext(Dispatchers.IO) {
93+
// Generate threadId upfront so we can pass it to the RAG adapter
94+
val threadId = UUIDv7.generateString()
9395
val prompt = WELCOME_PROMPT_TEMPLATE.format(displayName)
9496
val welcomeMessage = ragAdapter.sendMessage(
97+
threadId = threadId,
9598
message = prompt,
96-
fromUserId = ownerId, // Use the new user's ID for the chat session
97-
onEvent = { } // No status updates needed for welcome message
99+
fromUserId = ownerId,
100+
priorMessages = emptyList(), // No prior context for welcome thread
101+
onEvent = { } // No status updates needed for welcome message
98102
)
99103

100-
createThread(
104+
threadRepository.createWithMessage(
105+
threadId = threadId,
101106
ownerId = ownerId,
102107
title = "Welcome",
103108
message = welcomeMessage,

‎src/main/kotlin/com/embabel/hub/HubService.kt‎

Lines changed: 7 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,17 @@
11
package com.embabel.hub
22

3-
import com.embabel.guide.chat.model.DeliveredMessage
4-
import com.embabel.guide.chat.service.ChatService
5-
import com.embabel.guide.chat.service.ThreadService
63
import com.embabel.guide.domain.GuideUser
74
import com.embabel.guide.domain.GuideUserService
85
import com.embabel.guide.domain.WebUserData
96
import com.embabel.guide.util.UUIDv7
10-
import kotlinx.coroutines.CoroutineScope
11-
import kotlinx.coroutines.Dispatchers
12-
import kotlinx.coroutines.launch
137
import org.springframework.security.crypto.bcrypt.BCryptPasswordEncoder
148
import org.springframework.stereotype.Service
159

1610
@Service
1711
class HubService(
1812
private val guideUserService: GuideUserService,
1913
private val jwtTokenService: JwtTokenService,
20-
private val threadService: ThreadService,
21-
private val chatService: ChatService
14+
private val welcomeGreeter: WelcomeGreeter
2215
) {
2316

2417
private val passwordEncoder = BCryptPasswordEncoder()
@@ -74,21 +67,12 @@ class HubService(
7467
// Save the user through GuideUserService
7568
val guideUser = guideUserService.saveFromWebUser(webUser)
7669

77-
// Create a welcome thread for the new user with AI-generated greeting (fire-and-forget)
78-
// After creation, send the welcome message to the user via WebSocket
79-
CoroutineScope(Dispatchers.IO).launch {
80-
val timeline = threadService.createWelcomeThread(
81-
ownerId = guideUser.core.id,
82-
displayName = request.userDisplayName
83-
)
84-
// Send the welcome message to the user via WebSocket
85-
// Use WebUser ID (userId) as that's the principal in the JWT/WebSocket session
86-
val welcomeMessage = timeline.messages.firstOrNull()
87-
if (welcomeMessage != null) {
88-
val delivered = DeliveredMessage.createFrom(welcomeMessage, timeline.thread.threadId)
89-
chatService.sendToUser(userId, delivered)
90-
}
91-
}
70+
// Greet the new user with a welcome message (fire-and-forget)
71+
welcomeGreeter.greetNewUser(
72+
guideUserId = guideUser.core.id,
73+
webUserId = userId,
74+
displayName = request.userDisplayName
75+
)
9276

9377
return guideUser
9478
}

‎src/main/kotlin/com/embabel/hub/WelcomeGreeter.kt‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package com.embabel.hub
2+
3+
import com.embabel.guide.chat.model.DeliveredMessage
4+
import com.embabel.guide.chat.service.ChatService
5+
import com.embabel.guide.chat.service.ThreadService
6+
import kotlinx.coroutines.CoroutineScope
7+
import kotlinx.coroutines.Dispatchers
8+
import kotlinx.coroutines.launch
9+
import org.springframework.stereotype.Component
10+
11+
/**
12+
* Interface for greeting new users with a welcome message.
13+
*/
14+
interface WelcomeGreeter {
15+
/**
16+
* Greet a new user by creating a welcome thread and sending the message.
17+
* This is a fire-and-forget operation that runs asynchronously.
18+
*
19+
* @param guideUserId the GuideUser's core ID (owner of the thread)
20+
* @param webUserId the WebUser's ID (for WebSocket delivery)
21+
* @param displayName the user's display name for personalized greeting
22+
*/
23+
fun greetNewUser(guideUserId: String, webUserId: String, displayName: String)
24+
}
25+
26+
/**
27+
* Production implementation that creates AI-generated welcome threads.
28+
*/
29+
@Component
30+
class WelcomeGreeterImpl(
31+
private val threadService: ThreadService,
32+
private val chatService: ChatService
33+
) : WelcomeGreeter {
34+
35+
override fun greetNewUser(guideUserId: String, webUserId: String, displayName: String) {
36+
CoroutineScope(Dispatchers.IO).launch {
37+
val timeline = threadService.createWelcomeThread(
38+
ownerId = guideUserId,
39+
displayName = displayName
40+
)
41+
// Send the welcome message to the user via WebSocket
42+
val welcomeMessage = timeline.messages.firstOrNull()
43+
if (welcomeMessage != null) {
44+
val delivered = DeliveredMessage.createFrom(welcomeMessage, timeline.thread.threadId)
45+
chatService.sendToUser(webUserId, delivered)
46+
}
47+
}
48+
}
49+
}

‎src/test/kotlin/com/embabel/guide/GuideTestConfig.kt‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import com.embabel.common.ai.model.EmbeddingService
2020
import com.embabel.common.ai.model.SpringEmbeddingService
2121
import com.embabel.common.ai.model.Llm
2222
import com.embabel.common.ai.model.PricingModel
23+
import com.embabel.hub.WelcomeGreeter
2324
import org.springframework.ai.chat.messages.AssistantMessage
2425
import org.springframework.ai.chat.model.ChatModel
2526
import org.springframework.ai.chat.model.ChatResponse
@@ -63,6 +64,20 @@ class GuideTestConfig {
6364
provider = "test",
6465
)
6566
}
67+
68+
/**
69+
* No-op WelcomeGreeter for tests to avoid fire-and-forget coroutines
70+
* that can interfere with transactional test rollback.
71+
*/
72+
@Bean
73+
@org.springframework.context.annotation.Primary
74+
fun testWelcomeGreeter(): WelcomeGreeter {
75+
return object : WelcomeGreeter {
76+
override fun greetNewUser(guideUserId: String, webUserId: String, displayName: String) {
77+
// No-op for tests
78+
}
79+
}
80+
}
6681
}
6782

6883
/**

‎src/test/resources/application-test.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@ test:
1515

1616
OPENAI_API_KEY: test-fake-key
1717

18+
# Use fake RAG adapter in tests (avoids 2-minute timeout waiting for real chatbot)
19+
rag:
20+
adapter:
21+
type: fake
22+
1823
# Drivine datasource configuration for tests
1924
database:
2025
datasources:

0 commit comments

Comments
 (0)