feat(agents): WIP Amazon Bedrock AgentCore Memory as LongTermMemory

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

@@ -0,0 +1,3 @@
+# Module features-longterm-memory-aws
+
+[//]: # (TODO)

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

@@ -0,0 +1,54 @@
+import ai.koog.gradle.publish.maven.Publishing.publishToMaven
+import org.gradle.api.tasks.testing.Test
+
+group = rootProject.group
+version = rootProject.version
+
+plugins {
+    id("ai.kotlin.multiplatform")
+    alias(libs.plugins.kotlin.serialization)
+}
+
+kotlin {
+    sourceSets {
+        commonMain {
+            dependencies {
+                api(project(":agents:agents-features:agents-features-longterm-memory"))
+
+                api(libs.kotlinx.serialization.json)
+            }
+        }
+
+        commonTest {
+            dependencies {
+                implementation(kotlin("test"))
+                implementation(libs.kotlinx.coroutines.test)
+            }
+        }
+
+        jvmMain {
+            dependencies {
+                api(libs.aws.sdk.kotlin.bedrockagentcore)
+            }
+        }
+
+        jvmTest {
+            dependencies {
+                implementation(kotlin("test-junit5"))
+                implementation(project(":test-utils"))
+                implementation(libs.mockk)
+            }
+        }
+    }
+
+    explicitApi()
+}
+
+// Disable JUnit5 parallel execution for this module.
+// The AWS SDK BedrockAgentCoreClient relaxed mock is expensive to initialize via reflection,
+// and parallel execution causes thread contention that makes the test suite hang.
+tasks.withType<Test>().configureEach {
+    systemProperty("junit.jupiter.execution.parallel.enabled", "false")
+}
+
+publishToMaven()

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

@@ -0,0 +1,7 @@
+package ai.koog.agents.features.longtermmemory.aws
+
+/**
+ * This class is required for publishing iOS target when there's no commonMain set.
+ */
+@Suppress("unused")
+private class Stub

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

@@ -0,0 +1,32 @@
+package ai.koog.agents.features.longtermmemory.aws
+
+/**
+ * AgentCore Memory built-in strategies
+ */
+public enum class AgentcoreLongTermStrategyType {
+    /**
+     * Designed to identify and extract key pieces of factual information and contextual knowledge from conversational data.
+     * Default namespace: /strategies/{memoryStrategyId}/actors/{actorId}/
+     */
+    SEMANTIC,
+
+    /**
+     * Designed to automatically identify and extract user preferences, choices, and styles from conversational data.
+     * Default namespace: /strategies/{memoryStrategyId}/actors/{actorId}/
+     */
+    USER_PREFERENCE,
+
+    /**
+     * Responsible for generating condensed, real-time summaries of conversations within a single session.
+     * Default namespace: /strategies/{memoryStrategyId}/actors/{actorId}/sessions/{sessionId}/
+     */
+    SUMMARY,
+
+    /**
+     * Captures meaningful slices of user and system interaction
+     * so applications can recall context in a way that feels focused and relevant.
+     * Default namespace for extracted memories: /strategies/{memoryStrategyId}/actors/{actorId}/sessions/{sessionId}/
+     * Default namespace for reflection: /strategies/{memoryStrategyId}/actors/{actorId}/
+     */
+    EPISODIC
+}

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

@@ -0,0 +1,31 @@
+package ai.koog.agents.features.longtermmemory.aws
+
+/**
+ * Base exception for AgentCore Memory operations.
+ *
+ * Wraps AWS SDK failures so that callers of [AgentcoreSearchStorage]
+ * do not need to depend on AWS-specific exception types.
+ */
+public open class AgentcoreMemoryException : RuntimeException {
+    /**
+     * Creates an exception with the given error [message].
+     */
+    public constructor(message: String) : super(message)
+
+    /**
+     * Creates an exception with the given error [message] and [cause].
+     */
+    public constructor(message: String, cause: Throwable) : super(message, cause)
+
+    /**
+     * Thrown when a memory retrieve operation fails.
+     */
+    public class RetrieveException(message: String, cause: Throwable) :
+        AgentcoreMemoryException(message, cause)
+
+    /**
+     * Thrown when memory configuration is invalid.
+     */
+    public class ConfigurationException(message: String) :
+        AgentcoreMemoryException(message)
+}

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

@@ -0,0 +1,34 @@
+package ai.koog.agents.features.longtermmemory.aws
+
+import ai.koog.agents.longtermmemory.model.MemoryRecord
+import ai.koog.rag.base.TextDocument
+import ai.koog.rag.base.storage.search.Score
+import ai.koog.rag.base.storage.search.ScoreMetric
+import ai.koog.rag.base.storage.search.SearchResult
+import aws.sdk.kotlin.services.bedrockagentcore.model.MemoryRecordSummary
+import aws.sdk.kotlin.services.bedrockagentcore.model.MetadataValue
+
+/**
+ * Converts AWS Bedrock AgentCore memory record types to the framework's internal representations.
+ *
+ * Provides utilities to transform [MemoryRecordSummary] objects returned by the Bedrock AgentCore API
+ * into [SearchResult] instances wrapping [TextDocument], including score and metadata mapping.
+ */
+public object AgentcoreMemoryRecordConverter {
+
+    internal fun memoryRecordSummaryToSearchResult(memoryRecordSummary: MemoryRecordSummary): SearchResult<TextDocument> {
+        return SearchResult(
+            MemoryRecord(
+                memoryRecordSummary.content?.asTextOrNull() ?: "",
+                memoryRecordSummary.memoryRecordId,
+                mapMetadata(memoryRecordSummary.metadata)
+            ),
+            Score(memoryRecordSummary.score ?: 0.0, ScoreMetric.COSINE_SIMILARITY)
+        )
+    }
+
+    private fun mapMetadata(agentcoreMetadata: Map<String, MetadataValue>?): Map<String, Any> {
+        if (agentcoreMetadata.isNullOrEmpty()) return emptyMap()
+        return agentcoreMetadata.mapValues { (_, v) -> v.asStringValueOrNull() ?: "" }
+    }
+}

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

@@ -0,0 +1,39 @@
+package ai.koog.agents.features.longtermmemory.aws
+
+/**
+ * Builds AgentCore memory namespace strings.
+ *
+ * Two scopes are supported:
+ * - [actorScoped] — `/strategies/{memoryStrategyId}/actors/{actorId}/`
+ * - [sessionScoped] — `/strategies/{memoryStrategyId}/actors/{actorId}/sessions/{sessionId}/`
+ *
+ * Usage:
+ * ```kotlin
+ * val actorNs = AgentcoreNamespace.actorScoped("myStrategy", "alice")
+ * // "/strategies/myStrategy/actors/alice/"
+ *
+ * val sessionNs = AgentcoreNamespace.sessionScoped("myStrategy", "alice", "session-1")
+ * // "/strategies/myStrategy/actors/alice/sessions/session-1/"
+ * ```
+ */
+public object AgentcoreNamespace {
+
+    /**
+     * Returns an actor-scoped namespace: `/strategies/{memoryStrategyId}/actors/{actorId}/`
+     */
+    public fun actorScoped(memoryStrategyId: String, actorId: String): String {
+        require(memoryStrategyId.isNotBlank()) { "memoryStrategyId must not be blank" }
+        require(actorId.isNotBlank()) { "actorId must not be blank" }
+        return "/strategies/$memoryStrategyId/actors/$actorId/"
+    }
+
+    /**
+     * Returns a session-scoped namespace: `/strategies/{memoryStrategyId}/actors/{actorId}/sessions/{sessionId}/`
+     */
+    public fun sessionScoped(memoryStrategyId: String, actorId: String, sessionId: String): String {
+        require(memoryStrategyId.isNotBlank()) { "memoryStrategyId must not be blank" }
+        require(actorId.isNotBlank()) { "actorId must not be blank" }
+        require(sessionId.isNotBlank()) { "sessionId must not be blank" }
+        return "/strategies/$memoryStrategyId/actors/$actorId/sessions/$sessionId/"
+    }
+}

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

@@ -0,0 +1,97 @@
+package ai.koog.agents.features.longtermmemory.aws
+
+import ai.koog.agents.features.longtermmemory.aws.request.ListingSearchRequest
+import ai.koog.rag.base.TextDocument
+import ai.koog.rag.base.storage.SearchStorage
+import ai.koog.rag.base.storage.search.SearchRequest
+import ai.koog.rag.base.storage.search.SearchResult
+import ai.koog.rag.base.storage.search.SimilaritySearchRequest
+import aws.sdk.kotlin.services.bedrockagentcore.BedrockAgentCoreClient
+import aws.sdk.kotlin.services.bedrockagentcore.model.ListMemoryRecordsRequest
+import aws.sdk.kotlin.services.bedrockagentcore.model.MemoryRecordSummary
+import aws.sdk.kotlin.services.bedrockagentcore.model.RetrieveMemoryRecordsRequest
+import aws.sdk.kotlin.services.bedrockagentcore.model.SearchCriteria
+import org.slf4j.LoggerFactory
+
+/**
+ * A [SearchStorage] implementation backed by AWS Bedrock AgentCore memory.
+ *
+ * Supports two search strategies:
+ * - [SimilaritySearchRequest]: performs semantic similarity search using the Bedrock `RetrieveMemoryRecords` API.
+ * - [ListingSearchRequest]: lists memory records using the Bedrock `ListMemoryRecords` API.
+ *
+ * @param client the [BedrockAgentCoreClient] used to communicate with the AWS Bedrock AgentCore service.
+ * @param agentcoreMemoryId the identifier of the AgentCore memory store to search.
+ * @param agentcoreMemoryStrategyId the identifier of the memory strategy applied during search and listing.
+ */
+public class AgentcoreSearchStorage(
+    public val client: BedrockAgentCoreClient,
+    public val agentcoreMemoryId: String,
+    public val agentcoreMemoryStrategyId: String, // TODO: should be passed in AgentcoreSimilaritySearchRequest and AgentcoreListingSearchRequest
+) : SearchStorage<TextDocument, SearchRequest> {
+
+    private val logger = LoggerFactory.getLogger("AgentcoreSearchStorage")
+
+    override suspend fun search(
+        request: SearchRequest,
+        namespace: String?
+    ): List<SearchResult<TextDocument>> {
+        try {
+            val summaries = when (request) {
+                is SimilaritySearchRequest -> {
+                    retrieveMemoryRecords(request.limit, request.queryText, namespace)
+                        .map { AgentcoreMemoryRecordConverter.memoryRecordSummaryToSearchResult(it) }
+                        .filter { it.score.value >= (request.minScore ?: 0.0) }
+                }
+                is ListingSearchRequest -> {
+                    listMemoryRecords(request.limit, namespace)
+                        .map { AgentcoreMemoryRecordConverter.memoryRecordSummaryToSearchResult(it) }
+                }
+                else -> {
+                    throw IllegalArgumentException("Unsupported search request type: ${request::class.simpleName}")
+                }
+            }
+
+            return summaries
+        } catch (e: Exception) {
+            throw AgentcoreMemoryException.RetrieveException(
+                "Failed to search memory records: memoryId=$agentcoreMemoryId, namespace=$namespace",
+                e
+            )
+        }
+    }
+
+    private suspend fun retrieveMemoryRecords(topK: Int, searchQuery: String?, namespace: String?): List<MemoryRecordSummary> {
+        val request = RetrieveMemoryRecordsRequest {
+            memoryId = agentcoreMemoryId
+            this.namespace = namespace
+            searchCriteria = SearchCriteria {
+                memoryStrategyId = agentcoreMemoryStrategyId
+//                metadataFilters = null // TODO: pass the filterExpression to metadataFilters
+                this.searchQuery = searchQuery
+                this.topK = topK
+            }
+        }
+
+        logger.debug("Retrieving memory records for searchQuery $searchQuery and namespace $namespace")
+        val memoryRecordSummaries = client.retrieveMemoryRecords(request).memoryRecordSummaries
+        logger.debug("Retrieved ${memoryRecordSummaries.size} memory records")
+
+        return memoryRecordSummaries
+    }
+
+    private suspend fun listMemoryRecords(maxResults: Int?, namespace: String?): List<MemoryRecordSummary> {
+        val request = ListMemoryRecordsRequest {
+            memoryId = agentcoreMemoryId
+            memoryStrategyId = agentcoreMemoryStrategyId
+            this.namespace = namespace
+            this.maxResults = maxResults
+        }
+
+        logger.debug("Listing memory records for namespace $namespace")
+        val memoryRecordSummaries = client.listMemoryRecords(request).memoryRecordSummaries
+        logger.debug("Listed ${memoryRecordSummaries.size} memory records")
+
+        return memoryRecordSummaries
+    }
+}

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

@@ -0,0 +1,27 @@
+package ai.koog.agents.features.longtermmemory.aws
+
+import ai.koog.agents.longtermmemory.retrieval.SearchStrategy
+import ai.koog.rag.base.storage.search.SimilaritySearchRequest
+
+public class AgentcoreSimilaritySearchStrategy(
+    public val strategyType: AgentcoreLongTermStrategyType,
+    public val limit: Int = 10,
+    public val offset: Int = 0,
+    public val minScore: Double? = null,
+) : SearchStrategy {
+    override fun create(query: String): SimilaritySearchRequest = when (strategyType) {
+//        AgentcoreLongTermStrategyType.USER_PREFERENCE -> ListingSearchRequest(
+//            limit = limit,
+//            offset = offset,
+//        )
+        AgentcoreLongTermStrategyType.USER_PREFERENCE, // fixme: it won't work with SimilaritySearchRequest
+        AgentcoreLongTermStrategyType.SEMANTIC,
+        AgentcoreLongTermStrategyType.SUMMARY,
+        AgentcoreLongTermStrategyType.EPISODIC -> SimilaritySearchRequest(
+            queryText = query,
+            limit = limit,
+            offset = offset,
+            minScore = minScore,
+        )
+    }
+}

agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/request/AgentcoreListingSearchRequest.kt

@@ -0,0 +1,12 @@
+package ai.koog.agents.features.longtermmemory.aws.request
+
+import ai.koog.rag.base.storage.search.SearchRequest
+
+/**
+ * TODO: unused, keep it
+ */
+public data class AgentcoreListingSearchRequest(
+    override val limit: Int = 10,
+    override val offset: Int = 0,
+    val memoryStrategyId: String? = null,
+) : SearchRequest