Adding RAG to a Koog agent

In this topic, we will extend the customer support agent from the previous topic with long-term knowledge retrieval. By the end, the agent will automatically consult a document store of product policies and FAQs before responding — without being explicitly asked to, and without us writing any retrieval logic by hand.

Why ChatMemory alone isn't enough

The agent we built in the previous topic remembers what was said within a conversation via ChatMemory. But what happens when a customer asks about the return policy, or how a particular product works? These answers aren't in the conversation history — they live in documents that the team updates separately: policy PDFs, FAQ pages, product manuals.

We could put all of this in the system prompt, but that approach doesn't scale. A growing knowledge base would quickly consume the context window, increase costs on every single request, and make the system prompt impossible to maintain.

LongTermMemory solves this by storing documents in a vector store and automatically retrieving the most relevant chunks before each LLM call — exactly the RAG pattern we built from scratch in the "Building a RAG pipeline in Kotlin" topic, but wired as a single install() call on the agent.

Setup

Add the long-term memory and Qdrant dependencies alongside what was already present:

dependencies {
    implementation("ai.koog:koog-agents-jvm:[LATEST_VERSION]")
    implementation("ai.koog:agents-features-memory-jvm:[LATEST_VERSION]")
    // LongTermMemory feature — includes the RAG abstractions
    implementation("ai.koog:agents-features-longterm-memory-jvm:[LATEST_VERSION]")
    // Qdrant adapter for Koog's KoogVectorStore
    implementation("ai.koog:rag-vector-store-qdrant-jvm:[LATEST_VERSION]")
}

Make sure Qdrant is running locally:

docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant

Koog uses its own KoogVectorStore abstraction — the same concept as EmbeddingStore in LangChain4j and VectorStore in Spring AI, just in Koog's type system. We connect a QdrantKoogVectorStore to the local Qdrant instance:

import ai.koog.rag.vector.store.qdrant.QdrantKoogVectorStore
import io.qdrant.client.QdrantClient
import io.qdrant.client.QdrantGrpcClient
import io.qdrant.client.grpc.Collections.Distance
import io.qdrant.client.grpc.Collections.VectorParams

const val COLLECTION = "shopfast-knowledge"
const val VECTOR_DIM = 384 // matches AllMiniLmL6V2EmbeddingModel

// Create the Qdrant collection if it doesn't already exist
val qdrantClient = QdrantClient(
    QdrantGrpcClient.newBuilder("localhost", 6334, false).build()
)
if (!qdrantClient.collectionExistsAsync(COLLECTION).get()) {
    qdrantClient.createCollectionAsync(
        COLLECTION,
        VectorParams.newBuilder().setSize(VECTOR_DIM).setDistance(Distance.Cosine).build()
    ).get()
}

// Wrap it as a Koog vector store
val knowledgeBase = QdrantKoogVectorStore(
    client = qdrantClient,
    collectionName = COLLECTION
)

Ingest the knowledge base

Before the agent can retrieve anything, we need to load documents into the vector store. This is a one-time (or periodic) ingestion step — it runs when you deploy or update your knowledge base, not on every agent request.

Koog's LongTermMemory works with any embedding model that fits the LLMEmbeddingProvider interface. For the same reason we used AllMiniLmL6V2EmbeddingModel throughout the RAG topics — no API key, no network call, runs inside the JVM — we'll use it here too, wrapped in a Koog adapter:

import ai.koog.agents.features.longtermmemory.ingestion.ingestDocuments
import ai.koog.rag.base.TextDocument
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel

// In-process embedding model — no API key needed
val embeddingModel = AllMiniLmL6V2EmbeddingModel()

// The knowledge base documents — in production these would be loaded
// from files, a database, or a document management system
val policyDocuments = listOf(
    TextDocument(
        id = "returns-policy",
        content = """
            ShopFast Return Policy:
            Customers may return any item within 30 days of delivery for a full refund.
            Items must be unused and in original packaging.
            Electronics must be returned within 14 days.
            To initiate a return, contact support with your order ID.
            Refunds are processed within 3-5 business days after the item is received.
        """.trimIndent()
    ),
    TextDocument(
        id = "shipping-policy",
        content = """
            ShopFast Shipping Information:
            Standard shipping takes 3-5 business days.
            Express shipping (1-2 days) is available for an additional fee.
            Free standard shipping is available on orders over $50.
            International shipping is available to 50+ countries.
            Once dispatched, orders can be tracked via the link in your confirmation email.
        """.trimIndent()
    ),
    TextDocument(
        id = "damaged-items",
        content = """
            Damaged or Defective Items:
            If your item arrives damaged or defective, contact us within 7 days of delivery.
            We will arrange a free return and send a replacement at no charge.
            Photos of the damage are helpful but not required to start a claim.
        """.trimIndent()
    )
)

// Ingest documents — embed each one and store it in Qdrant
knowledgeBase.ingestDocuments(
    documents = policyDocuments,
    embeddingProvider = embeddingModel.asKoogEmbeddingProvider()
)

println("Ingested ${policyDocuments.size} documents into Qdrant.")

In a real application you'd run this ingestion step from a separate script or a scheduled job, not from the same process that runs the agent. The agent only reads from the vector store — it never writes to it.

Install LongTermMemory on the agent

Now, we will add install(LongTermMemory) to the agent's configuration block. This is the only change to the agent code — the strategy, tools, and ChatMemory from before all stay exactly as they were:

import ai.koog.agents.core.agent.AIAgent
import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.core.tools.reflect.asTools
import ai.koog.agents.features.longtermmemory.LongTermMemory
import ai.koog.agents.features.longtermmemory.SimilaritySearchStrategy
import ai.koog.agents.features.longtermmemory.UserPromptAugmenter
import ai.koog.agents.features.memory.ChatMemory
import ai.koog.prompt.executor.clients.openai.OpenAIModels
import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor

val toolRegistry = ToolRegistry {
    tools(SupportTools().asTools()) // same tools as previous topic
}

val agent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
    systemPrompt = """
        You are a customer support assistant for ShopFast.
        Be concise and policy-aware.
        Never invent order data — always use the available tools to look it up.
        Only process refunds when the customer explicitly requests one and the order is delivered.
        Relevant policy information will be provided as context — use it to answer policy questions accurately.
    """.trimIndent(),
    llmModel = OpenAIModels.Chat.GPT4oMini,
    toolRegistry = toolRegistry,
    strategy = supportStrategy // same strategy as previous topic
) {
    // Short-term memory: remembers the current conversation
    install(ChatMemory) {
        windowSize(20)
    }

    // Long-term memory: retrieves relevant policy documents before each LLM call
    install(LongTermMemory) {
        retrieval {
            storage = knowledgeBase
            searchStrategy = SimilaritySearchStrategy(
                topK = 3,                  // retrieve top 3 most relevant chunks
                similarityThreshold = 0.6  // lower than RAG pipeline default —
                                           // policy questions can use looser matching
            )
            promptAugmenter = UserPromptAugmenter() // appends retrieved context
                                                    // to the user message
        }
    }
}

UserPromptAugmenter appends the retrieved document chunks to the user's message before it reaches the LLM. The LLM sees both the question and the relevant policy context, allowing it to give accurate, grounded answers even for topics it wasn't explicitly told about in the system prompt.

With both memory features installed, run a conversation that mixes order look-ups (tool use) with policy questions (RAG retrieval):

import kotlinx.coroutines.runBlocking

fun main() = runBlocking {
    val sessionId = "customer-session-bob"

    // Turn 1: order status — uses the getOrderStatus tool
    println(agent.run("What's the status of my order ORD-12345?", sessionId))
    // → Looks up order via tool: "In Transit. Estimated delivery: tomorrow by 5 PM."

    // Turn 2: policy question — LongTermMemory retrieves the returns-policy document
    println(agent.run("Can I return it if I change my mind?", sessionId))
    // → ChatMemory knows this is about ORD-12345.
    //   LongTermMemory retrieves the returns policy.
    //   Agent answers: "Yes, you can return it within 30 days of delivery,
    //   as long as it's unused and in original packaging."

    // Turn 3: specific policy edge case — retrieves the damaged-items document
    println(agent.run("What if it arrives damaged?", sessionId))
    // → LongTermMemory retrieves damaged-items policy.
    //   Agent answers: "Contact us within 7 days of delivery.
    //   We'll arrange a free return and send a replacement."
}

Notice what's happening behind the scenes on Turn 2:

• ChatMemory loads the session history → the LLM knows we're discussing order ORD-12345.

• LongTermMemory embeds the user's question ("Can I return it if I change my mind?"), searches Qdrant, and finds the returns-policy document — then appends it to the user message as context.

• The LLM receives both the conversation history and the policy text, so it can give an accurate, grounded answer.

Neither of these steps required any code in the agent's tools or strategy — they happen automatically as part of the feature pipeline.

The two-flow design

LongTermMemory has two independently configurable flows:

• retrieval {} — runs before each LLM call. Embeds the current user message, searches the vector store, and augments the prompt with the top-k results. This is what we configured above.

• ingestion {} (optional) — runs after each agent turn. Automatically extracts facts or summaries from the conversation and writes them back to the vector store, so the agent's knowledge base grows from real interactions over time. We didn't configure this here because we have a static knowledge base — but for a live system where the knowledge base should evolve with conversations, you'd add it alongside retrieval {}.

When to use which memory mechanism

In most production agents you want both: ChatMemory for the current conversation, LongTermMemory for the knowledge that lives between conversations.

Conclusion

Here's what we've learned in this topic:

• LongTermMemory is Koog's built-in RAG feature. It automatically embeds the user's message, retrieves semantically relevant documents from a KoogVectorStore, and augments the prompt before each LLM call — with no changes to the agent's strategy or tools.

• The feature has two independent flows: retrieval {} for reading from the knowledge base, and ingestion {} for writing back to it automatically from conversations.

• ChatMemory and LongTermMemory are complementary and can both be installed on the same agent. ChatMemory handles in-session context; LongTermMemory handles cross-session knowledge.