Skip to content

Commit 1dbf520

Browse files
jasperbluesjasperblues
committed
Prep integration into Embabel.
1 parent 2cd1996 commit 1dbf520

File tree

10 files changed

+910
-35
lines changed

10 files changed

+910
-35
lines changed

‎pom.xml‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
<java.version>21</java.version>
2222
<drivine.version>0.0.22</drivine.version>
2323
<kotlin.version>2.0.20</kotlin.version>
24+
<embabel-agent.version>0.3.3-SNAPSHOT</embabel-agent.version>
25+
<embabel-common.version>0.1.9-SNAPSHOT</embabel-common.version>
2426
</properties>
2527

2628
<dependencies>
@@ -58,6 +60,24 @@
5860
<groupId>org.jetbrains.kotlin</groupId>
5961
<artifactId>kotlin-reflect</artifactId>
6062
</dependency>
63+
<dependency>
64+
<groupId>org.jetbrains.kotlinx</groupId>
65+
<artifactId>kotlinx-coroutines-core</artifactId>
66+
</dependency>
67+
68+
<!-- Embabel Common Core (for Timestamped interface) -->
69+
<dependency>
70+
<groupId>com.embabel.common</groupId>
71+
<artifactId>embabel-common-core</artifactId>
72+
<version>${embabel-common.version}</version>
73+
</dependency>
74+
75+
<!-- Embabel Agent API -->
76+
<dependency>
77+
<groupId>com.embabel.agent</groupId>
78+
<artifactId>embabel-agent-api</artifactId>
79+
<version>${embabel-agent.version}</version>
80+
</dependency>
6181

6282
<!-- Logging -->
6383
<dependency>

‎src/main/kotlin/com/embabel/chat/store/adapter/MessageMapper.kt‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
/*
2+
* Copyright 2024-2026 Embabel Pty Ltd.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.chat.store.adapter
17+
18+
import com.embabel.chat.AssistantMessage
19+
import com.embabel.chat.Message
20+
import com.embabel.chat.Role
21+
import com.embabel.chat.SystemMessage
22+
import com.embabel.chat.UserMessage
23+
import com.embabel.chat.store.model.MessageData
24+
import com.embabel.chat.store.model.StoredMessage
25+
import java.util.UUID
26+
27+
/**
28+
* Maps between embabel-agent Message types and embabel-chat-store MessageData.
29+
*
30+
* This mapper handles conversion in both directions:
31+
* - [Message] -> [MessageData] for persistence
32+
* - [StoredMessage] -> [Message] for retrieval
33+
*
34+
* Note: Multimodal content (images) is not currently persisted - only text content
35+
* is stored. The metadata field can be used to store additional information if needed.
36+
*
37+
* ## Usage
38+
*
39+
* You can use either the object methods or the extension functions:
40+
* ```kotlin
41+
* // Object methods
42+
* val messageData = MessageMapper.toMessageData(message)
43+
* val message = MessageMapper.toMessage(storedMessage)
44+
*
45+
* // Extension functions (more ergonomic)
46+
* val messageData = message.toMessageData()
47+
* val message = storedMessage.toMessage()
48+
* ```
49+
*/
50+
object MessageMapper {
51+
52+
/**
53+
* Convert an embabel-agent [Message] to [MessageData] for persistence.
54+
*
55+
* @param message the message to convert
56+
* @param messageId optional ID for the message (generated if not provided)
57+
* @return MessageData ready for persistence
58+
*/
59+
@JvmStatic
60+
@JvmOverloads
61+
fun toMessageData(
62+
message: Message,
63+
messageId: String = UUID.randomUUID().toString()
64+
): MessageData {
65+
return MessageData(
66+
messageId = messageId,
67+
role = message.role,
68+
content = message.content,
69+
createdAt = message.timestamp,
70+
metadata = buildMetadata(message)
71+
)
72+
}
73+
74+
/**
75+
* Convert a [StoredMessage] back to an embabel-agent [Message].
76+
*
77+
* @param storedMessage the stored message to convert
78+
* @return the appropriate Message subtype based on role
79+
*/
80+
@JvmStatic
81+
fun toMessage(storedMessage: StoredMessage): Message {
82+
return toMessage(storedMessage.message, storedMessage.author?.displayName)
83+
}
84+
85+
/**
86+
* Convert [MessageData] back to an embabel-agent [Message].
87+
*
88+
* @param messageData the message data to convert
89+
* @param authorName optional author name for the message
90+
* @return the appropriate Message subtype based on role
91+
*/
92+
@JvmStatic
93+
@JvmOverloads
94+
fun toMessage(messageData: MessageData, authorName: String? = null): Message {
95+
return when (messageData.role) {
96+
Role.USER -> UserMessage(
97+
content = messageData.content,
98+
name = authorName,
99+
timestamp = messageData.createdAt
100+
)
101+
Role.ASSISTANT -> AssistantMessage(
102+
content = messageData.content,
103+
name = authorName,
104+
timestamp = messageData.createdAt
105+
)
106+
Role.SYSTEM -> SystemMessage(
107+
content = messageData.content,
108+
timestamp = messageData.createdAt
109+
)
110+
}
111+
}
112+
113+
/**
114+
* Build metadata map from message properties.
115+
* Can be extended to store additional information.
116+
*/
117+
private fun buildMetadata(message: Message): Map<String, Any>? {
118+
val metadata = mutableMapOf<String, Any>()
119+
120+
// Store original sender name if present
121+
message.name?.let { metadata["senderName"] = it }
122+
123+
// Flag multimodal messages (content not fully preserved)
124+
if (message.isMultimodal) {
125+
metadata["isMultimodal"] = true
126+
metadata["imageCount"] = message.imageParts.size
127+
}
128+
129+
return metadata.ifEmpty { null }
130+
}
131+
}
132+
133+
/**
134+
* Extension function to convert a [Message] to [MessageData].
135+
*/
136+
fun Message.toMessageData(messageId: String = UUID.randomUUID().toString()): MessageData =
137+
MessageMapper.toMessageData(this, messageId)
138+
139+
/**
140+
* Extension function to convert a [StoredMessage] to a [Message].
141+
*/
142+
fun StoredMessage.toMessage(): Message =
143+
MessageMapper.toMessage(this)
144+
145+
/**
146+
* Extension function to convert [MessageData] to a [Message].
147+
*/
148+
fun MessageData.toMessage(authorName: String? = null): Message =
149+
MessageMapper.toMessage(this, authorName)

‎src/main/kotlin/com/embabel/chat/store/adapter/StoredConversation.kt‎

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
1+
/*
2+
* Copyright 2024-2026 Embabel Pty Ltd.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.embabel.chat.store.adapter
17+
18+
import com.embabel.chat.AssetTracker
19+
import com.embabel.chat.Conversation
20+
import com.embabel.chat.Message
21+
import com.embabel.chat.MessageAuthor
22+
import com.embabel.chat.Role
23+
import com.embabel.chat.event.MessageEvent
24+
import com.embabel.chat.store.model.SessionUser
25+
import com.embabel.chat.store.repository.ChatSessionRepository
26+
import com.embabel.chat.support.InMemoryAssetTracker
27+
import com.embabel.chat.support.InMemoryConversation
28+
import kotlinx.coroutines.CoroutineScope
29+
import kotlinx.coroutines.Dispatchers
30+
import kotlinx.coroutines.SupervisorJob
31+
import kotlinx.coroutines.launch
32+
import org.slf4j.LoggerFactory
33+
import org.springframework.context.ApplicationEventPublisher
34+
import java.util.UUID
35+
36+
/**
37+
* A [Conversation] implementation that persists messages to Neo4j via [ChatSessionRepository].
38+
*
39+
* This adapter bridges embabel-agent's Conversation interface with embabel-chat-store's
40+
* persistence layer. Messages added via [addMessage] are persisted **asynchronously**:
41+
*
42+
* 1. `addMessage()` returns immediately (non-blocking)
43+
* 2. Message is converted to [MessageData] and persisted in background
44+
* 3. On success: [MessageEvent] with status PERSISTED is published
45+
* 4. On failure: [MessageEvent] with status PERSISTENCE_FAILED is published
46+
*
47+
* ## Auto Title Generation
48+
*
49+
* If a [TitleGenerator] is provided, the session title is automatically generated
50+
* from the first message (if the session doesn't already have a title).
51+
*
52+
* ## Usage
53+
*
54+
* ```kotlin
55+
* val conversation = StoredConversation(
56+
* id = "session-123",
57+
* repository = chatSessionRepository,
58+
* eventPublisher = applicationEventPublisher,
59+
* sessionUser = currentUser,
60+
* titleGenerator = LlmTitleGenerator { prompt -> llm.generate(prompt) }
61+
* )
62+
*
63+
* // Returns immediately - persistence happens async
64+
* conversation.addMessage(UserMessage("Hello!"))
65+
*
66+
* // Subscribe to events for persistence confirmation
67+
* @EventListener
68+
* fun onPersisted(event: MessageEvent) {
69+
* if (event.status == MessageStatus.PERSISTED) { ... }
70+
* }
71+
* ```
72+
*
73+
* @param id the chat session ID (must already exist in the repository)
74+
* @param repository the repository for persistence operations
75+
* @param eventPublisher Spring's event publisher for broadcasting events
76+
* @param sessionUser optional user for attributing user messages
77+
* @param titleGenerator optional generator for auto-generating session title from first message
78+
* @param assetTracker tracker for conversation assets (defaults to in-memory)
79+
* @param scope coroutine scope for async operations (defaults to IO dispatcher with SupervisorJob)
80+
*/
81+
class StoredConversation(
82+
override val id: String,
83+
private val repository: ChatSessionRepository,
84+
private val eventPublisher: ApplicationEventPublisher? = null,
85+
private val sessionUser: SessionUser? = null,
86+
private val titleGenerator: TitleGenerator? = null,
87+
override val assetTracker: AssetTracker = InMemoryAssetTracker(),
88+
private val scope: CoroutineScope = CoroutineScope(Dispatchers.IO + SupervisorJob()),
89+
) : Conversation {
90+
91+
private val logger = LoggerFactory.getLogger(StoredConversation::class.java)
92+
93+
/**
94+
* Messages loaded from the repository.
95+
* Lazily refreshed on access.
96+
*/
97+
override val messages: List<Message>
98+
get() = repository.getMessages(id).map { it.toMessage() }
99+
100+
/**
101+
* Add a message to the conversation asynchronously using the default session user.
102+
*
103+
* For USER role messages, the default [sessionUser] is used as the author.
104+
* For other roles (ASSISTANT, SYSTEM), no author is attributed.
105+
*
106+
* This method returns immediately. The message is persisted in the background,
107+
* and events are published on success or failure.
108+
*
109+
* @param message the message to add
110+
* @return the message (returned immediately, before persistence completes)
111+
*/
112+
override fun addMessage(message: Message): Message {
113+
val author = when (message.role) {
114+
Role.USER -> sessionUser
115+
else -> null
116+
}
117+
return addMessageInternal(message, author)
118+
}
119+
120+
/**
121+
* Add a message with explicit author attribution.
122+
*
123+
* Use this for group chats or when the author differs per message.
124+
* The provided author must be a [SessionUser] for persistence.
125+
*
126+
* @param message the message to add
127+
* @param author the author of this message (must be SessionUser for persistence, null for system/assistant)
128+
* @return the message (returned immediately, before persistence completes)
129+
* @throws IllegalArgumentException if author is not null and not a SessionUser
130+
*/
131+
override fun addMessageFrom(message: Message, author: MessageAuthor?): Message {
132+
val sessionUserAuthor = when {
133+
author == null -> null
134+
author is SessionUser -> author
135+
else -> throw IllegalArgumentException(
136+
"Author must be a SessionUser for persistence. Got: ${author::class.simpleName}"
137+
)
138+
}
139+
return addMessageInternal(message, sessionUserAuthor)
140+
}
141+
142+
private fun addMessageInternal(message: Message, author: SessionUser?): Message {
143+
val messageData = message.toMessageData(
144+
messageId = UUID.randomUUID().toString()
145+
)
146+
val isFirstMessage = messages.isEmpty()
147+
148+
scope.launch {
149+
try {
150+
val updatedSession = repository.addMessage(id, messageData, author)
151+
val persistedMessage = updatedSession.messages.last().toMessage()
152+
153+
// Generate title from first message if no title exists
154+
if (isFirstMessage && titleGenerator != null) {
155+
val session = repository.findBySessionId(id).orElse(null)
156+
if (session?.session?.title.isNullOrBlank()) {
157+
try {
158+
val title = titleGenerator.generate(message)
159+
repository.updateSessionTitle(id, title)
160+
logger.debug("Generated title '{}' for session {}", title, id)
161+
} catch (e: Exception) {
162+
logger.warn("Failed to generate title for session {}: {}", id, e.message)
163+
}
164+
}
165+
}
166+
167+
eventPublisher?.publishEvent(
168+
MessageEvent.persisted(id, persistedMessage)
169+
)
170+
logger.debug("Message {} persisted to session {}", messageData.messageId, id)
171+
} catch (e: Exception) {
172+
logger.error("Failed to persist message to session {}: {}", id, e.message, e)
173+
eventPublisher?.publishEvent(
174+
MessageEvent.persistenceFailed(
175+
conversationId = id,
176+
content = message.content,
177+
role = message.role,
178+
error = e
179+
)
180+
)
181+
}
182+
}
183+
184+
return message
185+
}
186+
187+
override fun persistent(): Boolean = true
188+
189+
/**
190+
* Create a non-persistent view of the last n messages.
191+
*
192+
* Returns an [InMemoryConversation] since the view doesn't need persistence.
193+
*/
194+
override fun last(n: Int): Conversation {
195+
return InMemoryConversation(
196+
messages = messages.takeLast(n),
197+
id = id,
198+
persistent = false,
199+
assets = assetTracker
200+
)
201+
}
202+
}

0 commit comments

Comments
 (0)