google__adk-docs/examples/kotlin/snippets/sessions/MemoryExample.kt

/*
 * Copyright 2026 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.adk.kt.examples.sessions

import com.google.adk.kt.agents.CallbackContext
import com.google.adk.kt.agents.Instruction
import com.google.adk.kt.agents.LlmAgent
import com.google.adk.kt.callbacks.AfterAgentCallback
import com.google.adk.kt.callbacks.CallbackChoice
import com.google.adk.kt.events.Event
import com.google.adk.kt.memory.InMemoryMemoryService
import com.google.adk.kt.memory.MemoryEntry
import com.google.adk.kt.memory.VertexAiMemoryBankService
import com.google.adk.kt.memory.VertexAiRagMemoryService
import com.google.adk.kt.models.Gemini
import com.google.adk.kt.runners.InMemoryRunner
import com.google.adk.kt.sessions.InMemorySessionService
import com.google.adk.kt.sessions.SessionKey
import com.google.adk.kt.tools.LoadMemoryTool
import com.google.adk.kt.tools.PreloadMemoryTool
import com.google.adk.kt.tools.ToolContext
import com.google.adk.kt.types.Content
import com.google.adk.kt.types.Role
import kotlinx.coroutines.flow.collect
import kotlinx.coroutines.runBlocking

// --8<-- [start:full_example]

/**
 * This example demonstrates the basic flow using the `InMemoryMemoryService` in Kotlin.
 * It shows how to capture information in one session, add it to memory, and recall it in another.
 */
fun main() =
    runBlocking {
        // --- Constants ---
        val appName = "memory_example_app"
        val userId = "mem_user"
        val model = Gemini(name = "gemini-flash-latest")

        // --- Agent Definitions ---

        // Agent 1: Simple agent to capture information
        val infoCaptureAgent =
            LlmAgent(
                name = "InfoCaptureAgent",
                model = model,
                instruction = Instruction("Acknowledge the user's statement."),
            )

        // Agent 2: Agent that can use memory
        val memoryRecallAgent =
            LlmAgent(
                name = "MemoryRecallAgent",
                model = model,
                instruction =
                    Instruction(
                        "Answer the user's question. Use the 'load_memory' tool " +
                            "if the answer might be in past conversations.",
                    ),
                // Give the agent the tool
                tools = listOf(LoadMemoryTool()),
            )

        // --- Services ---
        // Services must be shared across runners to share state and memory
        val sessionService = InMemorySessionService()
        val memoryService = InMemoryMemoryService()

        // --- Turn 1: Capturing Information ---
        println("--- Turn 1: Capturing Information ---")
        val runner1 =
            InMemoryRunner(
                agent = infoCaptureAgent,
                appName = appName,
                sessionService = sessionService,
                memoryService = memoryService,
            )
        val sessionId1 = "session_info"
        val userInput1 = Content.fromText(Role.USER, "My favorite project is Project Alpha.")

        // Run the agent
        runner1
            .runAsync(
                userId = userId,
                sessionId = sessionId1,
                newMessage = userInput1,
            ).collect { event ->
                event.content?.parts?.forEach { part ->
                    if (!part.text.isNullOrBlank()) {
                        println("Agent Response: ${part.text}")
                    }
                }
            }

        // Get the completed session using SessionKey
        val session1 = sessionService.getSession(SessionKey(appName, userId, sessionId1))

        // Add this session's content to the Memory Service
        println("\n--- Adding Session 1 to Memory ---")
        if (session1 != null) {
            memoryService.addSessionToMemory(session1)
            println("Session added to memory.")
        }

        // --- Turn 2: Recalling Information ---
        println("\n--- Turn 2: Recalling Information ---")
        val runner2 =
            InMemoryRunner(
                agent = memoryRecallAgent,
                appName = appName,
                // Reuse the same service
                sessionService = sessionService,
                // Reuse the same service
                memoryService = memoryService,
            )
        val sessionId2 = "session_recall"
        val userInput2 = Content.fromText(Role.USER, "What is my favorite project?")

        // Run the second agent
        runner2
            .runAsync(
                userId = userId,
                sessionId = sessionId2,
                newMessage = userInput2,
            ).collect { event ->
                event.content?.parts?.forEach { part ->
                    if (!part.text.isNullOrBlank()) {
                        println("Agent Response: ${part.text}")
                    }
                }
            }
    }
// --8<-- [end:full_example]

// --8<-- [start:instantiate_service]
fun instantiateMemoryService() {
    val memoryService = InMemoryMemoryService()
}
// --8<-- [end:instantiate_service]

// --8<-- [start:search_within_tool]
suspend fun searchWithinTool(
    context: ToolContext,
    args: Map<String, Any>,
): String {
    val query = args["query"] as String
    val response =
        context.invocationContext.memoryService?.searchMemory(
            appName = context.invocationContext.session.key.appName,
            userId = context.invocationContext.session.key.userId,
            query = query,
        )
    // process response
    return response?.memories?.joinToString("\n") {
        it.content.parts.joinToString(" ") { p -> p.text ?: "" }
    } ?: ""
}
// --8<-- [end:search_within_tool]

// --8<-- [start:preload_memory_agent]
fun preloadMemoryAgent(model: Gemini) {
    val agent =
        LlmAgent(
            model = model,
            name = "weather_sentiment_agent",
            instruction = Instruction("..."),
            tools = listOf(PreloadMemoryTool()),
        )
}
// --8<-- [end:preload_memory_agent]

// --8<-- [start:multi_memory]

/**
 * Example of using two memory services in Kotlin.
 */
suspend fun searchAllMemory(
    toolContext: ToolContext,
    query: String,
    docsMemory: InMemoryMemoryService,
): Map<String, List<String>> {
    // Search the conversational memory (configured in the runner)
    val conversational =
        toolContext.invocationContext.memoryService?.searchMemory(
            appName = toolContext.invocationContext.session.key.appName,
            userId = toolContext.invocationContext.session.key.userId,
            query = query,
        )

    // Search a separate docs knowledge base
    val docs =
        docsMemory.searchMemory(
            appName = "docs",
            userId = "shared",
            query = query,
        )

    return mapOf(
        "from_conversations" to
            (
                conversational?.memories?.map {
                    it.content.parts.joinToString(" ") { p -> p.text ?: "" }
                } ?: emptyList()
            ),
        "from_docs" to
            docs.memories.map {
                it.content.parts.joinToString(" ") { p -> p.text ?: "" }
            },
    )
}

fun multiMemoryAgent(model: Gemini) {
    // docs_memory could be any MemoryService implementation
    val docsMemory = InMemoryMemoryService()

    val agent =
        LlmAgent(
            model = model,
            name = "multi_memory_agent",
            instruction =
                Instruction(
                    "Answer questions using both your conversation history and the " +
                        "docs knowledge base. Use the search_all_memory tool.",
                ),
            // In a real app, you'd wrap searchAllMemory in a @Tool annotated class
            // and pass docsMemory to its constructor.
        )
}
// --8<-- [end:multi_memory]

// --8<-- [start:auto_save_callback]
suspend fun autoSaveSessionToMemoryCallback(
    context: CallbackContext,
): CallbackChoice<Unit, Content> {
    context.addSessionToMemory()
    return CallbackChoice.Continue(Unit)
}

fun agentWithCallback(model: Gemini) {
    val agent =
        LlmAgent(
            model = model,
            name = "Generic_QA_Agent",
            instruction = Instruction("Answer the user's questions"),
            tools = listOf(PreloadMemoryTool()),
            afterAgentCallbacks = listOf(AfterAgentCallback(::autoSaveSessionToMemoryCallback)),
        )
}
// --8<-- [end:auto_save_callback]

// --8<-- [start:callback_memory_writes]

/**
 * Saves a chosen set of events rather than the whole session, tagged so they can
 * be filtered later. The events come from the caller: CallbackContext does not
 * expose the session.
 */
suspend fun saveEventsToMemory(
    context: CallbackContext,
    events: List<Event>,
) {
    // appName, userId and sessionId are taken from the invocation. Throws
    // IllegalStateException if the runner has no memory service configured.
    context.addEventsToMemory(events, customMetadata = mapOf("source" to "turn_callback"))
}

/** Writes an explicit fact, instead of letting the service derive one from events. */
suspend fun rememberPreferenceCallback(context: CallbackContext): CallbackChoice<Unit, Content> {
    val preference = MemoryEntry(content = Content.fromText(Role.USER, "Prefers window seats."))
    context.addMemory(listOf(preference))
    return CallbackChoice.Continue(Unit)
}
// --8<-- [end:callback_memory_writes]

// --8<-- [start:memory_bank]

/** Memory Bank keeps LLM-extracted memories in a Vertex AI Agent Engine. */
fun memoryBankRunner(agent: LlmAgent): InMemoryRunner {
    val memoryService =
        VertexAiMemoryBankService(
            project = "PROJECT_ID",
            location = "LOCATION",
            agentEngineId = "AGENT_ENGINE_ID",
        )
    return InMemoryRunner(
        agent = agent,
        appName = "memory_bank_app",
        memoryService = memoryService,
    )
}
// --8<-- [end:memory_bank]

// --8<-- [start:rag_memory]

/**
 * RAG memory stores whole transcripts in a Knowledge Engine corpus and retrieves them by vector
 * similarity.
 */
fun ragMemoryRunner(agent: LlmAgent): InMemoryRunner {
    val memoryService =
        VertexAiRagMemoryService(
            project = "PROJECT_ID",
            location = "LOCATION",
            // A bare corpus id, NOT a full resource name. Kotlin expands it to
            // projects/{project}/locations/{location}/ragCorpora/{id} and rejects an
            // already-expanded name -- unlike the Python tab above, which takes the full name.
            ragCorpus = "CORPUS_ID",
            similarityTopK = 5,
            vectorDistanceThreshold = 0.6,
        )
    return InMemoryRunner(
        agent = agent,
        appName = "rag_memory_app",
        memoryService = memoryService,
    )
}
// --8<-- [end:rag_memory]