diff --git a/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/AgentcoreCompositeSearchStrategy.kt b/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/AgentcoreCompositeSearchStrategy.kt index 6c533515f3..5e4ed6d508 100644 --- a/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/AgentcoreCompositeSearchStrategy.kt +++ b/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/AgentcoreCompositeSearchStrategy.kt @@ -5,7 +5,7 @@ import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreCompositeSear import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreListingSearchRequest import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreSearchRequest import ai.koog.agents.features.longtermmemory.aws.request.AgentcoreSimilaritySearchRequest -import ai.koog.agents.longtermmemory.retrieval.SearchStrategy +import ai.koog.agents.longtermmemory.retrieval.search.SearchStrategy import ai.koog.rag.base.storage.search.SearchRequest /** @@ -18,7 +18,7 @@ import ai.koog.rag.base.storage.search.SearchRequest * REFLECTIONS (actor-scoped). * * The outer query string produced by - * [ai.koog.agents.longtermmemory.retrieval.QueryExtractor] is injected into each + * [ai.koog.agents.longtermmemory.retrieval.search.SearchQueryProvider] is injected into each * similarity subrequest at [create] time. Listing subrequests do not use the query. * * Example: diff --git a/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreLongTermMemoryDsl.kt b/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreLongTermMemoryDsl.kt index 6e97c46e1e..2a6fe0ba76 100644 --- a/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreLongTermMemoryDsl.kt +++ b/agents/agents-features/agents-features-longterm-memory-aws/src/jvmMain/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreLongTermMemoryDsl.kt @@ -1,6 +1,5 @@ package ai.koog.agents.features.longtermmemory.aws.dsl -import ai.koog.agents.core.annotation.ExperimentalAgentsApi import ai.koog.agents.features.longtermmemory.aws.AgentcoreCompositeSearchStrategy import ai.koog.agents.features.longtermmemory.aws.AgentcoreCompositeSearchStrategy.AgentcoreSearchSubrequest import ai.koog.agents.features.longtermmemory.aws.AgentcoreNamespaceResolver @@ -11,6 +10,7 @@ import ai.koog.agents.features.longtermmemory.aws.augmentation.AgentcorePromptAu import ai.koog.agents.longtermmemory.feature.LongTermMemory import ai.koog.agents.longtermmemory.retrieval.augmentation.PromptAugmenter import ai.koog.agents.longtermmemory.retrieval.augmentation.SystemPromptAugmenter +import ai.koog.agents.longtermmemory.retrieval.search.SearchStrategy import aws.sdk.kotlin.services.bedrockagentcore.BedrockAgentCoreClient /** @@ -48,7 +48,6 @@ public annotation class AgentcoreLtmDsl * @param memoryId the AgentCore memory store identifier. * @param block DSL block appending one or more retrieval subrequests. */ -@ExperimentalAgentsApi public fun LongTermMemory.RetrievalSettingsBuilder.agentcore( client: BedrockAgentCoreClient, memoryId: String, @@ -283,7 +282,7 @@ public class AgentcoreRetrievalBuilder internal constructor( internal data class Configured( val storage: AgentcoreSearchStorage, - val searchStrategy: ai.koog.agents.longtermmemory.retrieval.SearchStrategy, + val searchStrategy: SearchStrategy, val namespace: String?, val promptAugmenter: PromptAugmenter, ) diff --git a/agents/agents-features/agents-features-longterm-memory-aws/src/jvmTest/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreRetrievalDslTest.kt b/agents/agents-features/agents-features-longterm-memory-aws/src/jvmTest/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreRetrievalDslTest.kt index a9d50b75a3..ffa54297c1 100644 --- a/agents/agents-features/agents-features-longterm-memory-aws/src/jvmTest/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreRetrievalDslTest.kt +++ b/agents/agents-features/agents-features-longterm-memory-aws/src/jvmTest/kotlin/ai/koog/agents/features/longtermmemory/aws/dsl/AgentcoreRetrievalDslTest.kt @@ -1,6 +1,5 @@ package ai.koog.agents.features.longtermmemory.aws.dsl -import ai.koog.agents.core.annotation.ExperimentalAgentsApi import ai.koog.agents.features.longtermmemory.aws.AgentcoreCompositeSearchStrategy import ai.koog.agents.features.longtermmemory.aws.AgentcoreNamespaceResolver import ai.koog.agents.features.longtermmemory.aws.AgentcoreNamespaceScope @@ -30,7 +29,6 @@ import kotlin.test.assertTrue * DSL directly against a [LongTermMemory.RetrievalSettingsBuilder] and inspect the * resulting strategy / storage / augmenter / namespace without talking to AWS. */ -@OptIn(ExperimentalAgentsApi::class) class AgentcoreRetrievalDslTest { private val client = mockk(relaxed = true) diff --git a/agents/agents-features/agents-features-longterm-memory/Module.md b/agents/agents-features/agents-features-longterm-memory/Module.md index 5918c3c9e3..2c6a613a40 100644 --- a/agents/agents-features/agents-features-longterm-memory/Module.md +++ b/agents/agents-features/agents-features-longterm-memory/Module.md @@ -9,8 +9,8 @@ The agents-features-longterm-memory module adds long-term memory capabilities to - **Retrieval (RAG)**: Searches a memory store for context relevant to the user's query and augments the LLM prompt before each call - **Ingestion**: Extracts and persists conversation messages into a memory store for future retrieval - **Flexible storage**: Plug any backend via `SearchStorage` / `WriteStorage` interfaces from the `rag-base` module; an in-memory `InMemoryRecordStorage` is included for testing -- **Configurable timing**: Ingest per-LLM-call or on agent completion -- **Prompt augmentation modes**: System prompt or user prompt or custom implementation +- **On-completion ingestion**: The accumulated session prompt/history is ingested once when the agent run completes +- **Prompt augmentation modes**: System prompt, user prompt, or custom implementation ### Key Components @@ -19,10 +19,11 @@ The agents-features-longterm-memory module adds long-term memory capabilities to | [`LongTermMemory`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemory.kt) | Agent feature with DSL config for retrieval & ingestion | | [`SearchStorage`](../../../../../../rag/rag-base/src/commonMain/kotlin/ai/koog/rag/base/storage/SearchStorage.kt) | Interface for searching memory records (defined in `rag-base`) | | [`WriteStorage`](../../../../../../rag/rag-base/src/commonMain/kotlin/ai/koog/rag/base/storage/WriteStorage.kt) | Interface for adding memory records (defined in `rag-base`) | -| [`SearchStrategy`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/SearchStrategy.kt) | Converts user query into a `SearchRequest`; `SimilaritySearchStrategy` is the default implementation | -| [`QueryExtractor`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/QueryExtractor.kt) | Extracts the search query string from a `Prompt` for retrieval | -| [`LastUserMessageQueryExtractor`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/QueryExtractor.kt) | Default `QueryExtractor` that uses the last user message content | -| [`ExtractionStrategy`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/ExtractionStrategy.kt) | Transforms messages into `TextDocument`s for storage | +| [`SearchStrategy`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchStrategy.kt) | Converts user query into a `SearchRequest`; `SimilaritySearchStrategy` is the default implementation | +| [`SearchQueryProvider`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchQueryProvider.kt) | Provides the search query string from a `Prompt` for retrieval | +| [`LastUserMessageQueryProvider`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/LastUserMessageQueryProvider.kt) | Default `SearchQueryProvider` that uses the last user message content | +| [`DocumentExtractor`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/DocumentExtractor.kt) | Transforms messages into `TextDocument`s for storage | +| [`MessagePassingDocumentExtractor`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/MessagePassingDocumentExtractor.kt) | Default `DocumentExtractor`; filters messages by role (`User`/`Assistant` by default) | | [`PromptAugmenter`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/PromptAugmenter.kt) | Interface for augmenting prompts with relevant context | | [`SystemPromptAugmenter`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/SystemPromptAugmenter.kt) | Inserts retrieved context as a system message | | [`UserPromptAugmenter`](src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/UserPromptAugmenter.kt) | Inserts retrieved context as a user message | diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/FailurePolicy.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/FailurePolicy.kt new file mode 100644 index 0000000000..aae71850a8 --- /dev/null +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/FailurePolicy.kt @@ -0,0 +1,48 @@ +package ai.koog.agents.longtermmemory.feature + +/** + * Policy controlling how failures from the underlying memory storage are handled + * by the [LongTermMemory] feature. + * + * Failures caused by [kotlin.coroutines.cancellation.CancellationException] are always + * propagated regardless of the selected policy. + */ +public enum class FailurePolicy { + /** + * Re-throw the underlying error wrapped into a dedicated [LongTermMemory] exception + * (either [LongTermMemoryRetrievalException] or [LongTermMemoryIngestionException]). + * + * For retrieval, this stops the LLM call before it is executed, which is usually + * safer than answering without the required memory context. + * + * For ingestion, this aborts the agent run instead of silently dropping memory records, + * which is useful for durable audit/logging use cases where losing memory is worse + * than failing the run. + */ + FAIL_FAST, + + /** + * Log the error and continue. For retrieval this means proceeding to the LLM call + * with no augmentation (treated as "no relevant memories"). For ingestion this means + * the records are dropped. + */ + LOG_AND_CONTINUE, +} + +/** + * Thrown when retrieval from the [LongTermMemory] storage fails and the configured + * [FailurePolicy] is [FailurePolicy.FAIL_FAST]. + */ +public class LongTermMemoryRetrievalException( + message: String, + cause: Throwable, +) : Exception(message, cause) + +/** + * Thrown when ingestion into the [LongTermMemory] storage fails and the configured + * [FailurePolicy] is [FailurePolicy.FAIL_FAST]. + */ +public class LongTermMemoryIngestionException( + message: String, + cause: Throwable, +) : Exception(message, cause) diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemory.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemory.kt index c0a5ce5066..69a07a46d4 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemory.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemory.kt @@ -5,7 +5,6 @@ import ai.koog.agents.core.agent.context.AIAgentContext import ai.koog.agents.core.agent.context.featureOrThrow import ai.koog.agents.core.agent.entity.AIAgentStorageKey import ai.koog.agents.core.agent.entity.createStorageKey -import ai.koog.agents.core.annotation.ExperimentalAgentsApi import ai.koog.agents.core.feature.AIAgentFunctionalFeature import ai.koog.agents.core.feature.AIAgentGraphFeature import ai.koog.agents.core.feature.AIAgentPlannerFeature @@ -15,27 +14,22 @@ import ai.koog.agents.core.feature.pipeline.AIAgentGraphPipeline import ai.koog.agents.core.feature.pipeline.AIAgentPipeline import ai.koog.agents.core.feature.pipeline.AIAgentPlannerPipeline import ai.koog.agents.longtermmemory.ingestion.IngestionSettings -import ai.koog.agents.longtermmemory.ingestion.IngestionTiming -import ai.koog.agents.longtermmemory.ingestion.extraction.ExtractionStrategy -import ai.koog.agents.longtermmemory.ingestion.extraction.FilteringExtractionStrategy -import ai.koog.agents.longtermmemory.retrieval.LastUserMessageQueryExtractor -import ai.koog.agents.longtermmemory.retrieval.QueryExtractor +import ai.koog.agents.longtermmemory.ingestion.extraction.DocumentExtractor +import ai.koog.agents.longtermmemory.ingestion.extraction.MessagePassingDocumentExtractor import ai.koog.agents.longtermmemory.retrieval.RetrievalSettings -import ai.koog.agents.longtermmemory.retrieval.SearchStrategy -import ai.koog.agents.longtermmemory.retrieval.SimilaritySearchStrategy import ai.koog.agents.longtermmemory.retrieval.augmentation.PromptAugmenter import ai.koog.agents.longtermmemory.retrieval.augmentation.SystemPromptAugmenter +import ai.koog.agents.longtermmemory.retrieval.search.LastUserMessageQueryProvider +import ai.koog.agents.longtermmemory.retrieval.search.SearchQueryProvider +import ai.koog.agents.longtermmemory.retrieval.search.SearchStrategy +import ai.koog.agents.longtermmemory.retrieval.search.SimilaritySearchStrategy import ai.koog.prompt.dsl.Prompt import ai.koog.prompt.message.Message -import ai.koog.prompt.streaming.StreamFrame -import ai.koog.prompt.streaming.toMessageResponses import ai.koog.rag.base.TextDocument import ai.koog.rag.base.storage.SearchStorage import ai.koog.rag.base.storage.WriteStorage import ai.koog.rag.base.storage.search.SearchRequest import io.github.oshai.kotlinlogging.KotlinLogging -import kotlinx.coroutines.sync.Mutex -import kotlinx.coroutines.sync.withLock import kotlin.coroutines.cancellation.CancellationException /** @@ -51,18 +45,10 @@ import kotlin.coroutines.cancellation.CancellationException * @see RetrievalSettings * @see IngestionSettings */ -@ExperimentalAgentsApi public class LongTermMemory( private val retrievalSettings: RetrievalSettings? = null, private val ingestionSettings: IngestionSettings? = null, ) { - /** - * Buffer to accumulate streaming frames by runId. - * Frames are accumulated during streaming and saved when streaming completes. - */ - private val streamingFramesBuffer: MutableMap> = mutableMapOf() - private val streamingFramesBufferMutex = Mutex() - /** * Configuration for the LongTermMemory feature. * @@ -118,7 +104,7 @@ public class LongTermMemory( * Example usage: * ```kotlin * ingestion { - * extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.User)) + * documentExtractor = MessagePassingDocumentExtractor(setOf(Message.Role.User)) * } * ``` */ @@ -146,12 +132,12 @@ public class LongTermMemory( /** * The extractor that defines how to derive the search query from the prompt. - * Defaults to [LastUserMessageQueryExtractor]. + * Defaults to [LastUserMessageQueryProvider]. * - * @see QueryExtractor - * @see LastUserMessageQueryExtractor + * @see SearchQueryProvider + * @see LastUserMessageQueryProvider */ - public var queryExtractor: QueryExtractor = LastUserMessageQueryExtractor() + public var searchQueryProvider: SearchQueryProvider = LastUserMessageQueryProvider() /** * The search strategy that defines how to search the retrieval storage. @@ -181,16 +167,25 @@ public class LongTermMemory( */ public var namespace: String? = null + /** + * How to react to retrieval failures (e.g. storage outage, invalid search request). + * + * Defaults to [FailurePolicy.FAIL_FAST] so a retrieval error stops the LLM call instead + * of silently producing an answer without the required memory context. Switch to + * [FailurePolicy.LOG_AND_CONTINUE] to fall back to a non-augmented LLM call. + */ + public var failurePolicy: FailurePolicy = FailurePolicy.FAIL_FAST + /** * Fluent setter for [storage]. */ public fun withStorage(storage: SearchStorage): RetrievalSettingsBuilder = apply { this.storage = storage } /** - * Fluent setter for [queryExtractor]. + * Fluent setter for [searchQueryProvider]. */ - public fun withQueryExtractor(queryExtractor: QueryExtractor): RetrievalSettingsBuilder = - apply { this.queryExtractor = queryExtractor } + public fun withSearchQueryProvider(searchQueryProvider: SearchQueryProvider): RetrievalSettingsBuilder = + apply { this.searchQueryProvider = searchQueryProvider } /** * Fluent setter for [searchStrategy]. @@ -216,18 +211,25 @@ public class LongTermMemory( public fun withNamespace(namespace: String): RetrievalSettingsBuilder = apply { this.namespace = namespace } + /** + * Fluent setter for [failurePolicy]. + */ + public fun withFailurePolicy(failurePolicy: FailurePolicy): RetrievalSettingsBuilder = + apply { this.failurePolicy = failurePolicy } + /** * RetrievalSettings builder. */ public fun build(): RetrievalSettings { - val retrievalStorage = storage ?: error("storage must be set in retrieval { } block") + val retrievalStorage = requireNotNull(storage) { "storage must be set in retrieval { } block" } return RetrievalSettings( retrievalStorage, - queryExtractor, + searchQueryProvider, searchStrategy, promptAugmenter, enableAutomaticRetrieval, - namespace + namespace, + failurePolicy, ) } } @@ -246,39 +248,42 @@ public class LongTermMemory( * The extractor that defines how to transform messages into memory records. * * Pre-built ingesters are available: - * - [ai.koog.agents.longtermmemory.ingestion.extraction.FilteringExtractionStrategy] - Filters messages by role + * - [ai.koog.agents.longtermmemory.ingestion.extraction.MessagePassingDocumentExtractor] - Filters messages by role * * Example usage: * ```kotlin * // Use pre-built extractor with parameters - * extractionStrategy = FilteringExtractionStrategy( + * documentExtractor = MessagePassingDocumentExtractor( * messageRolesToExtract = setOf(Message.Role.User) * ) * * // Or use lambda for custom logic - * extractionStrategy = ExtractionStrategy { messages -> + * documentExtractor = DocumentExtractor { messages -> * messages.map { TextDocument(content = it.content) } * } * ``` */ - public var extractionStrategy: ExtractionStrategy = FilteringExtractionStrategy() + public var documentExtractor: DocumentExtractor = MessagePassingDocumentExtractor() /** - * When `true` (default), ingestion happens automatically after LLM calls or on agent - * completion (depending on [timing]). When `false`, the storage is still accessible - * for manual use inside graph strategy nodes. + * When `true` (default), ingestion happens automatically on agent completion. + * When `false`, the storage is still accessible for manual use inside graph strategy nodes. */ public var enableAutomaticIngestion: Boolean = true /** - * When to mapMessages messages. Defaults to [IngestionTiming.ON_LLM_CALL]. + * Namespace (table/collection name) for a request. */ - public var timing: IngestionTiming = IngestionTiming.ON_LLM_CALL + public var namespace: String? = null /** - * Namespace (table/collection name) for a request. + * How to react to ingestion failures (e.g. storage outage). + * + * Defaults to [FailurePolicy.LOG_AND_CONTINUE] so transient ingestion errors do not + * abort the agent run. Switch to [FailurePolicy.FAIL_FAST] for durable audit/logging + * use cases where losing memory records is worse than failing the run. */ - public var namespace: String? = null + public var failurePolicy: FailurePolicy = FailurePolicy.LOG_AND_CONTINUE /** * Fluent setter for [storage]. @@ -286,10 +291,10 @@ public class LongTermMemory( public fun withStorage(storage: WriteStorage): IngestionSettingsBuilder = apply { this.storage = storage } /** - * Fluent setter for [extractionStrategy]. + * Fluent setter for [documentExtractor]. */ - public fun withExtractionStrategy(extractionStrategy: ExtractionStrategy): IngestionSettingsBuilder = - apply { this.extractionStrategy = extractionStrategy } + public fun withDocumentExtractor(documentExtractor: DocumentExtractor): IngestionSettingsBuilder = + apply { this.documentExtractor = documentExtractor } /** * Fluent setter for [enableAutomaticIngestion]. @@ -297,23 +302,30 @@ public class LongTermMemory( public fun withEnableAutomaticIngestion(enable: Boolean): IngestionSettingsBuilder = apply { this.enableAutomaticIngestion = enable } - /** - * Fluent setter for [timing]. - */ - public fun withTiming(timing: IngestionTiming): IngestionSettingsBuilder = apply { this.timing = timing } - /** * Fluent setter for [namespace]. */ public fun withNamespace(namespace: String): IngestionSettingsBuilder = apply { this.namespace = namespace } + /** + * Fluent setter for [failurePolicy]. + */ + public fun withFailurePolicy(failurePolicy: FailurePolicy): IngestionSettingsBuilder = + apply { this.failurePolicy = failurePolicy } + /** * IngestionSettings builder. */ public fun build(): IngestionSettings { - val ingestionStorage = storage ?: error("storage must be set in ingestion { } block") - return IngestionSettings(ingestionStorage, extractionStrategy, timing, enableAutomaticIngestion, namespace) + val ingestionStorage = requireNotNull(storage) { "storage must be set in ingestion { } block" } + return IngestionSettings( + ingestionStorage, + documentExtractor, + enableAutomaticIngestion, + namespace, + failurePolicy, + ) } } @@ -352,15 +364,12 @@ public class LongTermMemory( return ltmFeature } - // Note: ingestion interceptors on "Starting" events must be registered before - // retrieval interceptors so that messages are ingested before prompt augmentation. if (enableIngestion) { installIngestionInterceptors(ltmFeature, pipeline) } if (enableRetrieval) { installRetrievalInterceptors(ltmFeature, pipeline) } - installCleanupInterceptors(ltmFeature, pipeline) return ltmFeature } @@ -368,9 +377,8 @@ public class LongTermMemory( /** * Install interceptors for ingesting messages into the memory record repository. * - * Depending on [IngestionTiming], interceptors are installed either on individual LLM - * calls/streams ([IngestionTiming.ON_LLM_CALL]) or on agent completion - * ([IngestionTiming.ON_AGENT_COMPLETION]). + * Ingestion happens once at agent completion: the final accumulated session + * prompt/history is passed to the configured extraction strategy as a single batch. */ private fun installIngestionInterceptors( ltmFeature: LongTermMemory, @@ -378,70 +386,6 @@ public class LongTermMemory( ) { val ingestion = ltmFeature.ingestionSettings ?: return - when (ingestion.timing) { - IngestionTiming.ON_LLM_CALL -> installOnLLMCallIngestion(ltmFeature, ingestion, pipeline) - IngestionTiming.ON_AGENT_COMPLETION -> installOnAgentCompletionIngestion(ingestion, pipeline) - } - } - - /** - * Install ingestion interceptors for [IngestionTiming.ON_LLM_CALL] mode. - * Covers both regular and streaming LLM calls. - */ - private fun installOnLLMCallIngestion( - ltmFeature: LongTermMemory, - ingestion: IngestionSettings, - pipeline: AIAgentPipeline, - ) { - // Ingest original (not yet augmented) prompt messages before regular LLM call - pipeline.interceptLLMCallStarting(this) { ctx -> - ingestMessages(ingestion, ctx.prompt.messages) - } - - // Ingest assistant responses after regular LLM call - pipeline.interceptLLMCallCompleted(this) { ctx -> - ingestMessages(ingestion, ctx.responses) - } - - // Ingest original (not yet augmented) prompt messages before streaming LLM call - pipeline.interceptLLMStreamingStarting(this) { ctx -> - ingestMessages(ingestion, ctx.prompt.messages) - } - - // Accumulate streaming frames in buffer - pipeline.interceptLLMStreamingFrameReceived(this) { ctx -> - ltmFeature.streamingFramesBufferMutex.withLock { - ltmFeature.streamingFramesBuffer.getOrPut(ctx.runId) { mutableListOf() } - .add(ctx.streamFrame) - } - } - - // Clean up the buffer on streaming failure - pipeline.interceptLLMStreamingFailed(this) { ctx -> - ltmFeature.streamingFramesBufferMutex.withLock { - ltmFeature.streamingFramesBuffer.remove(ctx.runId) - } - } - - // Ingest accumulated streaming response on streaming completion - pipeline.interceptLLMStreamingCompleted(this) { ctx -> - val frames = ltmFeature.streamingFramesBufferMutex.withLock { - ltmFeature.streamingFramesBuffer.remove(ctx.runId) - } - if (!frames.isNullOrEmpty()) { - ingestMessages(ingestion, frames.toMessageResponses()) - } - } - } - - /** - * Install ingestion interceptors for [IngestionTiming.ON_AGENT_COMPLETION] mode. - * All messages are saved at once when the agent completes. - */ - private fun installOnAgentCompletionIngestion( - ingestion: IngestionSettings, - pipeline: AIAgentPipeline, - ) { pipeline.interceptAgentCompleted(this) { ctx -> ctx.context.llm.readSession { ingestMessages(ingestion, prompt.messages) @@ -479,32 +423,11 @@ public class LongTermMemory( } } - /** - * Install safety-net interceptors to clean up any leaked streaming buffer entries - * when the agent completes or fails. - */ - private fun installCleanupInterceptors( - ltmFeature: LongTermMemory, - pipeline: AIAgentPipeline, - ) { - pipeline.interceptAgentCompleted(this) { ctx -> - ltmFeature.streamingFramesBufferMutex.withLock { - ltmFeature.streamingFramesBuffer.remove(ctx.runId) - } - } - - pipeline.interceptAgentExecutionFailed(this) { ctx -> - ltmFeature.streamingFramesBufferMutex.withLock { - ltmFeature.streamingFramesBuffer.remove(ctx.runId) - } - } - } - private suspend fun ingestMessages( ingestion: IngestionSettings, messages: List, ) { - val records = ingestion.extractionStrategy.extract(messages) + val records = ingestion.documentExtractor.extract(messages) if (records.isEmpty()) { return } @@ -514,18 +437,25 @@ public class LongTermMemory( } catch (e: CancellationException) { throw e } catch (e: Exception) { - logger.error(e) { "Failed to ingest ${records.size} memory records." } + when (ingestion.failurePolicy) { + FailurePolicy.FAIL_FAST -> throw LongTermMemoryIngestionException( + "Failed to ingest ${records.size} memory records.", + e, + ) + FailurePolicy.LOG_AND_CONTINUE -> + logger.error(e) { "Failed to ingest ${records.size} memory records." } + } } } /** - * Returns an augmented prompt only if there are relevant memory records for the query extracted by queryExtractor. + * Returns an augmented prompt only if there are relevant memory records for the query provided by searchQueryProvider. */ private suspend fun getAugmentedPromptOrNull( prompt: Prompt, retrieval: RetrievalSettings, ): Prompt? { - val query = retrieval.queryExtractor.extract(prompt) ?: return null + val query = retrieval.searchQueryProvider.provide(prompt) ?: return null val searchResults = try { val request = retrieval.searchStrategy.create(query) @@ -533,8 +463,16 @@ public class LongTermMemory( } catch (e: CancellationException) { throw e } catch (e: Exception) { - logger.error(e) { "Failed to search memory records for ${retrieval.searchStrategy}." } - emptyList() + when (retrieval.failurePolicy) { + FailurePolicy.FAIL_FAST -> throw LongTermMemoryRetrievalException( + "Failed to search memory records for ${retrieval.searchStrategy}.", + e, + ) + FailurePolicy.LOG_AND_CONTINUE -> { + logger.error(e) { "Failed to search memory records for ${retrieval.searchStrategy}." } + emptyList() + } + } } if (searchResults.isEmpty()) { return null @@ -581,7 +519,6 @@ public class LongTermMemory( * @throws IllegalStateException if the [LongTermMemory] feature is not installed. * @see withLongTermMemory */ -@OptIn(ExperimentalAgentsApi::class) public fun AIAgentContext.longTermMemory(): LongTermMemory = featureOrThrow(LongTermMemory) /** @@ -606,6 +543,5 @@ public fun AIAgentContext.longTermMemory(): LongTermMemory = featureOrThrow(Long * @throws IllegalStateException if the [LongTermMemory] feature is not installed. * @see longTermMemory */ -@OptIn(ExperimentalAgentsApi::class) public suspend fun AIAgentContext.withLongTermMemory(action: suspend LongTermMemory.() -> T): T = longTermMemory().action() diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/IngestionSettings.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/IngestionSettings.kt index 4beae6c2c9..84a63e9065 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/IngestionSettings.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/IngestionSettings.kt @@ -1,34 +1,35 @@ package ai.koog.agents.longtermmemory.ingestion -import ai.koog.agents.core.annotation.ExperimentalAgentsApi -import ai.koog.agents.longtermmemory.ingestion.extraction.ExtractionStrategy -import ai.koog.agents.longtermmemory.ingestion.extraction.FilteringExtractionStrategy +import ai.koog.agents.longtermmemory.feature.FailurePolicy +import ai.koog.agents.longtermmemory.ingestion.extraction.DocumentExtractor +import ai.koog.agents.longtermmemory.ingestion.extraction.MessagePassingDocumentExtractor import ai.koog.rag.base.TextDocument import ai.koog.rag.base.storage.WriteStorage /** * Settings controlling how messages are persisted (ingested) into the memory repository. * + * Ingestion happens once at agent completion: the final accumulated session prompt/history + * is passed to the configured [documentExtractor] as a single batch. + * * @param storage The ingestion storage where memory records will be persisted. - * @param extractionStrategy The extractor that defines how to transform messages into memory records. + * @param documentExtractor The extractor that defines how to transform messages into memory records. * Pre-built ingesters are available: - * - [ai.koog.agents.longtermmemory.ingestion.extraction.FilteringExtractionStrategy] - Filters messages by role - * Custom ingesters can be provided as lambdas via the [ai.koog.agents.longtermmemory.ingestion.extraction.ExtractionStrategy] SAM interface. - * @param timing When to ingest messages. Defaults to [IngestionTiming.ON_LLM_CALL]. - * - [IngestionTiming.ON_LLM_CALL]: prompt messages are ingested before each LLM call starts, - * and the assistant output is ingested after completion or stream completion. - * - [IngestionTiming.ON_AGENT_COMPLETION]: the final accumulated session prompt/history is - * ingested once at agent completion. - * @param enableAutomaticIngestion When `true` (default), ingestion happens automatically after LLM - * calls or on agent completion (depending on [timing]). When `false`, the storage is still - * accessible for manual use inside graph strategy nodes via [ai.koog.agents.longtermmemory.feature.withLongTermMemory]. + * - [ai.koog.agents.longtermmemory.ingestion.extraction.MessagePassingDocumentExtractor] - Filters messages by role + * Custom ingesters can be provided as lambdas via the [ai.koog.agents.longtermmemory.ingestion.extraction.DocumentExtractor] SAM interface. + * @param enableAutomaticIngestion When `true` (default), ingestion happens automatically on agent + * completion. When `false`, the storage is still accessible for manual use inside graph strategy + * nodes via [ai.koog.agents.longtermmemory.feature.withLongTermMemory]. * @param namespace Namespace (table/collection name) for a request + * @param failurePolicy How to react to failures from [storage] when persisting records. + * Defaults to [FailurePolicy.LOG_AND_CONTINUE] so transient ingestion errors do not abort + * the agent run. Set to [FailurePolicy.FAIL_FAST] for durable audit/logging use cases + * where losing memory records is worse than failing the run. */ -@ExperimentalAgentsApi public data class IngestionSettings( val storage: WriteStorage, - val extractionStrategy: ExtractionStrategy = FilteringExtractionStrategy(), - val timing: IngestionTiming = IngestionTiming.ON_LLM_CALL, + val documentExtractor: DocumentExtractor = MessagePassingDocumentExtractor(), val enableAutomaticIngestion: Boolean = true, - val namespace: String? = null + val namespace: String? = null, + val failurePolicy: FailurePolicy = FailurePolicy.LOG_AND_CONTINUE, ) diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/IngestionTiming.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/IngestionTiming.kt deleted file mode 100644 index 68f8793b6b..0000000000 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/IngestionTiming.kt +++ /dev/null @@ -1,22 +0,0 @@ -package ai.koog.agents.longtermmemory.ingestion - -import ai.koog.agents.core.annotation.ExperimentalAgentsApi - -/** - * Defines when messages should be extracted and ingested into the memory repository. - */ -@ExperimentalAgentsApi -public enum class IngestionTiming { - /** - * Prompt messages are ingested before each LLM call starts, and the assistant output - * is ingested after the call completes (or after stream completion for streaming calls). - * Enables intra-session RAG and provides crash resilience. - */ - ON_LLM_CALL, - - /** - * The final accumulated session prompt/history is ingested once at agent completion. - * Enables holistic extraction/summarization and avoids critical-path latency. - */ - ON_AGENT_COMPLETION, -} diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/ExtractionStrategy.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/DocumentExtractor.kt similarity index 53% rename from agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/ExtractionStrategy.kt rename to agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/DocumentExtractor.kt index 4c54fa261c..327446c115 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/ExtractionStrategy.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/DocumentExtractor.kt @@ -2,6 +2,7 @@ package ai.koog.agents.longtermmemory.ingestion.extraction import ai.koog.prompt.message.Message import ai.koog.rag.base.TextDocument +import kotlin.jvm.JvmStatic /** * Extractor of memory records during message ingestion. @@ -9,48 +10,48 @@ import ai.koog.rag.base.TextDocument * This is a functional interface (SAM) that defines how a list of messages * should be transformed into a list of [TextDocument]s for storage. * It provides flexibility in how messages are filtered, transformed, and - * converted into memory records while maintaining type safety. + * converted into [TextDocument]s while maintaining type safety. * * Pre-built implementations are available for common ingestion patterns: - * - [FilteringExtractionStrategy] - Filters messages by role + * - [MessagePassingDocumentExtractor] - Filters messages by role * * ### Usage Examples * * **Using pre-built extractors (Kotlin):** * ```kotlin * // Extract User and Assistant messages (default) - * val extractor = FilteringExtractionStrategy() + * val extractor = MessagePassingDocumentExtractor() * * // Extract only User messages - * val extractor = FilteringExtractionStrategy( + * val extractor = MessagePassingDocumentExtractor( * messageRolesToExtract = setOf(Message.Role.User) * ) * ``` * * **Custom implementation as lambda (Kotlin):** * ```kotlin - * val customExtractor = ExtractionStrategy { messages -> + * val customExtractor = DocumentExtractor { messages -> * messages * .filter { it.role == Message.Role.Assistant } - * .extract { MemoryRecord(content = summarize(it.content)) } + * .map { MemoryRecord(content = it.content) } * } * ``` * * **Custom implementation as lambda (Java):** * ```java - * ExtractionStrategy customExtractor = (messages) -> + * DocumentExtractor customExtractor = (messages) -> * messages.stream() * .filter(m -> m.getRole() == Message.Role.Assistant) - * .extract(m -> new MemoryRecord(m.getContent())) + * .map(m -> new MemoryRecord(m.getContent(), null, Collections.emptyMap())) * .collect(Collectors.toList()); * ``` */ -public fun interface ExtractionStrategy { +public fun interface DocumentExtractor { /** - * Transforms a list of messages into a list of memory records for storage. + * Transforms a list of messages into a list of [TextDocument]s for storage. * - * @param messages The messages to transform into memory records - * @return List of memory records to be stored + * @param messages The messages to transform into [TextDocument]s + * @return List of [TextDocument]s to be stored */ public suspend fun extract(messages: List): List @@ -59,28 +60,28 @@ public fun interface ExtractionStrategy { */ public companion object { /** - * Returns a builder that lets you choose a default [ExtractionStrategy] implementation. + * Returns a builder that lets you choose a default [DocumentExtractor] implementation. * * Example usage (Java): * ```java - * ExtractionStrategy.builder() + * DocumentExtractor.builder() * .filtering() * .withExtractRoles(new HashSet<>(Arrays.asList(Message.Role.User, Message.Role.Assistant))) * .build() * ``` */ - @kotlin.jvm.JvmStatic - public fun builder(): ExtractionStrategyBuilder = ExtractionStrategyBuilder() + @JvmStatic + public fun builder(): DocumentExtractorBuilder = DocumentExtractorBuilder() } } /** - * Intermediate builder that lets callers select a [ExtractionStrategy] implementation. + * Intermediate builder that lets callers select a [DocumentExtractor] implementation. */ -public class ExtractionStrategyBuilder { +public class DocumentExtractorBuilder { /** - * Select the [FilteringExtractionStrategy] implementation. - * Returns its [FilteringExtractionStrategy.Builder] for further configuration. + * Select the [MessagePassingDocumentExtractor] implementation. + * Returns its [MessagePassingDocumentExtractor.Builder] for further configuration. */ - public fun filtering(): FilteringExtractionStrategy.Builder = FilteringExtractionStrategy.Builder() + public fun filtering(): MessagePassingDocumentExtractor.Builder = MessagePassingDocumentExtractor.Builder() } diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/FilteringExtractionStrategy.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/MessagePassingDocumentExtractor.kt similarity index 52% rename from agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/FilteringExtractionStrategy.kt rename to agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/MessagePassingDocumentExtractor.kt index 52872ceec5..f966d5e341 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/FilteringExtractionStrategy.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/ingestion/extraction/MessagePassingDocumentExtractor.kt @@ -8,21 +8,14 @@ import ai.koog.rag.base.TextDocument * Default extractor that filters messages by role. * * This extractor filters messages to only include those with roles in - * [messageRolesToExtract], then converts each message's content into [MemoryRecord]s. - * - * When [lastMessageOnly] is `true`, only the last message for each role in [messageRolesToExtract] - * is extracted. This is useful with [ai.koog.agents.longtermmemory.ingestion.IngestionTiming.ON_LLM_CALL] to avoid re-ingesting - * messages that were already stored in previous calls. + * [messageRolesToExtract], then converts each message's content into [TextDocument]s. * * @property messageRolesToExtract The set of message roles to extract and persist. * Defaults to `setOf(Message.Role.User, Message.Role.Assistant)`. - * @property lastMessageOnly When `true`, only the last message for each matching role is extracted. - * Defaults to `false`. */ -public class FilteringExtractionStrategy( +public class MessagePassingDocumentExtractor( public val messageRolesToExtract: Set = setOf(Message.Role.User, Message.Role.Assistant), - public val lastMessageOnly: Boolean = false, -) : ExtractionStrategy { +) : DocumentExtractor { private companion object { private const val MESSAGE_ROLE_FIELD_NAME = "messageRole" @@ -30,16 +23,15 @@ public class FilteringExtractionStrategy( } /** - * Builder for [FilteringExtractionStrategy]. + * Builder for [MessagePassingDocumentExtractor]. * - * Provides a fluent API for constructing a [FilteringExtractionStrategy], + * Provides a fluent API for constructing a [MessagePassingDocumentExtractor], * which is convenient for Java users. * * Example usage (Java): * ```java - * new FilteringExtractionStrategy.Builder() + * new MessagePassingDocumentExtractor.Builder() * .withExtractRoles(new HashSet<>(Arrays.asList(Message.Role.User, Message.Role.Assistant))) - * .withLastMessageOnly(true) * .build() * ``` */ @@ -49,34 +41,19 @@ public class FilteringExtractionStrategy( */ public var extractRoles: Set = setOf(Message.Role.User, Message.Role.Assistant) - /** - * When true, only the last message for each matching role is extracted. - * Defaults to false. - */ - public var lastMessageOnly: Boolean = false - /** Fluent setter for [extractRoles]. */ public fun withExtractRoles(roles: Set): Builder = apply { this.extractRoles = roles } - /** Fluent setter for [lastMessageOnly]. */ - public fun withLastMessageOnly(lastMessageOnly: Boolean): Builder = - apply { this.lastMessageOnly = lastMessageOnly } - - /** Builds a [FilteringExtractionStrategy] from the current settings. */ - public fun build(): FilteringExtractionStrategy = - FilteringExtractionStrategy(extractRoles, lastMessageOnly) + /** Builds a [MessagePassingDocumentExtractor] from the current settings. */ + public fun build(): MessagePassingDocumentExtractor = + MessagePassingDocumentExtractor(extractRoles) } override suspend fun extract(messages: List): List { - val filtered: List = if (lastMessageOnly) { - messageRolesToExtract.mapNotNull { role -> - messages.lastOrNull { it.role == role } - } - } else { - messages.filter { it.role in messageRolesToExtract } - } - return filtered.map { messageToMemoryRecord(it) } + return messages + .filter { it.role in messageRolesToExtract } + .map { messageToMemoryRecord(it) } } private fun messageToMemoryRecord(message: Message): TextDocument = MemoryRecord( diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/model/MemoryRecord.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/model/MemoryRecord.kt index 24607f013b..cd4cbe6fbe 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/model/MemoryRecord.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/model/MemoryRecord.kt @@ -8,8 +8,6 @@ import ai.koog.rag.base.TextDocument * @property content The main textual content to be embedded and searched * @property id Unique identifier for the record * @property metadata Flexible key-value metadata for filtering and custom fields. - * Values must be primitive types (String, Number, or Boolean) when used with - * Spring AI vector store backends. */ public data class MemoryRecord( /** diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/LastUserMessageQueryExtractor.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/LastUserMessageQueryExtractor.kt deleted file mode 100644 index b0b6535683..0000000000 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/LastUserMessageQueryExtractor.kt +++ /dev/null @@ -1,15 +0,0 @@ -package ai.koog.agents.longtermmemory.retrieval - -import ai.koog.agents.core.annotation.ExperimentalAgentsApi -import ai.koog.prompt.dsl.Prompt -import ai.koog.prompt.message.Message - -/** - * Default [QueryExtractor] implementation that extracts the content of the last user message from the prompt. - */ -@ExperimentalAgentsApi -public class LastUserMessageQueryExtractor : QueryExtractor { - override fun extract(prompt: Prompt): String? { - return prompt.messages.lastOrNull { it.role == Message.Role.User }?.content - } -} diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/RetrievalSettings.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/RetrievalSettings.kt index 7d26a26fad..ac53e44f36 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/RetrievalSettings.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/RetrievalSettings.kt @@ -1,8 +1,12 @@ package ai.koog.agents.longtermmemory.retrieval -import ai.koog.agents.core.annotation.ExperimentalAgentsApi +import ai.koog.agents.longtermmemory.feature.FailurePolicy import ai.koog.agents.longtermmemory.retrieval.augmentation.PromptAugmenter import ai.koog.agents.longtermmemory.retrieval.augmentation.SystemPromptAugmenter +import ai.koog.agents.longtermmemory.retrieval.search.LastUserMessageQueryProvider +import ai.koog.agents.longtermmemory.retrieval.search.SearchQueryProvider +import ai.koog.agents.longtermmemory.retrieval.search.SearchStrategy +import ai.koog.agents.longtermmemory.retrieval.search.SimilaritySearchStrategy import ai.koog.rag.base.TextDocument import ai.koog.rag.base.storage.SearchStorage import ai.koog.rag.base.storage.search.SearchRequest @@ -11,21 +15,24 @@ import ai.koog.rag.base.storage.search.SearchRequest * Settings controlling how memory records are retrieved and injected into prompts (RAG). * * @param storage The retrieval storage to search for relevant memory records. - * @param queryExtractor The extractor that defines how to derive the search query from the prompt. - * Defaults to [LastUserMessageQueryExtractor], which uses the last user message content. + * @param searchQueryProvider The provider that defines how to derive the search query from the prompt. + * Defaults to [ai.koog.agents.longtermmemory.retrieval.search.LastUserMessageQueryProvider], which uses the last user message content. * @param searchStrategy The strategy that defines how to search the retrieval store. * @param promptAugmenter The augmenter that defines how retrieved context is inserted into the prompt. * @param enableAutomaticRetrieval When `true` (default), retrieval and prompt augmentation happen * automatically before each LLM call. When `false`, the storage and strategy are still accessible * for manual use inside graph strategy nodes via [ai.koog.agents.longtermmemory.feature.withLongTermMemory]. * @param namespace Namespace (table/collection name) for a request. + * @param failurePolicy How to react to failures from [storage] or [searchStrategy]. + * Defaults to [FailurePolicy.FAIL_FAST] so that retrieval errors stop the LLM call instead + * of silently producing an answer without the required memory context. */ -@ExperimentalAgentsApi public data class RetrievalSettings( val storage: SearchStorage, - val queryExtractor: QueryExtractor = LastUserMessageQueryExtractor(), + val searchQueryProvider: SearchQueryProvider = LastUserMessageQueryProvider(), val searchStrategy: SearchStrategy = SimilaritySearchStrategy(), val promptAugmenter: PromptAugmenter = SystemPromptAugmenter(), val enableAutomaticRetrieval: Boolean = true, - val namespace: String? = null + val namespace: String? = null, + val failurePolicy: FailurePolicy = FailurePolicy.FAIL_FAST, ) diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/PromptAugmenter.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/PromptAugmenter.kt index 601f9dbaea..3ddd7ef73e 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/PromptAugmenter.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/PromptAugmenter.kt @@ -3,6 +3,7 @@ package ai.koog.agents.longtermmemory.retrieval.augmentation import ai.koog.prompt.dsl.Prompt import ai.koog.rag.base.TextDocument import ai.koog.rag.base.storage.search.SearchResult +import kotlin.jvm.JvmStatic /** * Interface for augmenting prompts with relevant context retrieved from memory. @@ -50,7 +51,7 @@ public fun interface PromptAugmenter { * .build() * ``` */ - @kotlin.jvm.JvmStatic + @JvmStatic public fun builder(): PromptAugmenterBuilder = PromptAugmenterBuilder() /** diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/SystemPromptAugmenter.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/SystemPromptAugmenter.kt index 401b6c3910..e1332a2570 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/SystemPromptAugmenter.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/augmentation/SystemPromptAugmenter.kt @@ -9,8 +9,7 @@ import ai.koog.rag.base.storage.search.SearchResult /** * A [PromptAugmenter] that inserts retrieved context as a system message at the beginning of the prompt. * - * If an existing system message is present, the new context system message is inserted - * before it, keeping each system message focused on a single concern. + * The new context system message is prepended before all existing messages. * If there is no system message in the prompt, the prompt is returned unchanged. * * @param template The template for the system message. Use [PromptAugmenter.RELEVANT_CONTEXT_PLACEHOLDER] placeholder. diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/LastUserMessageQueryProvider.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/LastUserMessageQueryProvider.kt new file mode 100644 index 0000000000..8bd90724fa --- /dev/null +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/LastUserMessageQueryProvider.kt @@ -0,0 +1,13 @@ +package ai.koog.agents.longtermmemory.retrieval.search + +import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.message.Message + +/** + * Default [SearchQueryProvider] implementation that extracts the content of the last user message from the prompt. + */ +public class LastUserMessageQueryProvider : SearchQueryProvider { + override fun provide(prompt: Prompt): String? { + return prompt.messages.lastOrNull { it.role == Message.Role.User }?.content + } +} diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/QueryExtractor.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchQueryProvider.kt similarity index 55% rename from agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/QueryExtractor.kt rename to agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchQueryProvider.kt index 5f451b4945..c70ee6e731 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/QueryExtractor.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchQueryProvider.kt @@ -1,23 +1,21 @@ -package ai.koog.agents.longtermmemory.retrieval +package ai.koog.agents.longtermmemory.retrieval.search -import ai.koog.agents.core.annotation.ExperimentalAgentsApi import ai.koog.prompt.dsl.Prompt /** * Extracts a search query string from a [Prompt] to be used for memory retrieval. * * Implementations define how the search query is derived from the prompt messages. - * For example, the default [LastUserMessageQueryExtractor] uses the content of the last user message. + * For example, the default [LastUserMessageQueryProvider] uses the content of the last user message. * - * @see LastUserMessageQueryExtractor + * @see LastUserMessageQueryProvider */ -@ExperimentalAgentsApi -public fun interface QueryExtractor { +public fun interface SearchQueryProvider { /** * Extracts a search query string from the given [prompt]. * * @param prompt the prompt to extract the query from. * @return the extracted query string, or `null` if no query could be extracted. */ - public fun extract(prompt: Prompt): String? + public fun provide(prompt: Prompt): String? } diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/SearchStrategy.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchStrategy.kt similarity index 93% rename from agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/SearchStrategy.kt rename to agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchStrategy.kt index 10e182f586..10081b2764 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/SearchStrategy.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/retrieval/search/SearchStrategy.kt @@ -1,13 +1,14 @@ -package ai.koog.agents.longtermmemory.retrieval +package ai.koog.agents.longtermmemory.retrieval.search import ai.koog.rag.base.storage.search.SearchRequest import ai.koog.rag.base.storage.search.SimilaritySearchRequest +import kotlin.jvm.JvmStatic /** * Search strategy for creating search requests during prompt augmentation. * * This is a functional interface (SAM) that defines how a user query string - * should be transformed into a [SimilaritySearchRequest] for storage. + * should be transformed into a [SearchRequest] for retrieval. * * **[SimilaritySearchStrategy] is the default implementation.** * It uses vector embeddings for semantic search and works with all supported vector backends. @@ -32,10 +33,10 @@ import ai.koog.rag.base.storage.search.SimilaritySearchRequest */ public fun interface SearchStrategy { /** - * Maps a query string into a [SearchRequest] for the storage. + * Maps a query string into a [SearchRequest] for retrieval. * * @param query The user's query string (typically the last user message content) - * @return The similarity search request to be executed + * @return The search request to be executed */ public fun create(query: String): SearchRequest @@ -55,7 +56,7 @@ public fun interface SearchStrategy { * .build() * ``` */ - @kotlin.jvm.JvmStatic + @JvmStatic public fun builder(): SearchStrategyBuilder = SearchStrategyBuilder() } } diff --git a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorage.kt b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorage.kt index 1d331e5851..6c3c6d0458 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorage.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/commonMain/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorage.kt @@ -11,23 +11,25 @@ import ai.koog.rag.base.storage.search.Score import ai.koog.rag.base.storage.search.ScoreMetric 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 kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock import kotlin.uuid.ExperimentalUuidApi import kotlin.uuid.Uuid /** - * In-memory implementation of [SearchStorage] - * and [WriteStorage] that stores records in a map. + * In-memory implementation of [SearchStorage], [WriteStorage], [LookupStorage], + * and [DeletionStorage] that stores records in a map. * * This implementation is useful for testing, development, and scenarios where persistence * is not required. All data is stored in memory and will be lost when the application stops. * * ## Limitations: * - Data is not persisted and will be lost on application restart - * - Both [ai.koog.rag.base.storage.search.KeywordSearchRequest] and - * [ai.koog.rag.base.storage.search.SimilaritySearchRequest] are accepted, but both are - * implemented as simple case-insensitive substring matching; no vector embeddings are used + * - [ai.koog.rag.base.storage.search.KeywordSearchRequest] is implemented as a simple + * case-insensitive substring match. + * - [ai.koog.rag.base.storage.search.SimilaritySearchRequest] is implemented as a Jaccard + * coefficient over case-insensitive word sets; no vector embeddings are used. * - Filter expressions are ignored * * @param defaultNamespace The default namespace to use when none is specified in method calls. @@ -89,7 +91,17 @@ public open class InMemoryRecordStorage( namespace ) - else -> throw UnsupportedOperationException("InMemoryRecordStorage supports only KeywordSearchRequest.") + is SimilaritySearchRequest -> searchBySimilarity( + request.queryText, + request.limit, + request.offset, + request.minScore ?: 0.0, + namespace + ) + + else -> throw UnsupportedOperationException( + "InMemoryRecordStorage supports only KeywordSearchRequest and SimilaritySearchRequest." + ) } } @@ -129,6 +141,30 @@ public open class InMemoryRecordStorage( .take(limit) } + private suspend fun searchBySimilarity( + query: String, + limit: Int, + offset: Int, + minScore: Double, + namespace: String? + ): List> { + val allRecords = mutex.withLock { getRecordsForNamespace(namespace).values.toList() } + val queryWords = query.lowercase().split(Regex("\\W+")).filter { it.isNotEmpty() }.toSet() + + return allRecords + .map { record -> + val docWords = record.content.lowercase().split(Regex("\\W+")).filter { it.isNotEmpty() }.toSet() + val intersection = queryWords.intersect(docWords).size.toDouble() + val union = (queryWords + docWords).size.toDouble() + val score = if (union == 0.0) 0.0 else intersection / union + SearchResult(record, Score(score, ScoreMetric.COSINE_SIMILARITY)) + } + .filter { it.score.value > 0.0 && it.score.value >= minScore } + .sortedByDescending { it.score.value } + .drop(offset) + .take(limit) + } + /** * Returns the number of records in the repository for the specified namespace. * diff --git a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionJavaTest.java b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionJavaTest.java index 1c2f41fdda..2c196cf2c6 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionJavaTest.java +++ b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionJavaTest.java @@ -1,9 +1,7 @@ package ai.koog.agents.longtermmemory.feature; import ai.koog.agents.core.agent.AIAgent; -import ai.koog.agents.core.annotation.ExperimentalAgentsApi; -import ai.koog.agents.longtermmemory.ingestion.IngestionTiming; -import ai.koog.agents.longtermmemory.ingestion.extraction.ExtractionStrategy; +import ai.koog.agents.longtermmemory.ingestion.extraction.DocumentExtractor; import ai.koog.agents.longtermmemory.storage.InMemoryRecordStorage; import ai.koog.agents.testing.tools.MockPromptExecutor; import ai.koog.prompt.executor.clients.openai.OpenAIModels; @@ -23,47 +21,11 @@ * Java tests for configuring {@link LongTermMemory} ingestion settings from Java code. * Each test case demonstrates a different ingestion configuration using builders. */ -@ExperimentalAgentsApi public class LongTermMemoryIngestionJavaTest { private static final JSONSerializer serializer = new JacksonSerializer(); @Test - public void testIngestionWithFilteringExtractorAndOnLlmCallTiming() { - InMemoryRecordStorage storage = new InMemoryRecordStorage(); - - var agent = AIAgent.builder() - .promptExecutor( - MockPromptExecutor.builder(serializer) - .mockLLMAnswer("The capital of France is Paris.").asDefaultResponse() - .build() - ) - .llmModel(OpenAIModels.Chat.GPT4o) - .systemPrompt("You are a helpful assistant.") - .install(LongTermMemory.Feature, config -> - config.ingestion( - new LongTermMemory.IngestionSettingsBuilder() - .withStorage(storage) - .withExtractionStrategy( - ExtractionStrategy.builder() - .filtering() - .withExtractRoles(new HashSet<>(Arrays.asList(Message.Role.User, Message.Role.Assistant))) - .withLastMessageOnly(false) - .build() - ) - .withTiming(IngestionTiming.ON_LLM_CALL) - .build() - ) - ) - .build(); - - String result = (String) agent.run("What is the capital of France?"); - - assertNotNull(result); - assertFalse(result.isEmpty()); - } - - @Test - public void testIngestionWithFilteringExtractorAndOnAgentCompletionTiming() { + public void testIngestionWithFilteringExtractor() { InMemoryRecordStorage storage = new InMemoryRecordStorage(); var agent = AIAgent.builder() @@ -78,48 +40,12 @@ public void testIngestionWithFilteringExtractorAndOnAgentCompletionTiming() { config.ingestion( new LongTermMemory.IngestionSettingsBuilder() .withStorage(storage) - .withExtractionStrategy( - ExtractionStrategy.builder() + .withDocumentExtractor( + DocumentExtractor.builder() .filtering() .withExtractRoles(new HashSet<>(Arrays.asList(Message.Role.User, Message.Role.Assistant))) .build() ) - .withTiming(IngestionTiming.ON_AGENT_COMPLETION) - .build() - ) - ) - .build(); - - String result = (String) agent.run("Hello"); - - assertNotNull(result); - assertFalse(result.isEmpty()); - } - - @Test - public void testIngestionWithLastMessageOnlyExtractor() { - InMemoryRecordStorage storage = new InMemoryRecordStorage(); - - var agent = AIAgent.builder() - .promptExecutor( - MockPromptExecutor.builder(serializer) - .mockLLMAnswer("answer").asDefaultResponse() - .build() - ) - .llmModel(OpenAIModels.Chat.GPT4o) - .systemPrompt("You are a helpful assistant.") - .install(LongTermMemory.Feature, config -> - config.ingestion( - new LongTermMemory.IngestionSettingsBuilder() - .withStorage(storage) - .withExtractionStrategy( - ExtractionStrategy.builder() - .filtering() - .withExtractRoles(new HashSet<>(Arrays.asList(Message.Role.Assistant))) - .withLastMessageOnly(true) - .build() - ) - .withTiming(IngestionTiming.ON_LLM_CALL) .build() ) ) @@ -147,13 +73,12 @@ public void testFullConfigurationWithIngestionAndRetrieval() { config.ingestion( new LongTermMemory.IngestionSettingsBuilder() .withStorage(storage) - .withExtractionStrategy( - ExtractionStrategy.builder() + .withDocumentExtractor( + DocumentExtractor.builder() .filtering() .withExtractRoles(new HashSet<>(Arrays.asList(Message.Role.User, Message.Role.Assistant))) .build() ) - .withTiming(IngestionTiming.ON_AGENT_COMPLETION) .build() ); config.retrieval( @@ -162,6 +87,7 @@ public void testFullConfigurationWithIngestionAndRetrieval() { .withSearchStrategy(query -> new SimilaritySearchRequest(query, 15, 0, 0.5, null) ) + .withFailurePolicy(FailurePolicy.LOG_AND_CONTINUE) .build() ); }) diff --git a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalJavaTest.java b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalJavaTest.java index 5c699c5dc2..5a2ff6d98a 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalJavaTest.java +++ b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/java/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalJavaTest.java @@ -1,9 +1,8 @@ package ai.koog.agents.longtermmemory.feature; import ai.koog.agents.core.agent.AIAgent; -import ai.koog.agents.core.annotation.ExperimentalAgentsApi; import ai.koog.agents.longtermmemory.retrieval.RetrievalSettings; -import ai.koog.agents.longtermmemory.retrieval.SearchStrategy; +import ai.koog.agents.longtermmemory.retrieval.search.SearchStrategy; import ai.koog.agents.longtermmemory.retrieval.augmentation.PromptAugmenter; import ai.koog.agents.longtermmemory.storage.InMemoryRecordStorage; import ai.koog.agents.testing.tools.MockExecutorBuilder; @@ -20,7 +19,6 @@ * Java tests for configuring {@link LongTermMemory} retrieval settings from Java code. * Each test case demonstrates a different retrieval configuration using builders. */ -@ExperimentalAgentsApi public class LongTermMemoryRetrievalJavaTest { private static final JSONSerializer serializer = new JacksonSerializer(); @@ -46,6 +44,7 @@ public void testSimilaritySearchViaSearchStrategyDefaultTopK() { var retrievalSettings = new LongTermMemory.RetrievalSettingsBuilder() .withStorage(storage) .withSearchStrategy(SearchStrategy.builder().similarity().withTopK(10).build()) + .withFailurePolicy(FailurePolicy.LOG_AND_CONTINUE) .build(); var agent = buildAgentWithRetrieval(retrievalSettings); @@ -62,6 +61,7 @@ public void testCustomSearchViaLambda() { var retrievalSettings = new LongTermMemory.RetrievalSettingsBuilder() .withStorage(storage) .withSearchStrategy(query -> new SimilaritySearchRequest(query, 20, 0, 0.0, null)) + .withFailurePolicy(FailurePolicy.LOG_AND_CONTINUE) .build(); var agent = buildAgentWithRetrieval(retrievalSettings); @@ -81,6 +81,7 @@ public void testSimilaritySearchWithSystemPromptAugmenter() { SearchStrategy.builder().similarity().withTopK(10).withSimilarityThreshold(0.7).build() ) .withPromptAugmenter(PromptAugmenter.builder().system().build()) + .withFailurePolicy(FailurePolicy.LOG_AND_CONTINUE) .build(); var agent = buildAgentWithRetrieval(retrievalSettings); @@ -100,6 +101,7 @@ public void testSimilaritySearchWithUserPromptAugmenter() { SearchStrategy.builder().similarity().withTopK(5).withSimilarityThreshold(0.1).build() ) .withPromptAugmenter(PromptAugmenter.builder().user().build()) + .withFailurePolicy(FailurePolicy.LOG_AND_CONTINUE) .build(); var agent = buildAgentWithRetrieval(retrievalSettings); @@ -118,6 +120,7 @@ public void testSimilaritySearchViaSearchStrategy() { .withSearchStrategy( SearchStrategy.builder().similarity().withTopK(10).withSimilarityThreshold(0.7).build() ) + .withFailurePolicy(FailurePolicy.LOG_AND_CONTINUE) .build(); var agent = buildAgentWithRetrieval(retrievalSettings); @@ -136,6 +139,7 @@ public void testSimilaritySearchViaSearchStrategyWithThreshold() { .withSearchStrategy( SearchStrategy.builder().similarity().withTopK(5).withSimilarityThreshold(0.1).build() ) + .withFailurePolicy(FailurePolicy.LOG_AND_CONTINUE) .build(); var agent = buildAgentWithRetrieval(retrievalSettings); diff --git a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionTest.kt b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionTest.kt index 8bbf69fd92..d48c57b178 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionTest.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryIngestionTest.kt @@ -3,15 +3,13 @@ package ai.koog.agents.longtermmemory.feature import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.core.agent.config.AIAgentConfig import ai.koog.agents.core.agent.entity.ToolSelectionStrategy -import ai.koog.agents.core.annotation.ExperimentalAgentsApi import ai.koog.agents.core.dsl.builder.strategy import ai.koog.agents.core.dsl.extension.nodeLLMRequest import ai.koog.agents.core.dsl.extension.nodeLLMRequestStreaming import ai.koog.agents.core.tools.ToolDescriptor import ai.koog.agents.core.tools.ToolRegistry -import ai.koog.agents.longtermmemory.ingestion.IngestionTiming -import ai.koog.agents.longtermmemory.ingestion.extraction.ExtractionStrategy -import ai.koog.agents.longtermmemory.ingestion.extraction.FilteringExtractionStrategy +import ai.koog.agents.longtermmemory.ingestion.extraction.DocumentExtractor +import ai.koog.agents.longtermmemory.ingestion.extraction.MessagePassingDocumentExtractor import ai.koog.agents.longtermmemory.model.MemoryRecord import ai.koog.agents.longtermmemory.storage.InMemoryRecordStorage import ai.koog.agents.testing.tools.getMockExecutor @@ -36,7 +34,6 @@ import kotlin.test.assertTrue /** * Tests for LongTermMemory ingestion (IngestionSettings): persisting messages into memory storage. */ -@OptIn(ExperimentalAgentsApi::class) class LongTermMemoryIngestionTest { private val defaultNamespace = "default" @@ -53,6 +50,23 @@ class LongTermMemoryIngestionTest { edge(llmNode forwardTo nodeFinish transformed { it.content }) } + /** + * Strategy that performs two LLM calls within a single agent run. + * + * The first input is used as the first user message; a fixed second user message is + * appended before the second LLM call. The prompt for the second call therefore + * contains the first user message and the first assistant response — exactly the + * "prompt-history replay" scenario that must not lead to duplicate ingestion. + */ + private val twoCallNonStreamingStrategy = + strategy("ingestion-two-call-test", toolSelectionStrategy = ToolSelectionStrategy.NONE) { + val firstLlmNode by nodeLLMRequest(name = "llm-node-1", allowToolCalls = false) + val secondLlmNode by nodeLLMRequest(name = "llm-node-2", allowToolCalls = false) + edge(nodeStart forwardTo firstLlmNode) + edge(firstLlmNode forwardTo secondLlmNode transformed { "Follow-up question" }) + edge(secondLlmNode forwardTo nodeFinish transformed { it.content }) + } + private val streamingStrategy = strategy("ingestion-streaming-test", toolSelectionStrategy = ToolSelectionStrategy.NONE) { val llmNode by nodeLLMRequestStreaming(name = "llm-node") @@ -87,7 +101,7 @@ class LongTermMemoryIngestionTest { } // ========================================== - // Default FilteringExtractionStrategy (User + Assistant) + // Default MessagePassingDocumentExtractor (User + Assistant) // ========================================== @Test @@ -107,7 +121,7 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy() + documentExtractor = MessagePassingDocumentExtractor() } } } @@ -148,7 +162,7 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.Assistant)) + documentExtractor = MessagePassingDocumentExtractor(setOf(Message.Role.Assistant)) } } } @@ -180,8 +194,7 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.User)) - timing = IngestionTiming.ON_LLM_CALL + documentExtractor = MessagePassingDocumentExtractor(setOf(Message.Role.User)) } } } @@ -201,41 +214,6 @@ class LongTermMemoryIngestionTest { ) } - @Test - fun `assistant-only extractor excludes user messages`() = runTest { - val storage = InMemoryRecordStorage() - - val executor = getMockExecutor(defaultAgentConfig.serializer) { - mockLLMAnswer("Kotlin is a modern language").asDefaultResponse - } - - val agent = AIAgent( - promptExecutor = executor, - strategy = nonStreamingStrategy, - agentConfig = defaultAgentConfig, - toolRegistry = ToolRegistry.EMPTY - ) { - install(LongTermMemory.Feature) { - ingestion { - this.storage = storage - extractionStrategy = FilteringExtractionStrategy(messageRolesToExtract = setOf(Message.Role.Assistant)) - } - } - } - - agent.run("What is Kotlin?") - - val allResults = storage.search(KeywordSearchRequest(queryText = "Kotlin"), defaultNamespace) - assertTrue( - allResults.none { it.document.content.contains("What is Kotlin") }, - "User message should NOT be stored" - ) - assertTrue( - allResults.any { it.document.content.contains("Kotlin is a modern language") }, - "Assistant message should be stored" - ) - } - // ========================================== // No ingestion when not configured // ========================================== @@ -289,79 +267,11 @@ class LongTermMemoryIngestionTest { } // ========================================== - // Streaming ingestion - // ========================================== - - @Test - fun `streaming frames are ingested when configured`() = runTest { - val storage = InMemoryRecordStorage() - - val executor = streamingExecutor("Hello ", "world", "!") - - val agent = AIAgent( - promptExecutor = executor, - strategy = streamingStrategy, - agentConfig = defaultAgentConfig, - toolRegistry = ToolRegistry.EMPTY - ) { - install(LongTermMemory.Feature) { - ingestion { - this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.Assistant)) - } - } - } - - agent.run("Hello") - - assertEquals(1, storage.size(), "Streaming frames should be stored as a memory record") - val results = storage.search(KeywordSearchRequest(queryText = "Hello world"), defaultNamespace) - assertEquals(1, results.size) - assertTrue( - results.first().document.content.contains("Hello world!"), - "Concatenated streaming content should be stored" - ) - } - - @Test - fun `user messages are ingested during streaming with ON_LLM_CALL timing`() = runTest { - val storage = InMemoryRecordStorage() - - val executor = streamingExecutor("Streaming reply") - - val agent = AIAgent( - promptExecutor = executor, - strategy = streamingStrategy, - agentConfig = defaultAgentConfig, - toolRegistry = ToolRegistry.EMPTY - ) { - install(LongTermMemory.Feature) { - ingestion { - this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.User)) - timing = IngestionTiming.ON_LLM_CALL - } - } - } - - agent.run("User streaming question about Kotlin") - - assertTrue(storage.size() > 0, "User message should be stored during streaming") - val results = storage.search(KeywordSearchRequest(queryText = "Kotlin"), defaultNamespace) - assertTrue(results.any { it.document.content.contains("User streaming question about Kotlin") }) - val streamResults = storage.search(KeywordSearchRequest(queryText = "Streaming reply"), defaultNamespace) - assertTrue( - streamResults.none { it.document.content.contains("Streaming reply") }, - "Streaming assistant response should NOT be stored" - ) - } - - // ========================================== - // IngestionTiming.ON_AGENT_COMPLETION + // On-completion ingestion // ========================================== @Test - fun `ON_AGENT_COMPLETION stores messages after agent run completes`() = runTest { + fun `stores messages after agent run completes`() = runTest { val storage = InMemoryRecordStorage() val executor = getMockExecutor(defaultAgentConfig.serializer) { @@ -377,8 +287,7 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.Assistant)) - timing = IngestionTiming.ON_AGENT_COMPLETION + documentExtractor = MessagePassingDocumentExtractor(setOf(Message.Role.Assistant)) } } } @@ -393,7 +302,7 @@ class LongTermMemoryIngestionTest { @Test @Timeout(5) - fun `ON_AGENT_COMPLETION does not store messages during LLM call`() = runTest { + fun `does not store messages during LLM call`() = runTest { val storage = InMemoryRecordStorage() var storageSizeDuringLLMCall = -1 @@ -429,8 +338,7 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.Assistant)) - timing = IngestionTiming.ON_AGENT_COMPLETION + documentExtractor = MessagePassingDocumentExtractor(setOf(Message.Role.Assistant)) } } } @@ -440,17 +348,21 @@ class LongTermMemoryIngestionTest { assertEquals( 0, storageSizeDuringLLMCall, - "No records should be stored during LLM call with ON_COMPLETION timing" + "No records should be stored during LLM call; ingestion happens only on agent completion" ) assertTrue(storage.size() > 0, "Records should be stored after agent completion") } + // ========================================== + // Custom DocumentExtractor + // ========================================== + @Test - fun `ON_AGENT_COMPLETION stores user messages`() = runTest { + fun `custom extractor transforms content before storing`() = runTest { val storage = InMemoryRecordStorage() val executor = getMockExecutor(defaultAgentConfig.serializer) { - mockLLMAnswer("Assistant reply").asDefaultResponse + mockLLMAnswer("First sentence. Second sentence. Third sentence.").asDefaultResponse } val agent = AIAgent( @@ -462,29 +374,36 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.User)) - timing = IngestionTiming.ON_AGENT_COMPLETION + documentExtractor = DocumentExtractor { messages -> + messages.filter { it.role == Message.Role.Assistant } + .flatMap { it.content.split(". ") } + .map { it.trim().removeSuffix(".") } + .filter { it.isNotBlank() } + .map { MemoryRecord(content = it) } + } } } } - agent.run("User question about Kotlin") + agent.run("Hello") - assertTrue(storage.size() > 0, "User message should be stored on completion") - val results = storage.search(KeywordSearchRequest(queryText = "Kotlin"), defaultNamespace) - assertTrue(results.any { it.document.content.contains("User question about Kotlin") }) - assertTrue( - results.none { it.document.content.contains("Assistant reply") }, - "Assistant messages should NOT be stored" - ) + assertEquals(3, storage.size(), "Custom extractor should split into 3 separate records") + val results = storage.search(KeywordSearchRequest(queryText = "sentence"), defaultNamespace) + assertTrue(results.any { it.document.content.contains("First sentence") }) + assertTrue(results.any { it.document.content.contains("Second sentence") }) + assertTrue(results.any { it.document.content.contains("Third sentence") }) } + // ========================================== + // Edge case: extractor returns empty list + // ========================================== + @Test - fun `ON_AGENT_COMPLETION stores both user and assistant messages`() = runTest { + fun `extractor returning empty list stores nothing`() = runTest { val storage = InMemoryRecordStorage() val executor = getMockExecutor(defaultAgentConfig.serializer) { - mockLLMAnswer("Assistant response about Kotlin").asDefaultResponse + mockLLMAnswer("Some response").asDefaultResponse } val agent = AIAgent( @@ -496,113 +415,73 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.User, Message.Role.Assistant)) - timing = IngestionTiming.ON_AGENT_COMPLETION + documentExtractor = DocumentExtractor { emptyList() } } } } - agent.run("User question about Kotlin") + agent.run("Hello") - assertTrue(storage.size() >= 2, "Both user and assistant records should be stored") - val results = storage.search(KeywordSearchRequest(queryText = "Kotlin"), defaultNamespace) - assertTrue( - results.any { it.document.content.contains("User question about Kotlin") }, - "User message should be stored" - ) - assertTrue( - results.any { it.document.content.contains("Assistant response about Kotlin") }, - "Assistant message should be stored" - ) + assertEquals(0, storage.size(), "No records should be stored when extractor returns empty list") } // ========================================== - // Custom ExtractionStrategy + // Multi-call ingestion: prompt history is ingested once on completion // ========================================== @Test - fun `custom extractor transforms content before storing`() = runTest { + fun `two LLM calls with default extractor do not duplicate prompt-history messages`() = runTest { val storage = InMemoryRecordStorage() val executor = getMockExecutor(defaultAgentConfig.serializer) { - mockLLMAnswer("First sentence. Second sentence. Third sentence.").asDefaultResponse + mockLLMAnswer("Second answer about coroutines") onRequestContains "Follow-up question" + mockLLMAnswer("First answer about coroutines").asDefaultResponse } val agent = AIAgent( promptExecutor = executor, - strategy = nonStreamingStrategy, + strategy = twoCallNonStreamingStrategy, agentConfig = defaultAgentConfig, toolRegistry = ToolRegistry.EMPTY ) { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = ExtractionStrategy { messages -> - messages.filter { it.role == Message.Role.Assistant } - .flatMap { it.content.split(". ") } - .map { it.trim().removeSuffix(".") } - .filter { it.isNotBlank() } - .map { MemoryRecord(content = it) } - } + documentExtractor = MessagePassingDocumentExtractor() } } } - agent.run("Hello") + agent.run("First question") - assertEquals(3, storage.size(), "Custom extractor should split into 3 separate records") - val results = storage.search(KeywordSearchRequest(queryText = "sentence"), defaultNamespace) - assertTrue(results.any { it.document.content.contains("First sentence") }) - assertTrue(results.any { it.document.content.contains("Second sentence") }) - assertTrue(results.any { it.document.content.contains("Third sentence") }) - } - - @Test - fun `custom extractor that uppercases content`() = runTest { - val storage = InMemoryRecordStorage() + // Expected ingested records: User1, Assistant1, User2, Assistant2 → exactly 4. + // Ingestion runs once on agent completion over the final prompt history, + // so each message appears exactly once regardless of how many LLM calls happened. + assertEquals( + 4, + storage.size(), + "Each prompt-history message must be ingested exactly once on agent completion" + ) - val executor = getMockExecutor(defaultAgentConfig.serializer) { - mockLLMAnswer("The answer is 42").asDefaultResponse - } + val firstUser = storage.search(KeywordSearchRequest(queryText = "First question"), defaultNamespace) + assertEquals(1, firstUser.size, "First user message must be stored exactly once") - val agent = AIAgent( - promptExecutor = executor, - strategy = nonStreamingStrategy, - agentConfig = defaultAgentConfig, - toolRegistry = ToolRegistry.EMPTY - ) { - install(LongTermMemory.Feature) { - ingestion { - this.storage = storage - extractionStrategy = ExtractionStrategy { messages -> - messages - .filter { it.role == Message.Role.Assistant } - .map { MemoryRecord(content = it.content.uppercase()) } - } - } - } - } + val firstAssistant = storage.search(KeywordSearchRequest(queryText = "First answer"), defaultNamespace) + assertEquals(1, firstAssistant.size, "First assistant message must be stored exactly once") - agent.run("What is the answer?") + val secondUser = storage.search(KeywordSearchRequest(queryText = "Follow-up"), defaultNamespace) + assertEquals(1, secondUser.size, "Second user message must be stored exactly once") - assertTrue(storage.size() > 0, "At least one record should be stored") - val results = storage.search(KeywordSearchRequest(queryText = "ANSWER"), defaultNamespace) - assertTrue( - results.any { it.document.content == "THE ANSWER IS 42" }, - "Custom extractor should have uppercased the content" - ) + val secondAssistant = storage.search(KeywordSearchRequest(queryText = "Second answer"), defaultNamespace) + assertEquals(1, secondAssistant.size, "Second assistant message must be stored exactly once") } - // ========================================== - // Edge case: extractor returns empty list - // ========================================== - @Test - fun `extractor returning empty list stores nothing`() = runTest { + fun `consecutive agent runs each ingest their own messages on completion`() = runTest { val storage = InMemoryRecordStorage() val executor = getMockExecutor(defaultAgentConfig.serializer) { - mockLLMAnswer("Some response").asDefaultResponse + mockLLMAnswer("Stable assistant reply").asDefaultResponse } val agent = AIAgent( @@ -614,13 +493,21 @@ class LongTermMemoryIngestionTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = ExtractionStrategy { emptyList() } + documentExtractor = MessagePassingDocumentExtractor() } } } - agent.run("Hello") + agent.run("Question one") + val sizeAfterFirstRun = storage.size() + agent.run("Question two") + val sizeAfterSecondRun = storage.size() - assertEquals(0, storage.size(), "No records should be stored when extractor returns empty list") + // Each run is independent and ingests its own prompt history on completion. + assertTrue( + sizeAfterSecondRun > sizeAfterFirstRun, + "Second run must ingest its own messages on completion; " + + "first=$sizeAfterFirstRun second=$sizeAfterSecondRun" + ) } } diff --git a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalTest.kt b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalTest.kt index 50e6477c4c..c46ca04844 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalTest.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryRetrievalTest.kt @@ -3,20 +3,16 @@ package ai.koog.agents.longtermmemory.feature import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.core.agent.config.AIAgentConfig import ai.koog.agents.core.agent.entity.ToolSelectionStrategy -import ai.koog.agents.core.annotation.ExperimentalAgentsApi import ai.koog.agents.core.dsl.builder.strategy import ai.koog.agents.core.dsl.extension.nodeLLMRequest import ai.koog.agents.core.dsl.extension.nodeLLMRequestStreaming import ai.koog.agents.core.tools.ToolDescriptor import ai.koog.agents.core.tools.ToolRegistry -import ai.koog.agents.longtermmemory.ingestion.IngestionTiming -import ai.koog.agents.longtermmemory.ingestion.extraction.FilteringExtractionStrategy +import ai.koog.agents.longtermmemory.ingestion.extraction.MessagePassingDocumentExtractor import ai.koog.agents.longtermmemory.model.MemoryRecord -import ai.koog.agents.longtermmemory.retrieval.SearchStrategy -import ai.koog.agents.longtermmemory.retrieval.SimilaritySearchStrategy -import ai.koog.agents.longtermmemory.retrieval.augmentation.UserPromptAugmenter +import ai.koog.agents.longtermmemory.retrieval.search.SearchStrategy +import ai.koog.agents.longtermmemory.retrieval.search.SimilaritySearchStrategy import ai.koog.agents.longtermmemory.storage.InMemoryRecordStorage -import ai.koog.agents.longtermmemory.storage.InMemorySimilaritySearchStorage import ai.koog.prompt.dsl.Prompt import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.model.PromptExecutor @@ -25,7 +21,11 @@ import ai.koog.prompt.llm.LLModel import ai.koog.prompt.message.Message import ai.koog.prompt.message.ResponseMetaInfo import ai.koog.prompt.streaming.StreamFrame +import ai.koog.rag.base.TextDocument +import ai.koog.rag.base.storage.SearchStorage import ai.koog.rag.base.storage.search.KeywordSearchRequest +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 kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.flow @@ -34,13 +34,14 @@ import kotlinx.coroutines.test.runTest import org.junit.jupiter.api.Test import org.junit.jupiter.api.Timeout import kotlin.test.assertEquals +import kotlin.test.assertFailsWith import kotlin.test.assertFalse +import kotlin.test.assertNotNull import kotlin.test.assertTrue /** * Tests for LongTermMemory retrieval (RetrievalSettings): prompt augmentation via storage search. */ -@OptIn(ExperimentalAgentsApi::class) class LongTermMemoryRetrievalTest { private val defaultNamespace = "default" @@ -243,7 +244,7 @@ class LongTermMemoryRetrievalTest { fun `search request strategy receives the user query`() = runTest { var capturedQuery: String? = null - val storage = InMemorySimilaritySearchStorage() + val storage = InMemoryRecordStorage() storage.add( listOf( MemoryRecord(content = "The weather in Paris is sunny today"), @@ -283,7 +284,7 @@ class LongTermMemoryRetrievalTest { @Test @Timeout(5) fun `similaritySearch builder retrieves matching records`() = runTest { - val storage = InMemorySimilaritySearchStorage() + val storage = InMemoryRecordStorage() storage.add( listOf( MemoryRecord(content = "Kotlin was developed by JetBrains"), @@ -322,7 +323,7 @@ class LongTermMemoryRetrievalTest { @Test @Timeout(5) fun `similaritySearch builder returns no augmentation when query does not match`() = runTest { - val storage = InMemorySimilaritySearchStorage() + val storage = InMemoryRecordStorage() storage.add( listOf( MemoryRecord(content = "Kotlin was developed by JetBrains"), @@ -364,7 +365,7 @@ class LongTermMemoryRetrievalTest { @Test @Timeout(5) fun `empty storage produces no augmentation`() = runTest { - val storage = InMemorySimilaritySearchStorage() + val storage = InMemoryRecordStorage() var augmented = false val executor = promptCapturingExecutor { content -> @@ -399,7 +400,7 @@ class LongTermMemoryRetrievalTest { @Test @Timeout(5) fun `ingested data is retrievable in subsequent agent run`() = runTest { - val storage = InMemorySimilaritySearchStorage() + val storage = InMemoryRecordStorage() // First agent run: ingest data val ingestExecutor = promptCapturingExecutor { "Kotlin supports coroutines for async programming" } @@ -413,7 +414,7 @@ class LongTermMemoryRetrievalTest { install(LongTermMemory.Feature) { ingestion { this.storage = storage - extractionStrategy = FilteringExtractionStrategy() + documentExtractor = MessagePassingDocumentExtractor() } } } @@ -450,24 +451,23 @@ class LongTermMemoryRetrievalTest { } // ========================================== - // Ingestion + Retrieval: ingestion stores original (non-augmented) prompt + // FailurePolicy.FAIL_FAST (default) on retrieval // ========================================== @Test @Timeout(5) - fun `ingestion stores original prompt when both ingestion and retrieval are configured`() = runTest { - val retrievalStorage = InMemoryRecordStorage() - retrievalStorage.add( - listOf(MemoryRecord(content = "Context about Kotlin coroutines")), - defaultNamespace - ) - - val ingestionStorage = InMemoryRecordStorage() + fun `default FAIL_FAST policy throws LongTermMemoryRetrievalException when storage search fails`() = runTest { + val storageError = RuntimeException("storage exploded") + val failingStorage = object : SearchStorage { + override suspend fun search(request: SearchRequest, namespace: String?): List> { + throw storageError + } + } - var promptSeenByLLM: String? = null - val executor = promptCapturingExecutor { content -> - promptSeenByLLM = content - "LLM response" + var llmCalled = false + val executor = promptCapturingExecutor { + llmCalled = true + "SHOULD_NOT_BE_REACHED" } val agent = AIAgent( @@ -478,33 +478,19 @@ class LongTermMemoryRetrievalTest { ) { install(LongTermMemory.Feature) { retrieval { - storage = retrievalStorage - searchStrategy = SearchStrategy { _ -> - KeywordSearchRequest(queryText = "Kotlin") - } - promptAugmenter = UserPromptAugmenter() - } - ingestion { - storage = ingestionStorage - extractionStrategy = FilteringExtractionStrategy(setOf(Message.Role.User)) - timing = IngestionTiming.ON_LLM_CALL + this.storage = failingStorage + searchStrategy = SearchStrategy { KeywordSearchRequest(queryText = "anything") } + // failurePolicy left at default (FAIL_FAST) } } } - val originalUserMessage = "Tell me about Kotlin" - - agent.run(originalUserMessage) - - // Verify the LLM saw the augmented prompt (retrieval worked) - assertTrue( - promptSeenByLLM!!.contains("Context about Kotlin coroutines"), - "LLM should see augmented prompt with retrieved context" - ) + val thrown = assertFailsWith { + agent.run("Tell me about Kotlin") + } - // Verify ingestion stored the ORIGINAL user message, not the augmented one - val ingestedRecords = ingestionStorage.search(KeywordSearchRequest(queryText = "Kotlin"), defaultNamespace) - assertEquals(1, ingestedRecords.size) - assertEquals(originalUserMessage, ingestedRecords.first().document.content) + assertEquals(storageError, thrown.cause, "Original storage error should be preserved as cause") + assertNotNull(thrown.message, "Exception should carry a message") + assertFalse(llmCalled, "LLM must not be called when retrieval fails under FAIL_FAST policy") } } diff --git a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryStrategyTest.kt b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryStrategyTest.kt index b8fe249934..59c8666ac5 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryStrategyTest.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/feature/LongTermMemoryStrategyTest.kt @@ -3,7 +3,6 @@ package ai.koog.agents.longtermmemory.feature import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.core.agent.config.AIAgentConfig import ai.koog.agents.core.agent.entity.ToolSelectionStrategy -import ai.koog.agents.core.annotation.ExperimentalAgentsApi import ai.koog.agents.core.dsl.builder.node import ai.koog.agents.core.dsl.builder.strategy import ai.koog.agents.longtermmemory.model.MemoryRecord @@ -26,7 +25,6 @@ import kotlin.test.assertTrue * Verifies that users can install the feature, provide their own storage implementations, * and call `search` and `add` methods from those storages within strategy nodes. */ -@OptIn(ExperimentalAgentsApi::class) class LongTermMemoryStrategyTest { private val myNamespace = "ns" private val serializer = KotlinxSerializer() diff --git a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorageTest.kt b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorageTest.kt index 9384a74bd4..9150787159 100644 --- a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorageTest.kt +++ b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/storage/InMemoryRecordStorageTest.kt @@ -2,6 +2,7 @@ package ai.koog.agents.longtermmemory.storage import ai.koog.agents.longtermmemory.model.MemoryRecord import ai.koog.rag.base.storage.search.KeywordSearchRequest +import ai.koog.rag.base.storage.search.SimilaritySearchRequest import kotlinx.coroutines.test.runTest import kotlin.test.Test import kotlin.test.assertContains @@ -199,6 +200,102 @@ class InMemoryRecordStorageTest { assertEquals(listOf("id-1", "id-2"), ids) } + @Test + fun testSearchBySimilarity() = runTest { + val repository = InMemoryRecordStorage() + repository.add( + listOf( + MemoryRecord(id = "id-1", content = "Kotlin is a modern programming language"), + MemoryRecord(id = "id-2", content = "Java is also a programming language"), + MemoryRecord(id = "id-3", content = "Bananas are yellow fruits") + ), + defaultNamespace, + ) + + val results = repository.search( + SimilaritySearchRequest(queryText = "Kotlin programming language"), + defaultNamespace + ) + + assertEquals(2, results.size) + assertEquals("id-1", results[0].document.id) + assertEquals("id-2", results[1].document.id) + assertTrue(results[0].score.value > results[1].score.value) + } + + @Test + fun testSearchBySimilarityWithMinScore() = runTest { + val repository = InMemoryRecordStorage() + repository.add( + listOf( + MemoryRecord(id = "id-1", content = "Kotlin is a modern programming language"), + MemoryRecord(id = "id-2", content = "Java is also a programming language"), + ), + defaultNamespace, + ) + + val allResults = repository.search( + SimilaritySearchRequest(queryText = "Kotlin programming language"), + defaultNamespace + ) + val topScore = allResults.first().score.value + + val filtered = repository.search( + SimilaritySearchRequest(queryText = "Kotlin programming language", minScore = topScore), + defaultNamespace + ) + + assertEquals(1, filtered.size) + assertEquals("id-1", filtered[0].document.id) + } + + @Test + fun testSearchBySimilarityWithLimitAndOffset() = runTest { + val repository = InMemoryRecordStorage() + repository.add( + listOf( + MemoryRecord(id = "id-1", content = "Kotlin is a modern programming language"), + MemoryRecord(id = "id-2", content = "Java is also a programming language"), + MemoryRecord(id = "id-3", content = "Programming in general is fun"), + ), + defaultNamespace, + ) + + val limited = repository.search( + SimilaritySearchRequest(queryText = "Kotlin programming language", limit = 1), + defaultNamespace + ) + assertEquals(1, limited.size) + assertEquals("id-1", limited[0].document.id) + + val offset = repository.search( + SimilaritySearchRequest(queryText = "Kotlin programming language", limit = 1, offset = 1), + defaultNamespace + ) + assertEquals(1, offset.size) + assertEquals("id-2", offset[0].document.id) + } + + @Test + fun testSearchBySimilarityExcludesNonOverlapping() = runTest { + val repository = InMemoryRecordStorage() + repository.add( + listOf( + MemoryRecord(id = "id-1", content = "Kotlin programming language"), + MemoryRecord(id = "id-2", content = "Bananas are yellow fruits"), + ), + defaultNamespace, + ) + + val results = repository.search( + SimilaritySearchRequest(queryText = "Kotlin programming language"), + defaultNamespace + ) + + assertEquals(1, results.size) + assertEquals("id-1", results[0].document.id) + } + @Test fun testAddWithoutIdGeneratesId() = runTest { val repository = InMemoryRecordStorage() diff --git a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/storage/InMemorySimilaritySearchStorage.kt b/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/storage/InMemorySimilaritySearchStorage.kt deleted file mode 100644 index adac3ed116..0000000000 --- a/agents/agents-features/agents-features-longterm-memory/src/jvmTest/kotlin/ai/koog/agents/longtermmemory/storage/InMemorySimilaritySearchStorage.kt +++ /dev/null @@ -1,94 +0,0 @@ -package ai.koog.agents.longtermmemory.storage - -import ai.koog.agents.longtermmemory.model.MemoryRecord -import ai.koog.rag.base.TextDocument -import ai.koog.rag.base.storage.SearchStorage -import ai.koog.rag.base.storage.WriteStorage -import ai.koog.rag.base.storage.search.Score -import ai.koog.rag.base.storage.search.ScoreMetric -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 kotlinx.coroutines.sync.Mutex -import kotlinx.coroutines.sync.withLock -import kotlin.uuid.ExperimentalUuidApi -import kotlin.uuid.Uuid - -/** - * Simple in-memory [SearchStorage] that supports only [SimilaritySearchRequest]. - * - * Similarity is computed as the Jaccard coefficient of word sets (case-insensitive), - * which is sufficient for unit tests without requiring real vector embeddings. - * - * @param defaultNamespace The default namespace to use when none is specified. - */ -internal class InMemorySimilaritySearchStorage( - private val defaultNamespace: String = "default" -) : SearchStorage, - WriteStorage { - - private val mutex = Mutex() - private val namespaceRecords = mutableMapOf>() - - private fun getRecordsForNamespace(namespace: String?): MutableMap { - val ns = namespace ?: defaultNamespace - return namespaceRecords.getOrPut(ns) { mutableMapOf() } - } - - @OptIn(ExperimentalUuidApi::class) - override suspend fun add(documents: List, namespace: String?): List { - return mutex.withLock { - val nsRecords = getRecordsForNamespace(namespace) - documents.map { doc -> - val recordId = doc.id ?: Uuid.random().toString() - val recordWithId = MemoryRecord(doc.content, recordId, doc.metadata) - nsRecords[recordId] = recordWithId - recordId - } - } - } - - override suspend fun update(documents: Map, namespace: String?): List { - return mutex.withLock { - val nsRecords = getRecordsForNamespace(namespace) - val updated = mutableListOf() - for ((id, doc) in documents) { - if (nsRecords.containsKey(id)) { - nsRecords[id] = MemoryRecord(doc.content, doc.id, doc.metadata) - updated.add(id) - } - } - updated - } - } - - override suspend fun search( - request: SearchRequest, - namespace: String? - ): List> { - require(request is SimilaritySearchRequest) { - "InMemorySimilaritySearchStorage supports only SimilaritySearchRequest, got: ${request::class.simpleName}" - } - val allRecords = mutex.withLock { getRecordsForNamespace(namespace).values.toList() } - val queryWords = request.queryText.lowercase().split(Regex("\\W+")).filter { it.isNotEmpty() }.toSet() - val minScore = request.minScore ?: 0.0 - - return allRecords - .map { record -> - val docWords = record.content.lowercase().split(Regex("\\W+")).filter { it.isNotEmpty() }.toSet() - val intersection = queryWords.intersect(docWords).size.toDouble() - val union = (queryWords + docWords).size.toDouble() - val score = if (union == 0.0) 0.0 else intersection / union - SearchResult(record, Score(score, ScoreMetric.COSINE_SIMILARITY)) - } - .filter { it.score.value > 0.0 && it.score.value >= minScore } - .sortedByDescending { it.score.value } - .drop(request.offset) - .take(request.limit) - } - - /** - * Returns the number of records stored in the given namespace. - */ - fun size(namespace: String? = null): Int = namespaceRecords[namespace ?: defaultNamespace]?.size ?: 0 -} diff --git a/docs/docs/features/long-term-memory.md b/docs/docs/features/long-term-memory.md index 376ebb3496..a46266f40b 100644 --- a/docs/docs/features/long-term-memory.md +++ b/docs/docs/features/long-term-memory.md @@ -1,22 +1,17 @@ # Long-term memory -Feature (Experimental) - The `LongTermMemory` feature adds persistent memory to Koog AI agents via two independent group of settings: - **Retrieval** — augments LLM prompts with relevant context from a memory storage (Retrieval-Augmented Generation or RAG) - **Ingestion** — persists conversation messages into a memory storage for future retrieval ## Quick Start -> **Note:** `LongTermMemory` is an experimental API. Annotate your code with `@OptIn(ExperimentalAgentsApi::class)` or add `@file:OptIn(ExperimentalAgentsApi::class)` at the top of your file. === "Kotlin" ```kotlin - @OptIn(ExperimentalAgentsApi::class) val myStorage = InMemoryRecordStorage() // or your vector DB adapter - @OptIn(ExperimentalAgentsApi::class) val agent = AIAgent( promptExecutor = executor, strategy = singleRunStrategy(), @@ -65,7 +60,6 @@ Use retrieval without ingestion when you have a pre-populated knowledge base: === "Kotlin" ```kotlin - @OptIn(ExperimentalAgentsApi::class) install(LongTermMemory) { retrieval { storage = myVectorDbStorage @@ -96,23 +90,22 @@ Use retrieval without ingestion when you have a pre-populated knowledge base: | `UserPromptAugmenter()` | Inserts context as a separate user message before the last user message | | `PromptAugmenter { prompt, context -> ... }` | Custom augmentation via lambda | -### Query Extractors +### Search Query Providers -By default, the retrieval flow uses the last user message as the search query. You can customize this by providing a `QueryExtractor`: +By default, the retrieval flow uses the last user message as the search query. You can customize this by providing a `SearchQueryProvider`: -| Extractor | Behavior | +| Provider | Behavior | |---|---| -| `LastUserMessageQueryExtractor()` | Uses the content of the last user message (default) | -| `QueryExtractor { prompt -> ... }` | Custom extraction via lambda | +| `LastUserMessageQueryProvider()` | Uses the content of the last user message (default) | +| `SearchQueryProvider { prompt -> ... }` | Custom query derivation via lambda | === "Kotlin" ```kotlin - @OptIn(ExperimentalAgentsApi::class) install(LongTermMemory) { retrieval { storage = myStorage - queryExtractor = QueryExtractor { prompt -> + searchQueryProvider = SearchQueryProvider { prompt -> // Combine the last two user messages as the search query prompt.messages .filter { it.role == Message.Role.User } @@ -129,7 +122,7 @@ By default, the retrieval flow uses the last user message as the search query. Y ```java var retrievalSettings = new LongTermMemory.RetrievalSettingsBuilder() .withStorage(myStorage) - .withQueryExtractor(prompt -> { + .withSearchQueryProvider(prompt -> { var userMessages = prompt.getMessages().stream() .filter(m -> m.getRole() == Message.Role.User) .toList(); @@ -153,15 +146,13 @@ Use ingestion without retrieval to build up a memory storage over time: === "Kotlin" ```kotlin - @OptIn(ExperimentalAgentsApi::class) install(LongTermMemory) { ingestion { storage = myVectorDbStorage namespace = "my-collection" // optional: scope to a specific namespace/collection - extractionStrategy = FilteringExtractionStrategy( + documentExtractor = MessagePassingDocumentExtractor( messageRolesToExtract = setOf(Message.Role.User, Message.Role.Assistant) ) - timing = IngestionTiming.ON_LLM_CALL } } ``` @@ -171,32 +162,24 @@ Use ingestion without retrieval to build up a memory storage over time: ```java var ingestionSettings = new LongTermMemory.IngestionSettingsBuilder() .withStorage(myVectorDbStorage) - .withExtractionStrategy( - ExtractionStrategy.builder() + .withDocumentExtractor( + DocumentExtractor.builder() .filtering() .withExtractRoles(new HashSet<>(Arrays.asList(Message.Role.User, Message.Role.Assistant))) - .withLastMessageOnly(false) .build() ) - .withTiming(IngestionTiming.ON_LLM_CALL) .build(); ``` -### Ingestion Timing - -| Timing | Behavior | -|---|---| -| `ON_LLM_CALL` | Prompt messages are ingested before each LLM call starts; assistant output is ingested after completion or stream completion. Enables intra-session RAG. | -| `ON_AGENT_COMPLETION` | The final accumulated session prompt/history is ingested once at agent completion. | +Ingestion runs once when the agent run completes: the final accumulated session prompt/history is passed to the configured `documentExtractor` as a single batch. ## Disabling Automatic Behavior -By default, retrieval and ingestion run automatically (before and after LLM calls, respectively). You can disable automatic behavior while still having access to the configured storage and strategies from within strategy nodes: +By default, retrieval and ingestion run automatically (retrieval runs before each LLM call; ingestion runs once when the agent completes). You can disable automatic behavior while still having access to the configured storage and strategies from within strategy nodes: === "Kotlin" ```kotlin - @OptIn(ExperimentalAgentsApi::class) install(LongTermMemory) { retrieval { storage = myStorage @@ -237,7 +220,6 @@ This gives you three clean modes: Use `withLongTermMemory { }` inside a strategy node to directly search or add records: ```kotlin -@OptIn(ExperimentalAgentsApi::class) val myNode by node { withLongTermMemory { // Manually add records @@ -254,20 +236,18 @@ val myNode by node { Use `longTermMemory()` to get the feature instance directly: ```kotlin -@OptIn(ExperimentalAgentsApi::class) val myNode by node { val memory = longTermMemory() val storage = memory.ingestionStorage } ``` -## Custom Extraction Strategy +## Custom Document Extractor -Implement `ExtractionStrategy` to control how messages are transformed before storage: +Implement `DocumentExtractor` to control how messages are transformed before storage: ```kotlin -@OptIn(ExperimentalAgentsApi::class) -val summarizingExtractor = ExtractionStrategy { messages -> +val summarizingExtractor = DocumentExtractor { messages -> messages .filter { it.role == Message.Role.Assistant } .map { MemoryRecord(content = summarize(it.content)) } @@ -276,7 +256,7 @@ val summarizingExtractor = ExtractionStrategy { messages -> install(LongTermMemory) { ingestion { storage = myStorage - extractionStrategy = summarizingExtractor + documentExtractor = summarizingExtractor } } ``` @@ -295,10 +275,10 @@ class MyVectorDbStorage : SearchStorage, WriteStora override suspend fun add( records: List, namespace: String? - ) { - // Upsert into your vector DB + ): List { + // Upsert into your vector DB and return the IDs of added records } } ``` -For testing, use the built-in `InMemoryRecordStorage` which keeps records in memory. It accepts both `KeywordSearchRequest` and `SimilaritySearchRequest`, but implements both as simple case-insensitive substring matching (no vector embeddings). +For testing, use the built-in `InMemoryRecordStorage` which keeps records in memory. It supports both `KeywordSearchRequest` (implemented as case-insensitive substring matching) and `SimilaritySearchRequest` (implemented as a Jaccard coefficient over case-insensitive word sets); no vector embeddings are used.