RAG Overview

RAG Overview

Implement Retrieval-Augmented Generation (RAG) with Koog's embedding and vector storage modules. RAG augments LLM prompts with relevant context from a knowledge base, enabling agents to answer questions about domain-specific data.

Module Architecture

Dependencies

dependencies {
    implementation("ai.koog:koog-agents:0.8.0")
    implementation("ai.koog:rag-base:0.8.0")
    implementation("ai.koog:rag-vector:0.8.0")
    implementation("ai.koog:embeddings-base:0.8.0")
    implementation("ai.koog:embeddings-llm:0.8.0")
}

RAG Architecture

┌─────────────┐     ┌──────────────────┐     ┌──────────────────┐
│  Documents   │────▶│ DocumentEmbedder │────▶│ VectorStorage    │
│  (files,     │     │ (text → vector)  │     │ Backend          │
│   text)      │     └──────────────────┘     │ (persistence)    │
└─────────────┘                                └──────────────────┘
                                                      │
┌─────────────┐     ┌──────────────────┐              │
│ User Query   │────▶│ SimilaritySearch │◀─────────────┘
│              │     │ Request          │
└─────────────┘     └────────┬─────────┘
                             │
                      ┌──────▼──────┐
                      │  LLM with   │
                      │  Context    │
                      └─────────────┘

Core RAG Flow

The primary workflow uses EmbeddingStorage, which composes a DocumentEmbedder with a VectorStorageBackend:

Step 1: Create an Embedder

import ai.koog.embeddings.llm.LLMEmbedder
import ai.koog.clients.ollama.OllamaClient
import ai.koog.agents.ext.llm.OllamaModels

val client = OllamaClient()
val embedder = LLMEmbedder(client, OllamaModels.Embeddings.NOMIC_EMBED_TEXT)

Step 2: Create a Document Embedder

import ai.koog.rag.vector.JVMTextDocumentEmbedder

val documentEmbedder = JVMTextDocumentEmbedder(embedder)

Step 3: Instantiate Embedding Storage

import ai.koog.rag.vector.EmbeddingStorage
import ai.koog.rag.vector.InMemoryVectorStorageBackend

val storage = EmbeddingStorage(
    documentEmbedder = documentEmbedder,
    vectorStorageBackend = InMemoryVectorStorageBackend()
)

Step 4: Add Documents

storage.add(listOf(
    Path("/path/to/document1.txt"),
    Path("/path/to/document2.md"),
    Path("/path/to/document3.pdf")
))

Step 5: Search

import ai.koog.rag.base.SimilaritySearchRequest

val request = SimilaritySearchRequest(
    queryText = "How does the authentication system work?",
    limit = 5,
    minScore = 0.7
)

val results = storage.search(request)
for (result in results) {
    println("Score: ${result.score}, Content: ${result.document.content}")
}

Vector Storage Backends

Document Embedders

Combined Storage Implementations

Agentic RAG (Tool-Based Approach)

Rather than injecting retrieved documents into prompts upfront, expose the storage as an agent tool. This grants the agent autonomy to decide when knowledge base lookups are needed:

import ai.koog.agents.core.tools.annotations.Tool
import ai.koog.agents.core.tools.annotations.LLMDescription

class KnowledgeBaseTools(
    private val storage: SearchStorage<TextDocument, SimilaritySearchRequest>
) {
    @Tool
    @LLMDescription("Search the knowledge base for information about a topic")
    suspend fun searchKnowledgeBase(
        @LLMDescription("The search query")
        query: String,
        @LLMDescription("Maximum number of results to return")
        limit: Int = 5
    ): String {
        val request = SimilaritySearchRequest(
            queryText = query,
            limit = limit,
            minScore = 0.5
        )
        val results = storage.search(request)
        return if (results.isEmpty()) {
            "No relevant information found."
        } else {
            results.joinToString("\n---\n") { it.document.content }
        }
    }
}

// Register the tool
val toolRegistry = ToolRegistry {
    tool(KnowledgeBaseTools(storage))
}

val agent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(apiKey),
    llmModel = OpenAIModels.Chat.GPT4o,
    systemPrompt = "You are a helpful assistant. Use the knowledge base tool when users ask about specific topics.",
    toolRegistry = toolRegistry
)

This pattern is useful when the agent handles diverse requests — some requiring knowledge base lookups and others not.

When to Use RAG vs Memory

Current Limitations

For custom backends, build from rag-base abstractions and implement your own storage adapter.

Troubleshooting