Introduction to Koog

In this topic, we will look at Koog — JetBrains' open-source framework for building AI agents on the JVM. We will cover what Koog is, how its agent loop works, how to define tools, how to add memory, and how to consume streaming responses.

What is Koog?

Koog is JetBrains' own production agent framework, built from the experience of shipping AI features inside their IDE products (AI Assistant and Junie). It is open-source, written in Kotlin, and targets the JVM — but because it is built on Kotlin Multiplatform, the same agent code also runs on Android, iOS, JavaScript, and WebAssembly targets.

The framework's main strengths compared to using a provider SDK directly or wiring LangChain4j by hand are:

• A graph-based agent loop — execution is an explicit state machine of nodes and edges rather than an opaque while-loop, which makes behavior predictable, testable, and resumable after a crash.

• Built-in features installed on the agent with a single install() call — chat memory, long-term memory, persistence/checkpointing, and OpenTelemetry tracing all plug in without custom plumbing.

• Provider flexibility — OpenAI, Anthropic, Google, DeepSeek, Bedrock, and local Ollama models are all supported through the same PromptExecutor interface, and the model can be switched mid-conversation without losing history.

• Kotlin-first DSL — tools are plain Kotlin functions annotated with @Tool, the strategy graph is described in a type-safe DSL, and coroutines are used throughout.

Add the dependency as follows:

dependencies {
    implementation("ai.koog:koog-agents-jvm:[LATEST_VERSION]")
}

JDK 17 or higher and Kotlin 2.3.10 or higher are required.

The simplest agent

The minimal entry point is AIAgent. You give it a PromptExecutor (the connection to an LLM provider), a model, and a system prompt, then call agent.run() with the user's input inside a coroutine:

import ai.koog.agents.core.agent.AIAgent
import ai.koog.prompt.executor.clients.openai.OpenAIModels
import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor
import kotlinx.coroutines.runBlocking

fun main() = runBlocking {
    val agent = AIAgent(
        promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
        systemPrompt = "You are a helpful assistant. Answer concisely.",
        llmModel = OpenAIModels.Chat.GPT4o
    )
    val result = agent.run("What is Kotlin Multiplatform?")
    println(result)
}

simpleOpenAIExecutor() is a convenience function that wraps OpenAILLMClient in the simplest single-provider PromptExecutor. For multi-provider setups, use MultiLLMPromptExecutor and pass multiple clients.

The agent loop model

By default, Koog runs the standard Thought → Act → Observe loop. Under the hood, this is a graph of nodes connected by typed edges:

nodeStart
    → nodeLLMRequest     (Thought: send prompt to LLM)
    → nodeExecuteTools   (Act: run tool calls returned by LLM)
    → nodeLLMSendToolResults (Observe: feed results back)
    → ... (loop until LLM responds with text, not a tool call)
    → nodeFinish

Koog calls this the strategy. The default strategy (singleRunStrategy()) implements the loop above. You can replace it with a custom graph strategy for more complex workflows — branching, parallel tool calls, sequential sub-tasks with different tool sets.

Tool definitions

Tools are how the agent interacts with the outside world. In Koog, the simplest way to define a tool is with the @Tool and @LLMDescription annotations on a regular Kotlin function inside a class that implements ToolSet. Koog generates the JSON schema for the LLM automatically from the function signature:

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

class WeatherTools : ToolSet {

    @Tool
    @LLMDescription("Returns the current weather conditions for a given city.")
    fun getCurrentWeather(
        @LLMDescription("The name of the city, e.g. 'London' or 'Tokyo'")
        city: String
    ): String {
        // Real implementation would call a weather API here
        return "Sunny, 22°C in $city"
    }
}

The @LLMDescription annotation on both the function and each parameter is what the LLM reads to decide when and how to call the tool — so clear, specific descriptions are as important here as anywhere else in the course. Vague descriptions lead to wrong tool invocations in exactly the same way vague system prompts lead to poor responses.

Register the tool with the agent using ToolRegistry:

import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.core.tools.reflect.asTools

val toolRegistry = ToolRegistry {
    tools(WeatherTools().asTools())
}

val agent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
    systemPrompt = "You are a weather assistant.",
    llmModel = OpenAIModels.Chat.GPT4o,
    toolRegistry = toolRegistry
)

Koog also supports tool(::myTopLevelFunction) for standalone Kotlin functions and class-based tools that extend SimpleTool<Args> directly for more complex cases — see the Koog docs for details.

Memory

Koog provides two memory mechanisms, both installed on the agent with install() inside the agent's configuration block.

ChatMemory is short-term memory: it persists and restores conversation history across agent.run() calls, so the agent remembers what was said earlier in the same session. By default it uses an in-memory store, but any ChatHistoryProvider implementation — a database, Redis, or a Spring AI ChatMemoryRepository — can be substituted. A windowSize limit prevents the history from growing unboundedly and consuming excessive tokens:

import ai.koog.agents.features.memory.ChatMemory

val agent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
    systemPrompt = "You are a helpful assistant.",
    llmModel = OpenAIModels.Chat.GPT4oMini,
    toolRegistry = toolRegistry
) {
    install(ChatMemory) {
        // Replace with a database-backed provider for production
        windowSize(20) // keep the last 20 messages per session
    }
}

// Pass a session ID to associate history with a specific conversation
agent.run("What is Kotlin Multiplatform?", "session-alice")
agent.run("Can I use it on iOS?", "session-alice") // agent remembers the first question

LongTermMemory is for facts that should persist across sessions — user preferences, domain knowledge, or conversation summaries. It is backed by a vector store and augments the agent's prompt automatically using RAG at each turn. The Koog Spring AI integration wires this to your existing Spring AI VectorStore automatically:

import ai.koog.agents.features.longtermmemory.LongTermMemory
import ai.koog.agents.features.longtermmemory.SimilaritySearchStrategy
import ai.koog.agents.features.longtermmemory.UserPromptAugmenter

val agent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
    systemPrompt = "You are a helpful assistant.",
    llmModel = OpenAIModels.Chat.GPT4oMini,
    toolRegistry = toolRegistry
) {
    install(ChatMemory) {
        windowSize(20)
    }
    install(LongTermMemory) {
        retrieval {
            storage = myVectorStore        // any KoogVectorStore implementation
            searchStrategy = SimilaritySearchStrategy(
                topK = 5,
                similarityThreshold = 0.7
            )
            promptAugmenter = UserPromptAugmenter()
        }
    }
}

Both features can be installed on the same agent — ChatMemory handles the recent conversation window while LongTermMemory handles persistent knowledge across sessions.

Streaming responses

Koog's Streaming API lets you consume LLM output incrementally as a Flow<StreamFrame> instead of waiting for the full response. This is useful for real-time UIs, processing structured data as it arrives, or detecting tool calls live without buffering.

Streaming is used inside a custom node in a strategy graph, via llm.writeSession and requestLLMStreaming():

import ai.koog.agents.core.dsl.builder.strategy
import ai.koog.agents.core.dsl.builder.forwardTo
import ai.koog.agents.core.dsl.extension.node
import ai.koog.agents.core.streaming.StreamFrame
import ai.koog.agents.core.streaming.filterTextOnly

val streamingStrategy = strategy<String, String>("streaming-assistant") {
    val streamNode by node<String, String> { userInput ->
        llm.writeSession {
            appendPrompt { user(userInput) }

            val stream = requestLLMStreaming() // Flow<StreamFrame>

            // Print each token as it arrives
            stream.filterTextOnly().collect { token ->
                print(token)
            }

            // Return the full assembled text
            stream.collectText()
        }
    }

    edge(nodeStart forwardTo streamNode)
    edge(streamNode forwardTo nodeFinish)
}

The StreamFrame type is sealed — TextDelta carries incremental text tokens, ToolCallComplete carries a fully assembled tool call once all its chunks have arrived, and End marks the end of the stream with metadata. filterTextOnly() and collectText() are convenience extensions that filter the stream down to just text output.

For most simple agents you won't write streaming nodes by hand — streaming is primarily relevant when you need real-time token display in a UI or when you want to parse structured output (like a list of objects in Markdown format) as it arrives rather than waiting for the full response.

Event handlers

Every significant thing that happens inside the agent loop fires an event. You can subscribe to these in a handleEvents block when creating the agent — useful for logging, debugging, or forwarding streaming tokens to a websocket:

val agent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")),
    systemPrompt = "You are a helpful assistant.",
    llmModel = OpenAIModels.Chat.GPT4o,
    toolRegistry = toolRegistry
) {
    handleEvents {
        onToolCallStarting { context ->
            println("Calling tool: ${context.toolName}")
        }
        onLLMStreamingFrameReceived { context ->
            if (context.streamFrame is StreamFrame.TextDelta) {
                print((context.streamFrame as StreamFrame.TextDelta).text)
            }
        }
        onAgentFinished { _, result ->
            println("\nAgent finished with: $result")
        }
    }
}

Conclusion

Let's recap what we've covered in this topic:

• Koog is JetBrains' open-source, Kotlin-native agent framework. It runs on JVM, Android, iOS, JS, and WebAssembly, and supports all major LLM providers through a single PromptExecutor interface.

• You define tools with @Tool and @LLMDescription annotations on Kotlin functions inside a ToolSet class. Clear @LLMDescription strings are essential — the LLM uses them to decide when and how to call each tool.

• You install memory with install(ChatMemory) for short-term session history and install(LongTermMemory) for vector-backed persistent knowledge — both on the same agent if needed.

• The Streaming API exposes LLM output as Flow<StreamFrame>, with typed frames for text deltas, tool call completions, and end-of-stream markers. filterTextOnly() and collectText() simplify the common case of consuming plain text output token by token.