Long Term Memory

Long-Term Memory

Enable persistent memory for Koog AI agents with the LongTermMemory feature. This experimental feature provides two core capabilities: retrieval (augmenting LLM prompts with relevant context) and ingestion (persisting conversation messages for future retrieval).

Dependency

dependencies {
    implementation("ai.koog:koog-agents:0.8.0")
    implementation("ai.koog:agents-features-longterm-memory:0.8.0")
}

Maven

<dependency>
    <groupId>ai.koog</groupId>
    <artifactId>agents-features-longterm-memory-jvm</artifactId>
    <version>0.8.0</version>
</dependency>

Opt-In Requirement

LongTermMemory is experimental. Annotate usage with @OptIn:

@OptIn(ExperimentalAgentsApi::class)
fun setupAgent() { ... }

Or use a file-level opt-in:

@file:OptIn(ExperimentalAgentsApi::class)

Installation

Install LongTermMemory inside the agent configuration block with retrieval and/or ingestion settings:

import ai.koog.agents.features.memory.LongTermMemory

val agent = AIAgent(
    promptExecutor = executor,
    llmModel = OpenAIModels.Chat.GPT4o,
    systemPrompt = "You are a helpful assistant.",
    toolRegistry = ToolRegistry.EMPTY
) {
    install(LongTermMemory) {
        retrieval {
            storage = myStorage
            searchStrategy = SimilaritySearchStrategy(topK = 5)
        }
    }
}

Java

AIAgent agent = AIAgent.builder()
    .promptExecutor(executor)
    .llmModel(OpenAIModels.Chat.GPT4o)
    .systemPrompt("You are a helpful assistant.")
    .install(LongTermMemory.Feature, config -> {
        config.retrieval(
            new LongTermMemory.RetrievalSettingsBuilder()
                .withStorage(myStorage)
                .withSearchStrategy(
                    SearchStrategy.builder().similarity().withTopK(5).build()
                )
                .build()
        );
    })
    .build();

Retrieval Configuration (RAG)

Used when you have a pre-populated knowledge base. The retrieval system augments LLM prompts with relevant context from storage.

Prompt Augmenters

Control how retrieved context is injected into the prompt:

install(LongTermMemory) {
    retrieval {
        storage = myStorage
        promptAugmenter = SystemPromptAugmenter()
        searchStrategy = SimilaritySearchStrategy(topK = 5)
    }
}

Query Extractors

Control how the search query is derived from the prompt:

install(LongTermMemory) {
    retrieval {
        storage = myStorage
        queryExtractor = QueryExtractor { prompt ->
            prompt.messages
                .filter { it.role == Message.Role.User }
                .takeLast(2)
                .joinToString(" ") { it.content }
                .ifEmpty { null }
        }
    }
}

Search Strategies

install(LongTermMemory) {
    retrieval {
        storage = myStorage
        searchStrategy = SimilaritySearchStrategy(
            topK = 5,
            similarityThreshold = 0.7
        )
        namespace = "my-knowledge-base"
    }
}

Ingestion Configuration

Used to build up a memory storage over time by persisting conversation messages.

Extraction Strategies

Control which message roles are extracted:

install(LongTermMemory) {
    ingestion {
        storage = myStorage
        extractionStrategy = FilteringExtractionStrategy(
            messageRolesToExtract = setOf(Message.Role.User, Message.Role.Assistant)
        )
    }
}

Ingestion Timing

install(LongTermMemory) {
    ingestion {
        storage = myStorage
        timing = IngestionTiming.ON_LLM_CALL
    }
}

Custom Extraction Strategy

Implement ExtractionStrategy for full control over message-to-record transformation:

val summarizingExtractor = ExtractionStrategy { messages ->
    messages
        .filter { it.role == Message.Role.Assistant }
        .map { MemoryRecord(content = summarize(it.content)) }
}

install(LongTermMemory) {
    ingestion {
        storage = myStorage
        extractionStrategy = summarizingExtractor
    }
}

Disabling Automatic Behavior

By default, retrieval and ingestion run automatically. Both can be disabled for manual control:

install(LongTermMemory) {
    retrieval {
        storage = myStorage
        enableAutomaticRetrieval = false // disable auto-augmentation
    }
    ingestion {
        storage = myStorage
        enableAutomaticIngestion = false // disable auto-persistence
    }
}

Operating Modes

Accessing from Strategy Nodes

Use withLongTermMemory { } inside a node for direct search and add operations:

val searchNode by node<String, String> { input ->
    withLongTermMemory {
        val record = MemoryRecord(content = "important fact")
        ingestionStorage?.add(listOf(record), namespace = "my-namespace")

        val request = SimilaritySearchRequest(queryText = input, limit = 5)
        val results = retrievalStorage?.search(request, namespace = "my-namespace")
        results?.joinToString("\n") { it.content } ?: "No results found"
    }
}

Or use longTermMemory() to get the feature instance directly:

val node by node<String, Unit> { input ->
    val ltm = longTermMemory()
    val storage = ltm.ingestionStorage
    // ... use storage directly
}

Custom Storage Implementation

Implement SearchStorage and/or WriteStorage interfaces to connect to any vector database:

class MyVectorDbStorage(
    private val client: MyVectorDbClient
) : SearchStorage<TextDocument, SimilaritySearchRequest>,
    WriteStorage<TextDocument> {

    override suspend fun search(
        request: SimilaritySearchRequest,
        namespace: String?
    ): List<SearchResult<TextDocument>> {
        val embedding = client.embed(request.queryText)
        val results = client.search(embedding, request.limit, namespace)
        return results.map { SearchResult(it.document, it.score) }
    }

    override suspend fun add(
        records: List<TextDocument>,
        namespace: String?
    ) {
        for (record in records) {
            val embedding = client.embed(record.content)
            client.upsert(record, embedding, namespace)
        }
    }
}

Built-in Testing Storage

InMemoryRecordStorage keeps records in memory for testing:

import ai.koog.agents.features.memory.storage.InMemoryRecordStorage

val storage = InMemoryRecordStorage()

install(LongTermMemory) {
    retrieval {
        storage = storage
        searchStrategy = SimilaritySearchStrategy(topK = 5)
    }
    ingestion {
        storage = storage
    }
}

Note: InMemoryRecordStorage implements both KeywordSearchRequest and SimilaritySearchRequest as simple case-insensitive substring matching (no vector embeddings).

Complete Example

@OptIn(ExperimentalAgentsApi::class)
suspend fun main() {
    val storage = InMemoryRecordStorage()

    val agent = AIAgent(
        promptExecutor = simpleOpenAIExecutor(apiKey),
        llmModel = OpenAIModels.Chat.GPT4o,
        systemPrompt = "You are a helpful assistant with long-term memory."
    ) {
        install(LongTermMemory) {
            retrieval {
                storage = storage
                promptAugmenter = SystemPromptAugmenter()
                searchStrategy = SimilaritySearchStrategy(topK = 5)
            }
            ingestion {
                storage = storage
                extractionStrategy = FilteringExtractionStrategy(
                    messageRolesToExtract = setOf(Message.Role.User, Message.Role.Assistant)
                )
                timing = IngestionTiming.ON_AGENT_COMPLETION
            }
        }
    }

    // First session
    agent.run("My favorite programming language is Kotlin")

    // Second session — agent can retrieve previous context
    agent.run("What programming language do I prefer?")
}

ChatMemory vs AgentMemory vs LongTermMemory

Troubleshooting