Skip to content

Commit 215acfd

Browse files
serge-p7vserge-p7v
authored
feat(agents): Amazon Bedrock AgentCore Memory as LongTermMemory (#1855)
The second and the last part of the integration with Amazon Bedrock AgentCore Memory. closes KG-603
1 parent 1e313ba commit 215acfd

25 files changed

Lines changed: 2661 additions & 13 deletions

File tree

‎agents/agents-features/agents-features-chat-history-aws/src/jvmMain/kotlin/ai/koog/agents/features/chathistory/aws/AgentcoreChatHistoryProvider.kt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ import org.slf4j.LoggerFactory
3939
* @param totalEventsLimit Optional cap on the total number of events to fetch during [load].
4040
* @param ignoreUnsupportedValues If `true`, non-conversational message/role types are silently skipped.
4141
* If `false`, they cause an [IllegalStateException].
42-
* @throws AgentcoreMemoryException.ConfigurationException if [memoryId] is blank.
42+
* @throws AgentcoreShortTermMemoryException.ConfigurationException if [memoryId] is blank.
4343
*/
4444
public class AgentcoreChatHistoryProvider @JvmOverloads constructor(
4545
public val client: BedrockAgentCoreClient,
@@ -56,7 +56,7 @@ public class AgentcoreChatHistoryProvider @JvmOverloads constructor(
5656

5757
init {
5858
if (memoryId.isBlank()) {
59-
throw AgentcoreMemoryException.ConfigurationException("memoryId cannot be null or empty")
59+
throw AgentcoreShortTermMemoryException.ConfigurationException("memoryId cannot be null or empty")
6060
}
6161
}
6262

@@ -91,7 +91,7 @@ public class AgentcoreChatHistoryProvider @JvmOverloads constructor(
9191
val response = client.createEvent(request)
9292
logger.debug("Created an event with id ${response.event?.eventId}")
9393
} catch (e: SdkBaseException) {
94-
throw AgentcoreMemoryException.WriteException(
94+
throw AgentcoreShortTermMemoryException.WriteException(
9595
"Failed to save messages for conversation: $conversationId",
9696
e
9797
)
@@ -185,7 +185,7 @@ public class AgentcoreChatHistoryProvider @JvmOverloads constructor(
185185
allEvents.reverse()
186186
return allEvents
187187
} catch (e: SdkBaseException) {
188-
throw AgentcoreMemoryException.ReadException(
188+
throw AgentcoreShortTermMemoryException.ReadException(
189189
"Failed to fetch events for actor: $actorId, session: $sessionId",
190190
e
191191
)

agents/agents-features/agents-features-chat-history-aws/src/jvmMain/kotlin/ai/koog/agents/features/chathistory/aws/AgentcoreMemoryException.kt renamed to agents/agents-features/agents-features-chat-history-aws/src/jvmMain/kotlin/ai/koog/agents/features/chathistory/aws/AgentcoreShortTermMemoryException.kt

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ package ai.koog.agents.features.chathistory.aws
66
* Wraps AWS SDK failures so that callers of [AgentcoreChatHistoryProvider]
77
* do not need to depend on AWS-specific exception types.
88
*/
9-
public open class AgentcoreMemoryException : RuntimeException {
9+
public open class AgentcoreShortTermMemoryException : Exception {
1010
/**
1111
* Creates an exception with the given error [message].
1212
*/
@@ -21,12 +21,12 @@ public open class AgentcoreMemoryException : RuntimeException {
2121
* Thrown when a memory read operation fails.
2222
*/
2323
public class ReadException(message: String, cause: Throwable) :
24-
AgentcoreMemoryException(message, cause)
24+
AgentcoreShortTermMemoryException(message, cause)
2525

2626
/**
2727
* Thrown when a memory write operation fails.
2828
*/
29-
public class WriteException : AgentcoreMemoryException {
29+
public class WriteException : AgentcoreShortTermMemoryException {
3030
public constructor(message: String, cause: Throwable) : super(message, cause)
3131
public constructor(message: String) : super(message)
3232
}
@@ -35,5 +35,5 @@ public open class AgentcoreMemoryException : RuntimeException {
3535
* Thrown when memory configuration is invalid.
3636
*/
3737
public class ConfigurationException(message: String) :
38-
AgentcoreMemoryException(message)
38+
AgentcoreShortTermMemoryException(message)
3939
}

‎agents/agents-features/agents-features-chat-history-aws/src/jvmTest/kotlin/ai/koog/agents/features/chathistory/aws/AgentcoreChatHistoryProviderTest.kt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,10 +43,10 @@ class AgentcoreChatHistoryProviderTest {
4343

4444
@Test
4545
fun testBlankMemoryIdThrows() {
46-
assertFailsWith<AgentcoreMemoryException.ConfigurationException> {
46+
assertFailsWith<AgentcoreShortTermMemoryException.ConfigurationException> {
4747
AgentcoreChatHistoryProvider(client, memoryId = "")
4848
}
49-
assertFailsWith<AgentcoreMemoryException.ConfigurationException> {
49+
assertFailsWith<AgentcoreShortTermMemoryException.ConfigurationException> {
5050
AgentcoreChatHistoryProvider(client, memoryId = " ")
5151
}
5252
}
@@ -511,7 +511,7 @@ class AgentcoreChatHistoryProviderTest {
511511
coEvery { client.createEvent(any<CreateEventRequest>()) } throws
512512
aws.smithy.kotlin.runtime.ServiceException("AWS error")
513513

514-
assertFailsWith<AgentcoreMemoryException.WriteException> {
514+
assertFailsWith<AgentcoreShortTermMemoryException.WriteException> {
515515
provider.store("actor:session", listOf(Message.User("hi", RequestMetaInfo.Empty)))
516516
}
517517
}
@@ -523,7 +523,7 @@ class AgentcoreChatHistoryProviderTest {
523523
coEvery { client.listEvents(any<ListEventsRequest>()) } throws
524524
aws.smithy.kotlin.runtime.ServiceException("AWS error")
525525

526-
assertFailsWith<AgentcoreMemoryException.ReadException> {
526+
assertFailsWith<AgentcoreShortTermMemoryException.ReadException> {
527527
provider.load("actor:session")
528528
}
529529
}

‎agents/agents-features/agents-features-longterm-memory-aws/Module.md‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Module features-longterm-memory-aws
2+
3+
AWS Bedrock AgentCore integration for the `LongTermMemory` feature.
4+
Provides storage, retrieval strategy, prompt augmentation, and namespace resolution
5+
components that wire the AgentCore memory service into the long-term memory pipeline.
6+
7+
## Package ai.koog.agents.features.longtermmemory.aws
8+
9+
### AgentcoreSearchStorage
10+
11+
A `SearchStorage` implementation backed by AWS Bedrock AgentCore memory. Dispatches three
12+
kinds of search requests to the AgentCore API: similarity search (`RetrieveMemoryRecords`),
13+
listing (`ListMemoryRecords`), and composite search that fans out multiple subrequests
14+
concurrently, isolating individual failures so that other subrequests still return results.
15+
16+
### AgentcoreCompositeSearchStrategy
17+
18+
A `SearchStrategy` that holds a fixed list of `AgentcoreSearchSubrequest` templates and
19+
produces an `AgentcoreCompositeSearchRequest` at retrieval time. Supports mixing different
20+
AgentCore strategy types (e.g. PREFERENCE listing together with SEMANTIC similarity) and
21+
different namespace scopes (e.g. session-scoped episodes together with actor-scoped
22+
reflections) in a single composite call.
23+
24+
### AgentcoreCompositeSearchStrategy.AgentcoreSearchSubrequest
25+
26+
A sealed template for one entry inside an `AgentcoreCompositeSearchStrategy`. Two variants
27+
exist: `Similarity` (injects the per-turn query into a `RetrieveMemoryRecords` subrequest)
28+
and `Listing` (produces a query-free `ListMemoryRecords` subrequest). Templates are resolved
29+
into concrete requests at strategy creation time.
30+
31+
### AgentcoreMemoryRecord
32+
33+
A `TextDocument` that carries a single memory record retrieved from AgentCore, including its
34+
textual content, optional unique identifier, key-value metadata, and the
35+
`AgentcoreMemoryStrategy` that governs how the record is injected into the prompt.
36+
37+
### AgentcoreMemoryStrategy
38+
39+
An enum that classifies AgentCore memory strategy kinds and drives the augmentation pathway
40+
used by `AgentcorePromptAugmenter`: `SEMANTIC` and `PREFERENCE` records are folded into the
41+
system message; `EPISODES` and `REFLECTIONS` records are placed under dedicated labelled
42+
sections in the system message; `SUMMARY` records rewrite the last user message to prepend
43+
retrieved context.
44+
45+
### AgentcorePromptAugmenter
46+
47+
A `PromptAugmenter` that routes each retrieved `AgentcoreMemoryRecord` to the correct
48+
injection point based on its `AgentcoreMemoryStrategy`. SEMANTIC and PREFERENCE content is
49+
appended to the system message; EPISODES and REFLECTIONS are rendered as distinct labelled
50+
sections in the system message; SUMMARY content rewrites the last user message. A system
51+
message is created automatically when none is present.
52+
53+
### AgentcoreNamespaceScope
54+
55+
A sealed descriptor passed to `AgentcoreNamespaceResolver` to identify the memory "folder"
56+
targeted by a retrieval or ingestion operation. `Actor` scope covers strategies that are
57+
actor-scoped (PREFERENCE, SEMANTIC, REFLECTIONS); `Session` scope additionally carries a
58+
session identifier for strategies that partition memory per conversation (SUMMARY, EPISODES).
59+
60+
### AgentcoreNamespaceResolver
61+
62+
A functional interface that converts an `AgentcoreNamespaceScope` into an AgentCore namespace
63+
string. The built-in `Default` implementation reproduces AWS's documented layout
64+
(`/strategies/{strategyId}/actors/{actorId}/` and the session-scoped variant). A
65+
`template(...)` factory allows overriding the layout via placeholder strings without
66+
implementing the interface directly.

‎agents/agents-features/agents-features-longterm-memory-aws/build.gradle.kts‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import ai.koog.gradle.publish.maven.Publishing.publishToMaven
2+
import org.gradle.api.tasks.testing.Test
3+
4+
group = rootProject.group
5+
version = rootProject.version
6+
7+
plugins {
8+
id("ai.kotlin.multiplatform")
9+
alias(libs.plugins.kotlin.serialization)
10+
}
11+
12+
kotlin {
13+
sourceSets {
14+
commonMain {
15+
dependencies {
16+
api(project(":agents:agents-features:agents-features-longterm-memory"))
17+
18+
api(libs.kotlinx.serialization.json)
19+
}
20+
}
21+
22+
commonTest {
23+
dependencies {
24+
implementation(kotlin("test"))
25+
implementation(libs.kotlinx.coroutines.test)
26+
}
27+
}
28+
29+
jvmMain {
30+
dependencies {
31+
api(libs.aws.sdk.kotlin.bedrockagentcore)
32+
}
33+
}
34+
35+
jvmTest {
36+
dependencies {
37+
implementation(kotlin("test-junit5"))
38+
implementation(project(":test-utils"))
39+
implementation(libs.mockk)
40+
}
41+
}
42+
}
43+
44+
explicitApi()
45+
}
46+
47+
// Disable JUnit5 parallel execution for this module.
48+
// The AWS SDK BedrockAgentCoreClient relaxed mock is expensive to initialize via reflection,
49+
// and parallel execution causes thread contention that makes the test suite hang.
50+
tasks.withType<Test>().configureEach {
51+
systemProperty("junit.jupiter.execution.parallel.enabled", "false")
52+
}
53+
54+
publishToMaven()

‎agents/agents-features/agents-features-longterm-memory-aws/src/commonMain/kotlin/ai/koog/agents/features/longtermmemory/aws/Stub.kt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package ai.koog.agents.features.longtermmemory.aws
2+
3+
/**
4+
* This class is required for publishing iOS target when there's no commonMain set.
5+
*/
6+
@Suppress("unused")
7+
private class Stub

‎agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/AgentcoreCompositeSearchStrategy.kt‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
package ai.koog.agents.features.longtermmemory.aws
2+
3+
import ai.koog.agents.features.longtermmemory.aws.augmentation.AgentcoreMemoryStrategy
4+
import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreCompositeSearchRequest
5+
import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreListingSearchRequest
6+
import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreSearchRequest
7+
import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreSimilaritySearchRequest
8+
import ai.koog.agents.longtermmemory.retrieval.SearchStrategy
9+
import ai.koog.rag.base.storage.search.SearchRequest
10+
11+
/**
12+
* A [SearchStrategy] that produces an [AgentcoreCompositeSearchRequest] from a fixed
13+
* list of subrequest templates at retrieval time.
14+
*
15+
* subrequests can target different AgentCore strategies (different `memoryStrategyId`) and
16+
* different namespace scopes — for example, a PREFERENCE listing merged with a
17+
* SEMANTIC similarity search, or EPISODES (session-scoped) merged with
18+
* REFLECTIONS (actor-scoped).
19+
*
20+
* The outer query string produced by
21+
* [ai.koog.agents.longtermmemory.retrieval.QueryExtractor] is injected into each
22+
* similarity subrequest at [create] time. Listing subrequests do not use the query.
23+
*
24+
* Example:
25+
* ```kotlin
26+
* val strategy = AgentcoreCompositeSearchStrategy(
27+
* listOf(
28+
* AgentcoreSearchSubrequest.similarity(
29+
* memoryStrategyId = "sem-1",
30+
* namespace = AgentcoreNamespaceResolver.Default.resolve(AgentcoreNamespaceScope.Actor("sem-1", "alice")),
31+
* limit = 5,
32+
* ),
33+
* AgentcoreSearchSubrequest.listing(
34+
* memoryStrategyId = "prefs-1",
35+
* namespace = AgentcoreNamespaceResolver.Default.resolve(AgentcoreNamespaceScope.Actor("prefs-1", "alice")),
36+
* limit = 20,
37+
* ),
38+
* )
39+
* )
40+
* ```
41+
*/
42+
public class AgentcoreCompositeSearchStrategy(
43+
public val subrequests: List<AgentcoreSearchSubrequest>,
44+
) : SearchStrategy {
45+
46+
init {
47+
require(subrequests.isNotEmpty()) { "AgentcoreCompositeSearchStrategy must contain at least one subrequest" }
48+
}
49+
50+
override fun create(query: String): SearchRequest = AgentcoreCompositeSearchRequest(
51+
entries = subrequests.map { template ->
52+
AgentcoreCompositeSearchRequest.Entry(
53+
request = template.buildRequest(query),
54+
namespace = template.namespace,
55+
)
56+
},
57+
)
58+
59+
/**
60+
* Declarative template for a single subrequest of an [AgentcoreCompositeSearchStrategy].
61+
*
62+
* Templates are resolved into concrete [AgentcoreSearchRequest]s at
63+
* [AgentcoreCompositeSearchStrategy.create] time; this lets the containing strategy
64+
* inject the per-turn query string into similarity subrequests while leaving listing subrequests
65+
* query-free.
66+
*/
67+
public sealed interface AgentcoreSearchSubrequest {
68+
/**
69+
* Namespace to which the produced subrequest request will be scoped.
70+
*/
71+
public val namespace: String
72+
73+
/**
74+
* Produces the concrete subrequest [AgentcoreSearchRequest] for the given [query].
75+
*/
76+
public fun buildRequest(query: String): AgentcoreSearchRequest
77+
78+
public companion object {
79+
/**
80+
* Build a similarity-search subrequest against [memoryStrategyId] in [namespace].
81+
*
82+
* The [query] passed to [AgentcoreCompositeSearchStrategy.create] is used
83+
* as `queryText`.
84+
*/
85+
public fun similarity(
86+
strategyType: AgentcoreMemoryStrategy,
87+
memoryStrategyId: String,
88+
namespace: String,
89+
limit: Int = 10,
90+
minScore: Double? = null,
91+
filterExpression: String? = null,
92+
): AgentcoreSearchSubrequest = Similarity(
93+
strategyType = strategyType,
94+
memoryStrategyId = memoryStrategyId,
95+
namespace = namespace,
96+
limit = limit,
97+
minScore = minScore,
98+
filterExpression = filterExpression,
99+
)
100+
101+
/**
102+
* Build a listing subrequest against [memoryStrategyId] in [namespace]. The query
103+
* is ignored.
104+
*/
105+
public fun listing(
106+
strategyType: AgentcoreMemoryStrategy,
107+
memoryStrategyId: String,
108+
namespace: String,
109+
limit: Int = 10,
110+
): AgentcoreSearchSubrequest = Listing(
111+
strategyType = strategyType,
112+
memoryStrategyId = memoryStrategyId,
113+
namespace = namespace,
114+
limit = limit,
115+
)
116+
}
117+
118+
private data class Similarity(
119+
val strategyType: AgentcoreMemoryStrategy,
120+
val memoryStrategyId: String,
121+
override val namespace: String,
122+
val limit: Int,
123+
val minScore: Double?,
124+
val filterExpression: String?,
125+
) : AgentcoreSearchSubrequest {
126+
override fun buildRequest(query: String): AgentcoreSearchRequest =
127+
AgentcoreSimilaritySearchRequest(
128+
strategyType = strategyType,
129+
memoryStrategyId = memoryStrategyId,
130+
queryText = query,
131+
limit = limit,
132+
minScore = minScore,
133+
filterExpression = filterExpression,
134+
)
135+
}
136+
137+
private data class Listing(
138+
val strategyType: AgentcoreMemoryStrategy,
139+
val memoryStrategyId: String,
140+
override val namespace: String,
141+
val limit: Int,
142+
) : AgentcoreSearchSubrequest {
143+
override fun buildRequest(query: String): AgentcoreSearchRequest =
144+
AgentcoreListingSearchRequest(
145+
strategyType = strategyType,
146+
memoryStrategyId = memoryStrategyId,
147+
limit = limit,
148+
)
149+
}
150+
}
151+
}

0 commit comments

Comments
 (0)