# Koog Koog is a Kotlin-based framework designed to build and run AI agents entirely in idiomatic Kotlin. # Documentation # Overview Koog is an open-source JetBrains framework for building AI agents designed specifically for the JVM ecosystem. It provides a first-class development experience for both Kotlin and Java developers, featuring an idiomatic, type-safe Kotlin DSL and fluent builder-style Java APIs. While Java developers can leverage the full power of Koog on the JVM using idiomatic APIs, Kotlin developers can also deploy agents across JS, WasmJS, Android, and iOS targets using Kotlin Multiplatform. - [**Quickstart**](quickstart/) ______________________________________________________________________ Build and run your first AI agent - [**Glossary**](glossary/) ______________________________________________________________________ Learn the essential terms - [**Module versioning**](module-versioning/) ______________________________________________________________________ Understand stable vs. beta modules and API guarantees ## Agents Learn about [agents in general](agents/) and how to create different types of agents using Koog: - [**Basic agents**](agents/basic-agents/) ______________________________________________________________________ Use a predefined strategy that works for most common use cases - [**Functional agents**](agents/functional-agents/) ______________________________________________________________________ Define custom logic as a lambda function in plain Kotlin or Java - [**Graph-based agents**](agents/graph-based-agents/) ______________________________________________________________________ Implement a custom workflow as a strategy graph - [**Planner agents**](agents/planner-agents/) beta ______________________________________________________________________ Iteratively build and execute a plan until the state matches the desired conditions ## Core components Learn about the core components of Koog agents in detail: - [**Prompts**](prompts/) ______________________________________________________________________ Create, manage, and run prompts that drive the agent's interaction with the LLM - [**Strategies**](predefined-agent-strategies/) ______________________________________________________________________ Design the agent's intended workflow as a directed graph - [**Tools**](tools/) ______________________________________________________________________ Enable the agent to interact with external data sources and services - [**Features**](features/) ______________________________________________________________________ Extend and enhance the functionality of AI agents ## Advanced usage - [**History compression**](history-compression/) ______________________________________________________________________ Optimize token usage while maintaining context in long-running conversations using advanced techniques - [**Agent persistence**](features/agent-persistence/) ______________________________________________________________________ Restore the agent state at specific points during execution - [**Structured output**](structured-output/) ______________________________________________________________________ Generate responses in structured formats - [**Streaming API**](streaming-api/) ______________________________________________________________________ Process responses in real-time with streaming support and parallel tool calls - [**Knowledge retrieval**](embeddings/) beta ______________________________________________________________________ Retain and retrieve knowledge across conversations using [vector embeddings](embeddings/) and [RAG](retrieval-augmented-generation/) - [**Tracing**](features/tracing/) ______________________________________________________________________ Debug and monitor agent execution with detailed, configurable tracing - [**Long Term Memory**](features/long-term-memory/) beta ______________________________________________________________________ Integrate vector databases and memory providers for RAG and persistent memory. ## Integrations - [**Model Context Protocol (MCP)**](model-context-protocol/) beta ______________________________________________________________________ Use MCP tools directly in AI agents - [**Spring Boot**](spring-boot/) beta ______________________________________________________________________ Add Koog to your Spring applications - [**Ktor**](ktor-plugin/) beta ______________________________________________________________________ Integrate Koog with Ktor servers - [**OpenTelemetry**](features/open-telemetry/) ______________________________________________________________________ Trace, log, and measure your agent with popular observability tools - [**A2A Protocol**](a2a/) beta ______________________________________________________________________ Connect agents and services over a shared protocol # Key features Key features of Koog include: - **Idiomatic Kotlin and Java support**: Choose between a type-safe Kotlin DSL or a dedicated, fluent Java builder API. The Java API is designed to feel natural to Java teams, using standard thread pool executors instead of exposing coroutines. - **Reliability and fault-tolerance**: Handle failures with built-in retries and restore the agent state at specific points during execution with the agent persistence feature. - **Intelligent history compression**: Optimize token usage while maintaining context in long-running conversations using advanced built-in history compression techniques. - **Enterprise-ready integrations**: Utilize integration with popular JVM frameworks such as Spring Boot and Ktor to embed Koog into your applications. - **Observability with OpenTelemetry exporters**: Monitor and debug applications with built-in support for popular observability providers (W&B Weave, Langfuse). - **LLM switching and seamless history adaptation**: Switch to a different LLM at any point without losing the existing conversation history or reroute between multiple LLM providers. - **Multiplatform development**: For agents written in Kotlin, deploy agents across JVM, JS, WasmJS, Android, and iOS targets using Kotlin Multiplatform. - **Model Context Protocol integration**: Use Model Context Protocol (MCP) tools in AI agents. - **Knowledge retrieval and memory**: Retain and retrieve knowledge across conversations using vector embeddings and RAG. - **Powerful Streaming API**: Process responses in real-time with streaming support and parallel tool calls. - **Modular feature system**: Customize agent capabilities through a composable architecture. - **Flexible graph workflows**: Design complex agent behaviors using intuitive graph-based workflows. - **Custom tool creation**: Enhance your agents with tools that access external systems and APIs. - **Comprehensive tracing**: Debug and monitor agent execution with detailed, configurable tracing. # Versioning Koog follows [Semantic Versioning](https://semver.org/) with the format `X.Y.Z` (e.g., `1.2.0`). The framework is API-stable: once a public API is released, it will not be broken without a major version bump. ## Version Components | Component | Name | Format | Meaning | | --------- | ------ | ------- | ---------------------------------------------------------------------- | | `X` | Major | `X.y.z` | Breaking changes to existing APIs | | `Y` | Minor | `x.Y.z` | New API additions and deprecations; all existing APIs continue to work | | `Z` | Bugfix | `x.y.Z` | Bug fixes only; no API changes | ### Major (`X`) - May introduce breaking changes to existing APIs. - Old APIs may be removed. - A migration guide will be provided. - Released at most once per year. ### Minor (`Y`) - May add new APIs. - May deprecate existing APIs (with replacements provided), but deprecated APIs remain functional. - No breaking changes — all code that compiled against the previous minor version continues to compile and work. - Released at most once per month. ### Bugfix (`Z`) - Contains bug fixes only. - No API additions, removals, or deprecations. - Released at most once per week. ## Deprecation Policy APIs deprecated in a minor release (`Y`) will remain available until at least the next major release (`X`). Deprecation warnings will indicate the recommended replacement. ## Stable and Beta Modules Some modules are considered experimental and published with a `-beta` version suffix (e.g., `1.2.0-beta`) rather than the standard `X.Y.Z`. A module may be beta for one of several reasons: - **External integrations** — the underlying LLM provider API or external framework (e.g., Spring AI) may itself be unstable or subject to frequent or expected change. - **Experimental functionality** — the feature area is still being explored and the API shape may evolve (e.g., GOAP planning strategies). - **Experimental protocols** — the module implements a protocol that is not yet stable itself (e.g., A2A, Kotlin MCP). While every effort is made to keep beta modules stable, some API changes may occur across minor releases. Beta changes will not affect any stable module. A stable module at version `X.Y.Z` is always compatible with a beta module at version `X.Y.Z-beta` (and vice versa). All modules can be updated in sync. ### Umbrella Modules | Module | Version | Contents | | ----------------------- | ------------ | ------------------------------------------------------------------------ | | `koog-agents` | `1.2.0` | All stable modules (transitive) — recommended starting point | | `koog-agents-additions` | `1.2.0-beta` | Most beta/experimental modules (except standalone external integrations) | ### Module Versions | Module | Version | | ------------------------------------ | ------- | | `agents` | `1.2.0` | | `agents-core` | `1.2.0` | | `agents-features` | `1.2.0` | | `agents-features-chat-history-jdbc` | `1.2.0` | | `agents-features-chat-memory-sql` | `1.2.0` | | `agents-features-event-handler` | `1.2.0` | | `agents-features-memory` | `1.2.0` | | `agents-features-opentelemetry` | `1.2.0` | | `agents-features-persistence-jdbc` | `1.2.0` | | `agents-features-snapshot` | `1.2.0` | | `agents-features-sql` | `1.2.0` | | `agents-features-tokenizer` | `1.2.0` | | `agents-features-trace` | `1.2.0` | | `agents-mcp-metadata` | `1.2.0` | | `agents-test` | `1.2.0` | | `agents-tools` | `1.2.0` | | `agents-utils` | `1.2.0` | | `embeddings` | `1.2.0` | | `embeddings-base` | `1.2.0` | | `embeddings-llm` | `1.2.0` | | `http-client` | `1.2.0` | | `http-client-core` | `1.2.0` | | `http-client-java` | `1.2.0` | | `http-client-ktor` | `1.2.0` | | `http-client-okhttp` | `1.2.0` | | `http-client-test` | `1.2.0` | | `koog-agents` | `1.2.0` | | `koog-spring-ai` | `1.2.0` | | `prompt` | `1.2.0` | | `prompt-cache` | `1.2.0` | | `prompt-cache-files` | `1.2.0` | | `prompt-cache-model` | `1.2.0` | | `prompt-executor` | `1.2.0` | | `prompt-executor-anthropic-client` | `1.2.0` | | `prompt-executor-bedrock-client` | `1.2.0` | | `prompt-executor-cached` | `1.2.0` | | `prompt-executor-clients` | `1.2.0` | | `prompt-executor-model` | `1.2.0` | | `prompt-executor-ollama-client` | `1.2.0` | | `prompt-executor-openai-client` | `1.2.0` | | `prompt-executor-openai-client-base` | `1.2.0` | | `prompt-executor-openrouter-client` | `1.2.0` | | `prompt-llm` | `1.2.0` | | `prompt-markdown` | `1.2.0` | | `prompt-model` | `1.2.0` | | `prompt-processor` | `1.2.0` | | `prompt-structure` | `1.2.0` | | `prompt-tokenizer` | `1.2.0` | | `prompt-xml` | `1.2.0` | | `rag-base` | `1.2.0` | | `serialization` | `1.2.0` | | `serialization-core` | `1.2.0` | | `serialization-jackson` | `1.2.0` | | `serialization-test` | `1.2.0` | | `test-tck` | `1.2.0` | | `test-utils` | `1.2.0` | | `utils` | `1.2.0` | | Module | Version | | ---------------------------------------- | ------------ | | `a2a-client` | `1.2.0-beta` | | `a2a-core` | `1.2.0-beta` | | `a2a-server` | `1.2.0-beta` | | `a2a-test` | `1.2.0-beta` | | `a2a-test-server-tck` | `1.2.0-beta` | | `a2a-transport-client-jsonrpc-http` | `1.2.0-beta` | | `a2a-transport-core-jsonrpc` | `1.2.0-beta` | | `a2a-transport-server-jsonrpc-http` | `1.2.0-beta` | | `agents-ext` | `1.2.0-beta` | | `agents-features-a2a-client` | `1.2.0-beta` | | `agents-features-a2a-core` | `1.2.0-beta` | | `agents-features-a2a-server` | `1.2.0-beta` | | `agents-features-acp` | `1.2.0-beta` | | `agents-features-chat-history-aws` | `1.2.0-beta` | | `agents-features-longterm-memory` | `1.2.0-beta` | | `agents-features-longterm-memory-aws` | `1.2.0-beta` | | `agents-mcp` | `1.2.0-beta` | | `agents-mcp-server` | `1.2.0-beta` | | `agents-planner` | `1.2.0-beta` | | `koog-agents-additions` | `1.2.0-beta` | | `koog-ktor` | `1.2.0-beta` | | `koog-spring-ai-common` | `1.2.0-beta` | | `koog-spring-ai-starter-chat-memory` | `1.2.0-beta` | | `koog-spring-ai-starter-model-chat` | `1.2.0-beta` | | `koog-spring-ai-starter-model-embedding` | `1.2.0-beta` | | `koog-spring-ai-starter-vector-store` | `1.2.0-beta` | | `koog-spring-boot-starter` | `1.2.0-beta` | | `prompt-cache-redis` | `1.2.0-beta` | | `prompt-executor-dashscope-client` | `1.2.0-beta` | | `prompt-executor-deepseek-client` | `1.2.0-beta` | | `prompt-executor-google-client` | `1.2.0-beta` | | `prompt-executor-litert-client` | `1.2.0-beta` | | `prompt-executor-llms-all` | `1.2.0-beta` | | `prompt-executor-mistralai-client` | `1.2.0-beta` | | `rag-vector` | `1.2.0-beta` | # LLM providers Koog works with major LLM providers and also supports local models using [Ollama](https://ollama.com/). The following providers are currently supported: | LLM provider | Choose for | | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------- | | [OpenAI](https://platform.openai.com/docs/overview) (including [Azure OpenAI Service](https://azure.microsoft.com/en-us/products/ai-foundry/models/openai)) | Advanced models with a wide range of capabilities. | | [Anthropic](https://www.anthropic.com/) | Long contexts and prompt caching. | | [Google](https://ai.google.dev/) β | Multimodal processing (audio, video), large contexts. | | [DeepSeek](https://www.deepseek.com/) β | Cost-effective reasoning and coding. | | [OpenRouter](https://openrouter.ai/) | One integration with an access to multiple models from multiple providers for flexibility, provider comparison, and unified API. | | [Amazon Bedrock](https://aws.amazon.com/bedrock/) | AWS-native environment, enterprise security and compliance, multi-provider access. | | [Mistral](https://mistral.ai/) β | European data hosting, GDPR compliance. | | [Alibaba](https://www.alibabacloud.com/en?_p_lc=1) β ([DashScope](https://dashscope.aliyun.com/) OpenAI-compatible client) | Large contexts and cost-efficient Qwen models. | | [Ollama](https://ollama.com/) | Privacy, local development, offline operation, and no API costs. | The table below shows the LLM capabilities that Koog supports and which providers offer these capabilities in their models. | LLM capability | OpenAI | Anthropic | Google β | DeepSeek β | OpenRouter | Amazon Bedrock | Mistral β | Alibaba β (DashScope OpenAI-compatible client) | Ollama (local models) | | ------------------------------- | ---------------------------- | ------------------------------- | --------------------------------------------- | ---------- | ---------------- | ---------------- | ------------------------------- | ---------------------------------------------- | --------------------- | | Supported input | Text, image, audio, document | Text, image, document[1](#fn:1) | Text, image, audio, video, document[1](#fn:1) | Text | Differs by model | Differs by model | Text, image, document[1](#fn:1) | Text, image, audio, video[1](#fn:1) | Text, image[1](#fn:1) | | Response streaming | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Tools | ✓ | ✓ | ✓ | ✓ | ✓ | ✓[1](#fn:1) | ✓ | ✓ | ✓ | | Tool choice | ✓ | ✓ | ✓ | ✓ | ✓ | ✓[1](#fn:1) | ✓ | ✓ | – | | Structured output (JSON Schema) | ✓ | ✓[1](#fn:1) | ✓ | ✓ | ✓[1](#fn:1) | – | ✓ | ✓[1](#fn:1) | ✓ | | Multiple choices | ✓ | – | ✓ | – | ✓[1](#fn:1) | ✓[1](#fn:1) | ✓ | ✓[1](#fn:1) | – | | Temperature | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Speculation | ✓[1](#fn:1) | – | – | – | ✓[1](#fn:1) | – | ✓[1](#fn:1) | ✓[1](#fn:1) | – | | Content moderation | ✓ | – | – | – | – | ✓ | ✓ | – | ✓ | | Embeddings | ✓ | – | – | – | – | ✓ | ✓ | – | ✓ | | Prompt caching | ✓[1](#fn:1) | ✓ | – | – | – | – | – | – | – | | Completion | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Local execution | – | – | – | – | – | – | – | – | ✓ | Note Koog supports the most commonly used capabilities for creating AI agents. LLMs from each provider may have additional features that Koog does not currently support. To learn more, refer to [Model capabilities](../model-capabilities/). ## Working with providers Koog lets you work with LLM providers on two levels: - Using an **LLM client** for direct interaction with a specific provider. Each client implements the `LLMClient` interface, handling authentication, request formatting, and response parsing for the provider. For details, see [LLM clients](../prompts/llm-clients/). - Using a **prompt executor** for a higher-level abstraction that wraps one or multiple LLM clients, manages their lifecycles, and unifies an interface across providers. It can switch between providers and optionally fall back to a configured provider and LLM using the corresponding client. You can either create your own executor or use a pre-defined prompt executor for a specific provider. For details, see [Prompt executors](../prompts/prompt-executors/). Using a prompt executor offers a higher‑level layer over one or more LLMClients. It manages client lifecycles and exposes a unified interface across providers. In multi‑provider setups, it can route requests between providers and optionally fall back to a designated client when needed for core requests. You can create your own executor or use pre‑defined ones—both single‑provider and multi‑provider options are available. ## Next steps - [Create and run an agent](../quickstart/) with a specific LLM provider. - Learn more about [prompts](../prompts/). ______________________________________________________________________ 1. Capability is supported only by some models of the provider. [↩](#fnref:1 "Jump back to footnote 1 in the text")[↩](#fnref2:1 "Jump back to footnote 1 in the text")[↩](#fnref3:1 "Jump back to footnote 1 in the text")[↩](#fnref4:1 "Jump back to footnote 1 in the text")[↩](#fnref5:1 "Jump back to footnote 1 in the text")[↩](#fnref6:1 "Jump back to footnote 1 in the text")[↩](#fnref7:1 "Jump back to footnote 1 in the text")[↩](#fnref8:1 "Jump back to footnote 1 in the text")[↩](#fnref9:1 "Jump back to footnote 1 in the text")[↩](#fnref10:1 "Jump back to footnote 1 in the text")[↩](#fnref11:1 "Jump back to footnote 1 in the text")[↩](#fnref12:1 "Jump back to footnote 1 in the text")[↩](#fnref13:1 "Jump back to footnote 1 in the text")[↩](#fnref14:1 "Jump back to footnote 1 in the text")[↩](#fnref15:1 "Jump back to footnote 1 in the text")[↩](#fnref16:1 "Jump back to footnote 1 in the text")[↩](#fnref17:1 "Jump back to footnote 1 in the text")[↩](#fnref18:1 "Jump back to footnote 1 in the text") # Glossary ## Agent - **Agent**: an AI entity that can interact with tools, handle complex workflows, and communicate with users. - **LLM (Large Language Model)**: the underlying AI model that powers agent capabilities. - **Message**: a unit of communication in the agent system that represents data passed from a user, assistant, or system. - **Prompt**: the conversation history provided to an LLM that consists of messages from a user, assistant, and system. - **System prompt**: instructions provided to an agent to guide its behavior, define its role, and supply key information necessary for its tasks. - **Context**: the environment in which LLM interactions occur, with access to the conversation history and tools. - **LLM session**: a structured way to interact with LLMs that includes the conversation history, available tools, and methods to make requests. ## Agent workflow - **Strategy**: a defined workflow for an agent that consists of sequential subgraphs. The strategy defines how the agent processes input, interacts with tools, and generates output. A strategy graph consists of nodes connected by edges that represent transitions between nodes. ### Strategy graphs - **Graph**: a structure of nodes connected by edges that defines an agent strategy workflow. - **Node**: a fundamental building block of an agent strategy workflow that represents a specific operation or transformation. - **Edge**: a connection between nodes in an agent graph that defines the flow of operations, often with conditions that specify when to follow each edge. - **Conditions**: rules that determine when to follow a particular edge. - **Subgraph**: a self-contained unit of processing within an agent strategy, with its own set of tools, context, and responsibilities. ## Tools - **Tool**: a function that an agent can use to perform specific tasks or access external systems. The agent is aware of the available tools and their arguments but lacks knowledge of their implementation details. - **Tool call**: a request from an LLM to run a specific tool using the provided arguments. It functions similarly to a function call. - **Tool descriptor**: tool metadata that includes its name, description, and parameters. - **Tool registry**: a list of tools available to an agent. The registry informs the agent about the available tools. - **Tool result**: an output produced by running a tool. For example, if the tool is a method, the result would be its return value. ## History compression - **History compression**: the process of reducing the size of the conversation history to manage token usage by applying various compression strategies. To learn more, see [History compression](../history-compression/). ## Features - **Feature**: a component that extends and enhances the functionality of AI agents. ### EventHandler feature - **EventHandler**: a feature that enables monitoring and responding to various agent events, providing hooks for tracking agent lifecycle, handling errors, and processing tool invocations throughout the workflow. # Quickstart # Quickstart This guide will help you start using Koog in your project. ## Prerequisites Ensure your environment and project meet the following requirements: - JDK 17+ - Kotlin 2.2.0+ - Gradle 8.0+ or Maven 3.8+ ## Install Koog Add the [Koog package](https://central.sonatype.com/artifact/ai.koog/koog-agents/) as a dependency: build.gradle.kts ``` dependencies { // Stable implementation("ai.koog:koog-agents:1.2.0") // Beta implementation("ai.koog:koog-agents-additions:1.2.0-beta") } ``` build.gradle ``` dependencies { // Stable implementation 'ai.koog:koog-agents:1.2.0' // Beta implementation 'ai.koog:koog-agents-additions:1.2.0-beta' } ``` pom.xml ``` ai.koog koog-agents-jvm 1.2.0 ai.koog koog-agents-additions-jvm 1.2.0-beta ``` Module Versioning Koog follows Semantic Versioning (`X.Y.Z`). Stable modules (e.g., `1.0.0`) have guaranteed APIs, while beta modules (e.g., `1.0.0-beta`) are experimental and may change between releases. See [Module versioning](../module-versioning/) for details. Nightly builds Nightly builds from the develop branch are published to the [JetBrains Grazie Maven](https://packages.jetbrains.team/maven/p/grazi/grazie-platform-public) repository. To use a nightly build, add the following repository to your build configuration: `https://packages.jetbrains.team/maven/p/grazi/grazie-platform-public`. Then update your Koog dependency to the desired nightly version. Nightly versions follow the pattern `[next-major-version]-develop-[date]-[time]`. You can browse the available nightly builds [here](https://packages.jetbrains.team/maven/p/grazi/grazie-platform-public/ai/koog/koog-agents/). ## Set up an API key Koog requires either an API key from a [supported LLM provider](../llm-providers/) or a locally running LLM. Warning Avoid hardcoding API keys in the source code. Use environment variables to store API keys. Get your [OpenAI API key](https://platform.openai.com/api-keys) and assign it to the `OPENAI_API_KEY` environment variable. ``` export OPENAI_API_KEY=your-api-key ``` ``` setx OPENAI_API_KEY "your-api-key" ``` Get your [Anthropic API key](https://console.anthropic.com/settings/keys) and assign it to the `ANTHROPIC_API_KEY` environment variable. ``` export ANTHROPIC_API_KEY=your-api-key ``` ``` setx ANTHROPIC_API_KEY "your-api-key" ``` Get your [Gemini API key](https://aistudio.google.com/app/api-keys) and assign it to the `GOOGLE_API_KEY` environment variable. ``` export GOOGLE_API_KEY=your-api-key ``` ``` setx GOOGLE_API_KEY "your-api-key" ``` Get your [DeepSeek API key](https://platform.deepseek.com/api_keys) and assign it to the `DEEPSEEK_API_KEY` environment variable. ``` export DEEPSEEK_API_KEY=your-api-key ``` ``` setx DEEPSEEK_API_KEY "your-api-key" ``` Get your [OpenRouter API key](https://openrouter.ai/keys) and assign it to the `OPENROUTER_API_KEY` environment variable. ``` export OPENROUTER_API_KEY=your-api-key ``` ``` setx OPENROUTER_API_KEY "your-api-key" ``` [Generate an Amazon Bedrock API key](https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys.html) and assign it to the `BEDROCK_API_KEY` environment variable. ``` export BEDROCK_API_KEY=your-api-key ``` ``` setx BEDROCK_API_KEY "your-api-key" ``` Get your [Mistral API key](https://console.mistral.ai/api-keys) and assign it to the `MISTRAL_API_KEY` environment variable. ``` export MISTRAL_API_KEY=your-api-key ``` ``` setx MISTRAL_API_KEY "your-api-key" ``` Run a local LLM in Ollama as described in the [Ollama documentation](https://docs.ollama.com/quickstart). ## Create your first Koog agent The following example creates and runs a simple Koog agent using the [`GPT-4o`](https://platform.openai.com/docs/models/gpt-4o) model via the OpenAI API. ``` fun main() = runBlocking { // Get the OpenAI API key from the OPENAI_API_KEY environment variable val apiKey = System.getenv("OPENAI_API_KEY") ?: error("The API key is not set.") // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor(OpenAILLMClient(apiKey)), llmModel = OpenAIModels.Chat.GPT4o ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Get the OpenAI API key from the OPENAI_API_KEY environment variable String apiKey = System.getenv("OPENAI_API_KEY"); if (apiKey == null) { throw new RuntimeException("The API key is not set."); } // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(new MultiLLMPromptExecutor(openAIClient(apiKey))) .llmModel(OpenAIModels.Chat.GPT4o) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` Hello! I'm here to help you with whatever you need. Here are just a few things I can do: - Answer questions. - Explain concepts or topics you're curious about. - Provide step-by-step instructions for tasks. - Offer advice, notes, or ideas. - Help with research or summarize complex material. - Write or edit text, emails, or other documents. - Brainstorm creative projects or solutions. - Solve problems or calculations. Let me know what you need help with—I’m here for you! ``` The following example creates and runs a simple Koog agent using the [`Claude Opus 4.1`](https://www.anthropic.com/news/claude-opus-4-1) model via the Anthropic API. ``` fun main() = runBlocking { // Get the Anthropic API key from the ANTHROPIC_API_KEY environment variable val apiKey = System.getenv("ANTHROPIC_API_KEY") ?: error("The API key is not set.") // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor(AnthropicLLMClient(apiKey)), llmModel = AnthropicModels.Opus_4_1 ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Get the Anthropic API key from the ANTHROPIC_API_KEY environment variable String apiKey = System.getenv("ANTHROPIC_API_KEY"); if (apiKey == null) { throw new RuntimeException("The API key is not set."); } // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(new MultiLLMPromptExecutor(anthropicClient(apiKey))) .llmModel(AnthropicModels.Opus_4_1) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` Hello! I can help you with: - **Answering questions** and explaining topics - **Writing** - drafting, editing, proofreading - **Learning** - homework, math, study help - **Problem-solving** and brainstorming - **Research** and information finding - **General tasks** - instructions, planning, recommendations What do you need help with today? ``` The following example creates and runs a simple Koog agent using the [`Gemini 2.5 Pro`](https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/2-5-pro) model via the Gemini API. ``` fun main() = runBlocking { // Get the Gemini API key from the GOOGLE_API_KEY environment variable val apiKey = System.getenv("GOOGLE_API_KEY") ?: error("The API key is not set.") // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor(GoogleLLMClient(apiKey)), llmModel = GoogleModels.Gemini2_5Pro ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Get the Gemini API key from the GOOGLE_API_KEY environment variable String apiKey = System.getenv("GOOGLE_API_KEY"); if (apiKey == null) { throw new RuntimeException("The API key is not set."); } // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(new MultiLLMPromptExecutor(googleClient(apiKey))) .llmModel(GoogleModels.Gemini2_5Pro) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` I'm an AI that can help you with tasks involving language and information. You can ask me to: * **Answer questions** * **Write or edit text** (emails, stories, code, etc.) * **Brainstorm ideas** * **Summarize long documents** * **Plan things** (like trips or projects) * **Be a creative partner** Just tell me what you need ``` The following example creates and runs a simple Koog agent using the `deepseek-v4-flash` model via the DeepSeek API. ``` fun main() = runBlocking { // Get the DeepSeek API key from the DEEPSEEK_API_KEY environment variable val apiKey = System.getenv("DEEPSEEK_API_KEY") ?: error("The API key is not set.") // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor(DeepSeekLLMClient(apiKey)), llmModel = DeepSeekModels.DeepSeekV4Flash ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Get the DeepSeek API key from the DEEPSEEK_API_KEY environment variable String apiKey = System.getenv("DEEPSEEK_API_KEY"); if (apiKey == null) { throw new RuntimeException("The API key is not set."); } // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(new MultiLLMPromptExecutor(deepSeekClient(apiKey))) .llmModel(DeepSeekModels.DeepSeekV4Flash) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` Hello! I'm here to assist you with a wide range of tasks, including answering questions, providing information, helping with problem-solving, offering creative ideas, and even just chatting. Whether you need help with research, writing, learning something new, or simply want to discuss a topic, feel free to ask—I’m happy to help! 😊 ``` The following example creates and runs a simple Koog agent using the [`GPT-4o`](https://openrouter.ai/openai/gpt-4o) model via the OpenRouter API. ``` fun main() = runBlocking { // Get the OpenRouter API key from the OPENROUTER_API_KEY environment variable val apiKey = System.getenv("OPENROUTER_API_KEY") ?: error("The API key is not set.") // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor(OpenRouterLLMClient(apiKey)), llmModel = OpenRouterModels.GPT4o ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Get the OpenRouter API key from the OPENROUTER_API_KEY environment variable String apiKey = System.getenv("OPENROUTER_API_KEY"); if (apiKey == null) { throw new RuntimeException("The API key is not set."); } // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(new MultiLLMPromptExecutor(openRouterClient(apiKey))) .llmModel(OpenRouterModels.GPT4o) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` I can answer questions, help with writing, solve problems, organize tasks, and more—just let me know what you need! ``` The following example creates and runs a simple Koog agent using the [`Claude Sonnet 4.5`](https://www.anthropic.com/news/claude-sonnet-4-5) model via the Bedrock API. ``` fun main() = runBlocking { // Get the Bedrock API key from the BEDROCK_API_KEY environment variable val apiKey = System.getenv("BEDROCK_API_KEY") ?: error("The API key is not set.") // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor( BedrockLLMClient( StaticBearerTokenProvider(apiKey), BedrockClientSettings() ) ), llmModel = BedrockModels.AnthropicClaude4_5Sonnet ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Get the Bedrock API key from the BEDROCK_API_KEY environment variable String apiKey = System.getenv("BEDROCK_API_KEY"); if (apiKey == null) { throw new RuntimeException("The API key is not set."); } // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(simpleBedrockExecutorWithBearerToken(apiKey, new BedrockClientSettings())) .llmModel(BedrockModels.INSTANCE.getAnthropicClaude4_5Sonnet()) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` Hello! I'm a helpful assistant and I can assist you in many ways, including: - **Answering questions** on a wide range of topics (science, history, technology, etc.) - **Writing help** - drafting emails, essays, creative content, or editing text - **Problem-solving** - working through math problems, logic puzzles, or troubleshooting issues - **Learning support** - explaining concepts, providing study notes, or tutoring - **Planning & organizing** - helping with projects, schedules, or breaking down tasks - **Coding assistance** - explaining programming concepts or helping debug code - **Creative brainstorming** - generating ideas for projects, stories, or solutions - **General conversation** - discussing topics or just chatting What would you like help with today? ``` The following example creates and runs a simple Koog agent using the [`Mistral Medium 3.1`](https://docs.mistral.ai/models/mistral-medium-3-1-25-08) model via the Mistral AI API. ``` fun main() = runBlocking { // Get the Mistral AI API key from the MISTRAL_API_KEY environment variable val apiKey = System.getenv("MISTRAL_API_KEY") ?: error("The API key is not set.") // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor(MistralAILLMClient(apiKey)), llmModel = MistralAIModels.Chat.MistralMedium31 ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Get the Mistral AI API key from the MISTRAL_API_KEY environment variable String apiKey = System.getenv("MISTRAL_API_KEY"); if (apiKey == null) { throw new RuntimeException("The API key is not set."); } // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(new MultiLLMPromptExecutor(mistralAIClient(apiKey))) .llmModel(MistralAIModels.Chat.MistralMedium31) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` I can assist you with a wide range of topics and tasks. Here are some examples: 1. **Answering questions**: I can provide information on various subjects, including history, science, technology, literature, and more. 2. **Providing definitions**: If you're unsure about the meaning of a word or phrase, I can help define it for you. 3. **Generating text**: Whether it's writing an email, creating content for social media, or composing a story, I can help with text generation. 4. **Translation**: I can translate text from one language to another. 5. **Conversation**: We can have a chat about any topic that interests you, and I'll respond accordingly. 6. **Language practice**: If you're learning a new language, I can help with pronunciation, grammar, and vocabulary practice. 7. **Brainstorming**: If you're stuck on a problem or need ideas for a project, I can help brainstorm solutions. 8. **Summarization**: If you have a long piece of text and want a summary, I can condense it for you. What's on your mind? Is there something specific you'd like help with? ``` The following example creates and runs a simple Koog agent using the [`llama3.2`](https://ollama.com/library/llama3.2) model running locally via Ollama. ``` fun main() = runBlocking { // Create an agent val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor(OllamaClient()), llmModel = OllamaModels.Meta.LLAMA_3_2 ) // Run the agent val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` // Create an agent AIAgent agent = AIAgent.builder() .promptExecutor(new MultiLLMPromptExecutor(ollamaClient("http://localhost:11434"))) .llmModel(OllamaModels.Meta.LLAMA_3_2) .build(); // Run the agent String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The example can produce the following output: ``` I can assist with various tasks such as answering questions, providing information, and even helping with language-related tasks like proofreading or writing suggestions. What's on your mind today? ``` ## Next steps - Learn more about [agent types](../agents/) # Agents # Basic agents A basic agent uses a predefined strategy with a simple execution flow that works for most common use cases. It accepts a string input (a question, request, or task description) and sends this input to the configured LLM. The LLM may decide to call provided tools. The agent will execute the tools and send the results back to the LLM. This repeats until the LLM does not request any more tool calls and returns a string response. The agent then outputs this response. In [Graph-based agents](../graph-based-agents/), you can see how to re-create the predefined strategy graph used by basic agents. Prerequisites Ensure your environment and project meet the following requirements: - JDK 17+ - Kotlin 2.2.0+ - Gradle 8.0+ or Maven 3.8+ Add the [Koog package](https://central.sonatype.com/artifact/ai.koog/koog-agents/) as a dependency: build.gradle.kts ``` dependencies { // Stable implementation("ai.koog:koog-agents:1.2.0") // Beta implementation("ai.koog:koog-agents-additions:1.2.0-beta") } ``` build.gradle ``` dependencies { // Stable implementation 'ai.koog:koog-agents:1.2.0' // Beta implementation 'ai.koog:koog-agents-additions:1.2.0-beta' } ``` pom.xml ``` ai.koog koog-agents-jvm 1.2.0 ai.koog koog-agents-additions-jvm 1.2.0-beta ``` Get an API key from an LLM provider or run a local LLM via Ollama. For more information, see [Quickstart](../../quickstart/). Examples on this page assume that you have set the `OPENAI_API_KEY` environment variable. ## Create a minimal agent To create the most basic agent, instantiate [`AIAgent`](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent/-a-i-agent/index.html) and provide a [prompt executor](../../prompts/prompt-executors/) with a [language model](../../model-capabilities/#creating-a-model-llmodel-configuration): ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4o ) ``` This agent will expect a string as input and return a string as output. To run the agent, use the `run()` function with some user input: ``` fun main() = runBlocking { val result = agent.run("Hello! How can you help me?") println(result) } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .llmModel(OpenAIModels.Chat.GPT4o) .build(); ``` This agent expects a string as input and returns a string as output. To run the agent, use the `run()` method with some user input: ``` String result = agent.run("Hello! How can you help me?"); System.out.println(result); ``` The agent will return a generic answer, such as: ``` I can assist with a wide range of topics and tasks. Here are some examples: 1. **Answering questions**: I can provide information on various subjects, from science and history to entertainment and culture. 2. **Generating text**: I can help with writing tasks, such as suggesting alternative phrases, providing definitions, or even creating entire articles or stories. 3. **Translation**: I can translate text from one language to another, including popular languages such as Spanish, French, German, Chinese, and many more. 4. **Conversation**: I can engage in natural-sounding conversations, using context and understanding to respond to questions and statements. 5. **Brainstorming**: I can help generate ideas for creative projects, such as writing stories, composing music, or coming up with business ideas. 6. **Learning**: I can help with language learning, explaining grammar rules, vocabulary, and pronunciation. 7. **Calculations**: I can perform mathematical calculations, including basic arithmetic, algebra, and more advanced math concepts. What's on your mind? Do you have a specific question, topic, or task you'd like to tackle? ``` ## Add a system prompt Provide a [system message](../../prompts/prompt-creation/#system-message) to define the agent's role as well as the purpose, context, and instructions related to the task. ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("YOUR_API_KEY")), systemPrompt = "You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.", llmModel = OpenAIModels.Chat.GPT4o ) ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.") .llmModel(OpenAIModels.Chat.GPT4o) .build(); ``` The instructions in the system prompt will guide the agent's response: ``` I'm here to help you navigate the wild world of internet memes! What's on your mind? Are you trying to understand a specific meme, need help finding a popular joke, or perhaps want some recommendations for trending memes? Let me know, and I'll do my best to provide you with some LOLs! ``` ## Configure LLM output You can provide some [LLM parameters](../../llm-parameters/#llm-parameter-reference) directly to the agent constructor (Kotlin) or via the builder methods (Java) to customize the behavior of the LLM. For example, use the `temperature` parameter to adjust the randomness of the generated responses: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("YOUR_API_KEY")), systemPrompt = "You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.", llmModel = OpenAIModels.Chat.GPT4o, temperature = 0.7 ) ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.") .llmModel(OpenAIModels.Chat.GPT4o) .temperature(0.7) .build(); ``` Here are some response examples with different temperature values: ``` I'm here to help you navigate the wild world of internet memes! Whether you're looking for explanations, examples, or just want to share a meme with someone, I'm your go-to expert. What's on your mind? Got a specific meme in mind that's got you curious? Or maybe you need some meme-related advice? Fire away! ``` ``` I'm here to help you navigate the wild world of internet memes! What's on your mind? Need help understanding a specific meme, finding a popular joke or trend, or maybe even creating your own meme? Let's get this meme party started! ``` ``` I'd be happy to help you navigate the wild world of internet memes! Whether you're looking for explanations of classic memes, suggestions for new ones to try out, or just want to discuss your favorite meme culture trends, I'm here to assist. What's on your mind? Do you have a specific question about memes (e.g., "What does this meme mean?"), or are you looking for some meme-related recommendations (e.g., "Can you recommend a funny meme to share with friends?"). Let me know how I can help! ``` ## Add tools Agents can use [tools](../../tools/) to perform specific tasks. First, create a tool by annotating a function (Kotlin) or method (Java) with the [`@Tool`](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools.annotations/-tool/index.html) annotation: ``` @Tool @LLMDescription("Ask the user a question by sending it to stdout and return the answer from stdin") fun askUser( @LLMDescription("Question from the agent") question: String ): String { println(question) return readln() } ``` Then, use the [`ToolRegistry`](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool-registry/index.html) to make this tool available to the agent: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("YOUR_API_KEY")), systemPrompt = "You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.", llmModel = OpenAIModels.Chat.GPT4o, temperature = 0.7, toolRegistry = ToolRegistry { tool(::askUser) } ) ``` In the example, `askUser` is a tool that helps the agent maintain a conversation with the user via printing and reading from the console. If the agent decides to ask the user a question, it can call this tool that writes to `stdout` via `println()` and reads from `stdin` via `readln()`. ``` // Create a ToolSet class class UserConversationTools implements ToolSet { @Tool @LLMDescription("Ask the user a question by sending it to stdout and return the answer from stdin") public String askUser( @LLMDescription("Question from the agent") String question ) { System.out.println(question); Scanner scanner = new Scanner(System.in); return scanner.nextLine(); } } ``` Then, use the [`ToolRegistry`](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool-registry/index.html) to make this tool available to the agent: ``` UserConversationTools askUser = new UserConversationTools(); ToolRegistry toolRegistry = ToolRegistry.builder() .tools(askUser) .build(); AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.") .llmModel(OpenAIModels.Chat.GPT4o) .temperature(0.7) .toolRegistry(toolRegistry) .build(); ``` In the example, `askUser` is a tool that helps the agent maintain a conversation with the user via printing and reading from the console. Here is an example interaction with the agent: ``` Agent: Which meme would you like me to explain? Please choose from: Grumpy Cat, Success Kid, or Doge. User: Explain Doge Agent: **Doge** Doge is a popular internet meme that originated on the imageboard website 4chan in 2013. The meme features a picture of a Shiba Inu dog with a caption written in a playful, exaggerated style. The typical format of a Doge meme includes: * A photo of a Shiba Inu dog * A caption written in all capital letters, using a intentionally simplistic and childlike tone * Exaggerated or made-up words or phrases, often used to convey a humorous or nonsensical idea Examples of Doge memes might include: * "Such wow. Such happy." * "I had fun today!" * "Wow, I am good at napping." The meme is known for its lighthearted and playful tone, and is often used to express excitement, happiness, or silliness. The meme has since become a cultural phenomenon, with countless variations and parodies emerging online. ``` ## Adjust agent iterations To avoid infinite loops, Koog allows any agent to take a limited number of steps (50 by default). Use the `maxIterations` parameter to either increase this limit if you expect the agent to require more steps (such as tool calls and LLM requests) or decrease it for agents that require only a few steps. For example, a simple agent described here is not likely to require more than 10 steps: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("YOUR_API_KEY")), systemPrompt = "You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.", llmModel = OpenAIModels.Chat.GPT4o, temperature = 0.7, toolRegistry = ToolRegistry { tool(::askUser) }, maxIterations = 10 ) ``` ``` // Create a ToolSet class class UserConversationTools implements ToolSet { @Tool @LLMDescription("Ask the user a question by sending it to stdout and return the answer from stdin") public String askUser( @LLMDescription("Question from the agent") String question ) { System.out.println(question); Scanner scanner = new Scanner(System.in); return scanner.nextLine(); } } // In main method: UserConversationTools askUser = new UserConversationTools(); ToolRegistry toolRegistry = ToolRegistry.builder() .tools(askUser) .build(); AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.") .llmModel(OpenAIModels.Chat.GPT4o) .temperature(0.7) .toolRegistry(toolRegistry) .maxIterations(10) .build(); ``` Tip Instead of passing the model, temperature, max iterations, and other parameters directly to the Kotlin constructor or Java builder, you can also define and pass them as a separate configuration object. For more information, see [Agent configuration](../#agent-configuration). ## Handle events during agent runtime To assist with testing and debugging, as well as making hooks for chained agent interactions, Koog provides the [EventHandler](https://api.koog.ai/agents/agents-features/agents-features-event-handler/ai.koog.agents.features.eventHandler.feature/-event-handler/index.html) feature. Call the `handleEvents()` function inside the agent constructor lambda to install the feature and register event handlers: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("YOUR_API_KEY")), systemPrompt = "You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.", llmModel = OpenAIModels.Chat.GPT4o, temperature = 0.7, toolRegistry = ToolRegistry { tool(::askUser) }, maxIterations = 10 ){ handleEvents { // Handle tool calls onToolCallStarting { eventContext -> println("Tool called: ${eventContext.toolName} with args ${eventContext.toolArgs}") } } } ``` Use the `.install()` method on the agent builder to register event handlers with `EventHandler.Feature`: ``` // Create a ToolSet class class UserConversationTools implements ToolSet { @Tool @LLMDescription("Ask the user a question by sending it to stdout and return the answer from stdin") public String askUser( @LLMDescription("Question from the agent") String question ) { System.out.println(question); Scanner scanner = new Scanner(System.in); return scanner.nextLine(); } } // In main method: UserConversationTools askUser = new UserConversationTools(); ToolRegistry toolRegistry = ToolRegistry.builder() .tools(askUser) .build(); AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("You are an expert in internet memes. Be helpful, friendly, and answer user questions concisely, showing your knowledge of memes.") .llmModel(OpenAIModels.Chat.GPT4o) .temperature(0.7) .toolRegistry(toolRegistry) .maxIterations(10) .install(EventHandler.Feature, config -> { config.onToolCallStarting(eventContext -> { System.out.println("Tool called: " + eventContext.getToolName() + " with args " + eventContext.getToolArgs()); }); }) .build(); ``` The agent will now output something similar to the following when it calls the `askUser` tool: ``` Tool called: askUser with args {"question":"Which meme would you like me to explain?"} ``` For more information about Koog agent features, see [Features](../../features/). ## Next steps - Learn more about building [graph-based agents](../graph-based-agents/) and [functional agents](../functional-agents/) # Graph-based agents With graph-based agents, you model the behavior as an explicit state machine: nodes of a graph strategy represent actions (LLM calls, tool execution) and edges represent data flow between nodes. The main advantages of graph-based agents are: - Easy to visualize - State persistence - Composable architecture Prerequisites Ensure your environment and project meet the following requirements: - JDK 17+ - Kotlin 2.2.0+ - Gradle 8.0+ or Maven 3.8+ Add the [Koog package](https://central.sonatype.com/artifact/ai.koog/koog-agents/) as a dependency: build.gradle.kts ``` dependencies { // Stable implementation("ai.koog:koog-agents:1.2.0") // Beta implementation("ai.koog:koog-agents-additions:1.2.0-beta") } ``` build.gradle ``` dependencies { // Stable implementation 'ai.koog:koog-agents:1.2.0' // Beta implementation 'ai.koog:koog-agents-additions:1.2.0-beta' } ``` pom.xml ``` ai.koog koog-agents-jvm 1.2.0 ai.koog koog-agents-additions-jvm 1.2.0-beta ``` Get an API key from an LLM provider or run a local LLM via Ollama. For more information, see [Quickstart](../../quickstart/). Examples on this page assume that you are running Llama 3.2 locally via Ollama. This page describes how to re-create the strategy graph used by [basic agents](../basic-agents/). It sends a request to an LLM and then either outputs the response (if the LLM responded with an assistant message) or executes a tool (if the LLM requested a tool call). In case of a tool call, the agent sends the tool result to the LLM and then either outputs the response or executes a tool. Here is an illustration of the strategy graph: ``` --- config: flowchart: defaultRenderer: "elk" --- graph TB subgraph nodeStart Input end subgraph nodeFinish Output end subgraph nodeSendInput llmRequest(Request LLM) end subgraph nodeExecuteTool executeTool(Execute tool call) end subgraph nodeSendToolResult sendToolResult(Request LLM) end Input --String--> llmRequest llmRequest --Message.Assistant--> onToolCalls{{onToolCalls}} llmRequest --Message.Assistant--> onTextMessage{{onTextMessage}} onTextMessage --String--> Output onToolCalls --ToolCalls--> executeTool --ReceivedToolResults--> sendToolResult sendToolResult --Message.Assistant--> onToolCalls sendToolResult --Message.Assistant--> onTextMessage ``` ## Build a strategy graph In Koog, you implement a strategy using [`AIAgentGraphStrategyBuilder`](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.builder/-a-i-agent-graph-strategy-builder/index.html). Just like every node has an input and output type, the strategy as a whole also defines some input and output type. This example assumes that the input and output types are strings, which means the agent implementing this strategy will expect a string and return a string. To create a strategy, use the [`strategy()`](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.builder/strategy.html) function with two generics as the input and output types, provide a unique identifier for the strategy, and define the nodes and edges. ``` val calculatorAgentStrategy = strategy("Simple calculator") { val nodeSendInput by nodeLLMRequest() val nodeExecuteTool by nodeExecuteTools() val nodeSendToolResult by nodeLLMSendToolResults() edge(nodeStart forwardTo nodeSendInput) edge(nodeSendInput forwardTo nodeFinish onTextMessage { true }) edge(nodeSendInput forwardTo nodeExecuteTool onToolCalls { true }) edge(nodeExecuteTool forwardTo nodeSendToolResult) edge(nodeSendToolResult forwardTo nodeFinish onTextMessage { true }) edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCalls { true }) } ``` ``` var calculatorAgentStrategy = AIAgentGraphStrategy.builder("Simple calculator") .withInput(String.class) .withOutput(String.class); var nodeSendInput = AIAgentNode.llmRequest("nodeSendInput"); var nodeExecuteTool = AIAgentNode.executeTools("nodeExecuteTool"); var nodeSendToolResult = AIAgentNode.llmSendToolResults("nodeSendToolResult"); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(calculatorAgentStrategy.nodeStart) .to(nodeSendInput) .build()); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendInput) .to(calculatorAgentStrategy.nodeFinish) .onTextMessage() .build()); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendInput) .to(nodeExecuteTool) .onToolCalls() .build()); calculatorAgentStrategy.edge(nodeExecuteTool, nodeSendToolResult); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(calculatorAgentStrategy.nodeFinish) .onTextMessage() .build()); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(nodeExecuteTool) .onToolCalls() .build()); ``` This example uses only [predefined nodes](../../nodes-and-components/), but you can also create [custom nodes](../../custom-nodes/). Every strategy graph must have a path from `nodeStart` to `nodeFinish` connected by [edges](../../custom-strategy-graphs/#edges). Edges can have conditions to determine when to follow a particular edge. Edges can also transform the output of the previous node before passing it to the next one. This is necessary to connect nodes that have non-matching output and input types. In the previous example, `onToolCalls { true }` means that the edge will follow only if the previous node returned an assistant message containing at least one tool call (`MessagePart.Tool.Call`). When using `onTextMessage { true }`, the edge will follow only if the previous node returned an assistant message containing text parts (`MessagePart.Text`). This function also extracts and joins the text content of those parts, effectively transforming `Message.Assistant` to `String`, because `nodeFinish` expects a string. Tip Instead of `onTextMessage { true }`, you can do the following: ``` onMessageParts(MessagePart.Text::class) transformed { it.joinToString("\n") { part -> part.text } } ``` Or: ``` onCondition { it is Message.Assistant } transformed { (it as Message.Assistant).parts.filterIsInstance().joinToString("\n") { part -> part.text } } ``` ## Create and run the agent Let's create an agent instance with this strategy and run it: ``` val calculatorAgentStrategy = strategy("Simple calculator") { val nodeSendInput by nodeLLMRequest() val nodeExecuteTool by nodeExecuteTools() val nodeSendToolResult by nodeLLMSendToolResults() edge(nodeStart forwardTo nodeSendInput) edge(nodeSendInput forwardTo nodeFinish onTextMessage { true }) edge(nodeSendInput forwardTo nodeExecuteTool onToolCalls { true }) edge(nodeExecuteTool forwardTo nodeSendToolResult) edge(nodeSendToolResult forwardTo nodeFinish onTextMessage { true }) edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCalls { true }) } val mathAgent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, strategy = calculatorAgentStrategy ) fun main() = runBlocking { val result = mathAgent.run("Multiply 3 by 4, then multiply the result by 5, then add 10, then add 123.") println(result) } ``` ``` var calculatorAgentStrategy = AIAgentGraphStrategy.builder("Simple calculator") .withInput(String.class) .withOutput(String.class); var nodeSendInput = AIAgentNode.llmRequest("nodeSendInput"); var nodeExecuteTool = AIAgentNode.executeTools("nodeExecuteTool"); var nodeSendToolResult = AIAgentNode.llmSendToolResults("nodeSendToolResult"); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(calculatorAgentStrategy.nodeStart) .to(nodeSendInput) .build()); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendInput) .to(calculatorAgentStrategy.nodeFinish) .onTextMessage() .build()); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendInput) .to(nodeExecuteTool) .onToolCalls() .build()); calculatorAgentStrategy.edge(nodeExecuteTool, nodeSendToolResult); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(calculatorAgentStrategy.nodeFinish) .onTextMessage() .build()); calculatorAgentStrategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(nodeExecuteTool) .onToolCalls() .build()); var promptExecutor = PromptExecutor.builder() .ollama("http://localhost:11434") .build(); AIAgent mathAgent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .graphStrategy(calculatorAgentStrategy.build()) .build(); String result = mathAgent.run("Multiply 3 by 4, then multiply the result by 5, then add 10, then add 123.", null); System.out.println(result); ``` When you run this agent, it will respond with something like this: ``` To calculate this, I'll follow the order of operations: 1. Multiply 3 by 4: 3 * 4 = 12 2. Multiply the result by 5: 12 * 5 = 60 3. Add 10: 60 + 10 = 70 4. Add 123: 70 + 123 = 193 The final answer is 193. ``` However, since this agent doesn't have any tools, the LLM never returns a tool call and simply generates the whole answer. This is what effectively happens: ``` --- config: flowchart: defaultRenderer: "elk" --- graph LR subgraph nodeStart Input end subgraph nodeFinish Output end subgraph nodeSendInput llmRequest(Request LLM) end Input --String--> llmRequest --Message.Assistant--> onTextMessage{{onTextMessage}} --String--> Output ``` Even though it is correct in this case, the answer will depend on the arithmetic abilities of the underlying LLM. To make sure the calculations are correct, we should provide the agent with math tools. Then the LLM will be able to decide to call tools that perform the calculations deterministically. ## Add tools Define [tools](../../tools/) for performing math operations and add them to a [ToolRegistry](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool-registry/index.html): ``` @LLMDescription("Tools for performing math operations") class MathTools : ToolSet { @Tool @LLMDescription("Adds two numbers and returns the result") fun add(a: Int, b: Int): Int { // This is not necessary, but it helps to see the tool call in the console output println("Adding $a and $b...") return a + b } @Tool @LLMDescription("Multiplies two numbers and returns the result") fun multiply(a: Int, b: Int): Int { // This is not necessary, but it helps to see the tool call in the console output println("Multiplying $a and $b...") return a * b } } val toolRegistry = ToolRegistry { tools(MathTools()) } ``` ``` @LLMDescription("Tools for performing math operations") public static class MathTools implements ToolSet { @Tool @LLMDescription("Adds two numbers and returns the result") public int add(int a, int b) { // This is not necessary, but it helps to see the tool call in the console output System.out.println("Adding " + a + " and " + b + "..."); return a + b; } @Tool @LLMDescription("Multiplies two numbers and returns the result") public int multiply(int a, int b) { // This is not necessary, but it helps to see the tool call in the console output System.out.println("Multiplying " + a + " and " + b + "..."); return a * b; } } public static void main(String[] args) { ToolRegistry toolRegistry = ToolRegistry.builder() .tools(new MathTools()) .build(); } ``` Add the tool registry to the agent configuration: ``` val mathAgent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, strategy = calculatorAgentStrategy, toolRegistry = toolRegistry ) fun main() = runBlocking { val result = mathAgent.run("Multiply 3 by 4, then multiply the result by 5, then add 10, then add 123.") println(result) } ``` ``` AIAgent mathAgent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .graphStrategy(calculatorAgentStrategy.build()) .toolRegistry(toolRegistry) .build(); String result = mathAgent.run("Multiply 3 by 4, then multiply the result by 5, then add 10, then add 123.", null); System.out.println(result); ``` When you run the agent now, it will respond with something like this: ``` Multiplying 3 and 4... The output from the first operation was multiplied by 5: 5 * 12 = 60 Then, 10 was added to the result: 60 + 10 = 70 Finally, 123 was added to the result: 70 + 123 = 193 ``` According to this output, the agent correctly performed the calculations, but it only called the `multiply` tool once instead of calling the corresponding tool for every operation. We can help the agent by describing its role and providing instructions for using appropriate tools in the system prompt. ## Provide a system prompt A [system prompt](../../prompts/prompt-creation/#system-message) defines the agent's role and instructions for performing tasks. In our example, it is important to describe how the agent should process complex multistep calculations: ``` val mathAgent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, systemPrompt = """ You are a simple calculator assistant. You can add and multiply two numbers using the 'add' and 'multiply' tools. When the user provides input, extract the numbers and operations they requested. Use the appropriate tool for the first operation, then the next one, and so on, until you calculate the result. Always respond with a clear, friendly message showing the calculation and result. """.trimIndent(), toolRegistry = toolRegistry, strategy = calculatorAgentStrategy ) fun main() = runBlocking { val result = mathAgent.run("Multiply 3 by 4, then multiply the result by 5, then add 10, then add 123.") println(result) } ``` ``` AIAgent mathAgent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .systemPrompt("You are a simple calculator assistant. You can add and multiply two numbers using the 'add' and 'multiply' tools. When the user provides input, extract the numbers and operations they requested. Use the appropriate tool for the first operation, then the next one, and so on, until you calculate the result. Always respond with a clear, friendly message showing the calculation and result.") .graphStrategy(calculatorAgentStrategy.build()) .toolRegistry(toolRegistry) .build(); String result = mathAgent.run("Multiply 3 by 4, then multiply the result by 5, then add 10, then add 123.", null); System.out.println(result); ``` When you run the agent now, it will respond with something like this: ``` Multiplying 3 and 4... Multiplying 12 and 5... Adding 60 and 10... Adding 70 and 123... The final result is: 193 ``` As you can see, the agent now correctly calls the appropriate tool for each operation, ensuring that it performs the calculations deterministically instead of risking a hallucinated result. ## Next steps - Compare to [functional agents](../functional-agents/) and [planner agents](../planner-agents/) - Enhance your agent by [installing features](../../features/) - Improve the predictability and reliability with [structured output](../../structured-output/) # Functional agents With functional agents, you implement the logic as a function that handles user input, interacts with LLMs, calls tools if necessary, and produces the final output. Compared to [graph-based agents](../graph-based-agents/), this usually means faster prototyping with the following downsides: - Not easy to visualize - No state persistence Prerequisites Ensure your environment and project meet the following requirements: - JDK 17+ - Kotlin 2.2.0+ - Gradle 8.0+ or Maven 3.8+ Add the [Koog package](https://central.sonatype.com/artifact/ai.koog/koog-agents/) as a dependency: build.gradle.kts ``` dependencies { // Stable implementation("ai.koog:koog-agents:1.2.0") // Beta implementation("ai.koog:koog-agents-additions:1.2.0-beta") } ``` build.gradle ``` dependencies { // Stable implementation 'ai.koog:koog-agents:1.2.0' // Beta implementation 'ai.koog:koog-agents-additions:1.2.0-beta' } ``` pom.xml ``` ai.koog koog-agents-jvm 1.2.0 ai.koog koog-agents-additions-jvm 1.2.0-beta ``` Get an API key from an LLM provider or run a local LLM via Ollama. For more information, see [Quickstart](../../quickstart/). Examples on this page assume that you are running Llama 3.2 locally via Ollama. This page describes how to implement a functional strategy to quickly prototype some custom logic for your agent. ## Create a minimal functional agent To create a minimal functional agent, use the same [`AIAgent`](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent/-a-i-agent/index.html) interface as for a [basic agent](../basic-agents/) and pass an instance of [`AIAgentFunctionalStrategy`](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent/-a-i-agent-functional-strategy/index.html) to it. You can define a functional strategy that expects an input and returns an output, makes one LLM call, then returns the content of the assistant message from the response. In Kotlin, the most convenient way is to use the `functionalStrategy {...}` DSL method. In Java, you can use the `functionalStrategy` method on the `AIAgent` builder. ``` val strategy = functionalStrategy { input -> val response = requestLLM(input) response.parts.filterIsInstance().joinToString("\n") { it.text } } val mathAgent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, strategy = strategy ) fun main() = runBlocking { val result = mathAgent.run("What is 12 × 9?") println(result) } ``` ``` AIAgent mathAgent = AIAgent.builder() .promptExecutor(SimpleLLMExecutorsKt.simpleOllamaAIExecutor("http://localhost:11434")) .llmModel(OllamaModels.Meta.LLAMA_3_2) .functionalStrategy("mathStrategy", (AIAgentFunctionalContext context, String input) -> { Message.Response response = context.requestLLM(input); if (response instanceof Message.Assistant) { return ((Message.Assistant) response).getContent(); } return ""; }) .build(); String result = mathAgent.run("What is 12 × 9?"); System.out.println(result); ``` The agent can produce the following output: ``` The answer to 12 × 9 is 108. ``` ## Make sequential LLM calls You can extend the previous strategy to make multiple sequential LLM calls: ``` fun Message.Assistant.text(): String = parts.filterIsInstance().joinToString("\n") { it.text } val strategy = functionalStrategy { input -> // The first LLM call produces an initial draft based on the user input val draft = requestLLM("Draft: $input").text() // The second LLM call improves the initial draft val improved = requestLLM("Improve and clarify.").text() // The final LLM call formats the improved text and returns the result requestLLM("Format the result as bold.").text() } ``` ``` AIAgent mathAgent = AIAgent.builder() .promptExecutor(simpleOllamaAIExecutor("http://localhost:11434")) .systemPrompt("You are a precise math assistant.") .llmModel(OllamaModels.Meta.LLAMA_3_2) .functionalStrategy((AIAgentFunctionalContext context, String input) -> { // The first LLM call produces an initial draft based on the user input Message.Response draftResponse = context.requestLLM("Draft: " + input); String draft = ""; if (draftResponse instanceof Message.Assistant) { draft = ((Message.Assistant) draftResponse).getContent(); } // The second LLM call improves the initial draft Message.Response improvedResponse = context.requestLLM("Improve and clarify."); String improved = ""; if (improvedResponse instanceof Message.Assistant) { improved = ((Message.Assistant) improvedResponse).getContent(); } // The final LLM call formats the improved text and returns the result Message.Response finalResponse = context.requestLLM("Format the result as bold."); if (finalResponse instanceof Message.Assistant) { return ((Message.Assistant) finalResponse).getContent(); } return ""; }) .build(); ``` The agent can produce the following output: ``` To calculate the product of 12 and 9, we multiply these two numbers together. 12 × 9 = **108** ``` ## Add tools In many cases, a functional agent needs to complete specific tasks, such as reading and writing data, calling APIs, or performing other deterministic operations. In Koog, you expose such capabilities as [tools](../../tools/) and let the LLM decide when to call them. Here is what you need to do: 1. Create an [annotation-based tool](../../tools/annotation-based-tools/). 1. Add it to a tool registry and pass the registry to the agent. 1. Make sure the agent strategy can identify tool calls in LLM responses, execute the requested tools, send their results back to the LLM, and repeat the process until there are no tool calls remaining. ``` @LLMDescription("Tools for performing math operations") class MathTools : ToolSet { @Tool @LLMDescription("Multiplies two numbers and returns the result") fun multiply(a: Int, b: Int): Int { // This is not necessary, but it helps to see the tool call in the console output println("Multiplying $a and $b...") return a * b } } val toolRegistry = ToolRegistry { tool(MathTools()::multiply) } val strategy = functionalStrategy { input -> // Send the user input to the LLM var response = requestLLM(input) // Only loop while the LLM requests tools var toolCalls = response.parts.filterIsInstance() while (toolCalls.isNotEmpty()) { // Execute the tools and return the results val results = executeTools(toolCalls) // Send the tool results back to the LLM. The LLM may call more tools or return a final output response = sendToolResults(results) toolCalls = response.parts.filterIsInstance() } // When no tool calls remain, extract and return the assistant message content from the response response.parts.filterIsInstance().joinToString("\n") { it.text } } val mathAgentWithTools = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, toolRegistry = toolRegistry, strategy = strategy ) fun main() = runBlocking { val result = mathAgentWithTools.run("Multiply 3 by 4, then multiply the result by 5.") println(result) } ``` ``` @LLMDescription(description = "Tools for performing math operations") public static class MathTools implements ToolSet { @Tool @LLMDescription(description = "Multiplies two numbers and returns the result") public int multiply(int a, int b) { // This is not necessary, but it helps to see the tool call in the console output System.out.println("Multiplying " + a + " and " + b + "..."); return a * b; } } public static void main(String[] args) { MathTools mathTools = new MathTools(); ToolRegistry toolRegistry = ToolRegistry.builder() .tools(mathTools) .build(); AIAgent mathAgentWithTools = AIAgent.builder() .promptExecutor(SimpleLLMExecutorsKt.simpleOllamaAIExecutor("http://localhost:11434")) .llmModel(OllamaModels.Meta.LLAMA_3_2) .toolRegistry(toolRegistry) .functionalStrategy("mathWithTools", (AIAgentFunctionalContext context, String input) -> { // Send the user input to the LLM List responses = context.requestLLMMultiple(input); // Only loop while the LLM requests tools while (context.containsToolCalls(responses)) { // Extract tool calls from the response List pendingCalls = context.extractToolCalls(responses); // Execute the tools and return the results List results = context.executeMultipleTools(pendingCalls, false); // Send the tool results back to the LLM responses = context.sendMultipleToolResults(results); } // Extract and return the assistant message content from the response Message.Response finalResponse = responses.get(0); if (finalResponse instanceof Message.Assistant) { return ((Message.Assistant) finalResponse).getContent(); } return ""; }) .build(); String result = mathAgentWithTools.run("Multiply 3 by 4, then multiply the result by 5."); System.out.println(result); } ``` The agent can produce the following output: ``` Multiplying 3 and 4... Multiplying 12 and 5... The result of multiplying 3 by 4 is 12. Multiplying 12 by 5 gives us a final answer of 60. ``` ## Next steps - Learn how to create [graph-based agents](../graph-based-agents/) # Prompts # Prompts Prompts are instructions for Large Language Models (LLMs) that guide them in generating responses. They define the content and structure of your interactions with LLMs. This section describes how to create and run prompts with Koog. ## Creating prompts In Koog, prompts are instances of the [**Prompt**](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.dsl/-prompt/index.html) data class with the following properties: - `id`: A unique identifier for the prompt. - `messages`: A list of messages that represent the conversation with the LLM. - `params`: Optional [LLM configuration parameters](prompt-creation/#prompt-parameters) (such as temperature, tool choice, and others). Although you can instantiate the `Prompt` class directly, the recommended way to create prompts is by using the [Kotlin DSL](prompt-creation/) or the Java builder API, which provide a structured way to define the conversation. Note Kotlin examples on this page use the Kotlin DSL. Java examples use the `Prompt.builder("id")` builder with explicit methods like `system(...)`, `user(...)`, `assistant(...)`, `toolCall(...)`, `toolResult(...)`, and `withOutput(Foo.class)` where applicable. ``` val myPrompt = prompt("hello-koog") { system("You are a helpful assistant.") user("What is Koog?") } ``` ``` var myPrompt = Prompt.builder("hello-koog") .system("You are a helpful assistant.") .user("What is Koog?") .build(); ``` Note AI agents can take a simple text prompt as input. They automatically convert the text prompt to the Prompt object and send it to the LLM for execution. This is useful for a [basic agent](../agents/basic-agents/) that only needs to run a single request and does not require complex conversation logic. ## Running prompts Koog provides two levels of abstraction for running prompts against LLMs: LLM clients and prompt executors. Both accept Prompt objects and can be used for direct prompt execution, without an AI agent. The execution flow is the same for both clients and executors: ``` flowchart TB A([Prompt built with Kotlin DSL or Java builder]) B{LLM client or prompt executor} C[LLM provider] D([Response to your application]) A -->|"passed to"| B B -->|"sends request"| C C -->|"returns response"| B B -->|"returns result"| D ``` - [**LLM clients**](llm-clients/) ______________________________________________________________________ Low‑level interfaces for direct interaction with specific LLM providers. Use them when you work with a single provider and do not need advanced lifecycle management. - [**Prompt executors**](prompt-executors/) ______________________________________________________________________ High-level abstractions that manage the lifecycles of one or multiple LLM clients. Use them when you need a unified API for running prompts across multiple providers, with dynamic switching between them and fallbacks. ## Optimizing performance and handling failures Koog allows you to optimize performance and handle failures when running prompts. - [**LLM response caching**](llm-response-caching/) ______________________________________________________________________ Cache LLM responses to optimize performance and reduce costs for repeated requests. - [**Handling failures**](handling-failures/) ______________________________________________________________________ Use built-in retries, timeouts, and other error handling mechanisms in your application. ## Prompts in AI agents In Koog, AI agents maintain and manage prompts during their lifecycle. While LLM clients or executors are used to run prompts, agents handle the flow of prompt updates, ensuring the conversation history remains relevant and consistent. The prompt lifecycle in an agent usually includes several stages: 1. Initial prompt setup. 1. Automatic prompt updates. 1. Context window management. 1. Manual prompt management. ### Initial prompt setup When you [initialize an agent](../quickstart/#create-your-first-koog-agent), you can define a [system message](prompt-creation/#system-message) that sets the agent's behavior. Then, when you call the agent's `run()` method, you typically provide an initial [user message](prompt-creation/#user-messages) as input. Together, these messages form the agent's initial prompt. For example: ``` // Create an agent val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(apiKey), systemPrompt = "You are a helpful assistant.", llmModel = OpenAIModels.Chat.GPT4o ) // Run the agent val result = agent.run("What is Koog?") ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("You are a helpful assistant. Answer user questions concisely.") .llmModel(OpenAIModels.Chat.GPT4o) .build(); var result = agent.run("What is Koog?"); ``` In the example, the agent automatically converts the text prompt to the Prompt object and sends it to the prompt executor: ``` flowchart TB A([Your application]) B{{Configured AI agent}} C["Text prompt"] D["Prompt object"] E{{Prompt executor}} F[LLM provider] A -->|"run() with text"| B B -->|"takes"| C C -->|"converted to"| D D -->|"sent via"| E E -->|"calls"| F F -->|"responds to"| E E -->|"result to"| B B -->|"result to"| A ``` For more advanced configurations, you can also use [AIAgentConfig](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.config/-a-i-agent-config/index.html) to define the agent's initial prompt. ### Automatic prompt updates As the agent runs its strategy, [predefined nodes](../nodes-and-components/) automatically update the prompt. For example: - [`nodeLLMRequest`](../nodes-and-components/#nodellmrequest): Appends a user message to the prompt and captures the LLM response. - [`nodeLLMSendToolResult`](../nodes-and-components/#nodellmsendtoolresult): Appends tool execution results to the conversation. - [`nodeAppendPrompt`](../nodes-and-components/#nodeappendprompt): Inserts specific messages into the prompt at any point in the workflow. ### Context window management To avoid exceeding the LLM context window in long-running interactions, agents can use the [history compression](../history-compression/) feature. ### Manual prompt management For complex workflows, you can manage the prompt manually using [LLM sessions](../sessions/). In an agent strategy or custom node, you can use `llm.writeSession` to access and change the `Prompt` object. This lets you add, remove, or reorder messages as needed. # Creating prompts Koog provides a structured way to create prompts with control over message types, their order, and content: - For **Kotlin** users, through a type-safe Kotlin DSL. - For **Java** users, through a fluent builder API. ## Basic structure The `prompt()` function in Kotlin or the `Prompt.builder()` in Java create a Prompt object with a unique ID and a list of messages: ``` val prompt = prompt("unique_prompt_id") { // List of messages } ``` ``` Prompt prompt = Prompt.builder("unique_prompt_id") // List of messages .build(); ``` ## Message types The Kotlin DSL and the Java builder API support the following types of messages, each of which corresponds to a specific role in a conversation: - **System message**: Provides the context, instructions, and constraints to the LLM, defining its behavior. - **User message**: Represents the user input. - **Assistant message**: Represents LLM responses that are used for few-shot learning or to continue the conversation. - **Tool message**: Represents tool calls and their results. ``` val prompt = prompt("unique_prompt_id") { // Add a system message to set the context system("You are a helpful assistant with access to tools.") // Add a user message user("What is 5 + 3 ?") // Add an assistant message assistant("The result is 8.") } ``` ``` Prompt prompt = Prompt.builder("unique_prompt_id") // Add a system message to set the context .system("You are a helpful assistant with access to tools.") // Add a user message .user("What is 5 + 3 ?") // Add an assistant message .assistant("The result is 8.") .build(); ``` ### System message A system message defines the LLM behavior and sets the context for the entire conversation. It can specify the model's role, tone, provide guidelines and constraints on responses, and provide response examples. To create the system message, provide a string as an argument to the `system()` Kotlin function or Java method: ``` val prompt = prompt("system_message") { system("You are a helpful assistant that explains technical concepts.") } ``` ``` Prompt prompt = Prompt.builder("system_message") .system("You are a helpful assistant that explains technical concepts.") .build(); ``` ### User messages A user message represents input from the user. To create the user message, provide a string as an argument to the `user()` Kotlin function or Java method: ``` val prompt = prompt("user_message") { system("You are a helpful assistant.") user("What is Koog?") } ``` ``` Prompt prompt = Prompt.builder("user_message") .system("You are a helpful assistant.") .user("What is Koog?") .build(); ``` Most user messages contain plain text, but they can also include multimodal content, such as images, audio, video, and documents. For details and examples, see [Multimodal content](multimodal-content/). ### Assistant messages An assistant message represents an LLM response, which can be used for few-shot learning in future similar interactions, to continue a conversation, or to demonstrate the expected output structure. To create the assistant message, provide a string as an argument to the `assistant()` Kotlin function or Java method: ``` val prompt = prompt("article_review") { system("Evaluate the article.") // Example 1 user("The article is clear and easy to understand.") assistant("positive") // Example 2 user("The article is hard to read but it's clear and useful.") assistant("neutral") // Example 3 user("The article is confusing and misleading.") assistant("negative") // New input to classify user("The article is interesting and helpful.") } ``` ``` Prompt prompt = Prompt.builder("article_review") .system("Evaluate the article.") // Example 1 .user("The article is clear and easy to understand.") .assistant("positive") // Example 2 .user("The article is hard to read but it's clear and useful.") .assistant("neutral") // Example 3 .user("The article is confusing and misleading.") .assistant("negative") // New input to classify .user("The article is interesting and helpful.") .build(); ``` ### Tool messages A tool message represents a tool call and its result, which can be used to pre-fill the history of tool calls. Tip An LLM generates tool calls during execution. Pre-filling them is helpful for few-shot learning or demonstrating how the tools are expected to be used. To create the tool message, call the `tool()` function in Kotlin or the `toolCall()` and `toolResult()` methods in Java: ``` val prompt = prompt("calculator_example") { system("You are a helpful assistant with access to tools.") user("What is 5 + 3?") // Tool call toolCall( id = "calculator_tool_id", tool = "calculator", args = """{"operation": "add", "a": 5, "b": 3}""" ) // Tool result toolResult( id = "calculator_tool_id", tool = "calculator", output = "8" ) // LLM response based on tool result assistant("The result of 5 + 3 is 8.") user("What is 4 + 5?") } ``` ``` Prompt prompt = Prompt.builder("calculator_example") .system("You are a helpful assistant with access to tools.") .user("What is 5 + 3?") // Tool call .toolCall("calculator_tool_id", "calculator", "{\"operation\": \"add\", \"a\": 5, \"b\": 3}") // Tool result .toolResult("calculator_tool_id", "calculator", "8") // LLM response based on tool result .assistant("The result of 5 + 3 is 8.") .user("What is 4 + 5?") .build(); ``` ## Text message builders Warning Text message builders are available only in Kotlin. When building a `system()`, `user()`, or `assistant()` message, you can use helper [text-building-functions](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.text/-text-content-builder/index.html) for rich text formatting. ``` val prompt = prompt("text_example") { user { +"Review the following code snippet:" +"fun greet(name: String) = println(\"Hello, \$name!\")" // Paragraph break br() text("Please include in your explanation:") // Indent content padding(" ") { +"1. What the function does." +"2. How string interpolation works." } } } ``` You can also use the [Markdown](https://api.koog.ai/prompt/prompt-markdown/ai.koog.prompt.markdown/markdown.html) and [XML](https://api.koog.ai/prompt/prompt-xml/ai.koog.prompt.xml/xml.html) builders to add the content in the corresponding format. ``` val prompt = prompt("markdown_xml_example") { // A user message in Markdown format user { markdown { h2("Evaluate the article using the following criteria:") bulleted { item { +"Clarity and readability" } item { +"Accuracy of information" } item { +"Usefulness to the reader" } } } } // An assistant message in XML format assistant { xml { xmlDeclaration() tag("review") { tag("clarity") { text("positive") } tag("accuracy") { text("neutral") } tag("usefulness") { text("positive") } } } } } ``` Tip You can mix the text building functions with the XML and Markdown builders. ## Prompt parameters Prompts can be customized by configuring parameters that control the LLM's behavior. ``` val prompt = prompt( id = "custom_params", params = LLMParams( temperature = 0.7, numberOfChoices = 1, toolChoice = LLMParams.ToolChoice.Auto ) ) { system("You are a creative writing assistant.") user("Write a song about winter.") } ``` ``` // Create params first LLMParams params = new LLMParams( 0.7, // temperature null, // maxTokens 1, // numberOfChoices null, // speculation null, // schema LLMParams.ToolChoice.Auto.INSTANCE, // toolChoice null, // user null // additionalProperties ); Prompt prompt = Prompt.builder("custom_params") .system("You are a creative writing assistant.") .user("Write a song about winter.") .build(); // Apply params to the built prompt prompt = prompt.withParams(params); ``` The following parameters are supported: - `temperature`: Controls randomness in the model's responses. - `toolChoice`: Controls tool calling behavior of the model. - `numberOfChoices`: Requests multiple alternative responses. - `schema`: Defines the structure for the model's response format. - `maxTokens`: Limits the number of tokens in the response. - `speculation`: Provides a hint about the expected response format (only supported by specific models). For more information, see [LLM parameters](../../llm-parameters/). ## Extending existing prompts You can extend an existing prompt by calling the `prompt()` function in Kotlin or the `Prompt.builder()` in Java with the existing prompt as an argument: ``` val basePrompt = prompt("base") { system("You are a helpful assistant.") user("Hello!") assistant("Hi! How can I help you?") } val extendedPrompt = prompt(basePrompt) { user("What's the weather like?") } ``` ``` Prompt basePrompt = Prompt.builder("base") .system("You are a helpful assistant.") .user("Hello!") .assistant("Hi! How can I help you?") .build(); Prompt extendedPrompt = Prompt.builder(String.valueOf(basePrompt)) .user("What's the weather like?") .build(); ``` This creates a new prompt that includes all messages from `basePrompt` and the new user message. ## Next steps - Learn how to work with [multimodal content](multimodal-content/). - Run prompts with [LLM clients](../llm-clients/) if you work with a single LLM provider. - Run prompts with [prompt executors](../prompt-executors/) if you work with multiple LLM providers. - Learn how to use llm cache with [cache control](cache-control/). # Multimodal content Multimodal content refers to content of different types, such as text, images, audio, video, and files. Koog lets you send images, audio, video, and files to LLMs within the `user` message along with text. You can add them to the `user` message by using the corresponding functions in Kotlin or methods in Java: - `image()`: Attaches images (JPG, PNG, WebP, GIF). - `audio()`: Attaches audio files (MP3, WAV, FLAC). - `video()`: Attaches video files (MP4, AVI, MOV). - `file()` / `binaryFile()` / `textFile()`: Attaches documents (PDF, TXT, MD, etc.). Each function or method supports two ways of configuring attachment parameters, so you can: - Pass a URL or a file path to the function or method, and it automatically handles attachment parameters. For `file()`, `binaryFile()`, and `textFile()`, you must also provide the MIME type. - Create and pass a `ContentPart` object to the function or method for custom control over attachment parameters. Note Multimodal content support varies by [LLM provider](../../../llm-providers/). Check the provider documentation for supported content types. ### Auto-configured attachments If you pass a URL or a file path to the attachment functions or methods, Koog automatically constructs the corresponding attachment parameters based on the file extension. The general format of the `user` message that includes a text message and a list of auto-configured attachments is as follows: ``` user { +"Describe these images:" image("https://example.com/test.png") image(Path("/path/to/image.png")) +"Focus on the main subjects." } ``` ``` ContentPartsBuilder partsBuilder = new ContentPartsBuilder(); partsBuilder.text("Describe these images:"); partsBuilder.image("https://example.com/test.png"); partsBuilder.text("Focus on the main subjects."); Prompt prompt = Prompt.builder("image_analysis") .user(partsBuilder.build()) .build(); ``` In Kotlin, the `+` operator adds text content to the user message along with the attachments. In Java, use the `text()` method of `ContentPartsBuilder`. ### Custom-configured attachments The [`ContentPart`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-content-part/index.html) interface lets you configure parameters for each attachment individually. All attachments implement the `ContentPart.Attachment` interface. You can create an instance of a specific implementation for each attachment, configure its parameters, and pass it to the corresponding `image()`, `audio()`, `video()`, or `file()` functions in Kotlin or methods in Java. The general format of the `user` message that includes a text message and a list of custom-configured attachments is as follows: ``` user { +"Describe this image" image( AttachmentSource.Image( content = AttachmentContent.URL("https://example.com/capture.png"), format = "png", mimeType = "image/png", fileName = "capture.png" ) ) } ``` ``` Prompt prompt = Prompt.builder("custom_image") .user(List.of( new ContentPart.Text("Describe this image"), new ContentPart.Image( new AttachmentContent.URL("https://example.com/capture.png"), "png", "image/png", "capture.png" ) )) .build(); ``` Koog provides the following specialized classes for each media type that implement the `ContentPart.Attachment` interface: - [`ContentPart.Image`](api:prompt-model::ai.koog.prompt.message.ContentPart.Image): image attachments, such as JPG or PNG files. - [`ContentPart.Audio`](api:prompt-model::ai.koog.prompt.message.ContentPart.Audio): audio attachments, such as MP3 or WAV files. - [`ContentPart.Video`](api:prompt-model::ai.koog.prompt.message.ContentPart.Video): video attachments, such as MP4 or AVI files. - [`ContentPart.File`](api:prompt-model::ai.koog.prompt.message.ContentPart.File): file attachments, such as PDF or TXT files. All `ContentPart.Attachment` types accept the following parameters: | Name | Data type | Required | Description | | ---------- | ------------------------------------------------------------------------------------------------------------------ | --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `content` | [AttachmentContent](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-attachment-content/index.html) | Yes | The source of the provided file content. | | `format` | String | Yes | The format of the provided file. For example, `png`. | | `mimeType` | String | Only for `ContentPart.File` | The MIME Type of the provided file. For `ContentPart.Image`, `ContentPart.Audio`, and `ContentPart.Video`, it defaults to `/` (for example, `image/png`). For `ContentPart.File`, it must be explicitly provided. | | `fileName` | String? | No | The name of the provided file including the extension. For example, `screenshot.png`. | #### Attachment content Implementations of the AttachmentContent interface define the type and source of content that is provided as input to the LLM: - [`AttachmentContent.URL`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-attachment-content/-u-r-l/index.html) defines the URL of the provided content: ``` AttachmentContent.URL("https://example.com/image.png") ``` - [`AttachmentContent.Binary.Bytes`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-attachment-content/-binary/index.html) defines the file content as a byte array: ``` AttachmentContent.Binary.Bytes(byteArrayOf(/* ... */)) ``` - [`AttachmentContent.Binary.Base64`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-attachment-content/-binary/index.html) defines the file content as a Base64-encoded string containing file data: ``` AttachmentContent.Binary.Base64("iVBORw0KGgoAAAANS...") ``` - [`AttachmentContent.PlainText`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-attachment-content/-plain-text/index.html) defines the file content as plain text (for [`ContentPart.File`](api:prompt-model::ai.koog.prompt.message.ContentPart.File) only): ``` AttachmentContent.PlainText("This is the file content.") ``` ### Mixed attachments In addition to providing different types of attachments in separate prompts or messages, you can also provide multiple and mixed types of attachments in a single `user()` message: ``` val prompt = prompt("mixed_content") { system("You are a helpful assistant.") user { +"Compare the image with the document content." image(Path("/path/to/image.png")) binaryFile(Path("/path/to/page.pdf"), "application/pdf") +"Structure the result as a table" } } ``` ``` Prompt prompt = Prompt.builder("mixed_content_example") .system("You are a helpful assistant.") .user(List.of( new ContentPart.Text("Please analyze this image and the attached document."), new ContentPart.Image( new AttachmentContent.URL("https://example.com/image.png"), "png", "image/png", "image.png" ), new ContentPart.File( new AttachmentContent.URL("https://example.com/document.pdf"), "pdf", "application/pdf", "document.pdf" ), new ContentPart.Text("Summarize the differences.") )) .build(); ``` ## Next steps - Run prompts with [LLM clients](../../llm-clients/) if you work with a single LLM provider. - Run prompts with [prompt executors](../../prompt-executors/) if you work with multiple LLM providers. # Handling failures This page describes how to handle failures for LLM clients and prompt executors using the built-in retry and timeout mechanisms. ## Retry functionality When working with LLM providers, transient errors like rate limits or temporary service unavailability may occur. The `RetryingLLMClient` decorator adds automatic retry logic to any LLM client in both Kotlin and Java. ### Basic usage Wrap any existing client with the retry capability: ``` // Wrap any client with the retry capability val client = OpenAILLMClient(apiKey) val resilientClient = RetryingLLMClient(client) // Now all operations will automatically retry on transient errors val response = resilientClient.execute(prompt, OpenAIModels.Chat.GPT4o) ``` ``` OpenAILLMClient client = openAIClient(apiKey); RetryingLLMClient resilientClient = new RetryingLLMClient(client); // Now all operations will automatically retry on transient errors List response = resilientClient.execute(prompt, OpenAIModels.Chat.GPT4o); ``` ### Configuring retry behavior By default, `RetryingLLMClient` configures an LLM client with the maximum of 3 retry attempts, a 1-second initial delay, and a 30-second maximum delay. You can specify a different retry configuration using a `RetryConfig` passed to `RetryingLLMClient`. For example: ``` // Use the predefined configuration val conservativeClient = RetryingLLMClient( delegate = client, config = RetryConfig.CONSERVATIVE ) ``` ``` OpenAILLMClient client = openAIClient(apiKey); // Use the predefined configuration RetryingLLMClient conservativeClient = new RetryingLLMClient( client, RetryConfig.Companion.getCONSERVATIVE() ); ``` Koog provides several predefined retry configurations available via `RetryConfig` in Kotlin and `RetryConfig.Companion` in Java: | Configuration (Kotlin) | Max attempts | Initial delay | Max delay | Use case | | -------------------------- | ------------ | ------------- | --------- | -------------------------------------------------------------------------------------------------------- | | `RetryConfig.DISABLED` | 1 (no retry) | - | - | Development, testing, and debugging. | | `RetryConfig.CONSERVATIVE` | 3 | 2s | 30s | Background or scheduled tasks where reliability is more important than speed. | | `RetryConfig.AGGRESSIVE` | 5 | 500ms | 20s | Critical operations where fast recovery from transient errors is more important than reducing API calls. | | `RetryConfig.PRODUCTION` | 3 | 1s | 20s | General production use. | You can use them directly or create custom configurations: ``` // Or create a custom configuration val customClient = RetryingLLMClient( delegate = client, config = RetryConfig( maxAttempts = 5, initialDelay = 1.seconds, maxDelay = 30.seconds, backoffMultiplier = 2.0, jitterFactor = 0.2 ) ) ``` ### Retry error patterns By default, the `RetryingLLMClient` recognizes common transient errors. This behavior is controlled by the [`RetryConfig.retryablePatterns`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/ai.koog.prompt.executor.clients.retry/-retry-config/retryable-patterns.html) patterns. Each pattern is represented by [`RetryablePattern`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/ai.koog.prompt.executor.clients.retry/-retryable-pattern/index.html) that checks the error message from a failed request and determines whether it should be retried. Koog provides the predefined retry configurations and patterns that work across all the supported LLM providers. You can keep the defaults or customize them for your specific needs. #### Pattern types You can use the following pattern types and combine any number of them: - `RetryablePattern.Status`: Matches a specific HTTP status code in the error message (such as `429`, `500`,`502`, etc.). - `RetryablePattern.Keyword`: Matches a keyword in the error message (such as `rate limit` or `request timeout`). - `RetryablePattern.Regex`: Matches a regular expression in the error message. - `RetryablePattern.Custom`: Matches a custom logic using a lambda function. If any pattern returns `true`, the error is considered retryable, and the LLM client retries the request. #### Default patterns Unless you customize the retry configuration, the following patterns are used by default: - **HTTP status codes**: - `429`: Rate limit - `500`: Internal server error - `502`: Bad gateway - `503`: Service unavailable - `504`: Gateway timeout - `529`: Anthropic overloaded - **Error keywords**: - rate limit - too many requests - request timeout - connection timeout - read timeout - write timeout - connection reset by peer - connection refused - temporarily unavailable - service unavailable These default patterns are defined in Koog as [`RetryConfig.DEFAULT_PATTERNS`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/ai.koog.prompt.executor.clients.retry/-retry-config/-companion/-d-e-f-a-u-l-t_-p-a-t-t-e-r-n-s.html). #### Custom patterns You can define custom patterns for your specific needs: ``` val config = RetryConfig( retryablePatterns = listOf( RetryablePattern.Status(429), // Specific status code RetryablePattern.Keyword("quota"), // Keyword in error message RetryablePattern.Regex(Regex("ERR_\\d+")), // Custom regex pattern RetryablePattern.Custom { error -> // Custom logic error.contains("temporary") && error.length > 20 } ) ) ``` You can also append custom patterns to the default `RetryConfig.DEFAULT_PATTERNS`: ``` val config = RetryConfig( retryablePatterns = RetryConfig.DEFAULT_PATTERNS + listOf( RetryablePattern.Keyword("custom_error") ) ) ``` ### Streaming with retry Streaming operations can optionally be retried. This feature is disabled by default. ``` val config = RetryConfig( maxAttempts = 3 ) val client = RetryingLLMClient(baseClient, config) val stream = client.executeStreaming(prompt, OpenAIModels.Chat.GPT4o) ``` Note Streaming retries only apply to connection failures that occur before the first token is received. Once streaming has started, the retry logic is disabled. If an error occurs during streaming, the operation is terminated. ### Retry with prompt executors When working with prompt executors, you can wrap the underlying LLM client with a retry mechanism before creating the executor in both Kotlin and Java. To learn more about prompt executors, see [Prompt executors](../prompt-executors/). ``` // Single provider executor with retry val resilientClient = RetryingLLMClient( OpenAILLMClient(System.getenv("OPENAI_API_KEY")), RetryConfig.PRODUCTION ) val executor = MultiLLMPromptExecutor(resilientClient) // Multi-provider executor with flexible client configuration val multiExecutor = MultiLLMPromptExecutor( LLMProvider.OpenAI to RetryingLLMClient( OpenAILLMClient(System.getenv("OPENAI_API_KEY")), RetryConfig.CONSERVATIVE ), LLMProvider.Anthropic to RetryingLLMClient( AnthropicLLMClient(System.getenv("ANTHROPIC_API_KEY")), RetryConfig.AGGRESSIVE ), // The Bedrock client already has a built-in AWS SDK retry LLMProvider.Bedrock to BedrockLLMClient( identityProvider = StaticCredentialsProvider { accessKeyId = System.getenv("AWS_ACCESS_KEY_ID") secretAccessKey = System.getenv("AWS_SECRET_ACCESS_KEY") sessionToken = System.getenv("AWS_SESSION_TOKEN") }, ), ) ``` ``` // Single provider executor with retry (Java) RetryingLLMClient resilientClient = new RetryingLLMClient( openAIClient(System.getenv("OPENAI_API_KEY")), RetryConfig.Companion.getPRODUCTION() ); MultiLLMPromptExecutor executor = new MultiLLMPromptExecutor(resilientClient); // Multi-provider executor with flexible client configuration (Java) LLMClient openai = new RetryingLLMClient( openAIClient(System.getenv("OPENAI_API_KEY")), RetryConfig.Companion.getCONSERVATIVE() ); LLMClient anthropic = new RetryingLLMClient( anthropicClient(System.getenv("ANTHROPIC_API_KEY")), RetryConfig.Companion.getAGGRESSIVE() ); Map clients = Map.of( LLMProvider.OpenAI, openai, LLMProvider.Anthropic, anthropic ); MultiLLMPromptExecutor multiExecutor = new MultiLLMPromptExecutor(clients); ``` ## Timeout configuration All LLM clients support timeout configuration in both Kotlin and Java to prevent hanging requests. You can specify timeout values for network connections when creating the client using the [`ConnectionTimeoutConfig`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/ai.koog.prompt.executor.clients/-connection-timeout-config/index.html) class. `ConnectionTimeoutConfig` has the following properties: | Property | Default Value | Description | | ---------------------- | -------------------- | ------------------------------------------------------------- | | `connectTimeoutMillis` | 60 seconds (60,000) | Maximum time to establish a connection to the server. | | `requestTimeoutMillis` | 15 minutes (900,000) | Maximum time for the entire request to complete. | | `socketTimeoutMillis` | 15 minutes (900,000) | Maximum time to wait for data over an established connection. | You can customize these values for your specific needs. For example: ``` val client = OpenAILLMClient( apiKey = apiKey, settings = OpenAIClientSettings( timeoutConfig = ConnectionTimeoutConfig( connectTimeoutMillis = 5000, // 5 seconds to establish connection requestTimeoutMillis = 60000, // 60 seconds for the entire request socketTimeoutMillis = 120000 // 120 seconds for data on the socket ) ) ) ``` ``` String apiKey = System.getenv("OPENAI_API_KEY"); ConnectionTimeoutConfig timeouts = new ConnectionTimeoutConfig( 5000L, // connectTimeoutMillis 60000L, // requestTimeoutMillis 120000L // socketTimeoutMillis ); OpenAIClientSettings settings = new OpenAIClientSettings( "https://api.openai.com", // baseUrl timeouts, "v1/chat/completions", // chatCompletionsPath "v1/responses", // responsesAPIPath "v1/embeddings", // embeddingsPath "v1/moderations", // moderationsPath "v1/models" // modelsPath ); OpenAILLMClient client = openAIClient(apiKey, settings); ``` Tip For long-running or streaming calls, set higher values for `requestTimeoutMillis` and `socketTimeoutMillis`. ## Error handling When working with LLMs in production, you need to implement error handling, including: - **Try-catch blocks** to handle unexpected errors. - **Logging errors with context** for debugging. - **Fallbacks** for critical operations. - **Monitoring retry patterns** to identify recurring issues. Here is an example of error handling in Kotlin and Java: ``` val logger = LoggerFactory.getLogger("Example") val resilientClient = RetryingLLMClient( OpenAILLMClient(System.getenv("OPENAI_API_KEY")), RetryConfig.PRODUCTION ) val prompt = prompt("test") { user("Hello") } val model = OpenAIModels.Chat.GPT4o fun processResponse(response: Any) { /* implmenentation */ } fun scheduleRetryLater() { /* implmenentation */ } fun notifyAdministrator() { /* implmenentation */ } fun useDefaultResponse() { /* implmenentation */ } try { val response = resilientClient.execute(prompt, model) processResponse(response) } catch (e: Exception) { logger.error("LLM operation failed", e) when { e.message?.contains("rate limit") == true -> { // Handle rate limiting specifically scheduleRetryLater() } e.message?.contains("invalid api key") == true -> { // Handle authentication errors notifyAdministrator() } else -> { // Fall back to an alternative solution useDefaultResponse() } } } ``` ``` Logger logger = LoggerFactory.getLogger("Example"); RetryingLLMClient resilientClient = new RetryingLLMClient( openAIClient(System.getenv("OPENAI_API_KEY")), RetryConfig.PRODUCTION ); Prompt prompt = Prompt.builder("test") .user("Hello") .build(); MultiLLMPromptExecutor promptExecutor = new MultiLLMPromptExecutor(resilientClient); Consumer processResponse = (resp) -> { /* implementation */ }; Runnable scheduleRetryLater = () -> { /* implementation */ }; Runnable notifyAdministrator = () -> { /* implementation */ }; Runnable useDefaultResponse = () -> { /* implementation */ }; try { Message.Assistant response = promptExecutor.execute(prompt, OpenAIModels.Chat.GPT4o); processResponse.accept(response); } catch (Exception e) { logger.error("LLM operation failed", e); String msg = e.getMessage() == null ? "" : e.getMessage().toLowerCase(); if (msg.contains("rate limit")) { scheduleRetryLater.run(); } else if (msg.contains("invalid api key")) { notifyAdministrator.run(); } else { useDefaultResponse.run(); } } ``` # LLM response caching For repeated requests that you run with a prompt executor, you can cache LLM responses to optimize performance and reduce costs in both Kotlin and Java. In Koog, caching is available for all prompt executors through `CachedPromptExecutor`, which is a wrapper around `PromptExecutor` that adds caching functionality. It lets you store responses from previously executed prompts and retrieve them when the same prompts are run again. To create a cached prompt executor in Kotlin or Java, perform the following: 1. Create a prompt executor for which you want to cache responses. 1. Create a `CachedPromptExecutor` instance by providing the desired cache and the prompt executor you created. 1. Run the created `CachedPromptExecutor` with the desired prompt and model. Here is an example: ``` // Create a prompt executor val client = OpenAILLMClient(System.getenv("OPENAI_API_KEY")) val promptExecutor = MultiLLMPromptExecutor(client) // Create a cached prompt executor val cachedExecutor = CachedPromptExecutor( cache = FilePromptCache(Path("path/to/your/cache/directory")), nested = promptExecutor ) // Run cached prompt executor for the first time // This will perform an actual LLM request val firstTime = measureTimeMillis { val firstResponse = cachedExecutor.execute(prompt, OpenAIModels.Chat.GPT4o) val text = firstResponse.parts.filterIsInstance().joinToString("\n") { it.text } println("First response: $text") } println("First execution took: ${firstTime}ms") // Run cached prompt executor for the second time // This will return the result immediately from the cache val secondTime = measureTimeMillis { val secondResponse = cachedExecutor.execute(prompt, OpenAIModels.Chat.GPT4o) val text = secondResponse.parts.filterIsInstance().joinToString("\n") { it.text } println("Second response: $text") } println("Second execution took: ${secondTime}ms") ``` ``` // Create a prompt Prompt prompt = Prompt.builder("test") .user("Hello") .build(); // Create a prompt executor OpenAILLMClient client = openAIClient(System.getenv("OPENAI_API_KEY")); MultiLLMPromptExecutor promptExecutor = new MultiLLMPromptExecutor(client); // Create a cached prompt executor FilePromptCache cache = new FilePromptCache(Path.of("path/to/your/cache/directory"), null); CachedPromptExecutor cachedExecutor = new CachedPromptExecutor(cache, promptExecutor, Clock.System.INSTANCE); // Run cached prompt executor for the first time // This will perform an actual LLM request long start1 = System.nanoTime(); List firstResponse = cachedExecutor.execute(prompt, OllamaModels.Meta.LLAMA_3_2); long firstTimeMs = (System.nanoTime() - start1) / 1_000_000L; System.out.println("First response: " + firstResponse.getFirst().getContent()); System.out.println("First execution took: " + firstTimeMs + "ms"); // Run cached prompt executor for the second time // This will return the result immediately from the cache long start2 = System.nanoTime(); List secondResponse = cachedExecutor.execute(prompt, OllamaModels.Meta.LLAMA_3_2); long secondTimeMs = (System.nanoTime() - start2) / 1_000_000L; System.out.println("Second response: " + secondResponse.getFirst().getContent()); System.out.println("Second execution took: " + secondTimeMs + "ms"); ``` The example produces the following output: ``` First response: Hello! It seems like we're starting a new conversation. What can I help you with today? First execution took: 48ms Second response: Hello! It seems like we're starting a new conversation. What can I help you with today? Second execution took: 1ms ``` The second response is retrieved from the cache, which took only 1ms. Note - If you call `executeStreaming()` in Kotlin or `executeStreamingWithPublisher()` in Java with the cached prompt executor, it produces a response as a single chunk. - If you call `moderate()` with the cached prompt executor in either Kotlin or Java, it forwards the request to the nested prompt executor and does not use the cache. - Caching of multiple choice responses (`executeMultipleChoices()`) is not supported in either Kotlin or Java. # Prompt caching control Prompt caching control lets you instruct a supported LLM provider to store a portion of your prompt server-side, so that subsequent requests that share the same prefix can be served from the cache instead of reprocessing the tokens. This reduces both latency and cost for repetitive workloads such as multi-turn conversations, large system prompts, or fixed tool definitions. Prompt caching vs. response caching Prompt caching control is a **provider-side** feature: the provider stores the prompt prefix, not the response. This is different from [`CachedPromptExecutor`](../../llm-response-caching/), which stores complete LLM responses locally so that identical prompts skip the network call entirely. Koog supports prompt caching control for **Anthropic** and **Amazon Bedrock**. ## Anthropic Anthropic supports two complementary approaches to prompt caching. ### Automatic caching (request-level) Set the `cacheControl` property on [`AnthropicParams`](../../../llm-parameters/) and pass it to your prompt. Anthropic will automatically place the cache breakpoint at the last cacheable block in the request, without you having to annotate individual messages. This is the recommended approach for multi-turn conversations. ``` // Enable automatic caching with the default 5-minute TTL val params = AnthropicParams(cacheControl = AnthropicCacheControl.Default) val prompt = prompt("assistant", params = params) { system("You are a helpful assistant with a very long system prompt...") user("What can you help me with?") } val response = client.execute(prompt, AnthropicModels.Sonnet_4) println(response) ``` ``` // Enable automatic caching with the default 5-minute TTL AnthropicParams params = new AnthropicParams( null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, AnthropicCacheControl.Default.INSTANCE ); Prompt prompt = Prompt.builder("assistant") .system("You are a helpful assistant with a very long system prompt...") .user("What can you help me with?") .build() .withParams(params); ``` ### Manual caching (block-level) Attach a `cacheControl` argument to individual messages or tool definitions to place the cache breakpoint at a specific position. Everything up to and including the annotated block is eligible for caching. #### System messages ``` val prompt = prompt("assistant") { // Cache the system prompt for 1 hour system("You are a knowledgeable assistant...", AnthropicCacheControl.OneHour) user("Summarize the latest AI research.") } val response = client.execute(prompt, AnthropicModels.Sonnet_4) println(response) ``` ``` Prompt prompt = Prompt.builder("assistant") // Cache the system prompt for 1 hour .system("You are a knowledgeable assistant...", AnthropicCacheControl.OneHour.INSTANCE) .user("Summarize the latest AI research.") .build(); ``` #### User and assistant messages ``` val prompt = prompt("conversation") { system("You are a helpful assistant.") // Cache after a large user message (e.g. document content) user(listOf(MessagePart.Text("Here is a long document: ...", cacheControl = AnthropicCacheControl.Default))) assistant(listOf(MessagePart.Text("I have read the document."))) user("Summarize it.") } val response = client.execute(prompt, AnthropicModels.Sonnet_4) println(response) ``` ``` Prompt prompt = Prompt.builder("conversation") .system("You are a helpful assistant.") // Cache after a large user message (e.g. document content) .user(List.of(new ContentPart.Text("Here is a long document: ...")), AnthropicCacheControl.Default.INSTANCE) .assistant("I have read the document.", AnthropicCacheControl.Default.INSTANCE) .user("Summarize it.") .build(); ``` #### Tool definitions When a tool list is fixed across many requests, caching the last tool definition means all tool schemas are cached together. ``` val searchTool = ToolDescriptor( name = "web_search", description = "Search the web for information.", requiredParameters = listOf( ToolParameterDescriptor("query", "Search query", ToolParameterType.String) ), // Cache all tool definitions up to and including this one cacheControl = AnthropicCacheControl.Default ) ``` ``` ToolDescriptor searchTool = new ToolDescriptor( "web_search", "Search the web for information.", List.of( new ToolParameterDescriptor("query", "Search query", ToolParameterType.String.INSTANCE) ), Collections.emptyList(), // Cache all tool definitions up to and including this one AnthropicCacheControl.Default.INSTANCE ); ``` ### Cache TTL options | Option | TTL | Price multiplier | | ------------------------------- | --------- | ---------------------- | | `AnthropicCacheControl.Default` | 5 minutes | 1.25× base input price | | `AnthropicCacheControl.OneHour` | 1 hour | 2× base input price | Cache writes are charged at a higher rate than regular input tokens, but cache reads are cheaper. See the [Anthropic prompt caching docs](https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching) for current pricing. ### Monitoring cache usage Anthropic reports cache statistics in the response usage. These are accessible via the raw API response and can be observed through tracing or logging features. | Field | Meaning | | -------------------------- | ---------------------------------------- | | `cacheReadInputTokens` | Tokens read from an existing cache entry | | `cacheCreationInputTokens` | Tokens written to a new cache entry | ### Combining automatic and block-level caching Both modes can be used simultaneously. Block-level `cacheControl` markers give you fine-grained control over breakpoint positions, while the request-level `cacheControl` in `AnthropicParams` handles the tail of the conversation automatically. ``` // Block-level: pin the system prompt in the 1-hour cache tier // Automatic: let Anthropic manage breakpoints for the conversation tail val params = AnthropicParams(cacheControl = AnthropicCacheControl.Default) val prompt = prompt("combined", params = params) { system("You are a helpful assistant...", AnthropicCacheControl.OneHour) user("Hello!") } ``` ``` // Block-level: pin the system prompt in the 1-hour cache tier // Automatic: let Anthropic manage breakpoints for the conversation tail AnthropicParams params = new AnthropicParams( AnthropicCacheControl.Default.INSTANCE ); Prompt prompt = Prompt.builder("combined") .system("You are a helpful assistant...", AnthropicCacheControl.OneHour.INSTANCE) .user("Hello!") .build() .withParams(params); ``` ______________________________________________________________________ ## Amazon Bedrock Amazon Bedrock uses a block-level caching model via the Converse API. When `cacheControl` is set on a message or tool, Bedrock inserts a `CachePoint` block immediately after the annotated element. Note Bedrock prompt caching is a JVM-only feature, as the Bedrock client itself is JVM-only. ### System messages ``` val prompt = prompt("assistant") { // Cache the system prompt using the default TTL system("You are a knowledgeable assistant...", BedrockCacheControl.Default) user("What is prompt caching?") } val response = client.execute(prompt, BedrockModels.AnthropicClaude4Sonnet) println(response) ``` ``` Prompt prompt = Prompt.builder("assistant") // Cache the system prompt using the default TTL .system("You are a knowledgeable assistant...", BedrockCacheControl.Default.INSTANCE) .user("What is prompt caching?") .build(); ``` ### User and assistant messages ``` val prompt = prompt("conversation") { system("You are a helpful assistant.") // Cache after the large context message user("Here is the document: ...", BedrockCacheControl.FiveMinutes) assistant(listOf(MessagePart.Text("I have read the document."))) user("Summarize it.") } val response = client.execute(prompt, BedrockModels.AnthropicClaude4Sonnet) println(response) ``` ``` Prompt prompt = Prompt.builder("conversation") .system("You are a helpful assistant.") // Cache after the large context message .user("Here is the document: ...", BedrockCacheControl.FiveMinutes.INSTANCE) .assistant("I have read the document.", BedrockCacheControl.Default.INSTANCE) .user("Summarize it.") .build(); ``` ### Tool definitions ``` val searchTool = ToolDescriptor( name = "web_search", description = "Search the web for information.", requiredParameters = listOf( ToolParameterDescriptor("query", "Search query", ToolParameterType.String) ), // Cache all tool definitions up to and including this one cacheControl = BedrockCacheControl.Default ) ``` ``` ToolDescriptor searchTool = new ToolDescriptor( "web_search", "Search the web for information.", List.of( new ToolParameterDescriptor("query", "Search query", ToolParameterType.String.INSTANCE) ), Collections.emptyList(), // Cache all tool definitions up to and including this one BedrockCacheControl.Default.INSTANCE ); ``` ### Cache TTL options | Option | TTL | | --------------------------------- | --------------------------------------- | | `BedrockCacheControl.Default` | Provider default (no explicit TTL sent) | | `BedrockCacheControl.FiveMinutes` | 5 minutes | | `BedrockCacheControl.OneHour` | 1 hour | See the [Amazon Bedrock prompt caching docs](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html) for supported models and pricing. ______________________________________________________________________ ## Choosing a caching strategy | Situation | Recommended approach | | ------------------------------------------------- | --------------------------------------------------------------- | | Multi-turn chat with a large, fixed system prompt | Anthropic automatic caching or Bedrock block-level on system | | Stable tool definitions reused across requests | Block-level `cacheControl` on the last tool definition | | Long document passed as user context | Block-level `cacheControl` on the user message | | Arbitrary multi-turn conversation (Anthropic) | Automatic caching via `AnthropicParams.cacheControl` | | Need 1-hour cache retention | `AnthropicCacheControl.OneHour` / `BedrockCacheControl.OneHour` | # Running prompts # LLM clients LLM clients are designed for direct interaction with LLM providers. Each client implements the [`LLMClient`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/ai.koog.prompt.executor.clients/-l-l-m-client/index.html) interface, which provides methods for executing prompts and streaming responses. You can use an LLM client when you work with a single LLM provider and don't need advanced lifecycle management. If you need to manage multiple LLM providers, use a [prompt executor](../prompt-executors/). The table below shows the available LLM clients and their capabilities. | LLM provider | LLMClient | Tool calling | Streaming | Multiple choices | Embeddings | Moderation | Model listing | Notes | | ---------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------ | --------- | ---------------- | ---------- | ----------- | ------------- | --------------------------------------------------------------------------------------------------------------------------- | | [OpenAI](https://platform.openai.com/docs/overview) | [OpenAILLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client/ai.koog.prompt.executor.clients.openai/-open-a-i-l-l-m-client/index.html) | ✓ | ✓ | ✓ | ✓ | ✓[1](#fn:1) | ✓ | | | [Anthropic](https://www.anthropic.com/) | [AnthropicLLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-anthropic-client/ai.koog.prompt.executor.clients.anthropic/-anthropic-l-l-m-client/index.html) | ✓ | ✓ | - | - | - | - | - | | [Google](https://ai.google.dev/) β | [GoogleLLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/ai.koog.prompt.executor.clients.google/-google-l-l-m-client/index.html) | ✓ | ✓ | ✓ | ✓ | - | ✓ | - | | [DeepSeek](https://www.deepseek.com/) β | [DeepSeekLLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-deepseek-client/ai.koog.prompt.executor.clients.deepseek/-deep-seek-l-l-m-client/index.html) | ✓ | ✓ | ✓ | - | - | ✓ | OpenAI-compatible chat client. | | [OpenRouter](https://openrouter.ai/) | [OpenRouterLLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openrouter-client/ai.koog.prompt.executor.clients.openrouter/-open-router-l-l-m-client/index.html) | ✓ | ✓ | ✓ | - | - | ✓ | OpenAI-compatible router client. | | [Amazon Bedrock](https://aws.amazon.com/bedrock/) | [BedrockLLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/ai.koog.prompt.executor.clients.bedrock/-bedrock-l-l-m-client/index.html) | ✓ | ✓ | - | ✓ | ✓[2](#fn:2) | - | JVM-only AWS SDK client that supports multiple model families. | | [Mistral](https://mistral.ai/) β | [MistralAILLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-mistralai-client/ai.koog.prompt.executor.clients.mistralai/-mistral-a-i-l-l-m-client/index.html) | ✓ | ✓ | ✓ | ✓ | ✓[3](#fn:3) | ✓ | OpenAI-compatible client. | | [Alibaba](https://www.alibabacloud.com/en?_p_lc=1) β | [DashScopeLLMClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-dashscope-client/ai.koog.prompt.executor.clients.dashscope/-dashscope-l-l-m-client/index.html) | ✓ | ✓ | ✓ | - | - | ✓ | OpenAI-compatible client that exposes provider-specific parameters (`enableSearch`, `parallelToolCalls`, `enableThinking`). | | [Ollama](https://ollama.com/) | [OllamaClient](https://api.koog.ai/prompt/prompt-executor/prompt-executor-clients/prompt-executor-ollama-client/ai.koog.prompt.executor.ollama.client/-ollama-client/index.html) | ✓ | ✓ | - | ✓ | ✓ | - | Local server client with model management APIs. | ## Running a prompt To run a prompt using an LLM client, perform the following: 1. Create an LLM client that handles the connection between your application and LLM providers. 1. Call the `execute()` method with the prompt and LLM as arguments. Here is an example that uses `OpenAILLMClient` to run prompts: ``` fun main() = runBlocking { // Create an OpenAI client val apiKey = System.getenv("OPENAI_API_KEY") val client = OpenAILLMClient(apiKey) // Create a prompt val prompt = prompt("prompt_name", LLMParams()) { // Add a system message to set the context system("You are a helpful assistant.") // Add a user message user("Tell me about Kotlin") // You can also add assistant messages for few-shot examples assistant("Kotlin is a modern programming language...") // Add another user message user("What are its key features?") } // Run the prompt val response = client.execute(prompt, OpenAIModels.Chat.GPT4o) // Print the response println(response) } ``` ``` // Create an OpenAI client String apiKey = System.getenv("OPENAI_API_KEY"); OpenAILLMClient client = openAIClient(apiKey); // Create a prompt Prompt prompt = Prompt.builder("prompt_name") // Add a system message to set the context .system("You are a helpful assistant.") // Add a user message .user("Tell me about Kotlin") // You can also add assistant messages for few-shot examples .assistant("Kotlin is a modern programming language...") // Add another user message .user("What are its key features?") .build(); // Run the prompt List response = client.execute(prompt, OpenAIModels.Chat.GPT4o, Collections.emptyList()); // Print the response System.out.println(response); client.close(); ``` ## Streaming responses Note Available for all LLM clients. When you need to process responses as they are generated, you can use the `executeStreaming()` method in Kotlin or `executeStreamingWithPublisher()` in Java to stream the model output. The streaming API provides different frame types: - **Delta frames** (`TextDelta`, `ReasoningDelta`, `ToolCallDelta`) — incremental content that arrives in chunks - **Complete frames** (`TextComplete`, `ReasoningComplete`, `ToolCallComplete`) — full content after all deltas are received - **End frame** (`End`) — signals stream completion with finish reason For models that support reasoning (such as Claude Sonnet 4.5 or GPT-o1), reasoning frames will be emitted during streaming. See the [Streaming API documentation](../../streaming-api/) for more details on working with frames. ``` // Set up the OpenAI client with your API key val token = System.getenv("OPENAI_API_KEY") val client = OpenAILLMClient(token) val response = client.executeStreaming( prompt = prompt("stream_demo") { user("Stream this response in short chunks.") }, model = OpenAIModels.Chat.GPT4_1 ) response.collect { frame -> when (frame) { is StreamFrame.TextDelta -> print(frame.text) is StreamFrame.ReasoningDelta -> print("[Reasoning] ${frame.text}") is StreamFrame.ToolCallComplete -> println("\nTool call: ${frame.name}") is StreamFrame.End -> println("\n[done] Reason: ${frame.finishReason}") else -> {} // Handle other frame types if needed } } ``` ``` // Set up the OpenAI client with your API key String token = System.getenv("OPENAI_API_KEY"); OpenAILLMClient client = openAIClient(token); Prompt prompt = Prompt.builder("stream_demo") .user("Stream this response in short chunks.") .build(); Publisher response = client.executeStreamingWithPublisher(prompt, OpenAIModels.Chat.GPT4_1); // Subscribe to the Publisher to consume frames response.subscribe(new Subscriber() { private Subscription subscription; @Override public void onSubscribe(Subscription s) { this.subscription = s; s.request(Long.MAX_VALUE); } @Override public void onNext(StreamFrame frame) { switch (frame) { case StreamFrame.TextDelta delta -> System.out.print(delta.getText()); case StreamFrame.ReasoningDelta reasoning -> System.out.print("[Reasoning] " + reasoning.getText()); case StreamFrame.ToolCallComplete toolCall -> System.out.println("\nTool call: " + toolCall.getName()); case StreamFrame.End end -> System.out.println("\n[done] Reason: " + end.getFinishReason()); default -> {} // Handle other frame types } } @Override public void onError(Throwable t) { t.printStackTrace(); } @Override public void onComplete() { } }); ``` ## Multiple choices Note Available for all LLM clients except `GoogleLLMClient`, `BedrockLLMClient`, and `OllamaClient` You can request multiple alternative responses from the model in a single call by using the `executeMultipleChoices()` method. It requires additionally specifying the [`numberOfChoices`](../prompt-creation/#prompt-parameters) LLM parameter in the prompt being executed. ``` fun main() = runBlocking { val apiKey = System.getenv("OPENAI_API_KEY") val client = OpenAILLMClient(apiKey) val choices = client.executeMultipleChoices( prompt = prompt("n_best", params = LLMParams(numberOfChoices = 3)) { system("You are a creative assistant.") user("Give me three different opening lines for a story.") }, model = OpenAIModels.Chat.GPT4o ) choices.forEachIndexed { i, choice -> val text = choice.parts.filterIsInstance().joinToString(" ") { it.text } println("Line #${i + 1}: $text") } } ``` ``` String apiKey = System.getenv("OPENAI_API_KEY"); OpenAILLMClient client = openAIClient(apiKey); // Configure parameters (LLMParams constructor requires all 8 arguments in Java) LLMParams params = new LLMParams( null, // temperature null, // maxTokens 3, // numberOfChoices null, // speculation null, // schema null, // toolChoice null, // user null // additionalProperties ); Prompt prompt = Prompt.builder("n_best") .system("You are a creative assistant.") .user("Give me three different opening lines for a story.") .build() .withParams(params); // LLMChoice is a type alias for List List> choices = client.executeMultipleChoices( prompt, OpenAIModels.Chat.GPT4o ); for (int i = 0; i < choices.size(); i++) { List choice = choices.get(i); StringBuilder text = new StringBuilder(); for (Message.Response msg : choice) { text.append(msg.getContent()).append(" "); } System.out.println("Line #" + (i + 1) + ": " + text.toString().trim()); } ``` ## Listing available models Note Available for all LLM clients except `AnthropicLLMClient`, `BedrockLLMClient`, and `OllamaClient`. To get a list of available model IDs supported by the LLM client, use the `models()` method: ``` fun main() = runBlocking { val apiKey = System.getenv("OPENAI_API_KEY") val client = OpenAILLMClient(apiKey) val models: List = client.models() models.forEach { println(it.id) } } ``` ``` String apiKey = System.getenv("OPENAI_API_KEY"); OpenAILLMClient client = openAIClient(apiKey); List models = client.models(); for (LLModel model : models) { System.out.println(model.getId()); } ``` ## Embeddings Note Available for `OpenAILLMClient`, `GoogleLLMClient`, `BedrockLLMClient`, `MistralAILLMClient`, and `OllamaClient`. You convert text into embedding vectors using the `embed()` method. Choose an embedding model and pass your text to this method: ``` fun main() = runBlocking { val apiKey = System.getenv("OPENAI_API_KEY") val client = OpenAILLMClient(apiKey) val embedding = client.embed( text = "This is a sample text for embedding", model = OpenAIModels.Embeddings.TextEmbedding3Large ) println("Embedding size: ${embedding.size}") } ``` ## Moderation Note Available for the following LLM clients: `OpenAILLMClient`, `BedrockLLMClient`, `MistralAILLMClient`, `OllamaClient`. You can use the `moderate()` method with a moderation model to check whether a prompt contains inappropriate content: ``` fun main() = runBlocking { val apiKey = System.getenv("OPENAI_API_KEY") val client = OpenAILLMClient(apiKey) val result = client.moderate( prompt = prompt("moderation") { user("This is a test message that may contain offensive content.") }, model = OpenAIModels.Moderation.Omni ) println(result) } ``` ``` String apiKey = System.getenv("OPENAI_API_KEY"); OpenAILLMClient client = openAIClient(apiKey); Prompt prompt = Prompt.builder("moderation") .user("This is a test message that may contain offensive content.") .build(); ModerationResult result = client.moderate(prompt, OpenAIModels.Moderation.Omni); System.out.println(result); ``` ## Integration with prompt executors [Prompt executors](../prompt-executors/) wrap LLM clients and provide additional functionality, such as routing, fallbacks, and unified usage across providers. They are recommended for production use, as they offer flexibility when working with multiple providers. ______________________________________________________________________ 1. Supports moderation via the OpenAI Moderation API. [↩](#fnref:1 "Jump back to footnote 1 in the text") 1. Moderation requires Guardrails configuration. [↩](#fnref:2 "Jump back to footnote 2 in the text") 1. Supports moderation via the Mistral `v1/moderations` endpoint. [↩](#fnref:3 "Jump back to footnote 3 in the text") # HTTP clients Every LLM client in Koog expects a [`KoogHttpClient`](api:http-client-core::ai.koog.http.client.KoogHttpClient) — the abstract HTTP contract the framework uses to talk to providers. You hand one in at construction. You can build that `KoogHttpClient` yourself, but it's real work: each provider has its own base URL, auth header shape, content-type, and SSE conventions. Getting all of that right per provider is exactly what [`KoogHttpClient.Factory`](api:http-client-core::ai.koog.http.client.KoogHttpClient.Factory) exists to spare you. You pass in a `Factory` and the provider client calls `Factory.create(...)` with the parameters that fit its API. Four backend factories ship out of the box — Ktor, the JDK `HttpClient`, OkHttp, and Spring's `WebClient` — and you can implement your own. ## How it works One factory works for any provider: pick a backend once and use it across clients. ``` fun main() { val factory = KtorKoogHttpClient.Factory() val openai = OpenAILLMClient( apiKey = System.getenv("OPENAI_API_KEY"), settings = OpenAIClientSettings(), httpClientFactory = factory, ) val anthropic = AnthropicLLMClient( apiKey = System.getenv("ANTHROPIC_API_KEY"), settings = AnthropicClientSettings(), httpClientFactory = factory, ) } ``` ``` import ai.koog.http.client.ktor.KtorKoogHttpClient; import ai.koog.prompt.executor.clients.anthropic.AnthropicClientSettings; import ai.koog.prompt.executor.clients.anthropic.AnthropicLLMClient; import ai.koog.prompt.executor.clients.openai.OpenAIClientSettings; import ai.koog.prompt.executor.clients.openai.OpenAILLMClient; KtorKoogHttpClient.Factory factory = new KtorKoogHttpClient.Factory(); OpenAILLMClient openai = new OpenAILLMClient( System.getenv("OPENAI_API_KEY"), new OpenAIClientSettings(), factory ); AnthropicLLMClient anthropic = new AnthropicLLMClient( System.getenv("ANTHROPIC_API_KEY"), new AnthropicClientSettings(), factory ); ``` ## Supported HTTP client flavors | Module | Notes | | ------------------------------------------------------------------- | --------------------------------------------- | | [`http-client-ktor`](api:http-client-ktor:) | The only backend usable from non-JVM targets. | | [`http-client-java`](api:http-client-java:) | Wraps the JDK 11+ `java.net.http.HttpClient`. | | [`http-client-okhttp`](api:http-client-okhttp:) | Backed by OkHttp. Android-friendly. | | [`http-client-spring-webclient`](api:http-client-spring-webclient:) | Backed by Spring `WebClient`. | ## Convenience APIs and factory auto-discovery On JVM and Android, you can construct each LLM client without passing a factory explicitly. Behind the scenes, [`HttpClientFactoryResolver`](api:http-client-core::ai.koog.http.client.HttpClientFactoryResolver) uses `java.util.ServiceLoader` to resolve `KoogHttpClient.Factory` from the runtime classpath: - Every backend module provides a `ServiceLoader` registration. - Resolution succeeds only when exactly one factory is visible on the runtime classpath. - `prompt-executor-llms-all` declares `http-client-ktor` as a `runtimeOnly` dependency, so you get Ktor by default without compile-time exposure to that module. - `simpleExecutor(apiKey)` and `PromptExecutorBuilder.(apiKey)` use the same resolution path. ``` fun main() { val apiKey = System.getenv("OPENAI_API_KEY") val client = OpenAILLMClient(apiKey) val executor = simpleOpenAIExecutor(apiKey) } ``` ``` import static ai.koog.prompt.executor.clients.openai.OpenAIClientFactory.openAIClient; import static ai.koog.prompt.executor.llms.all.SimplePromptExecutors.simpleOpenAIExecutor; String apiKey = System.getenv("OPENAI_API_KEY"); OpenAILLMClient client = openAIClient(apiKey); PromptExecutor executor = simpleOpenAIExecutor(apiKey); ``` Auto-discovery is not supported on KMP at the moment, so the convenience methods are not available outside the JVM either. From `commonMain`, pass a `Factory` explicitly. ### Auto-discovery gotchas - **Zero backends on the runtime classpath** → `IllegalStateException` on first resolution. Add a backend module to the runtime classpath, or pass a `Factory` explicitly. - **Two or more backends** → same exception; the message names the providers it found. Exclude all but one with Gradle (`exclude(module = "http-client-ktor")` on the offending dependency) or pass a `Factory` explicitly at the call site. ## Custom backends Any class implementing `KoogHttpClient.Factory` works. To make it auto-discoverable on the JVM, register it as a `ServiceLoader` provider: ``` src/main/resources/META-INF/services/ai.koog.http.client.KoogHttpClient$Factory ``` The file contains a single line: the fully qualified name of your factory class. The literal `$` (separator for the nested `Factory` class) is correct — the file is `KoogHttpClient$Factory`, not `KoogHttpClient.Factory`. If you don't want auto-discovery, skip the registration and pass your factory explicitly everywhere. # Prompt executors Prompt executors provide a higher-level abstraction that lets you manage the lifecycle of one or multiple LLM clients. You can work with multiple LLM providers through a unified interface, abstracting from provider-specific details, with dynamic switching between them and fallbacks. ## Executor types Koog provides three main types of prompt executors that implement the [`PromptExecutor`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.model/-prompt-executor/index.html) interface: | Type | Class | Description | | --------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | Single-provider | [`SingleLLMPromptExecutor`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-single-l-l-m-prompt-executor/index.html) | Wraps a single LLM client for one provider. Use this executor if your agent only requires switching between models within a single LLM provider. | | Multi-provider | [`MultiLLMPromptExecutor`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-multi-l-l-m-prompt-executor/index.html) | Wraps multiple LLM clients and routes calls based on the LLM provider. It can optionally use a configured fallback provider and LLM when the requested client is unavailable. Use this executor if your agent needs to switch between LLMs from different providers. | | Routing | [`RoutingLLMPromptExecutor`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-routing-l-l-m-prompt-executor/index.html) | Distributes requests to a given LLM model across multiple client instances using routing strategies. Use this executor to avoid rate limits, improve throughput, and implement failover strategies with load balancing. | ## Creating a single-provider executor To create a prompt executor for a specific LLM provider, perform the following: 1. Configure an LLM client for a specific provider with the corresponding API key. 1. Create a prompt executor using [`MultiLLMPromptExecutor`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-multi-l-l-m-prompt-executor/index.html). Here is an example: ``` val openAIClient = OpenAILLMClient(System.getenv("OPENAI_API_KEY")) val promptExecutor = MultiLLMPromptExecutor(openAIClient) ``` ``` OpenAILLMClient openAIClient = openAIClient(System.getenv("OPENAI_API_KEY")); MultiLLMPromptExecutor promptExecutor = new MultiLLMPromptExecutor(openAIClient); ``` ## Creating a multi-provider executor To create a prompt executor that works with multiple LLM providers, do the following: 1. Configure clients for the required LLM providers with the corresponding API keys. 1. Pass the configured clients to the [`MultiLLMPromptExecutor`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-multi-l-l-m-prompt-executor/index.html) class constructor to create a prompt executor with multiple LLM providers. ``` val openAIClient = OpenAILLMClient(System.getenv("OPENAI_API_KEY")) val ollamaClient = OllamaClient() val multiExecutor = MultiLLMPromptExecutor( LLMProvider.OpenAI to openAIClient, LLMProvider.Ollama to ollamaClient ) ``` ``` OpenAILLMClient openAIClient = openAIClient(System.getenv("OPENAI_API_KEY")); OllamaClient ollamaClient = ollamaClient(); MultiLLMPromptExecutor promptExecutor = new MultiLLMPromptExecutor(openAIClient, ollamaClient); ``` ## Creating a routing executor Experimental API Routing capabilities are experimental and may change in future releases. To use them, opt in with `@OptIn(ExperimentalRoutingApi::class)`. To create a prompt executor that distributes requests across multiple LLM client instances using routing strategies, do the following: 1. Configure multiple client instances (they can be for the same or different LLM providers) with the corresponding API keys. 1. Create a router using a routing strategy, such as [`RoundRobinRouter`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-round-robin-router/index.html). 1. Pass the router to the [`RoutingLLMPromptExecutor`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-routing-l-l-m-prompt-executor/index.html) class constructor. This is useful for avoiding rate limits, improving throughput, and implementing failover strategies. ``` // Create multiple client instances val openAI1 = OpenAILLMClient(apiKey = "openai-key-1") val openAI2 = OpenAILLMClient(apiKey = "openai-key-2") val anthropic = AnthropicLLMClient(apiKey = "anthropic-key") // Create router with round-robin strategy val router = RoundRobinRouter(openAI1, openAI2, anthropic) // Create routing executor val routingExecutor = RoutingLLMPromptExecutor(router) ``` ``` // Create multiple client instances OpenAILLMClient openAI1 = openAIClient("openai-key-1"); OpenAILLMClient openAI2 = openAIClient("openai-key-2"); AnthropicLLMClient anthropic = anthropicClient("anthropic-key"); // Create router with round-robin strategy RoundRobinRouter router = new RoundRobinRouter(openAI1, openAI2, anthropic); // Create routing executor RoutingLLMPromptExecutor routingExecutor = new RoutingLLMPromptExecutor(router); ``` When you execute prompts with this executor, requests to OpenAI models will alternate between `openAI1` and `openAI2` using the round-robin strategy. Requests to Anthropic models always go to the single `anthropic` client, as round-robin maintains an independent counter per provider. You can also implement custom routing strategies by creating a class that implements the [`LLMClientRouter`](https://api.koog.ai/prompt/prompt-executor/prompt-executor-model/ai.koog.prompt.executor.llms/-l-l-m-client-router/index.html) interface. ## Pre-defined prompt executors For faster setup, Koog provides ready-to-use executor implementations for common providers in both Kotlin and Java. The following table includes the **pre-defined single-provider executors** that return `SingleLLMPromptExecutor` configured with a specific LLM client. | LLM provider | Prompt executor | Description | | -------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- | | OpenAI | [simpleOpenAIExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-open-a-i-executor.html) | Wraps `OpenAILLMClient` that runs prompts with OpenAI models. | | OpenAI | [simpleAzureOpenAIExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-azure-open-a-i-executor.html) | Wraps `OpenAILLMClient` configured for using Azure OpenAI Service. | | Anthropic | [simpleAnthropicExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-anthropic-executor.html) | Wraps `AnthropicLLMClient` that runs prompts with Anthropic models. | | Google | [simpleGoogleAIExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-google-a-i-executor.html) | Wraps `GoogleLLMClient` that runs prompts with Google models. | | OpenRouter | [simpleOpenRouterExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-open-router-executor.html) | Wraps `OpenRouterLLMClient` that runs prompts with OpenRouter. | | Amazon Bedrock | [simpleBedrockExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-bedrock-executor.html) | Wraps `BedrockLLMClient` that runs prompts with AWS Bedrock. | | Amazon Bedrock | [simpleBedrockExecutorWithBearerToken](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-bedrock-executor-with-bearer-token.html) | Wraps `BedrockLLMClient` and uses the provided Bedrock API key to send requests. | | Mistral | [simpleMistralAIExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-mistral-a-i-executor.html) | Wraps `MistralAILLMClient` that runs prompts with Mistral models. | | Ollama | [simpleOllamaAIExecutor](https://api.koog.ai/prompt/prompt-executor/prompt-executor-llms-all/ai.koog.prompt.executor.llms.all/simple-ollama-a-i-executor.html) | Wraps `OllamaClient` that runs prompts with Ollama. | Here is an example of creating a pre-defined executor: ``` // Create an OpenAI executor val promptExecutor = simpleOpenAIExecutor("OPENAI_API_KEY") ``` ``` // Create an OpenAI executor PromptExecutor openAIExecutor = simpleOpenAIExecutor("OPENAI_API_KEY"); ``` ## Running a prompt To run a prompt using a prompt executor, do the following: 1. Create a prompt executor. 1. Run the prompt with the specific LLM using the `execute()` method. Here is an example: ``` // Create an OpenAI executor val promptExecutor = simpleOpenAIExecutor("OPENAI_API_KEY") // Execute a prompt val response = promptExecutor.execute( prompt = prompt("demo") { user("Summarize this.") }, model = OpenAIModels.Chat.GPT4o ) ``` ``` // Create an OpenAI executor PromptExecutor promptExecutor = simpleOpenAIExecutor("OPENAI_API_KEY"); // Create a prompt Prompt prompt = Prompt.builder("demo") .user("Summarize this.") .build(); // Run the prompt List response = promptExecutor.execute(prompt, OpenAIModels.Chat.GPT4o); ``` This will run the prompt with the `GPT4o` model and return the response. Note The prompt executors provide methods to run prompts using various capabilities, such as streaming, multiple choice generation, and content moderation. Since prompt executors wrap LLM clients, each executor supports the capabilities of the corresponding client. For details, refer to [LLM clients](../llm-clients/). ## Switching between providers When you work with multiple LLM providers using `MultiLLMPromptExecutor`, you can switch between them. The process is as follows: 1. Create an LLM client instance for each provider you want to use. 1. Create a `MultiLLMPromptExecutor` that maps LLM providers to LLM clients. 1. Run a prompt with a model from the corresponding client passed as an argument to the `execute()` method. The prompt executor will use the corresponding client based on the model provider to run the prompt. Here is an example of switching between providers: ``` // Create LLM clients for OpenAI, Anthropic, and Google providers val openAIClient = OpenAILLMClient("OPENAI_API_KEY") val anthropicClient = AnthropicLLMClient("ANTHROPIC_API_KEY") val googleClient = GoogleLLMClient("GOOGLE_API_KEY") // Create a MultiLLMPromptExecutor that maps LLM providers to LLM clients val executor = MultiLLMPromptExecutor( LLMProvider.OpenAI to openAIClient, LLMProvider.Anthropic to anthropicClient, LLMProvider.Google to googleClient ) // Create a prompt val p = prompt("demo") { user("Summarize this.") } // Run the prompt with an OpenAI model; the prompt executor automatically switches to the OpenAI client val openAIResult = executor.execute(p, OpenAIModels.Chat.GPT4o) // Run the prompt with an Anthropic model; the prompt executor automatically switches to the Anthropic client val anthropicResult = executor.execute(p, AnthropicModels.Sonnet_4_5) ``` ``` // Create LLM clients for OpenAI, Anthropic, and Google providers OpenAILLMClient openAIClient = openAIClient("OPENAI_API_KEY"); AnthropicLLMClient anthropicClient = anthropicClient("ANTHROPIC_API_KEY"); GoogleLLMClient googleClient = googleClient("GOOGLE_API_KEY"); // Create a MultiLLMPromptExecutor that maps LLM providers to LLM clients MultiLLMPromptExecutor promptExecutor = new MultiLLMPromptExecutor( Map.of( LLMProvider.OpenAI, openAIClient, LLMProvider.Anthropic, anthropicClient, LLMProvider.Google, googleClient ) ); // Create a prompt Prompt prompt = Prompt.builder("demo") .user("Summarize this.") .build(); // Run the prompt with an OpenAI model; the prompt executor automatically switches to the OpenAI client List openAIResult = promptExecutor.execute(prompt, OpenAIModels.Chat.GPT4o); // Run the prompt with an Anthropic model; the prompt executor automatically switches to the Anthropic client List anthropicResult = promptExecutor.execute(prompt, AnthropicModels.Sonnet_4_5); ``` You can optionally configure a fallback LLM provider and model to use when the requested client is unavailable. For details, refer to [Configuring fallbacks](#configuring-fallbacks). ## Configuring fallbacks Multi-provider and routing prompt executors can be configured to use a fallback LLM provider and model when the requested LLM client is unavailable. To configure the fallback mechanism, pass fallback settings when creating a `MultiLLMPromptExecutor` or `RoutingLLMPromptExecutor`: ``` val openAIClient = OpenAILLMClient(System.getenv("OPENAI_API_KEY")) val ollamaClient = OllamaClient() val multiExecutor = MultiLLMPromptExecutor( LLMProvider.OpenAI to openAIClient, LLMProvider.Ollama to ollamaClient, fallback = MultiLLMPromptExecutor.FallbackPromptExecutorSettings( fallbackProvider = LLMProvider.Ollama, fallbackModel = OllamaModels.Meta.LLAMA_3_2 ) ) ``` ``` OpenAILLMClient openAIClient = openAIClient(System.getenv("OPENAI_API_KEY")); OllamaClient ollamaClient = ollamaClient(); MultiLLMPromptExecutor multiExecutor = new MultiLLMPromptExecutor( Map.of( LLMProvider.OpenAI, openAIClient, LLMProvider.Ollama, ollamaClient ), new MultiLLMPromptExecutor.FallbackPromptExecutorSettings( LLMProvider.Ollama, OllamaModels.Meta.LLAMA_3_2 ) ); ``` If you pass a model from an LLM provider that is not included in the `MultiLLMPromptExecutor`, the prompt executor will use the fallback model: ``` // Create a prompt val p = prompt("demo") { user("Summarize this") } // If you pass a Google model, the prompt executor will use the fallback model, as the Google client is not included val response = multiExecutor.execute(p, GoogleModels.Gemini2_5Pro) ``` ``` // Create a prompt Prompt p = Prompt.builder("demo") .user("Summarize this") .build(); // If you pass a Google model, the prompt executor will use the fallback model, as the Google client is not included List response = multiExecutor.execute(p, GoogleModels.Gemini2_5Pro); ``` Note Fallbacks are available for the `execute()` and `executeMultipleChoices()` methods only. # Tools # Overview Agents use tools to perform specific tasks or access external systems. ## Tool workflow The Koog framework offers the following workflow for working with tools in Kotlin and Java: 1. Create a custom tool or use one of the built-in tools. 1. Add the tool to a tool registry. 1. Pass the tool registry to an agent. 1. Use the tool with the agent. ### Available tool types There are three types of tools in the Koog framework: - Built-in tools that provide functionality for agent-user interaction and conversation management. For details, see [Built-in tools](built-in-tools/). - Annotation-based custom tools that let you expose functions as tools to LLMs. For details, see [Annotation-based tools](annotation-based-tools/). - Custom tools that let you control tool parameters, metadata, execution logic, and how it is registered and invoked. For details, see [Class-based tools](class-based-tools/). ### Tool registry Before you can use a tool in an agent, you must add it to a tool registry. The tool registry manages all tools available to the agent. The key features of the tool registry: - Organizes tools. - Supports merging of multiple tool registries. - Provides methods to retrieve tools by name or type. To learn more, see [ToolRegistry](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool-registry/index.html). Here is an example of how to create the tool registry and add the tool to it: ``` val toolRegistry = ToolRegistry { tools(myTool) } ``` ``` // Create an instance of your ToolSet MyToolSet myTool = new MyToolSet(); // Build the ToolRegistry and register tools from the ToolSet ToolRegistry toolRegistry = ToolRegistry.builder() .tools(myTool) .build(); ``` To merge multiple tool registries, do the following: ``` val firstToolRegistry = ToolRegistry { tools(firstSampleTool) } val secondToolRegistry = ToolRegistry { tools(secondSampleTool) } val newRegistry = firstToolRegistry + secondToolRegistry ``` ``` // Create instances of your ToolSets FirstToolSet firstSampleTool = new FirstToolSet(); SecondToolSet secondSampleTool = new SecondToolSet(); // Build separate tool registries ToolRegistry firstToolRegistry = ToolRegistry.builder() .tools(firstSampleTool) .build(); ToolRegistry secondToolRegistry = ToolRegistry.builder() .tools(secondSampleTool) .build(); ToolRegistry newRegistry = firstToolRegistry.plus(secondToolRegistry); ``` ### Passing tools to an agent To enable an agent to use a tool, you need to provide a tool registry that contains this tool as an argument when creating the agent: ``` // Agent initialization val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), systemPrompt = "You are a helpful assistant with strong mathematical skills.", llmModel = OpenAIModels.Chat.GPT4o, // Pass your tool registry to the agent toolRegistry = toolRegistry ) ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("You are a helpful assistant with strong mathematical skills.") .llmModel(OpenAIModels.Chat.GPT4o) .toolRegistry(ToolRegistry.builder() .tools(secondSampleTool) .build() ) .build(); ``` ### Calling tools There are several ways to call tools within your agent code. The recommended approach is to use the provided methods in the agent context rather than calling tools directly, as this ensures proper handling of tool operation within the agent environment. Tip Ensure you have implemented proper [error handling](../features/agent-event-handlers/) in your tools to prevent agent failure. The tools are called within a specific session context represented by `AIAgentLLMWriteSession`. It provides several methods for calling tools so that you can: - Call a tool with the given arguments. - Call a tool by its name and the given arguments. - Call a tool by the provided tool class and arguments. - Call a tool of the specified type with the given arguments. - Call a tool that returns a raw string result. For more details, the API reference for [AIAgentLLMWriteSession](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.session/-a-i-agent-l-l-m-write-session/index.html). #### Parallel tool calls You can also call tools in parallel using the `toParallelToolCallsRaw` extension. For example: ``` @Serializable data class Book( val title: String, val author: String, val description: String ) class BookTool() : SimpleTool( argsType = typeToken(), name = NAME, description = "A tool to parse book information from Markdown" ) { companion object { const val NAME = "book" } override suspend fun execute(args: Book): String { println("${args.title} by ${args.author}:\n ${args.description}") return "Done" } } val strategy = strategy("strategy-name") { /*...*/ val myNode by node { _ -> llm.writeSession { flow { emit(Book("Book 1", "Author 1", "Description 1")) }.toParallelToolCallsRaw(BookTool::class).collect() } } } ``` ``` ``` #### Calling tools from nodes When building agent workflows with nodes, you can use special nodes to call tools: - **nodeExecuteTool**: calls a single tool call and returns its result. For details, see [API-reference](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-execute-tool.html). - **nodeExecuteSingleTool** that calls a specific tool with the provided arguments. For details, see [API-reference](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-execute-single-tool.html). - **nodeExecuteMultipleTools** that performs multiple tool calls and returns their results. For details, see [API-reference](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-execute-multiple-tools.html). - **nodeLLMSendToolResult** that sends a tool result to the LLM and gets a response. For details, see [API-reference](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-send-tool-result.html). - **nodeLLMSendMultipleToolResults** that sends multiple tool results to the LLM. For details, see [API-reference](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-send-multiple-tool-results.html). ## Using agents as tools The framework provides the capability to convert any AI agent into a tool that can be used by other agents. This powerful feature enables you to create hierarchical agent architectures where specialized agents can be called as tools by higher-level orchestrating agents. ### Converting agents to tools To convert an agent into a tool, use the `AIAgentService` and the `createAgentTool()` extension function: ``` // Create a specialized agent service, responsible for creating financial analysis agents. val analysisAgentService = AIAgentService( promptExecutor = simpleOpenAIExecutor(apiKey), llmModel = OpenAIModels.Chat.GPT4o, systemPrompt = "You are a financial analysis specialist.", toolRegistry = analysisToolRegistry ) // Create a tool that would run financial analysis agent once called. val analysisAgentTool = analysisAgentService.createAgentTool( agentName = "analyzeTransactions", agentDescription = "Performs financial transaction analysis", inputDescription = "Transaction analysis request", inputType = typeToken(), ) ``` ``` ``` ### Using agent tools in other agents Once converted to a tool, you can add the agent tool to another agent's tool registry: ``` // Create a coordinator agent that can use specialized agents as tools val coordinatorAgent = AIAgent( promptExecutor = simpleOpenAIExecutor(apiKey), llmModel = OpenAIModels.Chat.GPT4o, systemPrompt = "You coordinate different specialized services.", toolRegistry = ToolRegistry { tool(analysisAgentTool) // Add other tools as needed } ) ``` ``` ``` ### Agent tool execution When an agent tool is called: 1. The arguments are deserialized according to the input descriptor. 1. The wrapped agent is executed with the deserialized input. 1. The agent's output is serialized and returned as the tool result. ### Benefits of agents as tools - **Modularity**: Break complex workflows into specialized agents. - **Reusability**: Use the same specialized agent across multiple coordinator agents. - **Separation of concerns**: Each agent can focus on its specific domain. # Built-in tools Koog provides built-in tools for Kotlin and Java to help you quickly prototype and experiment with agent-user interaction. These tools are not intended for production use. To use them, add `ai.koog:agents-ext` to your dependencies. The following built-in tools are available: | Tool | Name | Description | | ----------------- | -------------------- | ------------------------------------------------------------------------------------------------------------------------ | | SayToUser | `__say_to_user__` | Lets the agent send a message to the user. It prints the agent message to the console with the `Agent says:` prefix. | | AskUser | `__ask_user__` | Lets the agent ask the user for input. It prints the agent message to the console and waits for user response. | | ExitTool | `__exit__` | Lets the agent finish the conversation and terminate the session. | | ReadFileTool | `__read_file__` | Reads text file with optional line range selection. Returns formatted content with metadata using 0-based line indexing. | | EditFileTool | `__edit_file__` | Makes a single, targeted text replacement in a file; can also create new files or fully replace contents. | | ListDirectoryTool | `__list_directory__` | Lists directory contents as a hierarchical tree with optional depth control and glob filtering. | | WriteFileTool | `__write_file__` | Writes text content to a file (creating parent directories if needed). | ## Registering built-in tools Like any other tool, a built-in tool must be added to the tool registry to become available for an agent. Here is an example: ``` // Create a tool registry with all built-in tools val toolRegistry = ToolRegistry { tool(SayToUser) tool(AskUser) tool(ExitTool) tool(ReadFileTool(JVMFileSystemProvider.ReadOnly)) tool(ListDirectoryTool(JVMFileSystemProvider.ReadOnly)) tool(WriteFileTool(JVMFileSystemProvider.ReadWrite)) } // Pass the registry when creating an agent val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(apiToken), systemPrompt = "You are a helpful assistant.", llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = toolRegistry ) ``` You can create a comprehensive set of capabilities for your agent by combining built-in tools and custom tools within the same registry in both Kotlin and Java. To learn more about custom tools, see [Annotation-based tools](../annotation-based-tools/) and [Class-based tools](../class-based-tools/). # Annotation-based tools Annotation-based tools provide a declarative way to expose functions and methods as tools for large language models (LLMs) in both Kotlin and Java. By using annotations, you can transform any function or method into a tool that LLMs can understand and use. This approach is useful when you need to expose existing functionality to LLMs in Kotlin or Java without implementing tool descriptions manually. Note Annotation-based tools are JVM-only and not available for other platforms. For multiplatform support, use the [class-based tool API](../class-based-tools/). ## Key annotations To start using annotation-based tools in your project, you need to understand the following key annotations: | Annotation | Description | | ----------------- | ----------------------------------------------------------------------- | | `@Tool` | Marks functions that should be exposed as tools to LLMs. | | `@LLMDescription` | Provides descriptive information about your tools and their components. | ## @Tool annotation The `@Tool` annotation is used to mark functions (Kotlin) or methods (Java) that should be exposed as tools to LLMs. The functions and methods annotated with `@Tool` are collected by reflection from objects that implement the `ToolSet` interface. For details, see [Implement the ToolSet interface](#1-implement-the-toolset-interface). ### Definition ``` @Target(AnnotationTarget.FUNCTION) public annotation class Tool(val customName: String = "") ``` ### Parameters | Name | Required | Description | | ------------ | -------- | ---------------------------------------------------------------------------------------- | | `customName` | No | Specifies a custom name for the tool. If not provided, the name of the function is used. | ### Usage To mark a function or method as a tool, apply the `@Tool` annotation to this function or method in a class that implements the `ToolSet` interface: ``` class MyToolSet : ToolSet { @Tool fun myTool(): String { // Tool implementation return "Result" } @Tool(customName = "customToolName") fun anotherTool(): String { // Tool implementation return "Result" } } ``` ``` public class MyToolSet implements ToolSet { @Tool public String myTool() { // Tool implementation return "Result"; } @Tool(customName = "customToolName") public String anotherTool() { // Tool implementation return "Result"; } } ``` ## @LLMDescription annotation The `@LLMDescription` annotation provides descriptive information about code elements (classes, functions, methods, parameters, and so on) to LLMs. This helps LLMs understand the purpose and usage of these elements. ### Definition ``` @Target( AnnotationTarget.PROPERTY, AnnotationTarget.CLASS, AnnotationTarget.TYPE, AnnotationTarget.VALUE_PARAMETER, AnnotationTarget.FUNCTION ) public annotation class LLMDescription(val description: String) ``` ### Parameters | Name | Required | Description | | ------------- | -------- | ---------------------------------------------- | | `description` | Yes | A string that describes the annotated element. | ### Usage The `@LLMDescription` annotation can be applied at various levels. For example: - Function level: ``` @Tool @LLMDescription("Performs a specific operation and returns the result") fun myTool(): String { // Function implementation return "Result" } ``` ``` @Tool @LLMDescription(description = "Performs a specific operation and returns the result") public String myTool() { // Function implementation return "Result"; } ``` - Parameter level: ``` @Tool @LLMDescription("Processes input data") fun processTool( @LLMDescription("The input data to process") input: String, @LLMDescription("Optional configuration parameters") config: String = "" ): String { // Function implementation return "Processed: $input with config: $config" } ``` ``` @Tool @LLMDescription(description = "Processes input data") public String processTool( @LLMDescription(description = "The input data to process") String input, @LLMDescription(description = "Optional configuration parameters") String config ) { // Function implementation return "Processed: " + input + " with config: " + config; } ``` ## Creating a tool ### 1. Implement the ToolSet interface Create a class that implements the [`ToolSet`](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools.reflect/-tool-set/index.html) interface. This interface marks your class as a container for tools. ``` class MyFirstToolSet : ToolSet { // Tools will go here } ``` ``` public class MyFirstToolSet implements ToolSet { // Tools will go here } ``` ### 2. Add tool functions Add functions or methods to your class and annotate them with `@Tool` to expose them as tools: ``` class MyFirstToolSet : ToolSet { @Tool fun getWeather(location: String): String { // In a real implementation, you would call a weather API return "The weather in $location is sunny and 72°F" } } ``` ``` public class MyFirstToolSet implements ToolSet { @Tool public String getWeather(String location) { // In a real implementation, you would call a weather API return "The weather in " + location + " is sunny and 72°F"; } } ``` ### 3. Add descriptions Add `@LLMDescription` annotations to provide context for the LLM: ``` @LLMDescription("Tools for getting weather information") class MyFirstToolSet : ToolSet { @Tool @LLMDescription("Get the current weather for a location") fun getWeather( @LLMDescription("The city and state/country") location: String ): String { // In a real implementation, you would call a weather API return "The weather in $location is sunny and 72°F" } } ``` ``` @LLMDescription(description = "Tools for getting weather information") public class MyFirstToolSet implements ToolSet { @Tool @LLMDescription(description = "Get the current weather for a location") public String getWeather( @LLMDescription(description = "The city and state/country") String location ) { // In a real implementation, you would call a weather API return "The weather in " + location + " is sunny and 72°F"; } } ``` ### 4. Use your tools with an agent Now you can use your tools with an agent: ``` fun main() { runBlocking { // Create your tool set val weatherTools = MyFirstToolSet() // Create an agent with your tools val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(apiToken), systemPrompt = "Provide weather information for a given location.", llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = ToolRegistry { tools(weatherTools) } ) // The agent can now use your weather tools agent.run("What's the weather like in New York?") } } ``` ``` String apiToken = System.getenv("OPENAI_API_KEY"); // Create your tool set MyFirstToolSet weatherTools = new MyFirstToolSet(); ToolRegistry toolRegistry = ToolRegistry.builder() .tools(weatherTools) .build(); // Create an agent with your tools AIAgent agent = AIAgent.builder() .promptExecutor(simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))) .systemPrompt("Provide weather information for a given location.") .llmModel(OpenAIModels.Chat.GPT4o) .toolRegistry(toolRegistry) .build(); // The agent can now use your weather tools String result = agent.run("What's the weather like in New York?"); System.out.println(result); ``` ## Usage examples Here are some real-world examples of tool annotations. ### Basic example: Switch controller This example shows a simple tool set for controlling a switch: ``` @LLMDescription("Tools for controlling a switch") class SwitchTools(val switch: Switch) : ToolSet { @Tool @LLMDescription("Switches the state of the switch") fun switch( @LLMDescription("The state to set (true for on, false for off)") state: Boolean ): String { switch.switch(state) return "Switched to ${if (state) "on" else "off"}" } @Tool @LLMDescription("Returns the current state of the switch") fun switchState(): String { return "Switch is ${if (switch.isOn()) "on" else "off"}" } } ``` ``` public class Switch { private boolean state; public Switch(boolean state) { this.state = state; } // "switch" is a reserved keyword in Java, so we use a different method name public void setState(boolean state) { this.state = state; } public boolean isOn() { return state; } } @LLMDescription(description = "Tools for controlling a switch") public class SwitchTools implements ToolSet { private final Switch sw; public SwitchTools(Switch sw) { this.sw = sw; } @Tool @LLMDescription(description = "Switches the state of the switch") public String switchStateTo( @LLMDescription(description = "The state to set (true for on, false for off)") boolean state ) { sw.setState(state); return "Switched to " + (state ? "on" : "off"); } @Tool @LLMDescription(description = "Returns the current state of the switch") public String switchState() { return "Switch is " + (sw.isOn() ? "on" : "off"); } } ``` When an LLM needs to control a switch, it can understand the following information from the provided description: - The purpose and functionality of the tools. - The required parameters for using the tools. - The acceptable values for each parameter. - The expected return values upon execution. ### Advanced example: Diagnostic tools This example shows a more complex tool set for device diagnostics: ``` @LLMDescription("Tools for performing diagnostics and troubleshooting on devices") class DiagnosticToolSet : ToolSet { @Tool @LLMDescription("Run diagnostic on a device to check its status and identify any issues") fun runDiagnostic( @LLMDescription("The ID of the device to diagnose") deviceId: String, @LLMDescription("Additional information for the diagnostic (optional)") additionalInfo: String = "" ): String { // Implementation return "Diagnostic results for device $deviceId" } @Tool @LLMDescription("Analyze an error code to determine its meaning and possible solutions") fun analyzeError( @LLMDescription("The error code to analyze (e.g., 'E1001')") errorCode: String ): String { // Implementation return "Analysis of error code $errorCode" } } ``` ``` @LLMDescription(description = "Tools for performing diagnostics and troubleshooting on devices") public class DiagnosticToolSet implements ToolSet { // Convenience overload (not exposed as a tool) public String runDiagnostic(String deviceId) { return runDiagnostic(deviceId, ""); } @Tool @LLMDescription(description = "Run diagnostic on a device to check its status and identify any issues") public String runDiagnostic( @LLMDescription(description = "The ID of the device to diagnose") String deviceId, @LLMDescription(description = "Additional information for the diagnostic (optional)") String additionalInfo ) { // Implementation return "Diagnostic results for device " + deviceId; } @Tool @LLMDescription(description = "Analyze an error code to determine its meaning and possible solutions") public String analyzeError( @LLMDescription(description = "The error code to analyze (e.g., 'E1001')") String errorCode ) { // Implementation return "Analysis of error code " + errorCode; } } ``` ## Best practices - **Provide clear descriptions**: write clear, concise descriptions that explain the purpose and behavior of tools, parameters, and return values. - **Describe all parameters**: add `@LLMDescription` to all parameters to help LLMs understand what each parameter is for. - **Use consistent naming**: use consistent naming conventions for tools and parameters to make them more intuitive. - **Group related tools**: group related tools in the same `ToolSet` implementation and provide a class-level description. - **Return informative results**: make sure tool return values provide clear information about the result of the operation. - **Handle errors gracefully**: include error handling in your tools and return informative error messages. - **Document default values**: when parameters have default values (Kotlin) or overloads (Java), document this in the description. - **Keep tools focused**: Each tool should perform a specific, well-defined task rather than trying to do too many things. ## Troubleshooting common issues When working with tool annotations, you might encounter some common issues. ### Tools not being recognized If the agent does not recognize your tools, check the following: - Your class implements the `ToolSet` interface. - All tool functions or methods are annotated with `@Tool`. - Tool functions or methods have appropriate return types (`String` is recommended for simplicity). - Your tools are properly registered with the agent. ### Unclear tool descriptions If the LLM does not use your tools correctly or misunderstands their purpose, try the following: - Use primitive parameter types when possible (`String`, `Boolean`, `Int` in Kotlin, or `String`, `boolean`, `int` in Java). - Clearly describe the expected format in the parameter description. - For complex types, consider using `String` parameters with a specific format and parse them in your tool. - Include examples of valid inputs in your parameter descriptions. - Note that Java doesn't support default parameters. Use method overloading instead. ### Parameter type issues If the LLM provides incorrect parameter types, try the following: - Use simple parameter types when possible (`String`, `Boolean`, `Int`). - Clearly describe the expected format in the parameter description. - For complex types, consider using `String` parameters with a specific format and parse them in your tool. - Include examples of valid inputs in your parameter descriptions. ### Performance issues If your tools cause performance problems, try the following: - Keep tool implementations lightweight. - For resource-intensive operations, consider implementing asynchronous processing. - Cache results when appropriate. - Log tool usage to identify bottlenecks. # Class-based tools This section explains the API designed for scenarios that require enhanced flexibility and customized behavior. With this approach in Kotlin, you have full control over a tool, including its parameters, metadata, execution logic, and how it is registered and invoked. In Java, tools are created using annotation-based methods with reflection-based registration. This level of control is ideal for creating sophisticated tools that extend basic use cases, enabling seamless integration into agent sessions and workflows. This page describes how to implement a tool in both Kotlin and Java, manage tools through registries, call them, and use within node-based agent architectures. Note The API is multiplatform for Kotlin. Java tools are implemented using annotation-based methods and registered via reflection. This lets you use the same tools across different platforms in Kotlin, while Java provides full JVM interoperability. ## Tool implementation The Koog framework provides the following approaches for implementing tools: For Kotlin: - Using the base class `Tool` for all tools. You should use this class when you need to return non-text results or require complete control over the tool behavior. - Using the `SimpleTool` class that extends the base `Tool` class and simplifies the creation of tools that return text results. You should use this approach for scenarios where the tool only needs to return a text. Both approaches use the same core components but differ in implementation and the results they return. For Java: - Using annotation-based methods (`@Tool` and `@LLMDescription`) with reflection-based registration. This is the recommended approach for Java interoperability, as subclassing Kotlin's `Tool` or `SimpleTool` from Java is not supported due to suspend function limitations. ### Tool class (Kotlin) The [`Tool`](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool/index.html) abstract class is the base class for creating tools in Kotlin. It lets you create tools that accept specific argument types (`Args`) and return results of various types (`Result`). Each tool consists of the following components: | Component | Description | | ------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `Args` | The serializable data class that defines arguments required for the tool. | | `Result` | The serializable type of result that the tool returns. If you want to present tool results in a custom format, please inherit [ToolResult.TextSerializable](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool-result/-text-serializable/index.html) class and implement `textForLLM(): String` method | | `argsSerializer` | The overridden variable that defines how the arguments for the tool are deserialized. See also [argsSerializer](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool/args-serializer.html). | | `resultSerializer` | The overridden variable that defines how the result of the tool is deserialized. See also [resultSerializer](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool/result-serializer.html). If you chose to inherit [ToolResult.TextSerializable](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool-result/-text-serializable/index.html) consider using `ToolResultUtils.toTextSerializer()` | | `descriptor` | The overridden variable that specifies tool metadata: - `name` - `description` - `requiredParameters` (empty by default) - `optionalParameters` (empty by default) See also [descriptor](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool/descriptor.html). | | `execute()` | The function that implements the logic of the tool. It takes arguments of type `Args` and returns a result of type `Result`. See also execute(). | Java Implementation In Java, instead of subclassing `Tool`, use annotation-based methods with `@Tool` and `@LLMDescription`. The framework handles serialization and registration automatically through reflection. For more details, see [Annotation-based methods](#annotation-based-methods-java) below. Tip Ensure your tools have clear descriptions and well-defined parameter names to make it easier for the LLM to understand and use them properly. In Kotlin, use the `descriptor` property; in Java, use `@LLMDescription` annotations. #### Usage example Here is an example of a custom tool implementation using the `Tool` class that returns a numeric result: ``` // Implement a simple calculator tool that adds two digits object CalculatorTool : Tool( argsType = typeToken(), resultType = typeToken(), name = "calculator", description = "A simple calculator that can add two digits (0-9)." ) { // Arguments for the calculator tool @Serializable data class Args( @property:LLMDescription("The first digit to add (0-9)") val digit1: Int, @property:LLMDescription("The second digit to add (0-9)") val digit2: Int ) { init { require(digit1 in 0..9) { "digit1 must be a single digit (0-9)" } require(digit2 in 0..9) { "digit2 must be a single digit (0-9)" } } } // Function to add two digits override suspend fun execute(args: Args): Int = args.digit1 + args.digit2 } ``` After implementing your tool, you need to add it to a tool registry and then use it with an agent. For details, see [Tool registry](../#tool-registry). For more details, see [API reference](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool/index.html). #### Reading the agent context from a tool Tools that need the agent's full state (LLM context, run id, configuration, storage, ...) extend `AgentContextAwareTool` instead of `Tool`. The framework injects the live `AIAgentContext` driving the call, and the tool receives it as a typed parameter rather than reading it out of the argument schema. ``` // A tool that reads the live AIAgentContext driving the call. object TracingCalculatorTool : AgentContextAwareTool( argsType = typeToken(), resultType = typeToken(), name = "tracing_calculator", description = "Adds two digits and emits a log line tagged with the agent run id." ) { @Serializable data class Args( @property:LLMDescription("The first digit to add (0-9)") val digit1: Int, @property:LLMDescription("The second digit to add (0-9)") val digit2: Int ) override suspend fun execute(args: Args, context: AIAgentContext): Int { val runId = context.runId // ... use runId for cross-cutting context (logging, tracing, correlation) return args.digit1 + args.digit2 } } ``` `AgentContextAwareTool` is dispatched by the framework via a per-call `ToolCallMetadata` side channel that the framework manages on the tool's behalf. Invoking such a tool outside an agent run throws `IllegalStateException` because no `AIAgentContext` was injected; production code should always go through `ContextualAgentEnvironment`, and unit tests can supply the context explicitly via `ToolCallMetadata.of(AgentContextAwareTool.AgentContextKey to context)`. #### Reading raw per-call metadata A small number of tools want to read caller- or feature-contributed entries that are *not* the agent context (for example a distributed-tracing span id contributed by an observability feature). These tools extend `ToolBase` directly, which exposes the full `ToolCallMetadata` bag: ``` object SpanAwareCalculatorTool : ToolBase( argsType = typeToken(), resultType = typeToken(), name = "span_aware_calculator", description = "Adds two digits, propagating a tracing span id from caller or feature metadata." ) { @Serializable data class Args( @property:LLMDescription("The first digit to add (0-9)") val digit1: Int, @property:LLMDescription("The second digit to add (0-9)") val digit2: Int ) override suspend fun execute(args: Args, metadata: ToolCallMetadata): Int { val traceSpanId = metadata["trace.span.id"] as? String // ... use traceSpanId for cross-cutting context (logging, tracing, correlation) return args.digit1 + args.digit2 } } ``` Callers can pass metadata through `SafeTool.execute(args, serializer, metadata)` or directly through `AIAgentEnvironment.executeTool(toolCall, metadata)`. Features can contribute metadata for every tool call during installation by calling `pipeline.provideToolCallMetadata(this) { eventContext -> mapOf(...) }`. Caller-supplied metadata always wins over feature contributions on key collision. Existing tools that extend `Tool` and override `execute(args)` continue to work unchanged: the framework dispatches them through the same path and discards any `ToolCallMetadata`. To opt in to metadata, switch to `AgentContextAwareTool` (typed context access) or `ToolBase` (raw bag access). ### SimpleTool class (Kotlin) The [`SimpleTool`](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-simple-tool/index.html) abstract class extends `Tool` and simplifies the creation of tools that return text results. Each simple tool consists of the following components: | Component | Description | | ---------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `Args` | The serializable data class that defines arguments required for the custom tool. | | `argsSerializer` | The overridden variable that defines how the arguments for the tool are serialized. See also [argsSerializer](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool/args-serializer.html). | | `descriptor` | The overridden variable that specifies tool metadata: - `name` - `description` - `requiredParameters` (empty by default) - `optionalParameters` (empty by default) See also [descriptor](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-tool/descriptor.html). | | `doExecute()` | The overridden function that describes the main action performed by the tool. It takes arguments of type `Args` and returns a `String`. See also [doExecute()](https://api.koog.ai/agents/agents-tools/ai.koog.agents.core.tools/-simple-tool/do-execute.html). | Java Implementation In Java, the equivalent approach is to use annotation-based methods that return `String`. The framework automatically handles the text result wrapping. For more details, see [Annotation-based methods](#annotation-based-methods-java) below. Tip Ensure your tools have clear descriptions and well-defined parameter names to make it easier for the LLM to understand and use them properly. In Kotlin, use the `descriptor` and constructor parameters; in Java, use `@Tool` and `@LLMDescription` annotations. #### Usage example Here is an example of a custom tool implementation using `SimpleTool` in Kotlin: ``` // Create a tool that casts a string expression to a double value object CastToDoubleTool : SimpleTool( argsType = typeToken(), name = "cast_to_double", description = "casts the passed expression to double or returns 0.0 if the expression is not castable" ) { // Define tool arguments @Serializable data class Args( @property:LLMDescription("An expression to case to double") val expression: String, @property:LLMDescription("A comment on how to process the expression") val comment: String ) // Function that executes the tool with the provided arguments override suspend fun execute(args: Args): String { return "Result: ${castToDouble(args.expression)}, " + "the comment was: ${args.comment}" } // Function to cast a string expression to a double value private fun castToDouble(expression: String): Double { return expression.toDoubleOrNull() ?: 0.0 } } ``` ### Annotation-based methods (Java) To implement tools in Java, instead of subclassing `Tool` or `SimpleTool`, use annotation-based methods with `@Tool` and `@LLMDescription`. Koog handles serialization and registration automatically through reflection. To learn more about the implementation, see Java examples below. #### Usage examples This is an example of a tool implementation in Java, equivalent to using the `Tool` class in Kotlin. ``` // Java equivalent: implement the tool as a Java method and register it via ToolRegistry.builder(). // This is the recommended Java interop path instead of subclassing the Kotlin Tool base class. public final class CalculatorTool { private CalculatorTool() {} @Tool(customName = "calculator") @LLMDescription(description = "A simple calculator that can add two digits (0-9).") public static int calculator( @LLMDescription(description = "The first digit to add (0-9)") int digit1, @LLMDescription(description = "The second digit to add (0-9)") int digit2 ) { if (digit1 < 0 || digit1 > 9) throw new IllegalArgumentException("digit1 must be a single digit (0-9)"); if (digit2 < 0 || digit2 > 9) throw new IllegalArgumentException("digit2 must be a single digit (0-9)"); return digit1 + digit2; } public static ToolRegistry registry() throws NoSuchMethodException { return ToolRegistry.builder() .tool(CalculatorTool.class.getMethod("calculator", int.class, int.class)) .build(); } } // Note: Subclassing the Kotlin Tool and overriding a suspend execute(...) from Java is not supported. // The Java interop uses reflection-based registration of Java methods as tools. ``` Here is an example of a tool implementation in Java, equivalent to using the `SimpleTool` class in Kotlin. This example implements a simple tool that returns a text result. ``` // Java equivalent of SimpleTool: provide a Java method and register it as a tool. public final class CastToDoubleTool { private CastToDoubleTool() {} @Tool(customName = "cast_to_double") @LLMDescription(description = "casts the passed expression to double or returns 0.0 if the expression is not castable") public static String castToDouble( @LLMDescription(description = "An expression to case to double") String expression, @LLMDescription(description = "A comment on how to process the expression") String comment ) { double value; try { value = Double.parseDouble(expression); } catch (Exception e) { value = 0.0; } return "Result: " + value + ", the comment was: " + comment; } public static ToolRegistry registry() throws NoSuchMethodException { return ToolRegistry.builder() .tool(CastToDoubleTool.class.getMethod("castToDouble", String.class, String.class)) .build(); } } // Note: Extending Kotlin SimpleTool from Java is not required; registering a Java method is the idiomatic approach. ``` ### Sending tool result to LLM in custom format For Kotlin: If you are not happy with JSON results sent to LLM (in some cases, LLMs can work better if tool output is structured as Markdown, for instance), you have to follow the following steps: 1. Implement `ToolResult.TextSerializable` interface, and override `textForLLM()` method 1. Override `resultSerializer` using `ToolResultUtils.toTextSerializer()` For Java: Return formatted text (such as Markdown) directly as a `String` from your annotated method. The framework handles this automatically. #### Example Here is an example showing custom formatted output in both Kotlin and Java: ``` // A tool that edits file object EditFile : Tool( argsType = typeToken(), resultType = typeToken(), name = "edit_file", description = "Edits the given file" ) { // Define tool arguments @Serializable public data class Args( val path: String, val original: String, val replacement: String ) @Serializable public data class Result( private val patchApplyResult: PatchApplyResult ) { @Serializable public sealed interface PatchApplyResult { @Serializable public data class Success(val updatedContent: String) : PatchApplyResult @Serializable public sealed class Failure(public val reason: String) : PatchApplyResult } // Textual output (in Markdown format) that will be visible to the LLM after the tool finishes. fun textForLLM(): String = markdown { if (patchApplyResult is PatchApplyResult.Success) { line { bold("Successfully").text(" edited file (patch applied)") } } else { line { text("File was ") .bold("not") .text(" modified (patch application failed: ${(patchApplyResult as PatchApplyResult.Failure).reason})") } } } override fun toString(): String = textForLLM() } // Function that executes the tool with the provided arguments override suspend fun execute(args: Args): Result { return TODO("Implement file edit") } } ``` ``` import ai.koog.agents.core.tools.ToolRegistry; import ai.koog.agents.core.tools.annotations.LLMDescription; import ai.koog.agents.core.tools.annotations.Tool; // Java equivalent: return Markdown text directly to the LLM from a Java method and register it as a tool. // This avoids needing a custom serializable Result type (which would require Kotlin serialization support). public final class EditFile { private EditFile() {} @Tool(customName = "edit_file") @LLMDescription(description = "Edits the given file") public static String editFile( String path, String original, String replacement ) { // TODO: Implement file edit logic; below is a placeholder illustrating Markdown output boolean success = false; if (success) { return "**Successfully** edited file (patch applied)"; } else { return "File was **not** modified (patch application failed: reason)"; } } public static ToolRegistry registry() throws NoSuchMethodException { return ToolRegistry.builder() .tool(EditFile.class.getMethod("editFile", String.class, String.class, String.class)) .build(); } } // Note: If you need a structured custom Result object from Java, you must expose a Kotlin @Serializable type // or another serializer-aware type. Returning String works out-of-the-box with Koog's Java interop. ``` After implementing your tool in Kotlin or Java, you need to add it to a tool registry and then use it with an agent. For details, see [Tool registry](../tools/index#tool-registry). # Events # Agent events Agent events are actions or interactions that occur as part of an agent workflow. They include: - Agent lifecycle events - Strategy events - Node execution events - LLM call events - LLM streaming events - Tool execution events Note: Feature events are defined in the agents-core module and live under the package `ai.koog.agents.core.feature.model.events`. Features such as `agents-features-trace`, and `agents-features-event-handler` consume these events to process and forward messages created during agent execution. ## Predefined event types Koog provides predefined event types that can be used in custom message processors. The predefined events can be classified into several categories, depending on the entity they relate to: - [Agent events](#agent-events) - [Strategy events](#strategy-events) - [Node events](#node-events) - [Subgraph events](#subgraph-events) - [LLM call events](#llm-call-events) - [LLM streaming events](#llm-streaming-events) - [Tool execution events](#tool-execution-events) ### Agent events #### AgentStartingEvent Represents the start of an agent run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `agentId` | String | Yes | | The unique identifier of the AI agent. | | `runId` | String | Yes | | The unique identifier of the AI agent run. | #### AgentCompletedEvent Represents the end of an agent run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `agentId` | String | Yes | | The unique identifier of the AI agent. | | `runId` | String | Yes | | The unique identifier of the AI agent run. | | `result` | String | Yes | | The result of the agent run. Can be `null` if there is no result. | #### AgentExecutionFailedEvent Represents the occurrence of an error during an agent run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | --------------------------------------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `agentId` | String | Yes | | The unique identifier of the AI agent. | | `runId` | String | Yes | | The unique identifier of the AI agent run. | | `error` | AIAgentError | Yes | | The specific error that occurred during the agent run. For more information, see [AIAgentError](#aiagenterror). | #### AgentClosingEvent Represents the closure or termination of an agent. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `agentId` | String | Yes | | The unique identifier of the AI agent. | The `AIAgentError` class provides more details about an error that occurred during an agent run. Includes the following fields: | Name | Data type | Required | Default | Description | | ------------ | --------- | -------- | ------- | ---------------------------------------------------------------- | | `message` | String | Yes | | The message that provides more details about the specific error. | | `stackTrace` | String | Yes | | The collection of stack records until the last executed code. | | `cause` | String | No | null | The cause of the error, if available. | The `AgentExecutionInfo` class provides contextual information about the execution path, enabling tracking of nested execution contexts within an agent run. Includes the following fields: | Name | Data type | Required | Default | Description | | ---------- | ------------------ | -------- | ------- | --------------------------------------------------------------------------------------------- | | `parent` | AgentExecutionInfo | No | null | Reference to the parent execution context. If null, this represents the root execution level. | | `partName` | String | Yes | | A string representing the name of the current part or segment of the execution. | ### Strategy events #### GraphStrategyStartingEvent Represents the start of a graph-based strategy run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `strategyName` | String | Yes | | The name of the strategy. | | `graph` | StrategyEventGraph | Yes | | The graph structure representing the strategy workflow. | #### FunctionalStrategyStartingEvent Represents the start of a functional strategy run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `strategyName` | String | Yes | | The name of the strategy. | #### StrategyCompletedEvent Represents the end of a strategy run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `strategyName` | String | Yes | | The name of the strategy. | | `result` | String | Yes | | The result of the run. Can be `null` if there is no result. | ### Node events #### NodeExecutionStartingEvent Represents the start of a node run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `nodeName` | String | Yes | | The name of the node whose run started. | | `input` | JsonElement | No | null | The input value for the node. | #### NodeExecutionCompletedEvent Represents the end of a node run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `nodeName` | String | Yes | | The name of the node whose run ended. | | `input` | JsonElement | No | null | The input value for the node. | | `output` | JsonElement | No | null | The output value produced by the node. | #### NodeExecutionFailedEvent Represents an error that occurred during a node run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | -------------------------------------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `nodeName` | String | Yes | | The name of the node where the error occurred. | | `input` | JsonElement | No | null | The input data provided to the node. | | `error` | AIAgentError | Yes | | The specific error that occurred during the node run. For more information, see [AIAgentError](#aiagenterror). | ### Subgraph events #### SubgraphExecutionStartingEvent Represents the start of a subgraph run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `subgraphName` | String | Yes | | The name of the subgraph whose run started. | | `input` | JsonElement | No | null | The input value for the subgraph. | #### SubgraphExecutionCompletedEvent Represents the end of a subgraph run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `subgraphName` | String | Yes | | The name of the subgraph whose run ended. | | `input` | JsonElement | No | null | The input value for the subgraph. | | `output` | JsonElement | No | null | The output value produced by the subgraph. | #### SubgraphExecutionFailedEvent Represents an error that occurred during a subgraph run. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------------------------------------------ | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy run. | | `subgraphName` | String | Yes | | The name of the subgraph where the error occurred. | | `input` | JsonElement | No | null | The input data provided to the subgraph. | | `error` | AIAgentError | Yes | | The specific error that occurred during the subgraph run. For more information, see [AIAgentError](#aiagenterror). | ### LLM call events #### LLMCallStartingEvent Represents the start of an LLM call. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ---------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the LLM run. | | `prompt` | Prompt | Yes | | The prompt that is sent to the model. For more information, see [Prompt](#prompt). | | `model` | ModelInfo | Yes | | The model information. See [ModelInfo](#modelinfo). | | `tools` | List | Yes | | The list of tools that the model can call. | The `Prompt` class represents a data structure for a prompt, consisting of a list of messages, a unique identifier, and optional parameters for language model settings. Includes the following fields: | Name | Data type | Required | Default | Description | | ---------- | --------- | -------- | ----------- | ------------------------------------------------------------ | | `messages` | List | Yes | | The list of messages that the prompt consists of. | | `id` | String | Yes | | The unique identifier for the prompt. | | `params` | LLMParams | No | LLMParams() | The settings that control the way the LLM generates content. | The `ModelInfo` class represents information about a language model, including its provider, model identifier, and characteristics. Includes the following fields: | Name | Data type | Required | Default | Description | | ----------------- | --------- | -------- | ------- | ---------------------------------------------------------------- | | `provider` | String | Yes | | The provider identifier (e.g., "openai", "google", "anthropic"). | | `model` | String | Yes | | The model identifier (e.g., "gpt-4", "claude-3"). | | `displayName` | String | No | null | Optional human-readable display name for the model. | | `contextLength` | Long | No | null | Maximum number of tokens the model can process. | | `maxOutputTokens` | Long | No | null | Maximum number of tokens the model can generate. | #### LLMCallCompletedEvent Represents the end of an LLM call. Includes the following fields: | Name | Data type | Required | Default | Description | | -------------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the LLM run. | | `prompt` | Prompt | Yes | | The prompt used in the call. | | `model` | ModelInfo | Yes | | The model information. See [ModelInfo](#modelinfo). | | `responses` | List | Yes | | One or more responses returned by the model. | | `moderationResponse` | ModerationResult | No | null | The moderation response, if any. | #### LLMCallFailedEvent Represents the occurrence of an error during an LLM call. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ---------------------------------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the LLM run. | | `prompt` | Prompt | Yes | | The prompt that was sent to the model. | | `model` | ModelInfo | Yes | | The model information. See [ModelInfo](#modelinfo). | | `tools` | List | Yes | | The list of tools that the model could call. | | `error` | AIAgentError | Yes | | The specific error that occurred during the call. For more information, see [AIAgentError](#aiagenterror). | ### LLM streaming events #### LLMStreamingStartingEvent Represents the start of an LLM streaming call. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the LLM run. | | `prompt` | Prompt | Yes | | The prompt that is sent to the model. | | `model` | ModelInfo | Yes | | The model information. See [ModelInfo](#modelinfo). | | `tools` | List | Yes | | The list of tools that the model can call. | #### LLMStreamingFrameReceivedEvent Represents a streaming frame received from the LLM. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the LLM run. | | `prompt` | Prompt | Yes | | The prompt that is sent to the model. | | `model` | ModelInfo | Yes | | The model information. See [ModelInfo](#modelinfo). | | `frame` | StreamFrame | Yes | | The frame received from the stream. | #### LLMStreamingFailedEvent Represents the occurrence of an error during an LLM streaming call. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ----------------------------------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the LLM run. | | `prompt` | Prompt | Yes | | The prompt that is sent to the model. | | `model` | ModelInfo | Yes | | The model information. See [ModelInfo](#modelinfo). | | `error` | AIAgentError | Yes | | The specific error that occurred during streaming. For more information, see [AIAgentError](#aiagenterror). | #### LLMStreamingCompletedEvent Represents the end of an LLM streaming call. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the LLM run. | | `prompt` | Prompt | Yes | | The prompt that is sent to the model. | | `model` | ModelInfo | Yes | | The model information. See [ModelInfo](#modelinfo). | | `tools` | List | Yes | | The list of tools that the model can call. | ### Tool execution events #### ToolCallStartingEvent Represents the event of a model calling a tool. Includes the following fields: | Name | Data type | Required | Default | Description | | --------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy/agent run. | | `toolCallId` | String | No | null | The identifier of the tool call, if available. | | `toolName` | String | Yes | | The name of the tool. | | `toolArgs` | JsonObject | Yes | | The arguments that are provided to the tool. | #### ToolValidationFailedEvent Represents the occurrence of a validation error during a tool call. Includes the following fields: | Name | Data type | Required | Default | Description | | ----------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------------------ | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy/agent run. | | `toolCallId` | String | No | null | The identifier of the tool call, if available. | | `toolName` | String | Yes | | The name of the tool for which validation failed. | | `toolArgs` | JsonObject | Yes | | The arguments that are provided to the tool. | | `toolDescription` | String | No | null | A description of the tool that encountered the validation error. | | `message` | String | No | null | A message describing the validation error. | | `error` | AIAgentError | Yes | | The specific error that occurred. For more information, see [AIAgentError](#aiagenterror). | #### ToolCallFailedEvent Represents a failure to execute a tool. Includes the following fields: | Name | Data type | Required | Default | Description | | ----------------- | ------------------ | -------- | ------- | --------------------------------------------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the strategy/agent run. | | `toolCallId` | String | No | null | The identifier of the tool call, if available. | | `toolName` | String | Yes | | The name of the tool. | | `toolArgs` | JsonObject | Yes | | The arguments that are provided to the tool. | | `toolDescription` | String | No | null | A description of the tool that failed. | | `error` | AIAgentError | Yes | | The specific error that occurred when trying to call a tool. For more information, see [AIAgentError](#aiagenterror). | #### ToolCallCompletedEvent Represents a successful tool call with the return of a result. Includes the following fields: | Name | Data type | Required | Default | Description | | ----------------- | ------------------ | -------- | ------- | ------------------------------------------------------------------------------- | | `eventId` | String | Yes | | A unique identifier for the event or a group of events. | | `executionInfo` | AgentExecutionInfo | Yes | | Provides contextual information about the execution associated with this event. | | `runId` | String | Yes | | The unique identifier of the run. | | `toolCallId` | String | No | null | The identifier of the tool call. | | `toolName` | String | Yes | | The name of the tool. | | `toolArgs` | JsonObject | Yes | | The arguments provided to the tool. | | `toolDescription` | String | No | null | A description of the tool that was executed. | | `result` | JsonElement | No | null | The result of the tool call. | ## FAQ and troubleshooting The following section includes commonly asked questions and answers related to the Tracing feature. ### How do I trace only specific parts of my agent's execution? Use the `messageFilter` property to filter events. For example, to trace only node execution: ``` install(Tracing) { val fileWriter = TraceFeatureMessageFileWriter.create(outputPath) addMessageProcessor(fileWriter) // Only trace LLM calls fileWriter.setMessageFilter { message -> message is LLMCallStartingEvent || message is LLMCallCompletedEvent } } ``` ``` .install(Tracing.Feature, config -> { var fileWriter = TraceFeatureMessageFileWriter.create(outputPath); config.addMessageProcessor(fileWriter); // Only trace LLM calls fileWriter.setMessageFilter(message -> message instanceof LLMCallStartingEvent || message instanceof LLMCallCompletedEvent ); }) ``` ### Can I use multiple message processors? Yes, you can add multiple message processors to trace to different destinations simultaneously: ``` install(Tracing) { addMessageProcessor(TraceFeatureMessageLogWriter(logger)) addMessageProcessor(TraceFeatureMessageFileWriter.create(outputPath)) addMessageProcessor(TraceFeatureMessageRemoteWriter(connectionConfig)) } ``` ``` .install(Tracing.Feature, config -> { config.addMessageProcessor(TraceFeatureMessageLogWriter.create(logger)); config.addMessageProcessor(TraceFeatureMessageFileWriter.create(outputPath)); config.addMessageProcessor(new TraceFeatureMessageRemoteWriter()); }) ``` ### How can I create a custom message processor? Implement the `FeatureMessageProcessor` interface: ``` class CustomTraceProcessor : FeatureMessageProcessor() { override suspend fun processMessage(message: FeatureMessage) { // Custom processing logic if (message is NodeExecutionStartingEvent) { // Process node start event } else if (message is LLMCallCompletedEvent) { // Process LLM call end event } else { // Handle other event types } } override suspend fun close() { // Close connections if established } } val agent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, ) { install(Tracing) { // Use your custom processor addMessageProcessor(CustomTraceProcessor()) } } ``` ``` class CustomTraceProcessor extends FeatureMessageProcessor { @Override protected void handleMessage(FeatureMessage message) { // Custom processing logic if (message instanceof NodeExecutionStartingEvent) { // Process node start event } else if (message instanceof LLMCallCompletedEvent) { // Process LLM call end event } else { // Handle other event types } } @Override public void handleClose() { // Close connections if established } } var agent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().ollama().build()) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Tracing.Feature, config -> { // Use your custom processor config.addMessageProcessor(new CustomTraceProcessor()); }) .build(); ``` For more information about existing event types that can be handled by message processors, see [Predefined event types](#predefined-event-types). # Strategies # Predefined nodes and components Nodes are the fundamental building blocks of agent workflows in the Koog framework. Each node represents a specific operation or transformation in the workflow, and they can be connected using edges to define the flow of execution. In general, nodes let you encapsulate complex logic into reusable components that can be easily integrated into different agent workflows. This guide will walk you through the existing nodes that can be used in your agent strategies. Each node is essentially a function (Kotlin) or action (Java) that takes an input of a specific type and returns an output of a specific type. ``` graph LR in:::hidden out:::hidden subgraph node ["Node"] execute(Do stuff) end in --Input--> execute --Output--> out classDef hidden display: none; ``` Here is how you can define a node that expects a string as input and returns the length of the string (an integer) as output: ``` val nodeLength by node { input -> input.length } ``` ``` var nodeLength = AIAgentNode.builder("nodeLength") .withInput(String.class) .withOutput(Integer.class) .withAction((input, ctx) -> input.length()) .build(); ``` For more information, see [node()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.builder/node.html) (Kotlin) or [AIAgentNode.builder()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/builder.html) for Java. ## Utility nodes ### Pass-through node A simple pass-through node that does nothing and returns the input as output. For details, see [nodeDoNothing](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-do-nothing.html) (Kotlin) or [AIAgentNode.doNothing()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/do-nothing.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Pass-through node"] execute(Do nothing) end in ---|T| execute --T--> out classDef hidden display: none; ``` You can use this node for the following purposes: - Create a placeholder node in your graph. - Create a connection point without modifying the data. Here is an example: ``` val passthrough by nodeDoNothing("passthrough") edge(nodeStart forwardTo passthrough) edge(passthrough forwardTo nodeFinish) ``` ``` var passthrough = AIAgentNode.builder("passthrough") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> input) .build(); strategy.edge(strategy.nodeStart, passthrough); strategy.edge(passthrough, strategy.nodeFinish); ``` ## LLM nodes ### Prompt preparation node **A node that adds messages to the LLM prompt using the provided prompt builder. This is useful for modifying the conversation context before making an actual LLM request.** For details, see [nodeAppendPrompt](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-append-prompt.html) (Kotlin) or [AIAgentNode.appendPrompt()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node-builder-with-input/append-prompt.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Prompt preparation node"] execute(Append prompt) end in ---|T| execute --T--> out classDef hidden display: none; ``` You can use this node for the following purposes: - Add system instructions to the prompt. - Insert user messages into the conversation. - Prepare the context for subsequent LLM requests. Here is an example: ``` val firstNode by node { // Transform input to output } val secondNode by node { // Transform output to output } // Node will get the value of type Output as input from the previous node and path through it to the next node val setupContext by nodeAppendPrompt("setupContext") { system("You are a helpful assistant specialized in Kotlin programming.") user("I need help with Kotlin coroutines.") } edge(firstNode forwardTo setupContext) edge(setupContext forwardTo secondNode) ``` ``` var firstNode = AIAgentNode.builder() .withInput(Input.class) .withOutput(Output.class) .withAction((input, ctx) -> { // Transform input to output return input; }) .build(); var secondNode = AIAgentNode.builder() .withInput(Output.class) .withOutput(Output.class) .withAction((output, ctx) -> { // Transform output to output return output; }) .build(); var setupContext = AIAgentNode.builder() .withInput(Output.class) .appendPrompt(prompt -> { prompt.system("You are a helpful assistant specialized in Kotlin programming."); prompt.user("I need help with Kotlin coroutines."); }); strategy.edge(firstNode, setupContext); strategy.edge(setupContext, secondNode); ``` ### Tool-only node A node that appends a user message to the LLM prompt and gets a response where the LLM can only call tools. For details, see [nodeLLMSendMessageOnlyCallingTools](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-send-message-only-calling-tools.html) (Kotlin) or [AIAgentNode.llmSendMessageOnlyCallingTools()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-send-message-only-calling-tools.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Tool-only node"] execute(Request LLM expecting only tool calls) end in --String--> execute --Message.Response--> out classDef hidden display: none; ``` ### Forced single tool use node A node that that appends a user message to the LLM prompt and forces the LLM to use a specific tool. For details, see [nodeLLMSendMessageForceOneTool](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-send-message-force-one-tool.html) (Kotlin) or [AIAgentNode.llmSendMessageForceOneTool()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-send-message-force-one-tool.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Forced single tool use node"] execute(Request LLM expecting a specific tool call) end in --String--> execute --Message.Response--> out classDef hidden display: none; ``` ### LLM request node A node that appends a user message to the LLM prompt and gets a response with optional tool usage. The node configuration determines whether tool calls are allowed during the processing of the message. For details, see [nodeLLMRequest](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-request.html) (Kotlin) or [AIAgentNode.llmRequest()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-request.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["LLM request node"] execute(Request LLM) end in --String--> execute --Message.Response--> out classDef hidden display: none; ``` You can use this node for the following purposes: - Generate LLM response for the current prompt, controlling if the LLM is allowed to generate tool calls. Here is an example: ``` val requestLLM by nodeLLMRequest("requestLLM") edge(getUserQuestion forwardTo requestLLM) ``` ``` var requestLLM = AIAgentNode.llmRequest("requestLLM"); strategy.edge(AIAgentEdge.builder() .from(getUserQuestion) .to(requestLLM) .build()); ``` ### LLM request node with structured response A node that appends a user message to the LLM prompt and requests structured data from the LLM with error correction capabilities. For details, see [nodeLLMRequestStructured](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-request-structured.html) (Kotlin) or [AIAgentNode.llmRequestStructured()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-request-structured.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["LLM request node, structured response"] execute(Request LLM structured) end in --String--> execute -- "Result<StructuredResponse>" --> out classDef hidden display: none; ``` ### LLM request node with streaming response A node that appends a user message to the LLM prompt and streams LLM response with or without stream data transformation. For details, see [nodeLLMRequestStreaming](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-request-streaming.html) (Kotlin) or [AIAgentNode.llmRequestStreaming()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-request-streaming.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["LLM request node, streaming response"] execute(Request LLM streaming) end in --String--> execute --Flow--> out classDef hidden display: none; ``` ### LLM request node with multiple responses A node that appends a user message to the LLM prompt and gets multiple LLM responses with tool calls enabled. For details, see [nodeLLMRequest](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-request.html) (Kotlin) or [AIAgentNode.llmRequest()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-request.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["LLM request node, multiple responses"] execute(Request LLM expecting multiple responses) end in --String--> execute -- "List<Message.Response>" --> out classDef hidden display: none; ``` You can use this node for the following purposes: - Handle complex queries that require multiple tool calls. - Generate multiple tool calls. - Implement a workflow that requires multiple parallel actions. Here is an example: ``` val requestLLMMultipleTools by nodeLLMRequest() edge(getComplexUserQuestion forwardTo requestLLMMultipleTools) ``` ``` var requestLLMMultipleTools = AIAgentNode.llmRequest("requestLLMMultipleTools"); strategy.edge(AIAgentEdge.builder() .from(getComplexUserQuestion) .to(requestLLMMultipleTools) .build()); ``` ### History compression node A node that compresses the current LLM prompt (message history) into a summary, replacing messages with a concise summary (TL;DR). This is useful for managing long conversations by compressing the history to reduce token usage. For details, see [nodeLLMCompressHistory](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-compress-history.html) (Kotlin) or [AIAgentNode.llmCompressHistory()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-compress-history.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["History compression node"] execute(Compress current prompt) end in ---|T| execute --T--> out classDef hidden display: none; ``` To learn more about history compression, see [History compression](../history-compression/). You can use this node for the following purposes: - Manage long conversations to reduce token usage. - Summarize conversation history to maintain context. - Implement memory management in long-running agents. Here is an example: ``` val compressHistory by nodeLLMCompressHistory( "compressHistory", strategy = HistoryCompressionStrategy.FromLastNMessages(10), preserveMemory = true ) edge(generateHugeHistory forwardTo compressHistory) ``` ``` var compressHistory = AIAgentNode.llmCompressHistory("compressHistory") .withInput(String.class) .build(); strategy.edge(generateHugeHistory, compressHistory); ``` ## Tool nodes ### Tool execution node A node that executes a single tool call and returns its result. This node is used to handle tool calls made by the LLM. For details, see [nodeExecuteTool](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-execute-tool.html) (Kotlin) or [AIAgentNode.executeTool()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/execute-tool.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Tool execution node"] execute(Execute tool call) end in --MessagePart.Tool.Call--> execute --ReceivedToolResult--> out classDef hidden display: none; ``` You can use this node for the following purposes: - Execute tools requested by the LLM. - Handle specific actions in response to LLM decisions. - Integrate external functionality into the agent workflow. Here is an example: ``` val requestLLM by nodeLLMRequest() val executeTool by nodeExecuteTools() edge(requestLLM forwardTo executeTool onToolCalls { true }) ``` ``` var requestLLM = AIAgentNode.llmRequest("requestLLM"); var executeTool = AIAgentNode.executeTools("executeTool"); strategy.edge(AIAgentEdge.builder() .from(requestLLM) .to(executeTool) .onToolCalls() .build()); ``` ### Tool result follow-up node A node that adds a tool result to the prompt and requests an LLM response. For details, see [nodeLLMSendToolResult](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-send-tool-result.html) (Kotlin) or [AIAgentNode.llmSendToolResult()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-send-tool-result.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Tool result follow-up node"] execute(Request LLM) end in --ReceivedToolResult--> execute --Message.Response--> out classDef hidden display: none; ``` You can use this node for the following purposes: - Process the results of tool executions. - Generate responses based on tool outputs. - Continue a conversation after tool execution. Here is an example: ``` val executeTool by nodeExecuteTools() val sendToolResultToLLM by nodeLLMSendToolResults() edge(executeTool forwardTo sendToolResultToLLM) ``` ``` var executeTool = AIAgentNode.executeTools("executeTool"); var sendToolResultToLLM = AIAgentNode.llmSendToolResults("sendToolResultToLLM"); strategy.edge(executeTool, sendToolResultToLLM); ``` ### Multi-tool execution node A node that executes multiple tool calls. These calls can optionally be executed in parallel. For details, see [nodeExecuteMultipleTools](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-execute-multiple-tools.html) (Kotlin) or [AIAgentNode.executeMultipleTools()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/execute-multiple-tools.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Multi-tool execution node"] execute(Execute multiple tool calls) end in -- "List<MessagePart.Tool.Call>" --> execute -- "List<ReceivedToolResult>" --> out classDef hidden display: none; ``` You can use this node for the following purposes: - Execute multiple tools in parallel. - Handle complex workflows that require multiple tool executions. - Optimize performance by batching tool calls. Here is an example: ``` val requestLLMMultipleTools by nodeLLMRequest() val executeMultipleTools by nodeExecuteTools(parallel = true) edge(requestLLMMultipleTools forwardTo executeMultipleTools onToolCalls { true }) ``` ``` var requestLLMMultipleTools = AIAgentNode.llmRequest("requestLLMMultipleTools"); var executeMultipleTools = AIAgentNode.executeTools("executeMultipleTools"); // Route tool calls from the assistant response to the tool-execution node strategy.edge(AIAgentEdge.builder() .from(requestLLMMultipleTools) .to(executeMultipleTools) .onToolCalls() .build()); ``` ### Multiple tool result follow-up node A node that adds multiple tool results to the prompt and gets multiple LLM responses. For details, see [nodeLLMSendMultipleToolResults](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.extension/node-l-l-m-send-multiple-tool-results.html) (Kotlin) or [AIAgentNode.llmSendMultipleToolResults()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-node/-companion/llm-send-multiple-tool-results.html) (Java). ``` graph LR in:::hidden out:::hidden subgraph node ["Multiple tool result follow-up node"] execute(Request LLM expecting multiple responses) end in -- "List<ReceivedToolResult>" --> execute -- "List<Message.Response>" --> out classDef hidden display: none; ``` You can use this node for the following purposes: - Process the results of multiple tool executions. - Generate multiple tool calls. - Implement complex workflows with multiple parallel actions. Here is an example: ``` val executeTools by nodeExecuteTools(parallel = true) val sendToolResultsToLLM by nodeLLMSendToolResults() edge(executeTools forwardTo sendToolResultsToLLM) ``` ``` var executeTools = AIAgentNode.executeTools("executeTools"); var sendToolResultsToLLM = AIAgentNode.llmSendToolResults("sendToolResultsToLLM"); strategy.edge(executeTools, sendToolResultsToLLM); ``` ## Node output transformation The framework provides the `transform` extension function in Kotlin that allows you to create transformed versions of nodes that apply transformations to their output. In Java, you achieve the same result by creating intermediate nodes with explicit transformations. This is useful when you need to convert the output of a node to a different type or format while preserving the original node's functionality. ``` graph LR in:::hidden out:::hidden subgraph nodeWithTransform [transformed node] subgraph node ["node"] execute(Do stuff) end transform end in --Input--> execute --> transform --Output--> out classDef hidden display: none; ``` ### Node transformation In Kotlin, the [transform()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.builder/-a-i-agent-node-delegate/transform.html) function creates a new `AIAgentNodeDelegate` that wraps the original node and applies a transformation function to its output. In Java, you need to manually compose nodes with transformation logic using `AIAgentNode.builder()` and explicit type parameters. ``` inline fun AIAgentNodeDelegate.transform( noinline transformation: suspend (Output) -> T ): AIAgentNodeDelegate ``` ``` // In Java, you need to manually compose nodes // with transformation logic using AIAgentNode.builder() and explicit type parameters. // See the examples below for the Java approach to node transformations. ``` #### Custom node transformation Transform the output of a custom node to a different data type: ``` val textNode by nodeDoNothing("textNode").transform { text -> text.split(" ").filter { it.isNotBlank() }.size } edge(nodeStart forwardTo textNode) edge(textNode forwardTo nodeFinish) ``` ``` var textNode = AIAgentNode.builder("textNode") .withInput(String.class) .withOutput(Integer.class) .withAction((text, ctx) -> { String[] words = text.split(" "); int count = 0; for (String word : words) { if (!word.isBlank()) { count++; } } return count; }) .build(); strategy.edge(strategy.nodeStart, textNode); strategy.edge(textNode, strategy.nodeFinish); ``` #### Built-in node transformation Transform the output of built-in nodes like `nodeLLMRequest` (Kotlin) or `AIAgentNode.llmRequest()` (Java): ``` val lengthNode by nodeLLMRequest("llmRequest").transform { assistantMessage -> assistantMessage.parts.filterIsInstance().joinToString("\n") { it.text }.length } edge(nodeStart forwardTo lengthNode) edge(lengthNode forwardTo nodeFinish) ``` ``` var llmRequest = AIAgentNode.llmRequest("llmRequest"); var lengthNode = AIAgentNode.builder("lengthNode") .withInput(Message.Assistant.class) .withOutput(Integer.class) .withAction((assistantMessage, ctx) -> { String text = assistantMessage.getParts().stream() .filter(p -> p instanceof MessagePart.Text) .map(p -> ((MessagePart.Text) p).getText()) .collect(Collectors.joining()); return text.length(); }) .build(); strategy.edge(AIAgentEdge.builder() .from(strategy.nodeStart) .to(llmRequest) .build()); strategy.edge(llmRequest, lengthNode); strategy.edge(lengthNode, strategy.nodeFinish); ``` ## Predefined subgraphs The framework provides predefined subgraphs that encapsulate commonly used patterns and workflows. These subgraphs simplify the development of complex agent strategies by handling the creation of base nodes and edges automatically. The API is consistent between Kotlin and Java, with Kotlin using DSL functions and Java using builder methods. By using the predefined subgraphs, you can implement various popular pipelines. Here is an example: 1. Prepare the data. 1. Run the task. 1. Validate the task results. If the results are incorrect, return to step 2 with a feedback message to make adjustments. ### Task execution subgraph A subgraph that performs a specific task using provided tools and returns a structured result. It supports multi-response LLM interactions (the assistant may produce several responses interleaved with tool calls) and lets you control how tool calls are executed. In Kotlin, use [subgraphWithTask()](https://api.koog.ai/agents/agents-core/ai.koog.agents.ext.agent/subgraph-with-task.html), and in Java, use [AIAgentSubgraph.builder().withTask()](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-typed-a-i-agent-subgraph-builder/with-task.html). You can use this subgraph for the following purposes: - Create special components that handle specific tasks within a larger workflow. - Encapsulate complex logic with clear input and output interfaces. - Configure task-specific tools, models, and prompts. - Manage conversation history with automatic compression. - Develop structured agent workflows and task execution pipelines. - Generate structured results from LLM task execution, including flows with multiple assistant responses and tool invocations. The API allows you to fine‑tune execution with optional parameters: - [runMode](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-subgraph-with-task-builder/run-mode.html): controls how tool calls are executed during the task (sequential by default). Use this to switch between different tool execution strategies when supported by the underlying model/executor. - [assistantResponseRepeatMax](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-subgraph-with-task-builder/assistant-response-repeat-max.html): limits how many assistant responses are allowed before concluding the task cannot be completed (defaults to a safe internal limit if not provided). You can provide a task to the subgraph as text, configure the LLM if needed, and provide the necessary tools, and the subgraph will process and solve the task. Here is an example: ``` val processQuery by subgraphWithTask( tools = listOf(searchTool, calculatorTool, weatherTool), llmModel = OpenAIModels.Chat.GPT4o, parallelTools = false, assistantResponseRepeatMax = 3, ) { userQuery -> """ You are a helpful assistant that can answer questions about various topics. Please help with the following query: $userQuery """ } ``` ``` var processQuery = AIAgentSubgraph.builder("processQuery") .limitedTools(List.of(searchTool, calculatorTool, weatherTool)) .withInput(String.class) .withOutput(String.class) .withTask(userQuery -> "You are a helpful assistant that can answer questions about various topics.\n" + "Please help with the following query:\n" + userQuery) .parallelTools(false) .assistantResponseRepeatMax(3) .build(); ``` ### Task execution subgraph with verification A special version of `subgraphWithTask` that verifies whether a task was performed correctly and provides details about any issues encountered. This subgraph is useful for workflows that require validation or quality checks. In Kotlin, use [subgraphWithVerification()](https://api.koog.ai/agents/agents-core/ai.koog.agents.ext.agent/subgraph-with-verification.html), and in Java, use `AIAgentSubgraph.builder().withVerification()`. You can use this subgraph for the following purposes: - Verify the correctness of task execution. - Implement quality control processes in your workflows. - Create self-validating components. - Generate structured verification results with success/failure status and detailed feedback. The subgraph ensures that the LLM calls a verification tool at the end of the workflow to check whether the task was successfully completed. It guarantees this verification is performed as the final step and returns a [CriticResult](https://api.koog.ai/agents/agents-core/ai.koog.agents.ext.agent/-critic-result/index.html) that indicates whether a task was completed successfully and provides detailed feedback. Here is an example: ``` val verifyCode by subgraphWithVerification( tools = listOf(runTestsTool, analyzeTool, readFileTool), llmModel = AnthropicModels.Opus_4_6, parallelTools = false, assistantResponseRepeatMax = 3, ) { codeToVerify -> """ You are a code reviewer. Please verify that the following code meets all requirements: 1. It compiles without errors 2. All tests pass 3. It follows the project's coding standards Code to verify: $codeToVerify """ } ``` ``` var verifyCode = AIAgentSubgraph.builder("verifyCode") .limitedTools(List.of(runTestsTool, analyzeTool, readFileTool)) .withInput(String.class) .withVerification(codeToVerify -> "You are a code reviewer. Please verify that the following code meets all requirements:\n" + "1. It compiles without errors\n" + "2. All tests pass\n" + "3. It follows the project's coding standards\n\n" + "Code to verify:\n" + codeToVerify) .parallelTools(false) .assistantResponseRepeatMax(3) .build(); ``` ## Predefined strategies and common strategy patterns Koog provides predefined strategies that combine various nodes. The nodes are connected using edges to define the flow of operations, with conditions that specify when to follow each edge. You can integrate these strategies into your agent workflows if needed. ### Single run strategy A single run strategy is designed for non-interactive use cases where the agent processes input once and returns a result. You can use this strategy when you need to run straightforward processes that do not require complex logic. ``` public fun singleRunStrategy(): AIAgentGraphStrategy = strategy("single_run") { val nodeCallLLM by nodeLLMRequest("sendInput") val nodeExecuteTool by nodeExecuteTools("nodeExecuteTool") val nodeSendToolResult by nodeLLMSendToolResults("nodeSendToolResult") edge(nodeStart forwardTo nodeCallLLM) edge(nodeCallLLM forwardTo nodeExecuteTool onToolCalls { true }) edge(nodeCallLLM forwardTo nodeFinish onTextMessage { true }) edge(nodeExecuteTool forwardTo nodeSendToolResult) edge(nodeSendToolResult forwardTo nodeFinish onTextMessage { true }) edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCalls { true }) } ``` ``` public static AIAgentGraphStrategy singleRunStrategy() { var strategy = AIAgentGraphStrategy.builder("single_run") .withInput(String.class) .withOutput(String.class); var nodeCallLLM = AIAgentNode.llmRequest("sendInput"); var nodeExecuteTool = AIAgentNode.executeTools("nodeExecuteTool"); var nodeSendToolResult = AIAgentNode.llmSendToolResults("nodeSendToolResult"); strategy.edge(AIAgentEdge.builder() .from(strategy.nodeStart) .to(nodeCallLLM) .build()); strategy.edge(AIAgentEdge.builder() .from(nodeCallLLM) .to(nodeExecuteTool) .onToolCalls() .build()); strategy.edge(AIAgentEdge.builder() .from(nodeCallLLM) .to(strategy.nodeFinish) .onTextMessage() .build()); strategy.edge(nodeExecuteTool, nodeSendToolResult); strategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(strategy.nodeFinish) .onTextMessage() .build()); strategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(nodeExecuteTool) .onToolCalls() .build()); return strategy.build(); } ``` ### Tool-based strategy A tool-based strategy is designed for workflows that heavily rely on tools to perform specific operations. It typically executes tools based on the LLM decisions and processes the results. ``` fun toolBasedStrategy(name: String, toolRegistry: ToolRegistry): AIAgentGraphStrategy { return strategy(name) { val nodeSendInput by nodeLLMRequest() val nodeExecuteTool by nodeExecuteTools() val nodeSendToolResult by nodeLLMSendToolResults() // Define the flow of the agent edge(nodeStart forwardTo nodeSendInput) // If the LLM responds with a message, finish edge( (nodeSendInput forwardTo nodeFinish) onTextMessage { true } ) // If the LLM calls a tool, execute it edge( (nodeSendInput forwardTo nodeExecuteTool) onToolCalls { true } ) // Send the tool result back to the LLM edge(nodeExecuteTool forwardTo nodeSendToolResult) // If the LLM calls another tool, execute it edge( (nodeSendToolResult forwardTo nodeExecuteTool) onToolCalls { true } ) // If the LLM responds with a message, finish edge( (nodeSendToolResult forwardTo nodeFinish) onTextMessage { true } ) } } ``` ``` public static AIAgentGraphStrategy toolBasedStrategy(String name, ToolRegistry toolRegistry) { var strategy = AIAgentGraphStrategy.builder(name) .withInput(String.class) .withOutput(String.class); var nodeSendInput = AIAgentNode.llmRequest("nodeSendInput"); var nodeExecuteTool = AIAgentNode.executeTools("nodeExecuteTool"); var nodeSendToolResult = AIAgentNode.llmSendToolResults("nodeSendToolResult"); // Define the flow of the agent strategy.edge(AIAgentEdge.builder() .from(strategy.nodeStart) .to(nodeSendInput) .build()); // If the LLM responds with a message, finish strategy.edge(AIAgentEdge.builder() .from(nodeSendInput) .to(strategy.nodeFinish) .onTextMessage() .build()); // If the LLM calls a tool, execute it strategy.edge(AIAgentEdge.builder() .from(nodeSendInput) .to(nodeExecuteTool) .onToolCalls(call -> true) .build()); // Send the tool result back to the LLM strategy.edge(nodeExecuteTool, nodeSendToolResult); // If the LLM calls another tool, execute it strategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(nodeExecuteTool) .onToolCalls() .build()); // If the LLM responds with a message, finish strategy.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(strategy.nodeFinish) .onTextMessage() .build()); return strategy.build(); } ``` ### Streaming data strategy A streaming data strategy is designed for processing streaming data from the LLM. It typically requests streaming data, processes it, and potentially calls tools with the processed data. ``` val agentStrategy = strategy>("library-assistant") { // Describe the node containing the output stream parsing val getMdOutput by node> { booksDescription -> val books = mutableListOf() val mdDefinition = markdownBookDefinition() llm.writeSession { appendPrompt { user(booksDescription) } // Initiate the response stream in the form of the definition `mdDefinition` val markdownStream = requestLLMStreaming(mdDefinition) // Call the parser with the result of the response stream and perform actions with the result parseMarkdownStreamToBooks(markdownStream).collect { book -> books.add(book) println("Parsed Book: ${book.title} by ${book.author}") } } books } // Describe the agent's graph making sure the node is accessible edge(nodeStart forwardTo getMdOutput) edge(getMdOutput forwardTo nodeFinish) } ``` ``` var strategy = AIAgentGraphStrategy.builder() .withInput(String.class) .withOutput(List.class); var getMdOutput = AIAgentNode.builder() .withInput(String.class) .>withOutput(TypeToken.of(new TypeCapture>() {})) .withAction((booksDescription, ctx) -> { var books = new ArrayList(); StructureDefinition mdDefinition = markdownBookDefinition(); ctx.getLlm().writeSession(session -> { session.appendPrompt(prompt -> { prompt.user(booksDescription); }); // Initiate the response stream in the form of the definition `mdDefinition` var markdownStream = session.requestLLMStreaming(mdDefinition); // Call the parser with the result of the response stream and perform actions with the result parseMarkdownStreamToBooks(markdownStream).subscribe(new Flow.Subscriber<>() { @Override public void onSubscribe(Flow.Subscription subscription) { } @Override public void onNext(Book book) { books.add(book); System.out.println("Parsed Book: " + book.getTitle() + " by " + book.getAuthor()); } @Override public void onError(Throwable throwable) { } @Override public void onComplete() { } }); return null; }); return books; }) .build(); strategy.edge(strategy.nodeStart, getMdOutput); strategy.edge(getMdOutput, strategy.nodeFinish); ``` # Predefined agent strategies To make agent implementations easier, Koog provides predefined agent strategies for common agent use cases. The following predefined strategies are available: - [Chat agent strategy](#chat-agent-strategy) - [ReAct strategy](#react-strategy) ## Chat agent strategy The Chat agent strategy is designed for executing a chat interaction process. It orchestrates interactions between different stages, nodes, and tools to handle user input, execute tools, and provide responses in a chat-like manner. ### Overview The Chat agent strategy implements a pattern where the agent: 1. Receives user input 1. Processes the input using an LLM 1. Either calls a tool or provides a direct response 1. Processes tool results and continues the conversation 1. Provides feedback if the LLM tries to respond with plain text instead of using tools This approach creates a conversational interface where the agent can use tools to fulfill user requests. ### Setup and dependencies The implementation of Chat agent strategy in Koog is done through the `chatAgentStrategy` function. To make the function available in your agent code, add the following dependency import: ``` ai.koog.agents.ext.agent.chatAgentStrategy ``` To use the strategy, create an AI agent following the pattern below: ``` val chatAgent = AIAgent( promptExecutor = promptExecutor, toolRegistry = toolRegistry, llmModel = model, // Set chatAgentStrategy as the agent strategy strategy = chatAgentStrategy() ) ``` ``` AIAgent chatAgent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().openAI("OPENAI_API_KEY").build()) .llmModel(OpenAIModels.Chat.O4Mini) .toolRegistry(ToolRegistry.builder().build()) // Set chatAgentStrategy as the agent strategy .graphStrategy(AIAgentStrategies.chatAgentStrategy()) .build(); ``` ### When to use the Chat agent strategy The Chat agent strategy is particularly useful for: - Building conversational agents that need to use tools - Creating assistants that can perform actions based on user requests - Implementing chatbots that need to access external systems or data - Scenarios where you want to enforce tool usage rather than plain text responses ### Example Here is a code sample of an AI agent that implements the predefined Chat agent strategy (`chatAgentStrategy`) and tools that the agent may use: ``` val chatAgent = AIAgent( promptExecutor = promptExecutor, llmModel = model, // Use chatAgentStrategy as the agent strategy strategy = chatAgentStrategy(), // Add tools the agent can use toolRegistry = ToolRegistry { tool(searchTool) tool(weatherTool) } ) suspend fun main() { // Run the agent with a user query val result = chatAgent.run("What's the weather like today and should I bring an umbrella?") } ``` ``` // Add tools the agent can use ToolRegistry toolRegistry = ToolRegistry.builder() .tools(new SearchAndWeatherTools()) .build(); AIAgent chatAgent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().openAI("OPENAI_API_KEY").build()) .llmModel(OpenAIModels.Chat.O4Mini) // Use chatAgentStrategy as the agent strategy .graphStrategy(AIAgentStrategies.chatAgentStrategy()) .toolRegistry(toolRegistry) .build(); // Run the agent with a user query String result = chatAgent.run("What's the weather like today and should I bring an umbrella?"); ``` ## ReAct strategy The ReAct (Reasoning and Acting) strategy is an AI agent strategy that alternates between reasoning and execution stages to dynamically process tasks and request output from a Large Language Model (LLM). ### Overview The ReAct strategy implements a pattern where the agent: 1. Reasons about the current state and plans the next steps 1. Takes actions based on that reasoning 1. Observes the results of those actions 1. Repeats the cycle This approach combines the strengths of reasoning (thinking through problems step by step) and acting (executing tools to gather information or perform operations). ### Flow diagram Here is the flow diagram of the ReAct strategy: ### Setup and dependencies The implementation of ReAct strategy in Koog is done through the `reActStrategy` function. To use the strategy, create an AI agent following the pattern below: ``` val reActAgent = AIAgent( promptExecutor = promptExecutor, toolRegistry = toolRegistry, llmModel = model, // Set reActStrategy as the agent strategy strategy = reActStrategy( // Set optional parameter values reasoningInterval = 1, name = "react_agent" ) ) ``` ``` AIAgent reActAgent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().openAI("OPENAI_API_KEY").build()) .llmModel(OpenAIModels.Chat.O4Mini) .toolRegistry(ToolRegistry.builder().build()) // Set reActStrategy as the agent strategy .graphStrategy(AIAgentStrategies.reActStrategy( // Set optional parameter values 1, // reasoningInterval "react_agent" // name )) .build(); ``` ### Parameters The `reActStrategy` function takes the following parameters: | Parameter | Type | Default | Description | | ------------------- | ------ | -------- | ------------------------------------------------------------------- | | `reasoningInterval` | Int | 1 | Specifies the interval for reasoning steps. Must be greater than 0. | | `name` | String | `re_act` | The name of the strategy. | ### Example use case Here is an example of how the ReAct strategy works with a simple banking agent: #### 1. User input The user sends the initial prompt. For example, this can be a question such as `How much did I spend last month?`. #### 2. Reasoning The agent performs the initial reasoning by taking the user input and the reasoning prompt. The reasoning can look as follows: ``` I need to follow these steps: 1. Get all transactions from last month 2. Filter out deposits (positive amounts) 3. Calculate total spending ``` #### 3. Action and execution, phase 1 Based on the action items that the agent defined in the previous step, it runs a tool to get all transactions from the previous month. In this case, the tool to run is `get_transactions`, along with the defined `startDate` and `endDate` arguments that match the request to get all transactions during the previous month: ``` {tool: "get_transactions", args: {startDate: "2025-05-19", endDate: "2025-06-18"}} ``` The tool returns a result that can look as follows: ``` [ {date: "2025-05-25", amount: -100.00, description: "Grocery Store"}, {date: "2025-05-31", amount: +1000.00, description: "Salary Deposit"}, {date: "2025-06-10", amount: -500.00, description: "Rent Payment"}, {date: "2025-06-13", amount: -200.00, description: "Utilities"} ] ``` #### 4. Reasoning With the result returned by the tool, the agent performs reasoning again to determine the next steps in its flow: ``` I have the transactions. Now I need to: 1. Remove the salary deposit of +1000.00 2. Sum up the remaining transactions ``` #### 5. Action and execution, phase 2 Based on the previous reasoning step, the agent calls the `calculate_sum` tool that sums up the amounts provided as tool arguments. As the reasoning also resulted in the action point of removing the positive amount from transactions, the amounts provided as tool arguments are only the negative ones: ``` {tool: "calculate_sum", args: {amounts: [-100.00, -500.00, -200.00]}} ``` The tool returns the final result: ``` -800.00 ``` #### 6. Final response The agent returns the final response (assistant message) that includes the calculated sum: ``` You spent $800.00 last month on groceries, rent, and utilities. ``` ### When to use the ReAct strategy The ReAct strategy is particularly useful for: - Complex tasks requiring multistep reasoning - Scenarios where the agent needs to gather information before providing a final answer - Problems that benefit from breaking down into smaller steps - Tasks requiring both analytical thinking and tool usage ### Example Here is a code sample of an AI agent that implements the predefined ReAct strategy (`reActStrategy`) and tools that the agent may use: ``` val bankingAgent = AIAgent( promptExecutor = promptExecutor, llmModel = model, // Use reActStrategy as the agent strategy strategy = reActStrategy( reasoningInterval = 1, name = "banking_agent" ), // Add tools the agent can use toolRegistry = ToolRegistry { tool(getTransactions) tool(calculateSum) } ) suspend fun main() { // Run the agent with a user query val result = bankingAgent.run("How much did I spend last month?") } ``` ``` // Add tools the agent can use ToolRegistry toolRegistry = ToolRegistry.builder() .tools(new BankingTools()) .build(); AIAgent bankingAgent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().openAI("OPENAI_API_KEY").build()) .llmModel(OpenAIModels.Chat.O4Mini) // Use reActStrategy as the agent strategy .graphStrategy(AIAgentStrategies.reActStrategy(1, "banking_agent")) .toolRegistry(toolRegistry) .build(); // Run the agent with a user query String result = bankingAgent.run("How much did I spend last month?"); ``` # Custom strategy graphs Strategy graphs are the backbone of agent workflows in the Koog framework. They define how the agent processes input, interacts with tools, and generates output. A strategy graph consists of nodes connected by edges, with conditions determining the flow of execution. Creating a strategy graph lets you tailor the behavior of an agent to your specific needs, whether you are building a simple chatbot, a complex data processing pipeline, or anything in between. ## Strategy graph architecture At a high level, a strategy graph consists of the following components: - **Strategy**: the top-level container for the graph, created using the `strategy` function with the specified input and output types using generic parameters. - **Subgraphs**: sections of the graph that can have their own set of tools and context. - **Nodes**: individual operations or transformations in the workflow. - **Edges**: connections between nodes that define transition conditions and transformations. The strategy graph begins at a special node called `nodeStart` and ends at `nodeFinish`. The path between these nodes is determined by the edges and conditions specified in the graph. ## Strategy graph components ### Nodes Nodes are building blocks of a strategy graph. Each node represents a specific operation. The Koog framework provides predefined nodes and also lets you create custom nodes by using the `node` function. For details, see [Predefined nodes and components](../nodes-and-components/) and [Custom nodes](../custom-nodes/). ### Edges Edges connect nodes and define the flow of operation in the strategy graph. An edge is created using the `edge` function and the `forwardTo` infix function: ``` edge(sourceNode forwardTo targetNode) ``` ``` strategy.edge(sourceNode, targetNode); ``` #### Conditions Conditions determine when to follow a particular edge in the strategy graph. There are several types of conditions, here are some common ones: | Condition type | Description | | --------------- | ---------------------------------------------------------------------------------------- | | onCondition | A general-purpose condition that takes a lambda expression that returns a boolean value. | | onToolCalls | A condition that matches when the LLM calls one or more tools. | | onTextMessage | A condition that matches when the LLM responds with a text message. | | onToolNotCalled | A condition that matches when the LLM does not call a tool. | You can transform the output before passing it to the target node by using the `transformed` function: ``` edge(sourceNode forwardTo targetNode onCondition { input -> input.length > 10 } transformed { input -> input.uppercase() } ) ``` ``` strategy.edge(AIAgentEdge.builder() .from(sourceNode) .to(targetNode) .onCondition(input -> input.length() > 10) .transformed(input -> input.toUpperCase()) .build()); ``` ### Subgraphs Subgraphs are sections of the strategy graph that operate with their own set of tools and context. The strategy graph can contain multiple subgraphs. Each subgraph is defined by using the `subgraph` function: ``` val strategy = strategy("strategy-name") { val firstSubgraph by subgraph("first") { // Define nodes and edges for this subgraph } val secondSubgraph by subgraph("second") { // Define nodes and edges for this subgraph } } ``` ``` var firstSubgraph = AIAgentSubgraph.builder("first") .withInput(FirstInput.class) .withOutput(FirstOutput.class) .define(subgraph -> { // Define nodes and edges for this subgraph }) .build(); var secondSubgraph = AIAgentSubgraph.builder("second") .withInput(SecondInput.class) .withOutput(SecondOutput.class) .define(subgraph -> { // Define nodes and edges for this subgraph }) .build(); ``` A subgraph can use any tool from a tool registry. However, you can specify a subset of tools from this registry that can be used in the subgraph and pass it as an argument to the `subgraph` function: ``` val strategy = strategy("strategy-name") { val firstSubgraph by subgraph( name = "first", tools = listOf(someTool) ) { // Define nodes and edges for this subgraph } // Define other subgraphs } ``` ``` var firstSubgraph = AIAgentSubgraph.builder("first") .withInput(FirstInput.class) .withOutput(FirstOutput.class) .limitedTools(someTools) .define(subgraph -> { // Define nodes and edges for this subgraph }) .build(); ``` ## Basic strategy graph creation The basic strategy graph operates as follows: 1. Sends the input to the LLM. 1. If the LLM responds with a message, finishes the process. 1. If the LLM calls a tool, runs the tool. 1. Sends the tool result back to the LLM. 1. If the LLM responds with a message, finishes the process. 1. If the LLM calls another tool, runs the tool, and the process repeats from step 4. Here is an example of a basic strategy graph: ``` val myStrategy = strategy("my-strategy") { val nodeCallLLM by nodeLLMRequest() val executeToolCall by nodeExecuteTools() val sendToolResult by nodeLLMSendToolResults() edge(nodeStart forwardTo nodeCallLLM) edge(nodeCallLLM forwardTo nodeFinish onTextMessage { true }) edge(nodeCallLLM forwardTo executeToolCall onToolCalls { true }) edge(executeToolCall forwardTo sendToolResult) edge(sendToolResult forwardTo nodeFinish onTextMessage { true }) edge(sendToolResult forwardTo executeToolCall onToolCalls { true }) } ``` ``` var graph = AIAgentGraphStrategy.builder("single_run") .withInput(String.class) .withOutput(String.class); var nodeCallLLM = AIAgentNode.llmRequest("sendInput"); var nodeExecuteTool = AIAgentNode.executeTools("nodeExecuteTool"); var nodeSendToolResult = AIAgentNode.llmSendToolResults("nodeSendToolResult"); graph.edge(AIAgentEdge.builder() .from(graph.nodeStart) .to(nodeCallLLM) .build()); graph.edge(AIAgentEdge.builder() .from(nodeCallLLM) .to(nodeExecuteTool) .onToolCalls() .build()); graph.edge(AIAgentEdge.builder() .from(nodeCallLLM) .to(graph.nodeFinish) .onTextMessage() .build()); graph.edge(nodeExecuteTool, nodeSendToolResult); graph.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(graph.nodeFinish) .onTextMessage() .build()); graph.edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(nodeExecuteTool) .onToolCalls() .build()); var strategy = graph.build(); ``` ## Visualizing strategy graph On JVM you may generate a [Mermaid state diagram](https://mermaid.js.org/syntax/stateDiagram.html) for the strategy graph. For the graph created in the previous example, you can run: ``` val mermaidDiagram: String = myStrategy.asMermaidDiagram() println(mermaidDiagram) ``` ``` var mermaidDiagram = MermaidDiagramGenerator.INSTANCE.generate(myStrategy); System.out.println(mermaidDiagram); ``` and the output will be: ``` --- title: my-strategy --- stateDiagram state "nodeCallLLM" as nodeCallLLM state "executeToolCall" as executeToolCall state "sendToolResult" as sendToolResult [*] --> nodeCallLLM nodeCallLLM --> [*] : transformed nodeCallLLM --> executeToolCall : onCondition executeToolCall --> sendToolResult sendToolResult --> [*] : transformed sendToolResult --> executeToolCall : onCondition ``` ## Advanced strategy techniques ### History compression For long-running conversations, the history can grow large and consume a lot of tokens. To learn how to compress the history, see [History compression](../history-compression/). ### Parallel tool execution For workflows that require executing multiple tools in parallel, you can use the `nodeExecuteTools` node with `parallel = true`: ``` val executeMultipleTools by nodeExecuteTools(parallel = true) val processMultipleResults by nodeLLMSendToolResults() edge(someNode forwardTo executeMultipleTools) edge(executeMultipleTools forwardTo processMultipleResults) ``` You can also use the `toParallelToolCallsRaw` extension function for streaming data: ``` parseMarkdownStreamToBooks(markdownStream).toParallelToolCallsRaw(BookTool::class).collect() ``` To learn more, see [Tools](../tools/#parallel-tool-calls). ### Parallel node execution Parallel node execution lets you run multiple nodes concurrently, improving performance and enabling complex workflows. To initiate parallel node runs, use the `parallel` method: ``` val calc by parallel( nodeCalcTokens, nodeCalcSymbols, nodeCalcWords, ) { selectByMax { it } } ``` The code above creates a node named `calc` that runs the `nodeCalcTokens`, `nodeCalcSymbols`, and `nodeCalcWords` nodes in parallel and returns the results as an instance of `AsyncParallelResult`. For more information related to parallel node execution and a detailed reference, see [Parallel node execution](../parallel-node-execution/). ### Conditional branching For complex workflows that require different paths based on certain conditions, you can use conditional branching: ``` val branchA by node { input -> // Logic for branch A "Branch A: $input" } val branchB by node { input -> // Logic for branch B "Branch B: $input" } edge( (someNode forwardTo branchA) onCondition { input -> input.contains("A") } ) edge( (someNode forwardTo branchB) onCondition { input -> input.contains("B") } ) ``` ## Best practices When you create custom strategy graphs, follow these best practices: - Keep it simple. Start with a simple graph and add complexity as needed. - Give your nodes and edges descriptive names to make the graph easier to understand. - Handle all possible paths and edge cases. - Test your graph with various inputs to ensure it behaves as expected. - Document the purpose and behavior of your graph for future reference. - Use predefined strategies or common patterns as a starting point. - For long-running conversations, use history compression to reduce token usage. - Use subgraphs to organize your graph and manage tool access. ## Usage examples ### Tone analysis strategy The tone analysis strategy is a good example of a tool-based strategy that includes history compression: ``` fun toneStrategy(name: String, toolRegistry: ToolRegistry): AIAgentGraphStrategy { return strategy(name) { val nodeSendInput by nodeLLMRequest() val nodeExecuteTool by nodeExecuteTools() val nodeSendToolResult by nodeLLMSendToolResults() val nodeCompressHistory by nodeLLMCompressHistory() // Define the flow of the agent edge(nodeStart forwardTo nodeSendInput) // If the LLM responds with a message, finish edge( (nodeSendInput forwardTo nodeFinish) onTextMessage { true } ) // If the LLM calls a tool, execute it edge( (nodeSendInput forwardTo nodeExecuteTool) onToolCalls { true } ) // If the history gets too large, compress it edge( (nodeExecuteTool forwardTo nodeCompressHistory) onCondition { _ -> llm.readSession { prompt.messages.size > 100 } } ) edge(nodeCompressHistory forwardTo nodeSendToolResult) // Otherwise, send the tool result directly edge( (nodeExecuteTool forwardTo nodeSendToolResult) onCondition { _ -> llm.readSession { prompt.messages.size <= 100 } } ) // If the LLM calls another tool, execute it edge( (nodeSendToolResult forwardTo nodeExecuteTool) onToolCalls { true } ) // If the LLM responds with a message, finish edge( (nodeSendToolResult forwardTo nodeFinish) onTextMessage { true } ) } } ``` This strategy does the following: 1. Sends the input to the LLM. 1. If the LLM responds with a message, the strategy finishes the process. 1. If the LLM calls a tool, the strategy runs the tool. 1. If the history is too large (more than 100 messages), the strategy compresses it before sending the tool result. 1. Otherwise, the strategy sends the tool result directly. 1. If the LLM calls another tool, the strategy runs it. 1. If the LLM responds with a message, the strategy finishes the process. ## Troubleshooting When creating custom strategy graphs, you might encounter some common issues. Here are some troubleshooting tips: ### Graph fails to reach the finish node If your graph does not reach the finish node, check the following: - All paths from the start node eventually lead to the finish node. - Your conditions are not too restrictive, preventing edges from being followed. - There are no cycles in the graph that do not have an exit condition. ### Tool calls are not running If tool calls are not running, check the following: - The tools are properly registered in the tool registry. - The edge from the LLM node to the tool execution node has the correct condition (`onToolCall { true }`). ### History gets too large If your history gets too large and consumes too many tokens, consider the following: - Add a history compression node. - Use a condition to check the size of the history and compress it when it gets too large. - Use a more aggressive compression strategy (e.g., `FromLastNMessages` with a smaller N value). ### Graph behaves unexpectedly If your graph takes unexpected branches, check the following: - Your conditions are correctly defined. - The conditions are evaluated in the expected order (edges are checked in the order they are defined). - You are not accidentally overriding conditions with more general ones. ### Performance issues occur If your graph has performance issues, consider the following: - Simplify the graph by removing unnecessary nodes and edges. - Use parallel tool execution for independent operations. - Compress history. - Use more efficient nodes and operations. ## Overview Parallel node execution lets you run multiple AI agent nodes concurrently, improving performance and enabling complex workflows. This feature is particularly useful when you need to: - Process the same input through different models or approaches simultaneously - Perform multiple independent operations in parallel - Implement competitive evaluation patterns where multiple solutions are generated and then compared ## Key components Parallel node execution in Koog consists of the methods and data structures described below. ### Methods - `parallel()`: executes multiple nodes in parallel and collects their results. ### Data structures - `ParallelResult`: represents the completed result of a parallel node execution. - `NodeExecutionResult`: contains the output and context of a node execution. ## Basic usage ### Running nodes in parallel To initiate parallel execution of nodes, use the `parallel` method in the following format: ``` val nodeName by parallel( firstNode, secondNode, thirdNode /* Add more nodes if needed */ ) { // Merge strategy goes here, for example: selectByMax { it.length } } ``` Here is an actual example of running three nodes in parallel and selecting the result with the maximum length: ``` val calc by parallel( nodeCalcTokens, nodeCalcSymbols, nodeCalcWords, ) { selectByMax { it } } ``` The code above runs the `nodeCalcTokens`, `nodeCalcSymbols`, and `nodeCalcWords` nodes in parallel and returns the result with the maximum value. ### Merge strategies After executing nodes in parallel, you need to specify how to merge the results. Koog provides the following merge strategies: - `selectBy()`: selects a result based on a predicate function. - `selectByMax()`: selects the result with the maximum value based on a comparison function. - `selectByIndex()`: selects a result based on an index returned by a selection function. - `fold()`: folds the results into a single value using an operation function. #### selectBy Selects a result based on a predicate function: ``` val nodeSelectJoke by parallel( nodeOpenAI, nodeAnthropicSonnet, nodeAnthropicOpus, ) { selectBy { it.contains("programmer") } } ``` This selects the first joke that contains the word "programmer". #### selectByMax Selects the result with the maximum value based on a comparison function: ``` val nodeLongestJoke by parallel( nodeOpenAI, nodeAnthropicSonnet, nodeAnthropicOpus, ) { selectByMax { it.length } } ``` This selects the joke with the maximum length. #### selectByIndex Selects a result based on an index returned by a selection function: ``` val nodeBestJoke by parallel( nodeOpenAI, nodeAnthropicSonnet, nodeAnthropicOpus, ) { selectByIndex { jokes -> // Use another LLM to determine the best joke llm.writeSession { model = OpenAIModels.Chat.GPT4o appendPrompt { system("You are a comedy critic. Select the best joke.") user("Here are three jokes: ${jokes.joinToString("\n\n")}") } val response = requestLLMStructured() response.getOrNull()!!.data.bestJokeIndex } } } ``` This uses another LLM call to determine the index of the best joke. #### fold Folds the results into a single value using an operation function: ``` val nodeAllJokes by parallel( nodeOpenAI, nodeAnthropicSonnet, nodeAnthropicOpus, ) { fold("Jokes:\n") { result, joke -> "$result\n$joke" } } ``` This combines all jokes into a single string. ## Example: Best joke agent Here is a complete example that uses parallel execution to generate jokes from different LLM models and select the best one: ``` val strategy = strategy("best-joke") { // Define nodes for different LLM models val nodeOpenAI by node { topic -> llm.writeSession { model = OpenAIModels.Chat.GPT4o appendPrompt { system("You are a comedian. Generate a funny joke about the given topic.") user("Tell me a joke about $topic.") } val response = requestLLMWithoutTools() response.parts.filterIsInstance().joinToString("\n") { it.text } } } val nodeAnthropicSonnet by node { topic -> llm.writeSession { model = AnthropicModels.Sonnet_4_5 appendPrompt { system("You are a comedian. Generate a funny joke about the given topic.") user("Tell me a joke about $topic.") } val response = requestLLMWithoutTools() response.parts.filterIsInstance().joinToString("\n") { it.text } } } val nodeAnthropicOpus by node { topic -> llm.writeSession { model = AnthropicModels.Opus_4_6 appendPrompt { system("You are a comedian. Generate a funny joke about the given topic.") user("Tell me a joke about $topic.") } val response = requestLLMWithoutTools() response.parts.filterIsInstance().joinToString("\n") { it.text } } } // Execute joke generation in parallel and select the best joke val nodeGenerateBestJoke by parallel( nodeOpenAI, nodeAnthropicSonnet, nodeAnthropicOpus, ) { selectByIndex { jokes -> // Another LLM (e.g., GPT4o) would find the funniest joke: llm.writeSession { model = OpenAIModels.Chat.GPT4o appendPrompt { prompt("best-joke-selector") { system("You are a comedy critic. Give a critique for the given joke.") user( """ Here are three jokes about the same topic: ${jokes.mapIndexed { index, joke -> "Joke $index:\n$joke" }.joinToString("\n\n")} Select the best joke and explain why it's the best. """.trimIndent() ) } } val response = requestLLMStructured() val bestJoke = response.getOrNull()!!.data bestJoke.bestJokeIndex } } } // Connect the nodes nodeStart then nodeGenerateBestJoke then nodeFinish } ``` ## Best practices 1. **Consider resource constraints**: Be mindful of resource usage when executing nodes in parallel, especially when making multiple LLM API calls simultaneously. 1. **Context management**: Each parallel execution creates a forked context. When merging results, choose which context to preserve or how to combine contexts from different executions. 1. **Optimize for your use case**: - For competitive evaluation (like the joke example), use `selectByIndex` to select the best result - For finding the maximum value, use `selectByMax` - For filtering based on a condition, use `selectBy` - For aggregation, use `fold` to combine all results into a composite output ## Performance considerations Parallel execution can significantly improve throughput, but it comes with some overhead: - Each parallel node creates a new coroutine - Context forking and merging add some computational cost - Resource contention may occur with many parallel executions For optimal performance, parallelize operations that: - Are independent of each other - Have significant execution time - Don't share mutable state ## Overview Koog provides a way to store and pass data using `AIAgentStorage`, which is a key-value storage system designed as a type-safe way to pass data between different nodes or even subgraphs. The storage is accessible through the `storage` property (`storage: AIAgentStorage`) available in agent nodes, allowing for seamless data sharing across different components of your AI agent system. ## Key and value structure The key-value data storage structure relies on the `AIAgentStorageKey` data class. For more information about `AIAgentStorageKey`, see the sections below. ### AIAgentStorageKey The storage uses a typed key system to provide type safety when storing and retrieving data. `AIAgentStorageKey` class represents a storage key used for identifying and accessing data. Here are the key features of this class: - The generic type parameter `T` specifies the type of data associated with this key, providing type safety. - Each key has a `name` property which is a string identifier for easier identification and debugging. - Each key instance is unique. `name` is not used to determine uniqueness, so it is acceptable to have several keys with the same name. This allows resuing existing strategy components without the risk of accidentally overwriting your data in the storage. ## Usage examples The following sections provide an actual example of creating a storage key and using it to store and retrieve data. ### Defining a class that represents your data The first step in storing data that you want to pass is creating a class that represents your data. Here is an example of a simple class with basic user data: ``` class UserData( val name: String, val age: Int ) ``` ``` record UserData( String name, int age ) {} ``` Once defined, use the class to create a storage key as described below. ### Creating a storage key Create a typed storage key for the defined data structure: ``` val userDataKey = createStorageKey("user-data") ``` ``` AIAgentStorageKey userDataKey = AIAgentStorage.createStorageKey("user-data", TypeToken.of(UserData.class)); ``` The `createStorageKey` function takes a string parameter used for identification and debugging purposes, and a `TypeToken` representing the value type (in Java; Kotlin uses reified generics automatically). ### Storing data To save data using a created storage key, use the `storage.set(key: AIAgentStorageKey, value: T)` method in a node: ``` val nodeSaveData by node { storage.set(userDataKey, UserData("John", 26)) } ``` ``` var nodeSaveData = AIAgentNode.builder("nodeSaveData") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { ctx.getStorage().set(userDataKey, new UserData("John", 26)); return ""; }) .build(); ``` ### Retrieving data To retrieve the data, use the `storage.get` method in a node: ``` val nodeRetrieveData by node { message -> storage.get(userDataKey)?.let { userFromStorage -> println("Hello dear $userFromStorage, here's a message for you: $message") } } ``` ``` var nodeRetrieveData = AIAgentNode.builder("nodeRetrieveData") .withInput(String.class) .withOutput(String.class) .withAction((message, ctx) -> { var userData = ctx.getStorage().get(userDataKey); System.out.println("Hello dear %s, here's a message for you: %s".formatted(userData, message)); return ""; }) .build(); ``` ## API documentation For a complete reference related to the `AIAgentStorage` class, see [AIAgentStorage](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/index.html). For individual functions available in the `AIAgentStorage` class, see the following API references: - [clear](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/clear.html) - [get](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/get.html) - [getValue](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/get-value.html) - [putAll](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/put-all.html) - [remove](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/remove.html) - [set](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/set.html) - [toMap](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.entity/-a-i-agent-storage/to-map.html) ## Additional information - `AIAgentStorage` is thread-safe, using a Mutex to ensure concurrent access is handled properly. - When retrieving values, type casting is handled automatically, ensuring type safety throughout your application. - For non-nullable access to values, use the `getValue` method which throws an exception if the key does not exist. - You can clear the storage entirely using the `clear` method, which removes all stored key-value pairs. # Features # Features Agent features provide a way to extend and enhance the functionality of AI agents. With features, you can: - Add new capabilities to agents - Intercept and modify agent behavior - Log and monitor agent execution - Register multiple handlers for the same event type within a single feature The Koog framework has the following features available out of the box: - [Event handling](agent-event-handlers/) ______________________________________________________________________ Monitor and respond to specific events during the agent execution - [Tracing](tracing/) ______________________________________________________________________ Capture detailed information about agent runs - [Chat memory](chat-memory/) ______________________________________________________________________ Store and retrieve chat message history between agent runs - [Long-term memory](long-term-memory/) ______________________________________________________________________ Add persistent memory to AI agents - [Agent persistence](agent-persistence/) ______________________________________________________________________ Save and restore the state of an agent at specific points during execution - [OpenTelemetry](open-telemetry/) ______________________________________________________________________ Generate, collect, and export telemetry data (traces) from your agent To learn how to implement your own features, see [Custom features](custom-features/). # Event handlers You can monitor and respond to specific events during the agent workflow by using event handlers for logging, testing, debugging, and extending agent behavior. ## Feature overview The EventHandler feature lets you hook into various agent events. It serves as an event delegation mechanism that: - Manages the lifecycle of AI agent operations. - Provides hooks for monitoring and responding to different stages of the workflow. - Enables error handling and recovery. - Facilitates tool invocation tracking and result processing. ### Installation and configuration The EventHandler feature integrates with the agent workflow through the `EventHandler` class, which provides a way to register callbacks for different agent events, and can be installed as a feature in the agent configuration. For details, see [API-reference](https://api.koog.ai/agents/agents-features/agents-features-event-handler/ai.koog.agents.features.eventHandler.feature/-event-handler/index.html). To install the feature and configure event handlers for the agent, do the following: ``` handleEvents { // Handle tool calls onToolCallStarting { eventContext -> println("Tool called: ${eventContext.toolName} with args ${eventContext.toolArgs}") } // Handle event triggered when the agent completes its execution onAgentCompleted { eventContext -> println("Agent finished with result: ${eventContext.result}") } // Other event handlers } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(simpleOllamaAIExecutor("http://localhost:11434")) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(EventHandler.Feature, cfg -> { // Handle tool calls cfg.onToolCallStarting(ctx -> { System.out.println("Tool called: " + ctx.getToolName() + " with args " + ctx.getToolArgs()); }); // Handle event triggered when the agent completes its execution cfg.onAgentCompleted(ctx -> { System.out.println("Agent finished with result: " + ctx.getResult()); }); }) .build(); ``` For more details about event handler configuration, see [API-reference](https://api.koog.ai/agents/agents-features/agents-features-event-handler/ai.koog.agents.features.eventHandler.feature/-event-handler-config/index.html). You can also set up event handlers using the `handleEvents` extension function when creating an agent. This function also installs the event handler feature and configures event handlers for the agent. Here is an example: ``` val agent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, ){ handleEvents { // Handle tool calls onToolCallStarting { eventContext -> println("Tool called: ${eventContext.toolName} with args ${eventContext.toolArgs}") } // Handle event triggered when the agent completes its execution onAgentCompleted { eventContext -> println("Agent finished with result: ${eventContext.result}") } // Other event handlers } } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(simpleOllamaAIExecutor("http://localhost:11434")) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(EventHandler.Feature, cfg -> { // Handle tool calls cfg.onToolCallStarting(ctx -> { System.out.println("Tool called: " + ctx.getToolName() + " with args " + ctx.getToolArgs()); }); // Handle event triggered when the agent completes its execution cfg.onAgentCompleted(ctx -> { System.out.println("Agent finished with result: " + ctx.getResult()); }); }) .build(); ``` # Tracing This page includes details about the Tracing feature, which provides comprehensive tracing capabilities for AI agents. ## Feature overview The Tracing feature is a powerful monitoring and debugging tool that captures detailed information about agent runs, including: - Strategy execution - LLM calls - LLM streaming (start, frames, completion, errors) - Tool calls - Node execution within the agent graph This feature operates by intercepting key events in the agent pipeline and forwarding them to configurable message processors. These processors can output the trace information to various destinations such as log files or other types of files in the filesystem, enabling developers to gain insights into agent behavior and troubleshoot issues effectively. ### Event flow 1. The Tracing feature intercepts events in the agent pipeline. 1. Events are filtered based on the configured message filter. 1. Filtered events are passed to registered message processors. 1. Message processors format and output the events to their respective destinations. ## Configuration and initialization ### Basic setup To use the Tracing feature, you need to: 1. Have one or more message processors (you can use the existing ones or create your own). 1. Install `Tracing` in your agent. 1. Configure the message filter (optional). 1. Add the message processors to the feature. ``` // Defining a logger/file that will be used as a destination of trace messages val logger = KotlinLogging.logger { } val outputPath = Path("/path/to/trace.log") // Creating an agent val agent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, ) { install(Tracing) { // Configure message processors to handle trace events addMessageProcessor(TraceFeatureMessageLogWriter(logger)) addMessageProcessor(TraceFeatureMessageFileWriter.create(outputPath)) } } ``` ``` // Defining a logger/file that will be used as a destination of trace messages var logger = LoggerFactory.getLogger("tracing"); var outputPath = Path.of("/path/to/trace.log"); // Creating an agent var agent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().ollama().build()) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Tracing.Feature, config -> { // Configure message processors to handle trace events config.addMessageProcessor(TraceFeatureMessageLogWriter.create(logger)); config.addMessageProcessor(TraceFeatureMessageFileWriter.create(outputPath)); }) .build(); ``` ### Message filtering You can process all existing events or select some of them based on specific criteria. The message filter lets you control which events are processed. This is useful for focusing on specific aspects of agent runs: ``` val fileWriter = TraceFeatureMessageFileWriter( outputPath, { path: Path -> SystemFileSystem.sink(path).buffered() } ) addMessageProcessor(fileWriter) // Filter for LLM-related events only fileWriter.setMessageFilter { message -> message is LLMCallStartingEvent || message is LLMCallCompletedEvent } // Filter for tool-related events only fileWriter.setMessageFilter { message -> message is ToolCallStartingEvent || message is ToolCallCompletedEvent || message is ToolValidationFailedEvent || message is ToolCallFailedEvent } // Filter for node execution events only fileWriter.setMessageFilter { message -> message is NodeExecutionStartingEvent || message is NodeExecutionCompletedEvent } ``` ``` var fileWriter = TraceFeatureMessageFileWriter.create( outputPath, path -> { try { return Files.newOutputStream(path); } catch (IOException e) { throw new UncheckedIOException(e); }} ); config.addMessageProcessor(fileWriter); // Filter for LLM-related events only fileWriter.setMessageFilter(message -> message instanceof LLMCallStartingEvent || message instanceof LLMCallCompletedEvent ); // Filter for tool-related events only fileWriter.setMessageFilter(message -> message instanceof ToolCallStartingEvent || message instanceof ToolCallCompletedEvent || message instanceof ToolValidationFailedEvent || message instanceof ToolCallFailedEvent ); // Filter for node execution events only fileWriter.setMessageFilter(message -> message instanceof NodeExecutionStartingEvent || message instanceof NodeExecutionCompletedEvent ); ``` ### Large trace volumes For agents with complex strategies or long-running executions, the volume of trace events can be substantial. Consider using the following methods to manage the volume of events: - Use specific message filters to reduce the number of events. - Implement custom message processors with buffering or sampling. - Use file rotation for log files to prevent them from growing too large. ### Dependency graph The Tracing feature has the following dependencies: ``` Tracing ├── AIAgentPipeline (for intercepting events) ├── TraceFeatureConfig │ └── FeatureConfig ├── Message Processors │ ├── TraceFeatureMessageLogWriter │ │ └── FeatureMessageLogWriter │ ├── TraceFeatureMessageFileWriter │ │ └── FeatureMessageFileWriter │ └── TraceFeatureMessageRemoteWriter │ └── FeatureMessageRemoteWriter └── Event Types (from ai.koog.agents.core.feature.model) ├── AgentStartingEvent ├── AgentCompletedEvent ├── AgentExecutionFailedEvent ├── AgentClosingEvent ├── GraphStrategyStartingEvent ├── FunctionalStrategyStartingEvent ├── StrategyCompletedEvent ├── NodeExecutionStartingEvent ├── NodeExecutionCompletedEvent ├── NodeExecutionFailedEvent ├── SubgraphExecutionStartingEvent ├── SubgraphExecutionCompletedEvent ├── SubgraphExecutionFailedEvent ├── LLMCallStartingEvent ├── LLMCallCompletedEvent ├── LLMCallFailedEvent ├── LLMStreamingStartingEvent ├── LLMStreamingFrameReceivedEvent ├── LLMStreamingFailedEvent ├── LLMStreamingCompletedEvent ├── ToolCallStartingEvent ├── ToolValidationFailedEvent ├── ToolCallFailedEvent └── ToolCallCompletedEvent ``` ## Examples and quickstarts ### Basic tracing to logger ``` // Create a logger val logger = KotlinLogging.logger { } // Create an agent with tracing val agent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, ) { install(Tracing) { addMessageProcessor(TraceFeatureMessageLogWriter(logger)) } } // Run the agent agent.run("Hello, agent!") ``` ``` // Create a logger var logger = LoggerFactory.getLogger("tracing"); // Create an agent with tracing var agent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().ollama().build()) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Tracing.Feature, config -> { config.addMessageProcessor(TraceFeatureMessageLogWriter.create(logger)); }) .build(); // Run the agent agent.run("Hello, agent!"); ``` ## Error handling and edge cases ### No message processors If no message processors are added to the Tracing feature, a warning will be logged: ``` Tracing Feature. No feature out stream providers are defined. Trace streaming has no target. ``` The feature will still intercept events, but they will not be processed or output anywhere. ### Resource management Message processors may hold resources (like file handles) that need to be properly released. Use the `use` extension function to ensure proper cleanup: ``` val writer = TraceFeatureMessageFileWriter( outputPath, { path: Path -> SystemFileSystem.sink(path).buffered() } ) // Creating an agent val agent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, ) { install(Tracing) { addMessageProcessor(writer) } } // Run the agent agent.run(input) // Writer will be automatically closed when the block exits ``` ``` var writer = TraceFeatureMessageFileWriter.create( outputPath, path -> { try { return Files.newOutputStream(path); } catch (IOException e) { throw new UncheckedIOException(e); }} ); // Creating an agent var agent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().ollama().build()) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Tracing.Feature, config -> { config.addMessageProcessor(writer); }) .build(); // Run the agent agent.run(input); // Writer will be automatically closed when the block exits ``` ### Tracing specific events to file ``` val fileWriter = TraceFeatureMessageFileWriter( outputPath, { path: Path -> SystemFileSystem.sink(path).buffered() } ) // Creating an agent val agent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, ) { install(Tracing) { addMessageProcessor(fileWriter) // Only trace LLM calls fileWriter.setMessageFilter { message -> message is LLMCallStartingEvent || message is LLMCallCompletedEvent } } } ``` ``` var fileWriter = TraceFeatureMessageFileWriter.create( outputPath, path -> { try { return Files.newOutputStream(path); } catch (IOException e) { throw new UncheckedIOException(e); }} ); // Creating an agent var agent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().ollama().build()) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Tracing.Feature, config -> { config.addMessageProcessor(fileWriter); // Only trace LLM calls fileWriter.setMessageFilter(message -> message instanceof LLMCallStartingEvent || message instanceof LLMCallCompletedEvent ); }) .build(); ``` ### Tracing specific events to remote endpoint You use tracing to remote endpoints when you need to send event data via the network. Once initiated, tracing to a remote endpoint launches a light server at the specified port number and sends events via Kotlin Server-Sent Events (SSE). ``` val connectionConfig = DefaultServerConnectionConfig(host = host, port = port) val writer = TraceFeatureMessageRemoteWriter(connectionConfig) // Creating an agent val agent = AIAgent( promptExecutor = simpleOllamaAIExecutor(), llmModel = OllamaModels.Meta.LLAMA_3_2, ) { install(Tracing) { addMessageProcessor(writer) } } // Run the agent agent.run(input) // Writer will be automatically closed when the block exits ``` ``` var connectionConfig = new DefaultServerConnectionConfig(host, port); var writer = new TraceFeatureMessageRemoteWriter(connectionConfig); // Creating an agent var agent = AIAgent.builder() .promptExecutor(PromptExecutor.builder().ollama().build()) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Tracing.Feature, config -> { config.addMessageProcessor(writer); }) .build(); // Run the agent agent.run(input); // Writer will be automatically closed when the block exits ``` On the client side, you can use `FeatureMessageRemoteClient` to receive events and deserialize them. ``` val clientConfig = DefaultClientConnectionConfig(host = host, port = port, protocol = URLProtocol.HTTP) val agentEvents = mutableListOf() val clientJob = launch { FeatureMessageRemoteClient(connectionConfig = clientConfig, scope = this).use { client -> val collectEventsJob = launch { client.receivedMessages.consumeAsFlow().collect { event -> // Collect events from server agentEvents.add(event as DefinedFeatureEvent) // Stop collecting events on agent finished if (event is AgentCompletedEvent) { cancel() } } } client.connect() collectEventsJob.join() client.healthCheck() } } listOf(clientJob).joinAll() ``` ## API documentation The Tracing feature follows a modular architecture with these key components: 1. [Tracing](https://api.koog.ai/agents/agents-features/agents-features-trace/ai.koog.agents.features.tracing.feature/-tracing/index.html): the main feature class that intercepts events in the agent pipeline. 1. [TraceFeatureConfig](https://api.koog.ai/agents/agents-features/agents-features-trace/ai.koog.agents.features.tracing.feature/-trace-feature-config/index.html): configuration class for customizing feature behavior. 1. Message Processors: components that process and output trace events: - [TraceFeatureMessageLogWriter](https://api.koog.ai/agents/agents-features/agents-features-trace/ai.koog.agents.features.tracing.writer/-trace-feature-message-log-writer/index.html): writes trace events to a logger. - [TraceFeatureMessageFileWriter](https://api.koog.ai/agents/agents-features/agents-features-trace/ai.koog.agents.features.tracing.writer/-trace-feature-message-file-writer/index.html): writes trace events to a file. - [TraceFeatureMessageRemoteWriter](https://api.koog.ai/agents/agents-features/agents-features-trace/ai.koog.agents.features.tracing.writer/-trace-feature-message-remote-writer/index.html): sends trace events to a remote server. # Chat memory The `ChatMemory` feature enagles AI agents to store conversation history and retrieve it across multiple runs. When installed, the agent automatically loads previous messages at the start of each run and stores the updated conversation when the run completes, enabling natural multi-turn chat. Key capabilities: - Automatic loading and storing of conversation history per session ID - Pluggable storage backend via `ChatHistoryProvider` - Built-in preprocessors to limit history size and filter messages - Custom preprocessor support for arbitrary message transformations ## Add dependencies Chat memory is an optional [feature](../) that is not available in Koog by default. To implement chat memory for your Koog agent, add a dependency for [`ai.koog:agents-features-memory`](https://mvnrepository.com/artifact/ai.koog/agents-features-memory): build.gradle.kts ``` dependencies { implementation("ai.koog:agents-features-memory:$koogVersion") } ``` build.gradle ``` dependencies { implementation 'ai.koog:agents-features-memory:$koogVersion' } ``` pom.xml ``` ai.koog agents-features-memory-jvm $koogVersion ``` Note The `ChatMemory` feature is available starting from Koog version **0.7.0**. ## Enable chat memory Install `ChatMemory` using the `install()` method when creating the agent: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4oMini ) { install(ChatMemory) } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(executor) .llmModel(OpenAIModels.Chat.GPT4oMini) .install(ChatMemory.Feature) .build(); ``` By default, it uses an in-memory [chat history provider](#history-providers) with no [preprocessors](#preprocessors). Configure the `ChatMemory` feature to use a custom chat history provider and preprocessors, for example: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4oMini ) { install(ChatMemory) { chatHistoryProvider = MyDatabaseChatHistoryProvider() windowSize(20) filterMessages { it is Message.User || it is Message.Assistant } } } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(executor) .llmModel(OpenAIModels.Chat.GPT4oMini) .install(ChatMemory.Feature, config -> config .chatHistoryProvider(new MyDatabaseChatHistoryProvider()) .windowSize(20) .filterMessages(msg -> msg instanceof Message.User || msg instanceof Message.Assistant)) .build(); ``` ## Session IDs Provide the session ID as the second argument to `agent.run()`. `ChatMemory` uses this ID to store and load conversations: ``` // First run - the agent saves the chat history at the end agent.run("What is the capital of France?", "session-1") // Second run — the agent loads the previous exchange agent.run("And what about Germany?", "session-1") ``` Different session IDs produce fully isolated histories. ## History providers The default `InMemoryChatHistoryProvider` is thread-safe but not persistent (history is lost on restart). For production, implement your own `ChatHistoryProvider` that stores messages persistently. ``` class MyDatabaseChatHistoryProvider(private val db: Database) : ChatHistoryProvider { override suspend fun store(conversationId: String, messages: List) { db.saveMessages(conversationId, messages) } override suspend fun load(conversationId: String): List { return db.loadMessages(conversationId) ?: emptyList() } } ``` ## Preprocessors Preprocessors transform the message list at both load time (before the agent sees it) and store time (before saving). They run sequentially in the order you add them to the `ChatMemory` feature configuration. ### Built-in preprocessors | Config method | Preprocessor class | Behavior | | ------------------------ | ---------------------------- | ------------------------------------- | | `windowSize(n)` | `WindowSizePreProcessor` | Keeps only the last `n` messages | | `filterMessages { ... }` | `FilterMessagesPreProcessor` | Keeps messages matching the predicate | ### Order of preprocessors Preprocessors run sequentially, with each output being the next input. This means that order matters. ``` // Effect: keep last 10 messages, then filter short ones from those 10 windowSize(10) filterMessages { it.content.length <= 100 } // Effect: filter short messages first, then keep last 10 of the survivors filterMessages { it.content.length <= 100 } windowSize(10) ``` ### Custom preprocessors To create a custom preprocessor, implement the `ChatMemoryPreProcessor` interface: ``` class RedactEmailsPreProcessor : ChatMemoryPreProcessor { override fun preprocess(messages: List): List { return messages.map { message -> // Replace email addresses in message content Message.User(message.content.replace(Regex("[\\w.]+@[\\w.]+"), "[REDACTED]")) } } } ``` Then add it to the config: ``` install(ChatMemory) { addPreProcessor(RedactEmailsPreProcessor()) windowSize(50) } ``` ## Chat memory vs agent persistence `ChatMemory` treats each `agent.run()` call as an atomic, self-contained loop. The agent loads chat history before running and stores it after a successful run. If the agent crashes during the run, it does not store the current chat messages, meaning that the chat history remains as it was before the run. [Persistence](../agent-persistence/) captures the agent's internal execution state (graph node, message history, inputs, and outputs) as checkpoints during the run. If the agent crashes, it can resume from the last checkpoint. | | ChatMemory | Persistence | | ------------------ | ------------------------------------------------ | ------------------------------------------------------------------ | | **What it saves** | Conversation messages | Execution state | | **When it saves** | After `agent.run()` completes | After each graph node or at manually defined points during the run | | **Crash behavior** | In-progress run is lost; previous history intact | Can resume from last checkpoint | | **Typical use** | Multi-turn chat continuity | Long-running agents with crash recovery | If your agent performs long-running tasks where a mid-execution crash would be costly, consider installing both features: ``` val agent = AIAgent( promptExecutor = executor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a helpful assistant.", ) { install(ChatMemory) { chatHistoryProvider = MyDatabaseProvider() windowSize(50) } install(Persistence) { storage = MyPersistenceStorageProvider() enableAutomaticPersistence = true } } ``` ## Best practices - **Always set a window size** to prevent unlimited conversation growth. - **Order preprocessors carefully**, as filtering before windowing and windowing before filtering produce different results. - **Use meaningful session IDs** for history isolation: user IDs, chat thread IDs, or UUIDs all work well. - **Implement a persistent provider for production** because the default `InMemoryChatHistoryProvider` loses history on restart. ## Next steps - Learn how to [build a simple CLI chat loop with memory](chat-agent-with-memory/) - See an example of a [chat endpoint with memory](chat-backend-with-memory/) # Agent persistence Agent Persistence is a feature that provides checkpoint functionality for AI agents in the Koog framework. It lets you save and restore the state of an agent at specific points during execution, enabling capabilities such as: - Resuming agent execution from a specific point - Rolling back to previous states - Persisting agent state across sessions ## Key concepts ### Checkpoints A checkpoint captures the complete state of an agent at a specific point in its execution, including: - Message history (all interactions between user, system, assistant, and tools) - Last node that was successfully executed - Output data from that node - Selected LLM - Common LLM parameters - Selected tools - `AIAgentStorage` contents (key-value data stored during execution) - Timestamp of creation Checkpoints are identified by unique IDs and are associated with a specific agent. ### `AIAgentStorage` persistence When a checkpoint is created, the framework serializes all values currently held in `AIAgentStorage` and includes them in the checkpoint. On restore, those values are deserialized and made available to the resumed agent exactly as they were at the time of the checkpoint. **Only serializable values are persisted.** The serializer used is the one configured in `AIAgentConfig` via its `serializer` property. Values that cannot be encoded by that serializer are silently skipped and will not be present in the restored storage. Non-serializable values are dropped silently when the checkpoint is written. If a value is missing after restoring from a checkpoint, verify that its type is serializable by the configured serializer. See [Serialization](../../serialization/) for more information. ## Installation To use the Agent Persistence feature, add it to your agent's configuration: ``` val agent = AIAgent( promptExecutor = executor, llmModel = OllamaModels.Meta.LLAMA_3_2, ) { install(Persistence) { // Use in-memory storage for snapshots storage = InMemoryPersistenceStorageProvider() } } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(executor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Persistence.Feature, cfg -> { // Use in-memory storage for snapshots cfg.setStorage(new InMemoryPersistenceStorageProvider()); }) .build(); ``` ## Configuration options The Agent Persistence feature has three main configuration options: - **Storage provider**: the provider used to save and retrieve checkpoints. - **Continuous persistence**: automatic creation of checkpoints after each node is run. - **Rollback strategy**: determines which state will be restored when rolling back to a checkpoint. ### Storage provider Set the storage provider that will be used to save and retrieve checkpoints: ``` install(Persistence) { storage = InMemoryPersistenceStorageProvider() } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(executor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Persistence.Feature, cfg -> { cfg.setStorage(new InMemoryPersistenceStorageProvider()); }) .build(); ``` The framework includes the following built-in providers: - `InMemoryPersistenceStorageProvider`: stores checkpoints in memory (lost when the application restarts). - `FilePersistenceStorageProvider`: persists checkpoints to the file system. - `NoPersistenceStorageProvider`: a no-op implementation that does not store checkpoints. This is the default provider. You can also implement custom storage providers by implementing the `PersistenceStorageProvider` interface. For more information, see [Custom storage providers](#custom-storage-providers). ### Continuous persistence Continuous persistence means that a checkpoint is automatically created after each node is run. To disable continuous persistence, use the code below: ``` install(Persistence) { enableAutomaticPersistence = false } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(executor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Persistence.Feature, cfg -> { cfg.setEnableAutomaticPersistence(true); }) .build(); ``` If continuous persistence is disabled, you can still create checkpoints manually. ## Basic usage ### Creating a checkpoint To learn how to create a checkpoint at a specific point in your agent's execution, see the code sample below: ``` suspend fun example(context: AIAgentContext) { // Create a checkpoint with the current state val checkpoint = context.persistence().createCheckpointAfterNode( agentContext = context, nodePath = context.executionInfo.path(), lastOutput = outputData, lastOutputType = outputType, checkpointId = context.runId, version = 0L ) // The checkpoint ID can be stored for later use val checkpointId = checkpoint?.checkpointId } ``` ``` // PersistenceKt.persistence() is the Java-accessible form of the Kotlin extension function Persistence persistence = PersistenceKt.persistence(context); // Create a checkpoint with the current state AgentCheckpointData checkpoint = persistence.createCheckpointAfterNode( context, context.getExecutionInfo().path(), outputData, TypeToken.of(String.class), 0L, context.getRunId() ); // The checkpoint ID can be stored for later use String checkpointId = checkpoint != null ? checkpoint.getCheckpointId() : null; ``` ### Restoring from a checkpoint To restore the state of an agent from a specific checkpoint, follow the code sample below: ``` suspend fun example(context: AIAgentContext, checkpointId: String) { // Roll back to a specific checkpoint context.persistence().rollbackToCheckpoint(checkpointId, context) // Or roll back to the latest checkpoint context.persistence().rollbackToLatestCheckpoint(context) } ``` ``` Persistence persistence = PersistenceKt.persistence(context); // Roll back to a specific checkpoint persistence.rollbackToCheckpoint(checkpointId, context); // Or roll back to the latest checkpoint persistence.rollbackToLatestCheckpoint(context); ``` #### Rolling back all side-effects produced by tools It's quite common for some tools to produce side-effects. Specifically, when you are running your agents on the backend, some of the tools would likely perform some database transactions. This makes it much harder for your agent to travel back in time. Imagine you have a tool `createUser` that creates a new user in your database. And your agent has populated multiple tool calls overtime: ``` tool call: createUser "Alex" ->>>> checkpoint-1 <<<<- tool call: createUser "Daniel" tool call: createUser "Maria" ``` And now you would like to roll back to a checkpoint. Restoring the agent's state (including message history, and strategy graph node) alone would not be sufficient to achieve the exact state of the world before the checkpoint. You should also restore the side-effects produced by your tool calls. In our example, this would mean removing `Maria` and `Daniel` from the database. With Koog Persistence you can achieve that by providing a `RollbackToolRegistry` to `Persistence` feature config: ``` install(Persistence) { enableAutomaticPersistence = true rollbackToolRegistry = RollbackToolRegistry { // For every `createUser` tool call there will be a `removeUser` invocation in the reverse order // when rolling back to the desired execution point. // Note: `removeUser` tool should take the same exact arguments as `createUser`. // It's the developer's responsibility to make sure that `removeUser` invocation rolls back all side-effects of `createUser`: registerRollback(::createUser, ::removeUser) } } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(executor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Persistence.Feature, cfg -> { cfg.setEnableAutomaticPersistence(true); cfg.setRollbackToolRegistry( RollbackToolRegistry.builder() // For every tool in UserToolSet there will be a corresponding rollback tool // in UserRollbackToolSet, invoked in reverse order when rolling back. // UserRollbackToolSet methods must be annotated with @Reverts to link // them to the corresponding tools in UserToolSet. .registerRollbacks(new UserToolSet(), new UserRollbackToolSet()) .build() ); }) .build(); ``` ### Using extension functions The Agent Persistence feature provides convenient extension functions for working with checkpoints: ``` suspend fun example(context: AIAgentContext) { // Access the checkpoint feature val checkpointFeature = context.persistence() // Or perform an action with the checkpoint feature context.withPersistence { ctx -> // 'this' is the checkpoint feature createCheckpointAfterNode( agentContext = ctx, nodePath = ctx.executionInfo.path(), lastOutput = outputData, lastOutputType = outputType, checkpointId = ctx.runId, version = 0L ) } } ``` ``` // Access the persistence feature via PersistenceKt (the Kotlin extension function) Persistence persistence = PersistenceKt.persistence(context); // Use the persistence feature directly to create a checkpoint persistence.createCheckpointAfterNode( context, context.getExecutionInfo().path(), outputData, TypeToken.of(String.class), 0L, context.getRunId() ); ``` ## Advanced usage ### Custom storage providers You can implement custom storage providers by implementing the `PersistenceStorageProvider` interface: ``` class MyCustomStorageProvider : PersistenceStorageProvider { override suspend fun getCheckpoints(sessionId: String, filter: MyFilterType?): List { TODO("Not yet implemented") } override suspend fun saveCheckpoint(sessionId: String, agentCheckpointData: AgentCheckpointData) { TODO("Not yet implemented") } override suspend fun getLatestCheckpoint(sessionId: String, filter: MyFilterType?): AgentCheckpointData? { TODO("Not yet implemented") } } ``` ``` class MyCustomStorageProvider extends AsyncPersistenceStorageProvider { @Override public CompletableFuture> getCheckpointsAsync( String agentId, Object filter) { throw new UnsupportedOperationException("Not yet implemented"); } @Override public CompletableFuture saveCheckpointAsync( String agentId, AgentCheckpointData checkpointData) { throw new UnsupportedOperationException("Not yet implemented"); } @Override public CompletableFuture getLatestCheckpointAsync( String agentId, Object filter) { throw new UnsupportedOperationException("Not yet implemented"); } } ``` To use your custom provider in the feature configuration, set it as the storage when configuring the Agent Persistence feature in your agent. ``` install(Persistence) { storage = MyCustomStorageProvider() } ``` ``` AIAgent agent = AIAgent.builder() .promptExecutor(executor) .llmModel(OllamaModels.Meta.LLAMA_3_2) .install(Persistence.Feature, cfg -> { cfg.setStorage(new MyCustomStorageProvider()); }) .build(); ``` ### Setting execution points For advanced control, you can directly set the execution point of an agent: ``` suspend fun example(context: AIAgentContext) { // You can set the execution point after some node and provide an output from the node: context.persistence().setExecutionPointAfterNode( agentContext = context, nodePath = context.executionInfo.path(), messageHistory = customMessageHistory, output = customOutput ) } ``` ``` Persistence persistence = PersistenceKt.persistence(context); // Set the execution point after a node and provide an output from the node: persistence.setExecutionPointAfterNode( context, context.getExecutionInfo().path(), customMessageHistory, customOutput ); ``` This allows for more fine-grained control over the agent's state beyond just restoring from checkpoints. # Custom features Features provide a way to extend and enhance the functionality of AI agents at runtime. They are designed to be modular and composable, allowing you to mix and match them according to your needs. In addition to [features](../) that are available in Koog out of the box, you can also implement your own features by extending a proper feature interface. This page presents the basic building blocks for your own feature using the current Koog APIs. ## Feature interfaces Koog provides the following interfaces that you can extend to implement custom features: - [AIAgentGraphFeature](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.feature/-a-i-agent-graph-feature/index.html): Represents a feature specific to [agents that have defined workflows](../../agents/graph-based-agents/) (graph-based agents). - [AIAgentFunctionalFeature](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.feature/-a-i-agent-functional-feature/index.html): Represents a feature that can be used with [functional agents](../../agents/functional-agents/). - [AIAgentPlannerFeature](https://api.koog.ai/agents/agents-planner/ai.koog.agents.planner/-a-i-agent-planner-feature/index.html): Represents a feature type that is specific to [planner agents](../../agents/planner-agents/). Note To create a custom feature that can be installed in graph-based, functional, and planner agents, you need to implement all interfaces. ## Implementing custom features To implement a custom feature, you need to create a feature structure according to the following steps: 1. Create a feature class. 1. Define a configuration class. The configuration class is an extension of the [FeatureConfig](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.feature.config/-feature-config/index.html) class. 1. Create a companion object that implements some or all of the following interfaces: [AIAgentGraphFeature](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.feature/-a-i-agent-graph-feature/index.html), [AIAgentFunctionalFeature](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.feature/-a-i-agent-functional-feature/index.html), [AIAgentPlannerFeature](https://api.koog.ai/agents/agents-planner/ai.koog.agents.planner/-a-i-agent-planner-feature/index.html). 1. Give your feature a unique storage key that is used for feature identification and retrieval in agent pipelines. The key is used inside the internal map in an agent pipeline that includes all registered features for an agent. When you run an agent, it needs to process all registered features, and the key is used to retrieve the feature from this map. 1. Implement the required methods. The code sample below shows the general pattern for implementing a custom feature that can be installed in graph-based, functional, and planner agents: ``` class MyFeature(val someProperty: String) { class Config : FeatureConfig() { var configProperty: String = "default" } companion object Feature : AIAgentGraphFeature, AIAgentFunctionalFeature, AIAgentPlannerFeature { // Unique storage key for retrieval in contexts override val key = createStorageKey("my-feature") override fun createInitialConfig(agentConfig: AIAgentConfig): Config = Config() // Feature installation for graph-based agents override fun install(config: Config, pipeline: AIAgentGraphPipeline) : MyFeature { val feature = MyFeature(config.configProperty) pipeline.interceptAgentStarting(this) { context -> // Event handler implementation } return feature } // Feature installation for functional agents override fun install(config: Config, pipeline: AIAgentFunctionalPipeline) : MyFeature { val feature = MyFeature(config.configProperty) pipeline.interceptAgentStarting(this) { context -> // Event handler implementation } return feature } // Feature installation for planner agents override fun install(config: Config, pipeline: AIAgentPlannerPipeline) : MyFeature { val feature = MyFeature(config.configProperty) pipeline.interceptAgentStarting(this) { context -> // Event handler implementation } return feature } } } ``` When creating an agent, install your feature using the `install` method: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), systemPrompt = "You are a helpful assistant. Answer user questions concisely.", llmModel = OpenAIModels.Chat.GPT4o ) { install(MyFeature) { configProperty = "value" } } ``` ### Pipeline interceptors Interceptors represent various points in the agent lifecycle where you can hook into the agent execution pipeline to implement your custom logic. Koog includes a range of predefined interceptors that you can use to observe various events. Below are the interceptors that you can register from your feature’s `install` method. The listed interceptors are grouped by type and apply to graph-based, functional, and planner agent pipelines. To reduce noise and optimize cost when developing actual features, register only the interceptors you need for the feature. Agent and environment lifecycle: - `interceptEnvironmentCreated`: Transforms the agent environment when it’s created. - `interceptAgentStarting`: Invoked before the start of an agent run. - `interceptAgentCompleted`: Invoked when an agent run completes successfully. - `interceptAgentExecutionFailed`: Invoked when an agent run fails. - `interceptAgentClosing`: Invoked just before the agent run closes (cleanup point). Strategy lifecycle: - `interceptStrategyStarting`: Invoked before the start of a strategy execution. - `interceptStrategyCompleted`: Invoked when a strategy execution completes successfully. LLM call lifecycle: - `interceptLLMCallStarting`: Invoked before an LLM call. - `interceptLLMCallFailed`: Invoked when an LLM call fails (the underlying prompt executor or moderation call throws). - `interceptLLMCallCompleted`: Invoked after an LLM call. LLM streaming lifecycle: - `interceptLLMStreamingStarting`: Invoked before streaming starts. - `interceptLLMStreamingFrameReceived`: Invoked for each received stream frame. - `interceptLLMStreamingFailed`: Invoked when streaming fails. - `interceptLLMStreamingCompleted`: Invoked after streaming completes. Tool call lifecycle: - `interceptToolCallStarting`: Invoked before a tool call. - `interceptToolValidationFailed`: Invoked when tool input validation fails. - `interceptToolCallFailed`: Invoked when tool execution fails. - `interceptToolCallCompleted`: Invoked after the tool completes (with a result). #### Interceptors specific to graph-based agents The following interceptors are available only on `AIAgentGraphPipeline` and let you observe node and subgraph lifecycle events. Node execution lifecycle: - `interceptNodeExecutionStarting`: Invoked before a node starts executing. - `interceptNodeExecutionCompleted`: Invoked after a node finishes executing. - `interceptNodeExecutionFailed`: Invoked when a node execution fails with an error. Subgraph execution lifecycle: - `interceptSubgraphExecutionStarting`: Invoked right before a subgraph starts executing. - `interceptSubgraphExecutionCompleted`: Invoked after a subgraph execution completes. - `interceptSubgraphExecutionFailed`: Invoked when a subgraph execution fails. For a feature to handle a specific type of event, it needs to register the corresponding pipeline interceptor. ### Filtering agent events When installing a feature in an agent, you may not want to handle all events that are registered in the feature. To filter out some events, you apply filters using the [FeatureConfig.setEventFilter](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.feature.config/-feature-config/set-event-filter.html) function. The following example shows how you can allow only LLM call start and end events for a feature: ``` install(MyFeature) { setEventFilter { context -> context.eventType is AgentLifecycleEventType.LLMCallStarting || context.eventType is AgentLifecycleEventType.LLMCallCompleted } } ``` #### Disabling event filtering for a feature If your feature logic relies on the complete agent event structure, event filtering can cause unexpected behavior. To prevent this, you need to disable event filtering when implementing the feature by overriding `setEventFilter` in your feature configuration to ignore any custom filters set when installing the feature. An example of a feature that relies on processing the entire agent event stream is [OpenTelemetry](../open-telemetry/), as it uses the complete agent event structure to compose an inherited structure of spans. Here is an example of how to disable event filtering for a feature: ``` class MyFeatureConfig : FeatureConfig() { override fun setEventFilter(filter: (AgentLifecycleEventContext) -> Boolean) { // Deactivate event filtering for the feature throw UnsupportedOperationException("Event filtering is not allowed.") } } ``` ## Example: A basic logging feature The following example shows how to implement a basic logging feature that logs agent lifecycle events. As the feature should be available graph-based, functional, and planner agents, interceptors that are common to all agent types are implemented in the `installCommon` method to avoid code duplication. The interceptors that are specific to individual agent types are implemented in the `installGraphPipeline`, `installFunctionalPipeline`, and `installPlannerPipeline` methods. ``` class LoggingFeature(val loggerName: String) { class Config : FeatureConfig() { var loggerName: String = "agent-logs" } companion object Feature : AIAgentGraphFeature, AIAgentFunctionalFeature, AIAgentPlannerFeature { override val key = createStorageKey("logging-feature") override fun createInitialConfig(agentConfig: AIAgentConfig): Config = Config() override fun install(config: Config, pipeline: AIAgentGraphPipeline) : LoggingFeature { val logging = LoggingFeature(config.loggerName) val logger = KotlinLogging.logger(config.loggerName) installGraphPipeline(pipeline, logger) return logging } override fun install(config: Config, pipeline: AIAgentFunctionalPipeline) : LoggingFeature { val logging = LoggingFeature(config.loggerName) val logger = KotlinLogging.logger(config.loggerName) installFunctionalPipeline(pipeline, logger) return logging } override fun install(config: Config, pipeline: AIAgentPlannerPipeline) : LoggingFeature { val logging = LoggingFeature(config.loggerName) val logger = KotlinLogging.logger(config.loggerName) installPlannerPipeline(pipeline, logger) return logging } private fun installCommon( pipeline: AIAgentPipeline, logger: KLogger, ) { pipeline.interceptAgentStarting(this) { e -> logger.info { "Agent starting: runId=${e.runId}" } } pipeline.interceptStrategyStarting(this) { e -> logger.info { "Strategy ${e.strategy.name} starting" } } pipeline.interceptLLMCallStarting(this) { e -> logger.info { "Making LLM call with ${e.tools.size} tools" } } pipeline.interceptLLMCallCompleted(this) { e -> logger.info { "Received response: ${e.response != null}" } } } private fun installGraphPipeline( pipeline: AIAgentGraphPipeline, logger: KLogger, ) { installCommon(pipeline, logger) pipeline.interceptNodeExecutionStarting(this) { e -> logger.info { "Node ${e.node.name} input: ${e.input}" } } pipeline.interceptNodeExecutionCompleted(this) { e -> logger.info { "Node ${e.node.name} output: ${e.output}" } } } private fun installFunctionalPipeline( pipeline: AIAgentFunctionalPipeline, logger: KLogger ) { installCommon(pipeline, logger) } private fun installPlannerPipeline( pipeline: AIAgentPlannerPipeline, logger: KLogger ) { installCommon(pipeline, logger) } } } ``` Here is an example of how to install the custom logging feature in an agent. The example shows a basic feature installation, along with the custom configuration property `loggerName` that lets you specify the name of the logger: ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), systemPrompt = "You are a helpful assistant. Answer user questions concisely.", llmModel = OpenAIModels.Chat.GPT4o ) { install(LoggingFeature) { loggerName = "my-custom-logger" } } agent.run("What is Kotlin?") ``` # OpenTelemetry support This page provides details about the support for OpenTelemetry with the Koog agentic framework for tracing and monitoring your AI agents. ## Overview OpenTelemetry is an observability framework that provides tools for generating, collecting, and exporting telemetry data (traces) from your applications. The Koog OpenTelemetry feature allows you to instrument your AI agents to collect telemetry data, which can help you: - Monitor agent performance and behavior - Debug issues in complex agent workflows - Visualize the execution flow of your agents - Track LLM calls and tool usage - Analyze agent behavior patterns ## Key OpenTelemetry concepts - **Spans**: spans represent individual units of work or operations within a distributed trace. They indicate the beginning and end of a specific activity in an application, such as an agent execution, a function call, an LLM call, or a tool call. - **Attributes**: attributes provide metadata about a telemetry-related item such as a span. Attributes are represented as key-value pairs. - **Events**: events are specific points in time during the lifetime of a span (span-related events) that represent something potentially noteworthy that happened. - **Exporters**: exporters are components responsible for sending the collected telemetry data to various backends or destinations. - **Collectors**: collectors receive, process, and export telemetry data. They act as intermediaries between your applications and your observability backend. - **Samplers**: samplers determine whether a trace should be recorded based on the sampling strategy. They are used to manage the volume of telemetry data. - **Resources**: resources represent entities that produce telemetry data. They are identified by resource attributes, which are key-value pairs that provide information about the resource. The OpenTelemetry feature in Koog automatically creates spans for various agent events, including: - Agent execution start and end - Node execution - LLM calls - Tool calls ## Installation To use OpenTelemetry with Koog, add the OpenTelemetry feature to your agent: ``` val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.GPT4o, systemPrompt = "You are a helpful assistant.", installFeatures = { install(OpenTelemetry) { // Configuration options go here } } ) ``` ``` var agent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OpenAIModels.Chat.GPT4o) .systemPrompt("You are a helpful assistant.") .install(OpenTelemetry.Feature, config -> { // Configuration options go here }) .build(); ``` ## Configuration ### Basic configuration Here is the full list of available properties that you set when configuring the OpenTelemetry feature in an agent: | Name | Data type | Default value | Description | | ---------------- | --------- | ---------------------------- | ---------------------------------------------------------------------------- | | `serviceName` | `String` | `ai.koog` | The name of the service being instrumented. | | `serviceVersion` | `String` | Current Koog library version | The version of the service being instrumented. | | `isVerbose` | `Boolean` | `false` | Whether to enable verbose logging for debugging OpenTelemetry configuration. | | `tracer` | `Tracer` | | The OpenTelemetry tracer instance used for creating spans. | Note The `tracer` property is a public property that you can access, but it is configured automatically based on the exporters and resource attributes you provide. The `OpenTelemetryConfig` class also includes methods that represent actions related to different configuration items. Here is an example of installing the OpenTelemetry feature with a basic set of configuration items: ``` install(OpenTelemetry) { // Set your service configuration setServiceInfo("my-agent-service", "1.0.0") // Add the Logging exporter addSpanExporter(LoggingSpanExporter.create()) } ``` ``` install(OpenTelemetry.Feature, config -> { // Set your service configuration config.setServiceInfo("my-agent-service", "1.0.0"); // Add the Logging exporter config.addSpanExporter(LoggingSpanExporter.create()); }) ``` For a reference of available methods, see the sections below. #### setServiceInfo Sets the service information including name and version. Takes the following arguments: | Name | Data type | Required | Default value | Description | | ---------------- | --------- | -------- | ------------- | ---------------------------------------------- | | `serviceName` | String | Yes | | The name of the service being instrumented. | | `serviceVersion` | String | Yes | | The version of the service being instrumented. | #### addSpanExporter Adds a span exporter to send telemetry data to external systems. Takes the following argument: | Name | Data type | Required | Default value | Description | | ---------- | -------------- | -------- | ------------- | ----------------------------------------------------------------------------- | | `exporter` | `SpanExporter` | Yes | | The `SpanExporter` instance to be added to the list of custom span exporters. | Both Kotlin SDK (`io.opentelemetry.kotlin.tracing.export.SpanExporter`) and Java SDK (`io.opentelemetry.sdk.trace.export.SpanExporter`) exporters are accepted. Java SDK exporters are automatically converted via the compat bridge. The exporter is registered behind a `batchSpanProcessor` — the OpenTelemetry-recommended default for production: spans are buffered and flushed on a worker so the agent never blocks on network I/O when a span ends. If you need full control over the processor (custom batching parameters, a simple processor for tests, or a composite processor), use [`addSpanProcessor`](#addspanprocessor) instead. #### addSpanProcessor Registers a `SpanProcessor` directly, bypassing the `batchSpanProcessor` wrapping done by [`addSpanExporter`](#addspanexporter). The factory runs inside the SDK's `TraceExportConfigDsl` scope, which exposes `batchSpanProcessor`, `simpleSpanProcessor`, and `compositeSpanProcessor`. Takes the following argument: | Name | Data type | Required | Default value | Description | | --------- | ------------------------------------------ | -------- | ------------- | ---------------------------------------------------- | | `factory` | `TraceExportConfigDsl.() -> SpanProcessor` | Yes | | Lambda that returns the `SpanProcessor` to register. | Reach for this when: - You want custom batching parameters: `addSpanProcessor { batchSpanProcessor(exporter, scheduleDelayMs = 500) }`. - You want spans flushed synchronously (useful in tests): `addSpanProcessor { simpleSpanProcessor(exporter) }`. - You want to fan out to several processors at once: `addSpanProcessor { compositeSpanProcessor(p1, p2) }`. For Java SDK exporters, wrap with `toOtelKotlinSpanExporter()` from the compat package first. #### addResourceAttributes Adds resource attributes to provide additional context about the service. Takes the following argument: | Name | Data type | Required | Default value | Description | | ------------ | ------------------ | -------- | ------------- | ------------------------------------------------------------------------------------------------------------------------------------ | | `attributes` | `Map` | Yes | | The key-value pairs that provide additional details about the service. Supported value types: `String`, `Long`, `Double`, `Boolean`. | #### setVerbose Enables or disables verbose logging. Takes the following argument: | Name | Data type | Required | Default value | Description | | --------- | --------- | -------- | ------------- | --------------------------------------------------------------- | | `verbose` | `Boolean` | Yes | `false` | If true, the application collects more detailed telemetry data. | Note Some content of OpenTelemetry spans is masked by default for security reasons. For example, LLM messages are masked as `HIDDEN:non-empty` instead of the actual message content. To get the content, set the value of the `verbose` argument to `true`. #### addMetricExporter Adds a metric exporter to send metric data to external systems. Takes the following arguments: | Name | Data type | Required | Default value | Description | | --------------- | ---------------- | -------- | ------------- | ---------------------------------------------------------------------------- | | `exporter` | `MetricExporter` | Yes | | The `MetricExporter` instance to register with a periodic metric reader. | | `meterInterval` | `Duration` | No | `1s` | The interval between metric reads. Also available as a `java.time.Duration`. | If no metric exporter is registered, metrics are disabled. Metrics are a JVM-only capability backed by the Java OpenTelemetry SDK; the Kotlin Multiplatform SDK 0.2.0 does not yet expose a metrics API. #### addMetricFilter Restricts the attribute keys that are reported for a specific metric instrument. This installs an OpenTelemetry `View` that drops any attribute not listed. Takes the following arguments: | Name | Data type | Required | Default value | Description | | -------------- | ------------- | -------- | ------------- | ----------------------------------------------------------- | | `metricName` | `String` | Yes | | The name of the metric instrument to apply the filter to. | | `keysToRetain` | `Set` | Yes | | The attribute keys that should be retained for this metric. | Use this to keep high-cardinality attributes (for example, request identifiers) from blowing up your metric backend while still exporting the metric itself. ### Advanced configuration For more advanced configuration, you can also customize resource attributes to add more information about the process that is producing telemetry data. ``` install(OpenTelemetry) { // Set your service configuration setServiceInfo("my-agent-service", "1.0.0") // Add the Logging exporter addSpanExporter(LoggingSpanExporter.create()) // Add resource attributes addResourceAttributes(mapOf( "custom.attribute" to "custom-value") ) } ``` ``` install(OpenTelemetry.Feature, config -> { // Set your service configuration config.setServiceInfo("my-agent-service", "1.0.0"); // Add the Logging exporter config.addSpanExporter(LoggingSpanExporter.create()); // Add resource attributes config.addResourceAttributes(Map.of( "custom.attribute", "custom-value" )); }) ``` #### Resource attributes Resource attributes represent additional information about a process producing telemetry data. Koog includes a set of resource attributes that are set by default: - `service.name` - `service.version` - `service.instance.time` - `os.type` - `os.version` - `os.arch` The default value of the `service.name` attribute is `ai.koog`, while the default `service.version` value is the currently used Koog library version. In addition to default resource attributes, you can also add custom attributes. To add a custom attribute to an OpenTelemetry configuration in Koog, use the `addResourceAttributes()` method in an OpenTelemetry configuration that takes a key and a value as its arguments. ``` addResourceAttributes(mapOf( "custom.attribute" to "custom-value") ) ``` ``` config.addResourceAttributes(Map.of( "custom.attribute", "custom-value" )); ``` ## What gets traced The OpenTelemetry feature captures the following agent activity: - **Agent lifecycle events**: agent start, stop, errors - **LLM interactions**: prompts, responses, token usage, latency, and failures (spans are marked with span status `ERROR` and `error.type` when an LLM call throws) - **Tool calls**: execution traces for tool invocations - **System context**: metadata such as model name, environment, Koog version By default, the contents of LLM prompts and responses are masked in exported spans to avoid exposing sensitive data. To include the full content, call [`setVerbose(true)`](#setverbose). For a detailed breakdown of individual span types and attributes, see [Span types and attributes](#span-types-and-attributes). ## Span types and attributes The OpenTelemetry feature automatically creates different types of spans to track various operations in your agent: - **CreateAgentSpan**: created when you run an agent, closed when the agent is closed or the process is terminated. - **InvokeAgentSpan**: the invocation of an agent. - **StrategySpan**: the execution of an agent's strategy (the top-level execution flow). - **NodeExecuteSpan**: the execution of a node in the agent's strategy. This is a custom, Koog-specific span. - **SubgraphExecuteSpan**: the execution of a subgraph within the agent strategy. This is a custom, Koog-specific span. - **InferenceSpan**: an LLM call. - **ExecuteToolSpan**: a tool call. - **McpClientSpan**: an MCP (Model Context Protocol) client operation. This span follows OpenTelemetry semantic conventions for MCP. Spans are organized in a nested, hierarchical structure. Here is an example of a span structure: ``` CreateAgentSpan InvokeAgentSpan StrategySpan NodeExecuteSpan InferenceSpan NodeExecuteSpan ExecuteToolSpan SubgraphExecuteSpan NodeExecuteSpan InferenceSpan ``` ### Span attributes Span attributes provide metadata related to a span. Each span has its set of attributes, while some spans can also repeat attributes. Koog supports a list of predefined attributes that follow OpenTelemetry's [Semantic conventions for generative AI spans](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/). For example, the conventions define an attribute named `gen_ai.conversation.id`, which is usually a required attribute for a span. In Koog, the value of this attribute is the unique identifier for an agent run, that is automatically set when you call the `agent.run()` method. In addition, Koog also includes custom, Koog-specific attributes. You can recognize most of these attributes by the `koog.` prefix. Here are the available custom attributes: - `koog.strategy.name`: the name of the agent strategy. A strategy is a Koog-related entity that describes the purpose of the agent. Used in the `StrategySpan` span. - `koog.node.id`: the identifier (name) of the node being executed. Used in the `NodeExecuteSpan` span. - `koog.node.input`: the input passed to the node at the beginning of execution. Present on `NodeExecuteSpan` when node starts. - `koog.node.output`: the output produced by the node upon completion. Present on `NodeExecuteSpan` when node completes successfully. - `koog.subgraph.id`: the identifier (name) of the subgraph being executed. Used in the `SubgraphExecuteSpan` span. - `koog.subgraph.input`: the input passed to the subgraph at the beginning of execution. Present on `SubgraphExecuteSpan` when subgraph starts. - `koog.subgraph.output`: the output produced by the subgraph upon completion. Present on `SubgraphExecuteSpan` when subgraph completes successfully. - `koog.moderation.result`: the JSON-encoded moderation outcome for the LLM call when one is available. Present on the `InferenceSpan` only when moderation was performed for the call. The OpenTelemetry GenAI semantic conventions do not define a moderation attribute, so Koog publishes this under the `koog.` namespace. ### Message content Per the OpenTelemetry GenAI semantic conventions, message content is carried on the `InferenceSpan` via two span attributes rather than per-message events: - `gen_ai.input.messages`: a JSON array of the messages sent to the model (system / user / assistant / tool roles). - `gen_ai.output.messages`: a JSON array of the messages returned by the model. Earlier versions of Koog emitted per-message OpenTelemetry events (`gen_ai.system.message`, `gen_ai.user.message`, `gen_ai.assistant.message`, `gen_ai.tool.message`, `gen_ai.choice`) to capture message content. Those events have been removed from the OpenTelemetry GenAI specification and are no longer emitted by Koog. Backends that still expect the indexed `gen_ai.prompt.{i}.*` / `gen_ai.completion.{i}.*` shape (Langfuse, Weave) continue to receive it through the corresponding span adapters. ## Metrics In addition to spans, the OpenTelemetry feature emits metrics that follow OpenTelemetry's [Semantic conventions for GenAI metrics](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/). Metrics are exported through the meter provider configured via [addMetricExporter](#addmetricexporter); if no exporter is registered, a console `LoggingMetricExporter` is used by default. The following instruments are registered: | Name | Instrument | Unit | Description | | ------------------------------------ | ---------- | --------- | -------------------------------------------------------------------------------------------------------- | | `gen_ai.client.token.usage` | Histogram | `{token}` | Token usage reported for each LLM call, split by `gen_ai.token.type` (`input`/`output`). | | `gen_ai.client.operation.duration` | Histogram | `s` | Duration of GenAI operations — both `text_completion` (LLM calls) and `execute_tool` (tool invocations). | | `koog.gen_ai.client.tool.call.count` | Counter | `{call}` | Koog-specific counter of tool calls performed by the agent, labelled by tool name and call status. | Explicit histogram bucket boundaries are provided as advice in line with the semantic conventions: - `gen_ai.client.token.usage`: `[1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864]` - `gen_ai.client.operation.duration`: `[0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, 40.96, 81.92]` ### gen_ai.provider.name Every data point carries a `gen_ai.provider.name` attribute: - For `text_completion` operations, the value is the LLM provider id (for example, `openai`, `anthropic`). - For `execute_tool` operations, the value is `koog`, because tool execution happens in-process rather than against a third-party provider. MCP tool executions keep this value and surface MCP-specific details through separate `mcp.*` attributes on the corresponding span, so tool metrics stay at low cardinality. ### error.type `error.type` is set only on failed `gen_ai.client.operation.duration` data points, per the GenAI semconv requirement. The value is the canonical Java class name of the error that caused the failure, so it is bounded by the exception hierarchy and safe to use as a metric dimension: - Subclasses of `AIAgentError` — for `execute_tool` failures and tool validation failures. - Any `Throwable` raised by the LLM client or agent runtime — for `text_completion` failures that surface through the agent-level failure hook. - `_OTHER` — fallback when an in-flight operation is flushed at agent close without an associated error. The attribute is not set on successful operations. ### restrictToolNameCardinality Tool metrics are labeled with `gen_ai.tool.name`. If you expose tools whose names are dynamic or user-generated, the tool-name cardinality can grow without bound. Use `restrictToolNameCardinality` to map any name outside an allow-list to a single fallback value. For metric-specific attribute filtering that applies to any instrument and any attribute key, use [addMetricFilter](#addmetricfilter). ## Exporters Exporters send collected telemetry data to an OpenTelemetry Collector or other types of destinations or backend implementations. To add an exporter, use the `addSpanExporter()` method when installing the OpenTelemetry feature. The method takes the following argument: | Name | Data type | Required | Default | Description | | ---------- | ------------ | -------- | ------- | --------------------------------------------------------------------------- | | `exporter` | SpanExporter | Yes | | The SpanExporter instance to be added to the list of custom span exporters. | The sections below provide information about some of the most commonly used exporters. Koog accepts both Kotlin SDK and Java SDK exporters — Java SDK exporters are automatically converted via the compat bridge. Note If you do not configure any custom exporters, Koog will use a console stdout exporter by default. This helps during local development and debugging. ### Logging exporter A logging exporter that outputs trace information to the console. `LoggingSpanExporter` (`io.opentelemetry.exporter.logging.LoggingSpanExporter`) is a part of the `opentelemetry-java` SDK. This type of export is useful for development and debugging purposes. ``` install(OpenTelemetry) { // Add the logging exporter addSpanExporter(LoggingSpanExporter.create()) // Add more exporters as needed } ``` ``` install(OpenTelemetry.Feature, config -> { // Add the logging exporter config.addSpanExporter(LoggingSpanExporter.create()); // Add more exporters as needed }) ``` ### OpenTelemetry HTTP exporter OpenTelemetry HTTP exporter (`OtlpHttpSpanExporter`) is a part of the `opentelemetry-java` SDK (`io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter`) and sends span data to a backend through HTTP. ``` install(OpenTelemetry) { // Add OpenTelemetry HTTP exporter addSpanExporter( OtlpHttpSpanExporter.builder() // Set the maximum time to wait for the collector to process an exported batch of spans .setTimeout(30, TimeUnit.SECONDS) // Set the OpenTelemetry endpoint to connect to .setEndpoint("http://localhost:3000/api/public/otel/v1/traces") // Add the authorization header .addHeader("Authorization", "Basic $AUTH_STRING") .build() ) } ``` ``` install(OpenTelemetry.Feature, config -> { // Add OpenTelemetry HTTP exporter config.addSpanExporter( OtlpHttpSpanExporter.builder() // Set the maximum time to wait for the collector to process an exported batch of spans .setTimeout(30, TimeUnit.SECONDS) // Set the OpenTelemetry endpoint to connect to .setEndpoint("http://localhost:3000/api/public/otel/v1/traces") // Add the authorization header .addHeader("Authorization", "Basic " + AUTH_STRING) .build() ); }) ``` ### OpenTelemetry gRPC exporter OpenTelemetry gRPC exporter (`OtlpGrpcSpanExporter`) is a part of the `opentelemetry-java` SDK (`io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter`). It exports telemetry data to a backend through gRPC and lets you define the host and port of the backend, collector, or endpoint that receives the data. The default port is `4317`. ``` install(OpenTelemetry) { // Add OpenTelemetry gRPC exporter addSpanExporter( OtlpGrpcSpanExporter.builder() // Set the host and the port .setEndpoint("http://localhost:4317") .build() ) } ``` ``` install(OpenTelemetry.Feature, config -> { // Add OpenTelemetry gRPC exporter config.addSpanExporter( OtlpGrpcSpanExporter.builder() // Set the host and the port .setEndpoint("http://localhost:4317") .build() ); }) ``` ## Integration with Langfuse Langfuse provides trace visualization and analytics for LLM/agent workloads. You can configure Koog to export OpenTelemetry traces directly to Langfuse using a helper function: ``` install(OpenTelemetry) { addLangfuseExporter( langfuseUrl = "https://cloud.langfuse.com", langfusePublicKey = "...", langfuseSecretKey = "..." ) } ``` ``` install(OpenTelemetry.Feature, config -> { LangfuseKt.addLangfuseExporter( config, "https://cloud.langfuse.com", "...", "..." ); }) ``` Please read the [full documentation](opentelemetry-langfuse-exporter/) about integration with Langfuse. ## Integration with W&B Weave W&B Weave provides trace visualization and analytics for LLM/agent workloads. Integration with W&B Weave can be configured via a predefined exporter: ``` install(OpenTelemetry) { addWeaveExporter( weaveOtelBaseUrl = "https://trace.wandb.ai", weaveEntity = "my-team", weaveProjectName = "my-project", weaveApiKey = "..." ) } ``` ``` install(OpenTelemetry.Feature, config -> { WeaveKt.addWeaveExporter( config, "https://trace.wandb.ai", "my-team", "my-project", "..." ); }) ``` Please read the [full documentation](opentelemetry-weave-exporter/) about integration with W&B Weave. ## Integration with Datadog Datadog provides monitoring, observability, and analytics for cloud-scale applications. Integration with Datadog can be configured via a predefined exporter: ``` install(OpenTelemetry) { addDatadogExporter( datadogApiKey = "...", url = "datadoghq.com" ) } ``` ``` install(OpenTelemetry.Feature, config -> { DatadogKt.addDatadogExporter( config, "...", // datadogApiKey "datadoghq.com" // url ); }) ``` Please read the [full documentation](opentelemetry-datadog-exporter/) about integration with Datadog. ## Integration with Jaeger Jaeger is a popular distributed tracing system that works with OpenTelemetry. The `opentelemetry` directory within `examples` in the Koog repository includes an example of using OpenTelemetry with Jaeger and Koog agents. ### Prerequisites To test OpenTelemetry with Koog and Jaeger, start the Jaeger OpenTelemetry all-in-one process using the provided `docker-compose.yaml` file, by running the following command: ``` docker compose up -d ``` The provided Docker Compose YAML file includes the following content: ``` # docker-compose.yaml services: jaeger-all-in-one: image: jaegertracing/all-in-one:1.39 container_name: jaeger-all-in-one environment: - COLLECTOR_OTLP_ENABLED=true ports: - "4317:4317" - "16686:16686" ``` To access the Jaeger UI and view your traces, open `http://localhost:16686`. ### Example To export telemetry data for use in Jaeger, the example uses `LoggingSpanExporter` (`io.opentelemetry.exporter.logging.LoggingSpanExporter`) and `OtlpGrpcSpanExporter` (`io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter`) from the `opentelemetry-java` SDK. Here is the full code sample: ``` fun main() = runBlocking { val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.O4Mini, systemPrompt = "You are a code assistant. Provide concise code examples." ) { install(OpenTelemetry) { // Add a console logger for local debugging addSpanExporter(LoggingSpanExporter.create()) // Send traces to OpenTelemetry collector addSpanExporter( OtlpGrpcSpanExporter.builder() .setEndpoint("http://localhost:4317") .build() ) } } agent.use { agent -> println("Running the agent with OpenTelemetry tracing...") val result = agent.run("Tell me a joke about programming") println("Agent run completed with result: '$result'." + "\nCheck Jaeger UI at http://localhost:16686 to view traces") } } ``` ``` public static void main(String[] args) { var agent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OpenAIModels.Chat.O4Mini) .systemPrompt("You are a code assistant. Provide concise code examples.") .install(OpenTelemetry.Feature, config -> { // Add a console logger for local debugging config.addSpanExporter(LoggingSpanExporter.create()); // Send traces to OpenTelemetry collector config.addSpanExporter( OtlpGrpcSpanExporter.builder() .setEndpoint("http://localhost:4317") .build() ); }) .build(); System.out.println("Running the agent with OpenTelemetry tracing..."); var result = agent.run("Tell me a joke about programming"); System.out.println( "Agent run completed with result: '" + result + "'." + "\nCheck Jaeger UI at http://localhost:16686 to view traces" ); } ``` ## Troubleshooting ### Common issues 1. **No traces appearing in the backend** - Confirm all required environment variables are set and exported in your shell. - Verify that your API key or secret is valid, has not been revoked, and has write/trace permissions. - Ensure the service is running and the OpenTelemetry port (4317) is accessible. - Check that the exporter is configured with the correct endpoint. - Wait a few seconds after agent execution — traces may not appear instantly. 1. **Connection issues** - Confirm your environment can reach the exporter's intake endpoint. - Check for firewall or proxy settings that block outbound HTTPS traffic. 1. **Missing spans or incomplete traces** - Verify that the agent execution completes successfully. - Ensure that you're not closing the application too quickly after agent execution. - Add a delay after agent execution to allow time for spans to be exported. 1. **Excessive number of spans** - Consider using a different sampling strategy by configuring the `sampler` property. - For example, use `Sampler.traceIdRatioBased(0.1)` to sample only 10% of traces. 1. **Span adapters override each other** - Currently, the OpenTelemetry agent feature does not support applying multiple span adapters [KG-265](https://youtrack.jetbrains.com/issue/KG-265/Adding-Weave-exporter-breaks-Langfuse-exporter). ## MCP (Model Context Protocol) telemetry support Koog provides comprehensive OpenTelemetry instrumentation for MCP operations following the [official OpenTelemetry semantic conventions for MCP](https://github.com/open-telemetry/semantic-conventions/pull/2083). ### Overview The MCP telemetry support includes: - **Automatic enrichment** of tool execution spans with MCP-specific attributes - **Client-side instrumentation** for MCP client operations (tools/call) - **Full semantic convention compliance** with all required, conditionally required, and recommended attributes ### MCP attributes MCP telemetry follows OpenTelemetry semantic conventions and includes the following attribute groups: **Required attributes:** - `mcp.method.name`: The MCP method name (e.g., "tools/call") **Conditionally required attributes:** - `gen_ai.tool.name`: When operation involves a tool - `gen_ai.prompt.name`: When operation involves a prompt - `jsonrpc.request.id`: When executing a request (not a notification) - `error.type`: When operation fails **Recommended attributes:** - `mcp.session.id`: Session identifier - `mcp.protocol.version`: MCP protocol version (e.g., "2025-06-18") - `network.transport`: Transport type ("pipe" for stdio, "tcp" for HTTP) - `server.address` and `server.port`: For client operations ### Span naming convention MCP spans follow the naming convention: `{mcp.method.name} {target}` Where `{target}` is the tool name or prompt name when applicable. Examples: - `"tools/call search"` - calling a tool named "search" ### Best practices - **Always set session IDs** when working with persistent MCP sessions to enable session tracking - **Propagate request IDs** from JSON-RPC requests for complete request tracing - **Monitor metrics** to identify performance bottlenecks in MCP operations ### Example: Full MCP client with telemetry ``` // Create MCP tools registry val toolRegistry = McpToolRegistryProvider.fromSseUrl("http://localhost:3000") // Create agent with OpenTelemetry enabled and pass the tool registry val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.GPT4o, systemPrompt = "You are a helpful assistant.", toolRegistry = toolRegistry ) { install(OpenTelemetry) { setServiceInfo("mcp-agent-service", "1.0.0") addSpanExporter(LoggingSpanExporter.create()) } } // Run agent - MCP tool calls will be automatically instrumented agent.use { it.run("Use the search tool to find information") } ``` This setup provides complete observability for MCP operations with minimal code changes, following OpenTelemetry best practices and semantic conventions. # Datadog exporter Koog emits agent traces using [OpenTelemetry](https://opentelemetry.io/), an open standard for observability data. To ship those traces to [Datadog](https://www.datadoghq.com/), Koog includes a built-in OpenTelemetry exporter — no manual instrumentation required. Once connected, Datadog's [OpenTelemetry support](https://docs.datadoghq.com/opentelemetry/) lets you visualize, analyze, and debug how your agents interact with LLMs, tools, and external APIs. ______________________________________________________________________ ## Setup instructions 1. Create a Datadog account at 1. Get your API key from [Organization Settings > API Keys](https://app.datadoghq.com/organization-settings/api-keys) 1. Provide your API key — either as a parameter to [`addDatadogExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.datadog/add-datadog-exporter.html), or via an environment variable: ``` export DD_API_KEY="" ``` 1. (Optional) To use a Datadog region other than US1 (`datadoghq.com`), pass the site as a parameter to [`addDatadogExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.datadog/add-datadog-exporter.html), or set an environment variable: ``` export DD_SITE="datadoghq.eu" ``` Supported sites: | Site | Region | | ------------------- | ------------- | | `datadoghq.com` | US1 (default) | | `datadoghq.eu` | EU1 | | `us3.datadoghq.com` | US3 | | `us5.datadoghq.com` | US5 | | `ap1.datadoghq.com` | AP1 (Japan) | ## Configuration To enable Datadog export, install the **OpenTelemetry feature** and call [`addDatadogExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.datadog/add-datadog-exporter.html). ### Basic example ``` fun main() = runBlocking { val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a code assistant. Provide concise code examples." ) { install(OpenTelemetry) { addDatadogExporter() } } println("Running agent with Datadog tracing") val result = agent.run("Tell me a joke about programming") println("Result: $result\nSee traces in Datadog LLM Observability") } ``` ``` public static void main(String[] args) { var agent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OpenAIModels.Chat.GPT4oMini) .systemPrompt("You are a code assistant. Provide concise code examples.") .install(OpenTelemetry.Feature, config -> DatadogKt.addDatadogExporter(config) ) .build(); System.out.println("Running agent with Datadog tracing"); var result = agent.run("Tell me a joke about programming"); System.out.println("Result: " + result + "\nSee traces in Datadog LLM Observability"); } ``` ## Trace attributes When Koog sends agent activity to Datadog, it does so as a series of *spans* — individual records of work, such as an LLM call or a tool execution. Related spans are grouped into a *trace*, which represents a complete agent run from start to finish. [`addDatadogExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.datadog/add-datadog-exporter.html) accepts a `resourceAttributes` parameter — a map of key-value pairs describing the application emitting the traces. These are attached to every span, making it easy to filter and group traces in Datadog by properties such as environment or version. Common attributes to include: - **env**: Environment name (for example, `production`, `staging`, or `development`) - **service.name**: Name of your service or application - **version**: Application version, useful for comparing behavior across deployments ### Example with trace attributes ``` fun main() = runBlocking { val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a helpful assistant." ) { install(OpenTelemetry) { addDatadogExporter( url = "datadoghq.eu", // Use EU region resourceAttributes = mapOf( "env" to "production", "service.name" to "my-agent", "version" to "1.0.0" ) ) } } println("Running agent with Datadog tracing") agent.run("What is Kotlin?") } ``` ``` public static void main(String[] args) { var agent = AIAgent.builder() .promptExecutor(promptExecutor) .systemPrompt("You are a helpful assistant.") .llmModel(OpenAIModels.Chat.GPT4oMini) .install(OpenTelemetry.Feature, config -> DatadogKt.addDatadogExporter( config, null, // datadogApiKey: use DD_API_KEY env var "datadoghq.eu" // url: use EU region )) .build(); System.out.println("Running agent with Datadog tracing"); agent.run("What is Kotlin?"); } ``` Note Setting `resourceAttributes` from Java is currently not supported because the underlying Kotlin function carries a [`kotlin.time.Duration`](https://kotlinlang.org/api/latest/jvm/stdlib/kotlin.time/-duration/) parameter (a value class) that causes JVM-name mangling on all overloads including parameters after it. Use the Kotlin example above when you need `resourceAttributes`. ## Sending to multiple backends To send traces to Datadog and another backend at the same time, register Datadog through [`addDatadogExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.datadog/add-datadog-exporter.html) and add the second exporter through [`addSpanExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.feature/-open-telemetry-config/add-span-exporter.html). Each call registers an independent batch span processor so the two backends are exported in parallel: ``` install(OpenTelemetry) { addDatadogExporter() addSpanExporter( OtlpHttpSpanExporter.builder() .setEndpoint("http://localhost:4318/v1/traces") .build() ) } ``` ## What gets traced The Datadog exporter captures the same activity as Koog's general OpenTelemetry integration. For the full list of captured spans and how to include LLM prompt and response content, see [What gets traced](../#what-gets-traced). For more details on Datadog's OpenTelemetry support, see [Datadog OTLP API Intake](https://docs.datadoghq.com/opentelemetry/guide/otlp_api/). ______________________________________________________________________ ## Troubleshooting - **No traces**: confirm `DD_API_KEY` and `DD_SITE` are set correctly (see [Setup instructions](#setup-instructions)). - **Authentication errors**: verify your key is active in [Organization Settings > API Keys](https://app.datadoghq.com/organization-settings/api-keys). - **Connection issues**: confirm your environment can reach `https://otlp./v1/traces` — for example, `https://otlp.datadoghq.com/v1/traces` for US1. For general troubleshooting, see [Troubleshooting](../#troubleshooting). # Langfuse exporter Koog emits agent traces using [OpenTelemetry](https://opentelemetry.io/), an open standard for observability data. To send those traces to [Langfuse](https://langfuse.com/), Koog includes a built-in OpenTelemetry exporter — no manual instrumentation required. Once connected, Langfuse's [OpenTelemetry support](https://langfuse.com/integrations/native/opentelemetry) lets you visualize, analyze, and debug how your agents interact with LLMs, tools, and external APIs. ______________________________________________________________________ ## Setup instructions 1. Create a Langfuse project using the [setup guide](https://langfuse.com/docs/get-started#create-new-project-in-langfuse). 1. Get your `public key` and `secret key` from [Organization Settings > API Keys](https://langfuse.com/faq/all/where-are-langfuse-api-keys). 1. Provide the host, public key, and secret key — either as parameters to [`addLangfuseExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.langfuse/add-langfuse-exporter.html), or via environment variables: ``` export LANGFUSE_HOST="https://cloud.langfuse.com" export LANGFUSE_PUBLIC_KEY="" export LANGFUSE_SECRET_KEY="" ``` ## Configuration Install the **OpenTelemetry feature** and call [`addLangfuseExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.langfuse/add-langfuse-exporter.html) to enable Langfuse export. ### Basic example ``` fun main() = runBlocking { val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a code assistant. Provide concise code examples." ) { install(OpenTelemetry) { addLangfuseExporter() } } println("Running agent with Langfuse tracing") val result = agent.run("Tell me a joke about programming") println("Result: $result\nSee traces on the Langfuse instance") } ``` ``` public static void main(String[] args) { var agent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OpenAIModels.Chat.GPT4oMini) .systemPrompt("You are a code assistant. Provide concise code examples.") .install(OpenTelemetry.Feature, config -> LangfuseKt.addLangfuseExporter(config) ) .build(); System.out.println("Running agent with Langfuse tracing"); var result = agent.run("Tell me a joke about programming"); System.out.println("Result: " + result + "\nSee traces on the Langfuse instance"); } ``` ## Trace attributes When Koog sends agent activity to Langfuse, it does so as a series of *spans* — individual records of work, such as an LLM call or a tool execution. Related spans are grouped into a *trace*, which represents a complete agent run from start to finish. [`addLangfuseExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.langfuse/add-langfuse-exporter.html) accepts a `traceAttributes` parameter — a list of key-value pairs attached to the root of each trace. These enable Langfuse-specific features such as sessions, environments, and tags, making it easy to filter and group traces in the Langfuse UI. For the full list of supported attributes, see [Langfuse OpenTelemetry docs](https://langfuse.com/integrations/native/opentelemetry#trace-level-attributes). Common attributes to include: - **Session ID** (`langfuse.session.id`): Groups related traces for aggregated metrics, cost analysis, and scoring - **Environment** (`langfuse.environment`): Isolates production traces from development and staging - **Tags** (`langfuse.trace.tags`): Labels traces with feature names, experiment IDs, or customer segments (array of strings) ### Example with session and tags ``` fun main() = runBlocking { val sessionId = UUID.randomUUID().toString() val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a helpful assistant." ) { install(OpenTelemetry) { addLangfuseExporter( traceAttributes = listOf( CustomAttribute("langfuse.session.id", sessionId), CustomAttribute("langfuse.trace.tags", listOf("chat", "kotlin", "production")) ) ) } } println("Running agent with Langfuse tracing") // Multiple runs with the same session ID will be grouped in Langfuse agent.run("What is Kotlin?") agent.run("Show me a coroutine example") } ``` Note Setting `traceAttributes` from Java is currently not supported because the underlying Kotlin function carries a [`kotlin.time.Duration`](https://kotlinlang.org/api/latest/jvm/stdlib/kotlin.time/-duration/) parameter (a value class) that causes JVM-name mangling on all overloads including parameters after it. Use the Kotlin example above when you need `traceAttributes`. ## What gets traced The Langfuse exporter captures the same activity as Koog’s general OpenTelemetry integration. It also captures span attributes required by Langfuse to show [Agent Graphs](https://langfuse.com/docs/observability/features/agent-graphs). For the full list of captured spans and how to include LLM prompt and response content, see [What gets traced](../#what-gets-traced). When visualized in Langfuse, the trace appears as follows: For more details on Langfuse OpenTelemetry tracing, see:\ [Langfuse OpenTelemetry Docs](https://langfuse.com/integrations/native/opentelemetry#opentelemetry-endpoint). ______________________________________________________________________ ## Troubleshooting - **No traces**: confirm `LANGFUSE_HOST`, `LANGFUSE_PUBLIC_KEY`, and `LANGFUSE_SECRET_KEY` are set, and that the key pair belongs to the correct project. - **Connection issues**: if running self-hosted Langfuse, confirm `LANGFUSE_HOST` is reachable from your environment. For general troubleshooting, see [Troubleshooting](../#troubleshooting). # W&B Weave exporter Koog emits agent traces using [OpenTelemetry](https://opentelemetry.io/), an open standard for observability data. To send those traces to [W&B Weave](https://wandb.ai/site/weave/), Koog includes a built-in OpenTelemetry exporter — no manual instrumentation required. Once connected, Weave’s [OpenTelemetry support](https://weave-docs.wandb.ai/guides/tracking/otel/) lets you visualize, analyze, and debug how your agents interact with LLMs, tools, and external APIs. ______________________________________________________________________ ## Setup instructions 1. Create a W&B account at . 1. Get your API key from . 1. Find your entity name at the [W&B Dashboard](https://wandb.ai/home) — it matches your username for personal accounts, or the team/organization name for shared workspaces. 1. Choose a project name. If the project doesn't exist yet, it will be created automatically when the first trace is sent. 1. Provide the entity, project name, and API key — either as parameters to [`addWeaveExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.weave/add-weave-exporter.html), or via environment variables: ``` export WEAVE_API_KEY="" export WEAVE_ENTITY="" export WEAVE_PROJECT_NAME="koog-tracing" ``` ## Configuration Install the **OpenTelemetry feature** and call [`addWeaveExporter()`](https://api.koog.ai/agents/agents-features/agents-features-opentelemetry/ai.koog.agents.features.opentelemetry.integration.weave/add-weave-exporter.html) to enable Weave export. ### Basic example ``` fun main() = runBlocking { val entity = System.getenv()["WEAVE_ENTITY"] ?: throw IllegalArgumentException("WEAVE_ENTITY is not set") val projectName = System.getenv()["WEAVE_PROJECT_NAME"] ?: "koog-tracing" val agent = AIAgent( promptExecutor = promptExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a code assistant. Provide concise code examples." ) { install(OpenTelemetry) { addWeaveExporter() } } println("Running agent with Weave tracing") val result = agent.run("Tell me a joke about programming") println("Result: $result\nSee traces on https://wandb.ai/$entity/$projectName/weave/traces") } ``` ``` public static void main(String[] args) { var entity = Optional.ofNullable(System.getenv("WEAVE_ENTITY")) .filter(env -> !env.isBlank()) .orElseThrow(() -> new IllegalArgumentException("WEAVE_ENTITY is not set")); var projectName = Optional.ofNullable(System.getenv("WEAVE_PROJECT_NAME")) .filter(env -> !env.isBlank()) .orElse("koog-tracing"); var agent = AIAgent.builder() .promptExecutor(promptExecutor) .llmModel(OpenAIModels.Chat.GPT4oMini) .systemPrompt("You are a helpful assistant.") .install(OpenTelemetry.Feature, config -> WeaveKt.addWeaveExporter( config, null, // weaveOtelBaseUrl: falls back to WEAVE_URL, defaults to https://trace.wandb.ai entity, projectName // remaining params (apiKey, timeout) use defaults ) ) .build(); System.out.println("Running agent with Weave tracing"); var result = agent.run("Tell me a joke about programming"); System.out.println("Result: " + result + "\nSee traces on https://wandb.ai/" + entity + "/" + projectName + "/weave/traces"); } ``` ## What gets traced The Weave exporter captures the same activity as Koog’s general OpenTelemetry integration. For the full list of captured spans and how to include LLM prompt and response content, see [What gets traced](../#what-gets-traced). When visualized in W&B Weave, the trace appears as follows: For more details, see the official [Weave OpenTelemetry Docs](https://weave-docs.wandb.ai/guides/tracking/otel/). ______________________________________________________________________ ## Troubleshooting - **No traces**: confirm `WEAVE_API_KEY`, `WEAVE_ENTITY`, and `WEAVE_PROJECT_NAME` are set, and that your W&B account has access to the specified entity and project. - **Authentication errors**: verify `WEAVE_API_KEY` is valid and has write permission for the selected entity. - **Connection issues**: confirm your environment can reach W&B’s OpenTelemetry ingestion endpoints. For general troubleshooting, see [Troubleshooting](../#troubleshooting). # History compression # History compression AI agents maintain a message history that includes user messages, assistant responses, tool calls, and tool responses. This history grows with each interaction as the agent follows its strategy. For long-running conversations, the history can become large and consume a lot of tokens. History compression helps reduce this by summarizing the full list of messages into one or several messages that contain only important information necessary for further agent operation. History compression addresses key challenges in agent systems: - Optimizes context usage. Focused and smaller contexts improve LLM performance and prevent failures from exceeding token limits. - Improves performance. Compressing history reduces the number of messages the LLM processes, resulting in faster responses. - Enhances accuracy. Focusing on relevant information helps the LLM remain focused and complete tasks without distractions. - Reduces costs. Reducing irrelevant messages lowers token usage, decreasing the overall cost of API calls. ## When to compress history History compression is performed at specific steps in the agent workflow: - Between logical steps (subgraphs) of the agent strategy. - When context becomes too long. ## History compression implementation There are two main approaches to implementing history compression in your agent: - In a strategy graph - In a custom node ### History compression in a strategy graph To compress the history in a strategy graph, you need to use the pre-defined node that compresses the current message history into a concise summary: - **Kotlin**: `nodeLLMCompressHistory` - **Java**: `AIAgentNode.llmCompressHistory()` For more information and specific examples, see [History compression node](../nodes-and-components/#history-compression-node). Depending on which step you decide to perform compression, the following scenarios are available: - To compress the history when it becomes too long, check the message count in your edge conditions and add a history compression node. To check the history length, do the following: - **Kotlin**: Define a helper extension. - **Java**: Use inline lambda expressions in `.onCondition()`. ``` // Define that the history is too long if there are more than 100 messages private suspend fun AIAgentContext.historyIsTooLong(): Boolean = llm.readSession { prompt.messages.size > 100 } val strategy = strategy("execute-with-history-compression") { val callLLM by nodeLLMRequest() val executeTool by nodeExecuteTools() val sendToolResult by nodeLLMSendToolResults() // Compress the LLM history and keep the current ReceivedToolResults for the next node val compressHistory by nodeLLMCompressHistory() edge(nodeStart forwardTo callLLM) edge(callLLM forwardTo nodeFinish onTextMessage { true }) edge(callLLM forwardTo executeTool onToolCalls { true }) // Compress history after executing any tool if the history is too long edge(executeTool forwardTo compressHistory onCondition { historyIsTooLong() }) edge(compressHistory forwardTo sendToolResult) // Otherwise, proceed to the next LLM request edge(executeTool forwardTo sendToolResult onCondition { !historyIsTooLong() }) edge(sendToolResult forwardTo executeTool onToolCalls { true }) edge(sendToolResult forwardTo nodeFinish onTextMessage { true }) } ``` ``` var graph = AIAgentGraphStrategy.builder("execute-with-history-compression") .withInput(String.class) .withOutput(String.class); var callLLM = AIAgentNode.llmRequest(null); var executeTool = AIAgentNode.executeTools(null); var sendToolResult = AIAgentNode.llmSendToolResults(null); // Compress the LLM history; the carried ReceivedToolResults flows into the next node. var compressHistory = AIAgentNode .llmCompressHistory("compressHistory") .withInput(ReceivedToolResults.class) .build(); // Edge from start to callLLM graph.edge(AIAgentEdge.builder() .from(graph.nodeStart) .to(callLLM) .build()); // Edge from callLLM to finish on text response graph.edge(AIAgentEdge.builder() .from(callLLM) .to(graph.nodeFinish) .onTextMessage() .build()); // Edge from callLLM to executeTool on tool call graph.edge(AIAgentEdge.builder() .from(callLLM) .to(executeTool) .onToolCalls() .build()); // Compress history after executing any tool if the history is too long graph.edge(AIAgentEdge.builder() .from(executeTool) .to(compressHistory) .onCondition((message, ctx) -> ctx.getLlm().readSession(session -> session.getPrompt().getMessages().size() > 100 ) ) .build()); graph.edge(compressHistory, sendToolResult); // Otherwise, proceed to the next LLM request graph.edge(AIAgentEdge.builder() .from(executeTool) .to(sendToolResult) .onCondition((message, ctx) -> ctx.getLlm().readSession(session -> session.getPrompt().getMessages().size() <= 100 ) ) .build()); // Edge from sendToolResult to executeTool on tool call graph.edge(AIAgentEdge.builder() .from(sendToolResult) .to(executeTool) .onToolCalls() .build()); // Edge from sendToolResult to finish on text response graph.edge(AIAgentEdge.builder() .from(sendToolResult) .to(graph.nodeFinish) .onTextMessage() .build()); ``` In this example, the strategy checks if the history is too long after each tool call. The history is compressed before sending the tool result back to the LLM. This prevents the context from growing during long conversations. - To compress the history between the logical steps (subgraphs) of your strategy, you can implement your strategy as follows: ``` val strategy = strategy("execute-with-history-compression") { val collectInformation by subgraph { // Some steps to collect the information } val compressHistory by nodeLLMCompressHistory() val makeTheDecision by subgraph { // Some steps to make the decision based on the current compressed history and collected information } nodeStart then collectInformation then compressHistory then makeTheDecision } ``` ``` var graph = AIAgentGraphStrategy.builder("execute-with-history-compression") .withInput(String.class) .withOutput(String.class); // Subgraph to collect information var collectInformation = AIAgentSubgraph.builder("collectInformation") .withInput(String.class) .withOutput(String.class) .limitedTools(Collections.emptyList()) .withTask(input -> "Collect information based on: " + input) .build(); // Compress history after collecting information var compressHistory = AIAgentNode .llmCompressHistory("compressHistory") .withInput(String.class) .build(); // Subgraph to make decision based on compressed history var makeTheDecision = AIAgentSubgraph.builder("makeTheDecision") .withInput(String.class) .withOutput(String.class) .limitedTools(Collections.emptyList()) .withTask(input -> "Make a decision based on the information") .build(); // Build the flow: start -> collectInformation -> compressHistory -> makeTheDecision -> finish graph.edge(graph.nodeStart, collectInformation); graph.edge(collectInformation, compressHistory); graph.edge(compressHistory, makeTheDecision); graph.edge(makeTheDecision, graph.nodeFinish); ``` In this example, the history is compressed after completing the information collection phase, but before proceeding to the decision-making phase. ### History compression in a custom node If you are implementing a custom node, you can compress history using the `replaceHistoryWithTLDR()` function (Kotlin) as follows: ``` llm.writeSession { replaceHistoryWithTLDR() } ``` This approach gives you more flexibility to implement compression at any point in your custom node logic, based on your specific requirements. To learn more about custom nodes, see [Custom nodes](../custom-nodes/). ## History compression strategies You can customize the compression process using the optional `strategy` parameter: - **Kotlin**: Pass the strategy to `nodeLLMCompressHistory(strategy=...)` or `replaceHistoryWithTLDR(strategy=...)`. - **Java**: Use the `.compressionStrategy()` builder method. The framework provides several built-in strategies. ### WholeHistory (Default) The default strategy that compresses the entire history into one TLDR message that summarizes what has been achieved so far. This strategy works well for most general use cases where you want to maintain awareness of the entire conversation context while reducing token usage. You can use it as follows: - In a strategy graph: ``` val compressHistory by nodeLLMCompressHistory( strategy = HistoryCompressionStrategy.WholeHistory ) ``` ``` // Using WholeHistory strategy in a compression node var compressHistory = AIAgentNode .llmCompressHistory("compressHistory") .withInput(String.class) .compressionStrategy(HistoryCompressionStrategy.WholeHistory) .build(); // Note: This example only shows the node creation. // You would need to add edges and other nodes to complete the graph. ``` - In a custom node: ``` llm.writeSession { replaceHistoryWithTLDR(strategy = HistoryCompressionStrategy.WholeHistory) } ``` ``` ctx.getLlm().writeSession(session -> { session.replaceHistoryWithTLDR(HistoryCompressionStrategy.WholeHistory); return null; }); ``` ### FromLastNMessages The strategy compresses only the last `n` messages into a TLDR message and completely discards earlier messages. This is useful when only the latest achievements of the agent (or the latest discovered facts, the latest context) are relevant for solving the problem. You can use it as follows: - In a strategy graph: ``` val compressHistory by nodeLLMCompressHistory( strategy = HistoryCompressionStrategy.FromLastNMessages(5) ) ``` ``` // Using FromLastNMessages strategy to compress only the last 5 messages var compressHistory = AIAgentNode .llmCompressHistory("compressHistory") .withInput(String.class) .compressionStrategy(HistoryCompressionStrategy.FromLastNMessages(5)) .build(); // Note: This example only shows the node creation. // You would need to add edges and other nodes to complete the graph. ``` - In a custom node: ``` llm.writeSession { replaceHistoryWithTLDR(strategy = HistoryCompressionStrategy.FromLastNMessages(5)) } ``` ``` ctx.getLlm().writeSession(session -> { session.replaceHistoryWithTLDR(HistoryCompressionStrategy.FromLastNMessages(5)); return null; }); ``` ### Chunked The strategy splits the whole message history into chunks of a fixed size and compresses each chunk independently into a TLDR message. This is useful when you need not only the concise TLDR of what has been done so far but also want to keep track of the overall progress, and some older information might also be important. You can use it as follows: - In a strategy graph: ``` val compressHistory by nodeLLMCompressHistory( strategy = HistoryCompressionStrategy.Chunked(10) ) ``` ``` // Using Chunked strategy to compress history in chunks of 10 messages var compressHistory = AIAgentNode .llmCompressHistory("compressHistory") .withInput(String.class) .compressionStrategy(HistoryCompressionStrategy.Chunked(10)) .build(); // Note: This example only shows the node creation. // You would need to add edges and other nodes to complete the graph. ``` - In a custom node: ``` llm.writeSession { replaceHistoryWithTLDR(strategy = HistoryCompressionStrategy.Chunked(10)) } ``` ``` ctx.getLlm().writeSession(session -> { session.replaceHistoryWithTLDR(HistoryCompressionStrategy.Chunked(10)); return null; }); ``` ### FactRetrievalHistoryCompressionStrategy The strategy searches for specific facts relevant to the provided list of concepts in the history and retrieves them. It changes the whole history to just these facts and leaves them as context for future LLM requests. This is useful when you have an idea of what exact facts will be relevant for the LLM to perform better on the task. You can use it as follows: - In a strategy graph: ``` val compressHistory by nodeLLMCompressHistory( strategy = FactRetrievalHistoryCompressionStrategy( Concept( keyword = "user_preferences", // Description to the LLM -- what specifically to search for description = "User's preferences for the recommendation system, including the preferred conversation style, theme in the application, etc.", // LLM would search for multiple relevant facts related to this concept: factType = FactType.MULTIPLE ), Concept( keyword = "product_details", // Description to the LLM -- what specifically to search for description = "Brief details about products in the catalog the user has been checking", // LLM would search for multiple relevant facts related to this concept: factType = FactType.MULTIPLE ), Concept( keyword = "issue_solved", // Description to the LLM -- what specifically to search for description = "Was the initial user's issue resolved?", // LLM would search for a single answer to the question: factType = FactType.SINGLE ) ) ) ``` ``` // Using FactRetrievalHistoryCompressionStrategy strategy to extract specific facts var compressHistory = AIAgentNode .llmCompressHistory("compressHistory") .withInput(ReceivedToolResult.class) .compressionStrategy(new FactRetrievalHistoryCompressionStrategy( new Concept( "user_preferences", "User's preferences for the recommendation system, including the preferred conversation style, theme in the application, etc.", FactType.MULTIPLE ), new Concept( "product_details", "Brief details about products in the catalog the user has been checking", FactType.MULTIPLE ), new Concept( "issue_solved", "Was the initial user's issue resolved?", FactType.SINGLE ) )) .build(); // Note: This example only shows the node creation. // You would need to add edges and other nodes to complete the graph. ``` - In a custom node: ``` llm.writeSession { replaceHistoryWithTLDR( strategy = FactRetrievalHistoryCompressionStrategy( Concept( keyword = "user_preferences", // Description to the LLM -- what specifically to search for description = "User's preferences for the recommendation system, including the preferred conversation style, theme in the application, etc.", // LLM would search for multiple relevant facts related to this concept: factType = FactType.MULTIPLE ), Concept( keyword = "product_details", // Description to the LLM -- what specifically to search for description = "Brief details about products in the catalog the user has been checking", // LLM would search for multiple relevant facts related to this concept: factType = FactType.MULTIPLE ), Concept( keyword = "issue_solved", // Description to the LLM -- what specifically to search for description = "Was the initial user's issue resolved?", // LLM would search for a single answer to the question: factType = FactType.SINGLE ) ) ) } ``` ``` ctx.getLlm().writeSession(session -> { session.replaceHistoryWithTLDR(new FactRetrievalHistoryCompressionStrategy( new Concept( "user_preferences", // Description to the LLM -- what specifically to search for "User's preferences for the recommendation system, including the preferred conversation style, theme in the application, etc.", // LLM would search for multiple relevant facts related to this concept: FactType.MULTIPLE ), new Concept( "product_details", // Description to the LLM -- what specifically to search for "Brief details about products in the catalog the user has been checking", // LLM would search for multiple relevant facts related to this concept: FactType.MULTIPLE ), new Concept( "issue_solved", // Description to the LLM -- what specifically to search for "Was the initial user's issue resolved?", // LLM would search for a single answer to the question: FactType.SINGLE ) )); return null; }); ``` ## Custom history compression strategy implementation Warning Custom history compression strategies are available only in Kotlin. You can create your own history compression strategy by extending the `HistoryCompressionStrategy` abstract class and implementing the `compress` method. Here is an example: ``` class MyCustomCompressionStrategy : HistoryCompressionStrategy() { override suspend fun compress( llmSession: AIAgentLLMWriteSession, memoryMessages: List ) { // 1. Process the current history in llmSession.prompt.messages // 2. Create new compressed messages // 3. Update the prompt with the compressed messages // Save original messages to preserve them val originalMessages = llmSession.prompt.messages // Example implementation: val importantMessages = llmSession.prompt.messages .filterIsInstance() .filter { message -> // Your custom filtering logic message.parts.filterIsInstance().any { it.text.contains("important") } } // Note: you can also make LLM requests using the `llmSession` and ask the LLM to do some job for you using, for example, `llmSession.requestLLMWithoutTools()` // Or you can change the current model: `llmSession.model = AnthropicModels.Opus_4_6` and ask some other LLM model -- but don't forget to change it back after // Compose the prompt with the filtered messages val compressedMessages = composeMessageHistory( originalMessages, importantMessages, memoryMessages ) } } ``` In this example, the custom strategy filters messages that contain the word "important" and keeps only those in the compressed history. Then you can use it as follows: - In a strategy graph: ``` val compressHistory by nodeLLMCompressHistory( strategy = MyCustomCompressionStrategy() ) ``` - In a custom node: ``` llm.writeSession { replaceHistoryWithTLDR(strategy = MyCustomCompressionStrategy()) } ``` ## Memory preservation during compression All history compression methods support memory preservation, which determines whether memory-related messages should be preserved during compression. In Kotlin, use the `preserveMemory` parameter. In Java, use the `.preserveMemory()` builder method. These are messages that contain facts retrieved from memory or indicate that the memory feature is not enabled. To enable memory preservation: - **Kotlin**: Use the `preserveMemory` parameter. - **Java**: Use the `.preserveMemory()` builder method. - In a strategy graph: ``` val compressHistory by nodeLLMCompressHistory( strategy = HistoryCompressionStrategy.WholeHistory, preserveMemory = true ) ``` ``` // Using WholeHistory strategy with preserveMemory=true var compressHistory = AIAgentNode .llmCompressHistory("compressHistory") .withInput(String.class) .compressionStrategy(HistoryCompressionStrategy.WholeHistory) .preserveMemory(true) .build(); // Note: This example only shows the node creation. // You would need to add edges and other nodes to complete the graph. ``` - In a custom node: ``` llm.writeSession { replaceHistoryWithTLDR( strategy = HistoryCompressionStrategy.WholeHistory, preserveMemory = true ) } ``` ``` ctx.getLlm().writeSession(session -> { session.replaceHistoryWithTLDR( /** strategy */ HistoryCompressionStrategy.WholeHistory, /** preserveMemory */ true ); return null; }); ``` # Model Context Protocol # Model Context Protocol Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../module-versioning/) for details. Model Context Protocol (MCP) is a standardized protocol that lets AI agents interact with external tools and services through a consistent interface. MCP exposes tools and prompts as API endpoints that AI agents can call. Each tool has a specific name and an input schema that describes its inputs and outputs using the JSON Schema format. The Koog framework provides integration with MCP servers, enabling you to incorporate MCP tools into your Koog agents. To learn more about the protocol, see the [Model Context Protocol](https://modelcontextprotocol.io) documentation. ## MCP servers MCP servers implement Model Context Protocol and provide a standardized way for AI agents to interact with tools and services. You can find ready-to-use MCP servers in the [MCP Marketplace](https://mcp.so/) or [MCP DockerHub](https://hub.docker.com/u/mcp). The MCP servers support the following transport protocols to communicate with agents: - Standard input/output (stdio) transport protocol used to communicate with the MCP servers running as separate processes. For example, a Docker container or a CLI tool. - Server-sent events (SSE) transport protocol (optional) used to communicate with the MCP servers over HTTP. ## Integration with Koog The Koog framework integrates with MCP using the [MCP SDK](https://github.com/modelcontextprotocol/kotlin-sdk) with the additional API extensions presented in the `agent-mcp` module. This integration lets the Koog agents perform the following: - Connect to MCP servers through various transport mechanisms (stdio, SSE). - Retrieve available tools from an MCP server. - Transform MCP tools into the Koog tool interface. - Register the transformed tools in a tool registry. - Call MCP tools with arguments provided by the LLM. ### Key components Here are the main components of the MCP integration in Koog: | Component | Description | | ---------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------- | | [`McpTool`](https://api.koog.ai/agents/agents-mcp/ai.koog.agents.mcp/-mcp-tool/index.html) | Serves as a bridge between the Koog tool interface and the MCP SDK. | | [`McpToolDescriptorParser`](https://api.koog.ai/agents/agents-mcp/ai.koog.agents.mcp/-mcp-tool-descriptor-parser/index.html) | Parses MCP tool definitions into the Koog tool descriptor format. | | [`McpToolRegistryProvider`](https://api.koog.ai/agents/agents-mcp/ai.koog.agents.mcp/-mcp-tool-registry-provider/index.html) | Creates MCP tool registries that connect to MCP servers through various transport mechanisms (stdio, SSE). | ## Getting started ### 1. Set up an MCP connection To use MCP with Koog, you need to set up a connection: 1. Start an MCP server (either as a process, Docker container, or web service). 1. Create a transport mechanism to communicate with the server. MCP servers support the stdio and SSE transport mechanisms to communicate with the agent, so you can connect using one of them. #### Connect with stdio This protocol is used when an MCP server runs as a separate process. Here is an example of setting up an MCP connection using the stdio transport: ``` // Start an MCP server (for example, as a process) val process = ProcessBuilder("path/to/mcp/server").start() // Create the stdio transport val transport = McpToolRegistryProvider.defaultStdioTransport(process) ``` #### Connect with SSE This protocol is used when an MCP server runs as a web service. Here is an example of setting up an MCP connection using the SSE transport: ``` // Create the SSE transport val transport = McpToolRegistryProvider.defaultSseTransport("http://localhost:8931") ``` ### 2. Create a tool registry Once you have the MCP connection, you can create a tool registry with tools from the MCP server in one of the following ways: - Using the provided transport mechanism for communication. For example: ``` // Create a tool registry with tools from the MCP server val toolRegistry = McpToolRegistryProvider.fromTransport( transport = transport, serverInfo = McpServerInfo(url = "http://localhost:8931", command = "path/to/mcp/server"), name = "my-client", version = "1.0.0" ) ``` - Using an MCP client connected to the MCP server. For example: ``` // Create a tool registry from an existing MCP client val toolRegistry = McpToolRegistryProvider.fromClient( mcpClient = existingMcpClient, serverInfo = McpServerInfo(url = "http://localhost:8931") ) ``` ### 3. Integrate with your agent To use MCP tools with your Koog agent, you need to register the tool registry with the agent: ``` // Create an agent with the tools val agent = AIAgent( promptExecutor = executor, strategy = strategy, llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = toolRegistry ) // Run the agent with a task that uses an MCP tool val result = agent.run("Use the MCP tool to perform a task") ``` ## Usage examples ### Google Maps MCP integration This example demonstrates how to connect to a [Google Maps](https://mcp.so/server/google-maps/modelcontextprotocol) server for geographic data using MCP: ``` // Start the Docker container with the Google Maps MCP server val process = ProcessBuilder( "docker", "run", "-i", "-e", "GOOGLE_MAPS_API_KEY=$googleMapsApiKey", "mcp/google-maps" ).start() // Create the ToolRegistry with tools from the MCP server val toolRegistry = McpToolRegistryProvider.fromProcess(process = process) // Create and run the agent val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(openAIApiToken), llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = toolRegistry, ) agent.run("Get elevation of the Jetbrains Office in Munich, Germany?") ``` ### Playwright MCP integration This example demonstrates how to connect to a [Playwright](https://mcp.so/server/playwright-mcp/microsoft) server for web automation using MCP: ``` // Start the Playwright MCP server val process = ProcessBuilder( "npx", "@playwright/mcp@latest", "--port", "8931" ).start() // Create the ToolRegistry with tools from the MCP server val toolRegistry = McpToolRegistryProvider.fromSseUrl("http://localhost:8931") // Create and run the agent val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(openAIApiToken), llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = toolRegistry, ) agent.run("Open a browser, navigate to jetbrains.com, accept all cookies, click AI in toolbar") ``` # A2A Protocol # A2A protocol Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../module-versioning/) for details. This page provides an overview of the A2A (Agent-to-Agent) protocol implementation in the Koog agentic framework. ## What is the A2A protocol? The A2A (Agent-to-Agent) protocol is a standardized communication protocol that enables AI agents to interact with each other and with client applications. It defines a set of methods, message formats, and behaviors that allow for consistent and interoperable agent communication. For more information and a detailed specification of the A2A protocol, see the official [A2A Protocol website](https://a2a-protocol.org/latest/). ## Getting Started **Important**: A2A dependencies are **not** included by default in the `koog-agents` meta-dependency. You must explicitly add the A2A modules you need to your project. To use A2A in your project, add dependencies based on your use case: - **For A2A client**: See [A2A Client documentation](a2a-client/#dependencies) - **For A2A server**: See [A2A Server documentation](a2a-server/#dependencies) - **For Koog integration**: See [A2A Koog Integration documentation](a2a-koog-integration/#dependencies) ## Key A2A components Koog provides full implementation of A2A protocol v0.3.0 for both client and server, as well as integration with the Koog agent framework: - [A2A Server](a2a-server/) is an agent or agentic system that exposes an endpoint implementing the A2A protocol. It receives requests from clients, processes tasks, and returns results or status updates. It can also be used independently of Koog agents. - [A2A Client](a2a-client/) is a client application or agent that initiates communication with an A2A server using the A2A protocol. It can also be used independently of Koog agents. - [A2A Koog Integration](a2a-koog-integration/) is a set of classes and utilities that simplify the integration of A2A with Koog Agents. It contains components (A2A features and nodes) for seamless A2A agent connections and communication within the Koog framework. For more examples, follow the [examples](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/a2a) # A2A Server Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../../module-versioning/) for details. The A2A server enables you to expose AI agents through the standardized A2A (Agent-to-Agent) protocol. It provides a complete implementation of the [A2A protocol specification](https://a2a-protocol.org/latest/specification/), handling client requests, executing agent logic, managing complex task lifecycles, and supporting real-time streaming responses. ## Dependencies To use the A2A server in your project, add the following dependencies to your `build.gradle.kts`: ``` dependencies { // Core A2A server library implementation("ai.koog:a2a-server:$koogVersion") // HTTP JSON-RPC transport (most common) implementation("ai.koog:a2a-transport-server-jsonrpc-http:$koogVersion") // Ktor server engine (choose one that fits your needs) implementation("io.ktor:ktor-server-netty:$ktorVersion") } ``` ## Overview The A2A server acts as a bridge between the A2A protocol transport layer and your custom agent logic. It orchestrates the entire request lifecycle while maintaining protocol compliance and providing robust session management. ## Core components ### A2AServer The main server class implementing the complete A2A protocol. It serves as the central coordinator that: - **Validates** incoming requests against protocol specifications - **Manages** concurrent sessions and task lifecycles - **Orchestrates** communication between transport, storage, and business logic layers - **Handles** all protocol operations: message sending, task querying, cancellation, push notifications The `A2AServer` accepts two required parameters: - `AgentExecutor` which defines business logic implementation of the agent - `AgentCard` which defines agent capabilities and metadata And a number of optional parameters that can be used to customize its storage and transport behavior. ### AgentExecutor The `AgentExecutor` interface is where you implement your agent's core business logic. It acts as the bridge between the A2A protocol and your specific AI agent capabilities. To start the execution of your agent, you must implement the `execute` method where define your agent's logic. To cancel the agent, you must implement the `cancel` method. ``` class MyAgentExecutor : AgentExecutor { override suspend fun execute( context: RequestContext, eventProcessor: SessionEventProcessor ) { // Agent logic here } override suspend fun cancel( context: RequestContext, eventProcessor: SessionEventProcessor, agentJob: Deferred? ) { // Cancel agent here, optional } } ``` The `RequestContext` provides rich information about the current request, including the `contextId` and `taskId` of the current session, the `message` sent, and the `params` of the request. The `SessionEventProcessor` communicates with clients: - **`sendMessage(message)`**: Send immediate responses (chat-style interactions) - **`sendTaskEvent(event)`**: Send task-related updates (long-running operations) ``` // For immediate responses (like chatbots) eventProcessor.sendMessage( Message( messageId = generateId(), role = Role.Agent, parts = listOf(TextPart("Here's your answer!")), contextId = context.contextId ) ) // For task-based operations eventProcessor.sendTaskEvent( TaskStatusUpdateEvent( contextId = context.contextId, taskId = context.taskId, status = TaskStatus( state = TaskState.Working, message = Message(/* progress update */), timestamp = Clock.System.now() ), final = false // More updates to come ) ) ``` ### AgentCard The `AgentCard` serves as your agent's self-describing manifest. It tells clients what your agent can do, how to communicate with it, and what security requirements it has. ``` val agentCard = AgentCard( // Basic Identity name = "Advanced Recipe Assistant", description = "AI agent specialized in cooking advice, recipe generation, and meal planning", version = "2.1.0", protocolVersion = "0.3.0", // Communication Settings url = "https://api.example.com/a2a", preferredTransport = TransportProtocol.JSONRPC, // Optional: Multiple transport support additionalInterfaces = listOf( AgentInterface("https://api.example.com/a2a", TransportProtocol.JSONRPC), ), // Capabilities Declaration capabilities = AgentCapabilities( streaming = true, // Support real-time responses pushNotifications = true, // Send async notifications stateTransitionHistory = true // Maintain task history ), // Content Type Support defaultInputModes = listOf("text/plain", "text/markdown", "image/jpeg"), defaultOutputModes = listOf("text/plain", "text/markdown", "application/json"), // Define available security schemes securitySchemes = mapOf( "bearer" to HTTPAuthSecurityScheme( scheme = "Bearer", bearerFormat = "JWT", description = "JWT token authentication" ), "api-key" to APIKeySecurityScheme( `in` = In.Header, name = "X-API-Key", description = "API key for service authentication" ) ), // Specify security requirements (logical OR of requirements) security = listOf( mapOf("bearer" to listOf("read", "write")), // Option 1: JWT with read/write scopes mapOf("api-key" to emptyList()) // Option 2: API key ), // Enable extended card for authenticated users supportsAuthenticatedExtendedCard = true, // Skills/Capabilities skills = listOf( AgentSkill( id = "recipe-generation", name = "Recipe Generation", description = "Generate custom recipes based on ingredients, dietary restrictions, and preferences", tags = listOf("cooking", "recipes", "nutrition"), examples = listOf( "Create a vegan pasta recipe with mushrooms", "I have chicken, rice, and vegetables. What can I make?" ) ), AgentSkill( id = "meal-planning", name = "Meal Planning", description = "Plan weekly meals and generate shopping lists", tags = listOf("meal-planning", "nutrition", "shopping") ) ), // Optional: Branding iconUrl = "https://example.com/agent-icon.png", documentationUrl = "https://docs.example.com/recipe-agent", provider = AgentProvider( organization = "CookingAI Inc.", url = "https://cookingai.com" ) ) ``` ### Transport Layer The A2A itself supports multiple transport protocols for communicating with clients. Currently, Koog provides implementations for JSON-RPC server transport over HTTP. #### HTTP JSON-RPC Transport ``` val transport = HttpJSONRPCServerTransport(server) transport.start( engineFactory = CIO, // Ktor engine (CIO, Netty, Jetty) port = 8080, // Server port path = "/a2a", // API endpoint path wait = true // Block until server stops ) ``` ### Storage The A2A server uses a pluggable storage architecture that separates different types of data. All storage implementations are optional and default to in-memory variants for development. - **TaskStorage**: Task lifecycle management - stores and manages task states, history, and artifacts - **MessageStorage**: Conversation history - manages message history within conversation contexts - **PushNotificationConfigStorage**: Webhook management - manages webhook configurations for asynchronous notifications ## Quickstart ### 1. Create AgentCard Define your agent's capabilities and metadata. ``` val agentCard = AgentCard( name = "IO Assistant", description = "AI agent specialized in input modification", version = "2.1.0", protocolVersion = "0.3.0", // Communication Settings url = "https://api.example.com/a2a", preferredTransport = TransportProtocol.JSONRPC, // Capabilities Declaration capabilities = AgentCapabilities( streaming = true, // Support real-time responses pushNotifications = true, // Send async notifications stateTransitionHistory = true // Maintain task history ), // Content Type Support defaultInputModes = listOf("text/plain", "text/markdown", "image/jpeg"), defaultOutputModes = listOf("text/plain", "text/markdown", "application/json"), // Skills/Capabilities skills = listOf( AgentSkill( id = "echo", name = "echo", description = "Echoes back user messages", tags = listOf("io"), ) ) ) ``` ### 2. Create an AgentExecutor In executor manages implement agent logic, handles incoming requests and sends responses. ``` class EchoAgentExecutor : AgentExecutor { override suspend fun execute( context: RequestContext, eventProcessor: SessionEventProcessor ) { val userMessage = context.params.message val userText = userMessage.parts .filterIsInstance() .joinToString(" ") { it.text } // Echo the user's message back val response = Message( messageId = UUID.randomUUID().toString(), role = Role.Agent, parts = listOf(TextPart("You said: $userText")), contextId = context.contextId, taskId = context.taskId ) eventProcessor.sendMessage(response) } } ``` ### 2. Create the Server Pass the agent executor and agent card to the server. ``` val server = A2AServer( agentExecutor = EchoAgentExecutor(), agentCard = agentCard ) ``` ### 3. Add Transport Layer Create a transport layer and start the server. ``` // HTTP JSON-RPC transport val transport = HttpJSONRPCServerTransport(server) transport.start( engineFactory = CIO, port = 8080, path = "/agent", wait = true ) ``` ## Agent Implementation Patterns ### Simple Response Agent If your agent only needs to respond to a single message, you can implement it as a simple agent. It can be also used if agent execution logic is not complex and time-consuming. ``` class SimpleAgentExecutor : AgentExecutor { override suspend fun execute( context: RequestContext, eventProcessor: SessionEventProcessor ) { val response = Message( messageId = UUID.randomUUID().toString(), role = Role.Agent, parts = listOf(TextPart("Hello from agent!")), contextId = context.contextId, taskId = context.taskId ) eventProcessor.sendMessage(response) } } ``` ### Task-Based Agent If the execution logic of your agent is complex and requires multiple steps, you can implement it as a task-based agent. It can be also used if agent execution logic is time-consuming and suspending. ``` class TaskAgentExecutor : AgentExecutor { override suspend fun execute( context: RequestContext, eventProcessor: SessionEventProcessor ) { // Send working status eventProcessor.sendTaskEvent( TaskStatusUpdateEvent( contextId = context.contextId, taskId = context.taskId, status = TaskStatus( state = TaskState.Working, timestamp = Clock.System.now() ), final = false ) ) // Do work... // Send completion eventProcessor.sendTaskEvent( TaskStatusUpdateEvent( contextId = context.contextId, taskId = context.taskId, status = TaskStatus( state = TaskState.Completed, timestamp = Clock.System.now() ), final = true ) ) } } ``` # A2A Client Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../../module-versioning/) for details. The A2A client enables you to communicate with A2A-compliant agents over the network. It provides a complete implementation of the [A2A protocol specification](https://a2a-protocol.org/latest/specification/), handling agent discovery, message exchange, task management, and real-time streaming responses. ## Dependencies To use the A2A client in your project, add the following dependencies to your `build.gradle.kts`: ``` dependencies { // Core A2A client library implementation("ai.koog:a2a-client:$koogVersion") // HTTP JSON-RPC transport (most common) implementation("ai.koog:a2a-transport-client-jsonrpc-http:$koogVersion") // Ktor client engine (choose one that fits your needs) implementation("io.ktor:ktor-client-cio:$ktorVersion") } ``` ## Overview The A2A client acts as a bridge between your application and A2A-compliant agents. It orchestrates the entire communication lifecycle while maintaining protocol compliance and providing robust session management. ## Core components ### A2AClient The main client class implementing the complete A2A protocol. It serves as the central coordinator that: - **Manages** connections and agent discovery through pluggable resolvers - **Orchestrates** message exchange and task operations with automatic protocol compliance - **Handles** streaming responses and real-time communication when supported by agents - **Provides** comprehensive error handling and fallback mechanisms for robust applications The `A2AClient` accepts two required parameters: - `ClientTransport` which handles network communication layer - `AgentCardResolver` which handles agent discovery and metadata retrieval The `A2AClient` interface provides several key methods for interacting with A2A agents: - `connect` method - To connect to the agent and retrieve its capabilities, which discovers what the agent can do and caches the AgentCard - `sendMessage` method - To send a message to the agent and receive a single response for simple request-response patterns - `sendMessageStreaming` method - To send a message with streaming support for real-time responses, which returns a Flow of events including partial messages and task updates - `getTask` method - To query the status and details of a specific task - `cancelTask` method - To cancel a running task if the agent supports cancellation - `cachedAgentCard` method - To get the cached agent card without making a network request, which returns null if connect hasn't been called yet ### ClientTransport The `ClientTransport` interface handles the low-level network communication while the A2A client manages the protocol logic. It abstracts away transport-specific details, allowing you to use different protocols seamlessly. #### HTTP JSON-RPC Transport The most common transport for A2A agents: ``` val transport = HttpJSONRPCClientTransport( url = "https://agent.example.com/a2a", // Agent endpoint URL httpClient = HttpClient(CIO) { // Optional: custom HTTP client install(ContentNegotiation) { json() } install(HttpTimeout) { requestTimeoutMillis = 30000 } } ) ``` ### AgentCardResolver The `AgentCardResolver` interface retrieves agent metadata and capabilities. It enables agent discovery from various sources and supports caching strategies for optimal performance. #### URL Agent Card Resolver Fetch agent cards from HTTP endpoints following A2A conventions: ``` val agentCardResolver = UrlAgentCardResolver( baseUrl = "https://agent.example.com", // Base URL of the agent service path = "/.well-known/agent-card.json", // Standard agent card location httpClient = HttpClient(CIO), // Optional: custom HTTP client ) ``` ## Quickstart ### 1. Create the Client Define the transport and agent card resolver and create the client. ``` // HTTP JSON-RPC transport val transport = HttpJSONRPCClientTransport( url = "https://agent.example.com/a2a" ) // Agent card resolver val agentCardResolver = UrlAgentCardResolver( baseUrl = "https://agent.example.com", path = "/.well-known/agent-card.json" ) // Create client val client = A2AClient(transport, agentCardResolver) ``` ### 2. Connect and Discover Connect to the agent and retrieve its card. Having agent's card enables you to query its capabilities and perform other operations, for example, check if it supports streaming. ``` // Connect and retrieve agent capabilities client.connect() val agentCard = client.cachedAgentCard() println("Connected to: ${agentCard.name}") println("Supports streaming: ${agentCard.capabilities.streaming}") ``` ### 3. Send Messages Send a message to the agent and receive a single response. The response can be either the message if the agent responded directly, or a task event if the agent is performing a task. ``` val message = Message( messageId = UUID.randomUUID().toString(), role = Role.User, parts = listOf(TextPart("Hello, agent!")), contextId = "conversation-1" ) val request = Request(data = MessageSendParams(message)) val response = client.sendMessage(request) // Handle response when (val event = response.data) { is Message -> { val text = event.parts .filterIsInstance() .joinToString { it.text } print(text) // Stream partial responses } is TaskEvent -> { if (event.final) { println("\nTask completed") } } } ``` ### 4. Send Messages Streaming The A2A client supports streaming responses for real-time communication. Instead of receiving a single response, it returns a `Flow` of events including messages and task updates. ``` // Check if agent supports streaming if (client.cachedAgentCard()?.capabilities?.streaming == true) { client.sendMessageStreaming(request).collect { response -> when (val event = response.data) { is Message -> { val text = event.parts .filterIsInstance() .joinToString { it.text } print(text) // Stream partial responses } is TaskStatusUpdateEvent -> { if (event.final) { println("\nTask completed") } } } } } else { // Fallback to non-streaming val response = client.sendMessage(request) // Handle single response } ``` ### 5. Manage Tasks A2A Client provides methods to control server tasks by asking for their status and cancelling them. ``` // Query task status val taskRequest = Request(data = TaskQueryParams(taskId = "task-123")) val taskResponse = client.getTask(taskRequest) val task = taskResponse.data println("Task state: ${task.status.state}") // Cancel running task if (task.status.state == TaskState.Working) { val cancelRequest = Request(data = TaskIdParams(taskId = "task-123")) val cancelledTask = client.cancelTask(cancelRequest).data println("Task cancelled: ${cancelledTask.status.state}") } ``` # A2A and Koog Integration Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../../module-versioning/) for details. Koog provides seamless integration with the A2A protocol, allowing you to expose Koog agents as A2A servers and connect Koog agents to other A2A-compliant agents. ## Dependencies A2A Koog integration requires specific feature modules depending on your use case: ### For Exposing Koog Agents as A2A Servers Add these dependencies to your `build.gradle.kts`: ``` dependencies { // Koog A2A server integration feature implementation("ai.koog:agents-features-a2a-server:$koogVersion") // HTTP JSON-RPC transport implementation("ai.koog:a2a-transport-server-jsonrpc-http:$koogVersion") // Ktor server engine (choose one that fits your needs) implementation("io.ktor:ktor-server-netty:$ktorVersion") } ``` ### For Connecting Koog Agents to A2A Agents Add these dependencies to your `build.gradle.kts`: ``` dependencies { // Koog A2A client integration feature implementation("ai.koog:agents-features-a2a-client:$koogVersion") // HTTP JSON-RPC transport implementation("ai.koog:a2a-transport-client-jsonrpc-http:$koogVersion") // Ktor client engine (choose one that fits your needs) implementation("io.ktor:ktor-client-cio:$ktorVersion") } ``` ## Overview The integration enables two main patterns: 1. **Expose Koog agents as A2A servers** - Make your Koog agents discoverable and accessible via the A2A protocol 1. **Connect Koog agents to A2A agents** - Let your Koog agents communicate with other A2A-compliant agents ## Exposing Koog Agents as A2A Servers ### Define Koog Agent with A2A feature Let's define a Koog agent first. The logic of the agent can vary, but here's an example basic single run agent with tools. The agent receives a message from the user, forwards it to the llm. If the llm response contains a tool call, the agent executes the tool and forwards the result to the llm. If the llm response contains an assistant message, the agent sends the assistant message to the user and finishes. On input resize, the agent sends a task submitted event to the A2A client with the input message. On each tool call, the agent sends a task working event to the A2A client with the tool call and result. On assistant message, the agent sends a task complete event to the A2A client with the assistant message. ``` /** * Create a Koog agent with A2A feature */ @OptIn(ExperimentalUuidApi::class) private fun createAgent( context: RequestContext, eventProcessor: SessionEventProcessor, ) = AIAgent( promptExecutor = MultiLLMPromptExecutor( LLMProvider.Google to GoogleLLMClient("api-key") ), toolRegistry = ToolRegistry { // Declare tools here }, strategy = strategy("test") { val nodeSetup by node { inputMessage -> // Convenience function to transform A2A message into Koog message val input = inputMessage.toKoogMessage() llm.writeSession { appendPrompt { message(input) } } // Send update event to A2A client withA2AAgentServer { sendTaskUpdate("Request submitted: ${input.content}", TaskState.Submitted) } } // Calling llm val nodeLLMRequest by node { llm.writeSession { requestLLM() } } // Executing tool val nodeProcessTool by node { toolCall -> withA2AAgentServer { sendTaskUpdate("Executing tool: ${toolCall.content}", TaskState.Working) } val toolResult = environment.executeTool(toolCall) llm.writeSession { appendPrompt { tool { result(toolResult) } } } withA2AAgentServer { sendTaskUpdate("Tool result: ${toolResult.content}", TaskState.Working) } } // Sending assistant message val nodeProcessAssistant by node { assistantMessage -> withA2AAgentServer { sendTaskUpdate(assistantMessage, TaskState.Completed) } } edge(nodeStart forwardTo nodeSetup) edge(nodeSetup forwardTo nodeLLMRequest) // If a tool call is returned from llm, forward to the tool processing node and then back to llm edge(nodeLLMRequest forwardTo nodeProcessTool onToolCall { true }) edge(nodeProcessTool forwardTo nodeLLMRequest) // If an assistant message is returned from llm, forward to the assistant processing node and then to finish edge(nodeLLMRequest forwardTo nodeProcessAssistant onAssistantMessage { true }) edge(nodeProcessAssistant forwardTo nodeFinish) }, agentConfig = AIAgentConfig( prompt = prompt("agent") { system("You are a helpful assistant.") }, model = GoogleModels.Gemini2_5Pro, maxAgentIterations = 10 ), ) { install(A2AAgentServer) { this.context = context this.eventProcessor = eventProcessor } } /** * Convenience function to send task update event to A2A client * @param content The message content * @param state The task state */ @OptIn(ExperimentalUuidApi::class) private suspend fun A2AAgentServer.sendTaskUpdate( content: String, state: TaskState, ) { val message = A2AMessage( messageId = Uuid.random().toString(), role = Role.Agent, parts = listOf( TextPart(content) ), contextId = context.contextId, taskId = context.taskId, ) val task = Task( id = context.taskId, contextId = context.contextId, status = TaskStatus( state = state, message = message, timestamp = Clock.System.now(), ) ) eventProcessor.sendTaskEvent(task) } ``` ## A2AAgentServer Feature Mechanism The `A2AAgentServer` is a Koog agent feature that enables seamless integration between Koog agents and the A2A protocol. The `A2AAgentServer` feature provides access to the `RequestContext` and `SessionEventProcessor` entities, which are used to communicate with the A2A client inside the Koog agent. To install the feature, call the `install` function on the agent and pass the `A2AAgentServer` feature along with the `RequestContext` and `SessionEventProcessor`: ``` // Install the feature install(A2AAgentServer) { this.context = context this.eventProcessor = eventProcessor } ``` To access these entities from Koog agent strategy, the feature provides a `withA2AAgentServer` function that allows agent nodes to access A2A server capabilities within their execution context. It retrieves the installed `A2AAgentServer` feature and provides it as the receiver for the action block. ``` // Usage within agent nodes withA2AAgentServer { // 'this' is now A2AAgentServer instance eventProcessor.sendTaskUpdate("Processing your request...", TaskState.Working) } ``` ### Start A2A Server After running the server Koog agent will be discoverable and accessible via the A2A protocol. ``` val agentCard = AgentCard( name = "Koog Agent", url = "http://localhost:9999/koog", description = "Simple universal agent powered by Koog", version = "1.0.0", protocolVersion = "0.3.0", preferredTransport = TransportProtocol.JSONRPC, capabilities = AgentCapabilities(streaming = true), defaultInputModes = listOf("text"), defaultOutputModes = listOf("text"), skills = listOf( AgentSkill( id = "koog", name = "Koog Agent", description = "Universal agent powered by Koog. Supports tool calling.", tags = listOf("chat", "tool"), ) ) ) // Server setup val server = A2AServer(agentExecutor = KoogAgentExecutor(), agentCard = agentCard) val transport = HttpJSONRPCServerTransport(server) transport.start(engineFactory = Netty, port = 8080, path = "/chat", wait = true) ``` ## Connecting Koog Agents to A2A Agents ### Create A2A Client and connect to the A2A Server ``` val transport = HttpJSONRPCClientTransport(url = "http://localhost:9999/koog") val agentCardResolver = UrlAgentCardResolver(baseUrl = "http://localhost:9999", path = "/koog") val client = A2AClient(transport = transport, agentCardResolver = agentCardResolver) val agentId = "koog" client.connect() ``` ### Create Koog Agent and add A2A Client to A2AAgentClient Feature To connect to A2A agent from your Koog Agent, you can use the A2AAgentClient feature, which provides a client API for connecting to A2A agents. The principle of the client is the same as the server: you install the feature and pass the `A2AAgentClient` feature along with the `RequestContext` and `SessionEventProcessor`. ``` val agent = AIAgent( promptExecutor = MultiLLMPromptExecutor( LLMProvider.Google to GoogleLLMClient("api-key") ), toolRegistry = ToolRegistry { // declare tools here }, strategy = strategy("test") { val nodeCheckStreaming by nodeA2AClientGetAgentCard().transform { it.capabilities.streaming } val nodeA2ASendMessageStreaming by nodeA2AClientSendMessageStreaming() val nodeA2ASendMessage by nodeA2AClientSendMessage() val nodeProcessStreaming by node>, Unit> { it.collect { response -> when (response.data) { is Task -> { // Process task } is A2AMessage -> { // Process message } is TaskStatusUpdateEvent -> { // Process task status update } is TaskArtifactUpdateEvent -> { // Process task artifact update } } } } val nodeProcessEvent by node { event -> when (event) { is Task -> { // Process task } is A2AMessage -> { // Process message } } } // If streaming is supported, send a message, process response and finish edge(nodeStart forwardTo nodeCheckStreaming transformed { agentId }) edge( nodeCheckStreaming forwardTo nodeA2ASendMessageStreaming onCondition { it == true } transformed { buildA2ARequest(agentId) } ) edge(nodeA2ASendMessageStreaming forwardTo nodeProcessStreaming) edge(nodeProcessStreaming forwardTo nodeFinish) // If streaming is not supported, send a message, process response and finish edge( nodeCheckStreaming forwardTo nodeA2ASendMessage onCondition { it == false } transformed { buildA2ARequest(agentId) } ) edge(nodeA2ASendMessage forwardTo nodeProcessEvent) edge(nodeProcessEvent forwardTo nodeFinish) // If streaming is not supported, send a message, process response and finish edge(nodeCheckStreaming forwardTo nodeFinish onCondition { it == null } transformed { println("Failed to get agents card") } ) }, agentConfig = AIAgentConfig( prompt = prompt("agent") { system("You are a helpful assistant.") }, model = GoogleModels.Gemini2_5Pro, maxAgentIterations = 10 ), ) { install(A2AAgentClient) { this.a2aClients = mapOf(agentId to client) } } @OptIn(ExperimentalUuidApi::class) private fun AIAgentGraphContextBase.buildA2ARequest(agentId: String): A2AClientRequest = A2AClientRequest( agentId = agentId, callContext = ClientCallContext.Default, params = MessageSendParams( message = A2AMessage( messageId = Uuid.random().toString(), role = Role.User, parts = listOf( TextPart(agentInput as String) ) ) ) ) ``` # Agent Client Protocol # Agent Client Protocol Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../module-versioning/) for details. Agent Client Protocol (ACP) is an open-source, standardized protocol that enables client applications to communicate with AI agents through a consistent, bidirectional interface. By implementing ACP in your Koog agent, you ensure it can easily integrate into any ACP-compliant environment, such as an IDE. For more information, see the [Agent Client Protocol](https://agentclientprotocol.com) documentation. ## Integration with Koog The Koog framework integrates with ACP using the [ACP Kotlin SDK](https://github.com/agentclientprotocol/kotlin-sdk) with additional API extensions. This integration provides: - Standardized communication for a Koog agent with ACP-compliant client applications - Automatic execution updates for tool calls, agent thoughts, and completions - Seamless message conversion between Koog's multimodal message formats and ACP's content blocks - Lifecycle mapping of Koog agent states to ACP session events Note Since [ACP Kotlin SDK](https://github.com/agentclientprotocol/kotlin-sdk) is JVM-specific, the ACP integration is currently available only on the JVM platform. ### Add dependencies ACP support is an optional [feature](../features/) that is not available in Koog by default. To implement ACP for your Koog agent, add a dependency for [ai.koog:agents-features-acp](https://mvnrepository.com/artifact/ai.koog/agents-features-acp), which itself has a dependency on [com.agentclientprotocol:acp](https://mvnrepository.com/artifact/com.agentclientprotocol/acp). For example, in case of `build.gradle.kts`: ``` dependencies { implementation("ai.koog:agents-features-acp:$koogVersion") } ``` ### Enable ACP for a Koog agent To bridge a Koog agent's internal [event system](../agent-events/) with the ACP protocol, install the `ai.koog.agents.features.acp.AcpAgent` feature. When installed, it listens for lifecycle events (like tool calls or LLM responses) and sends them to the ACP client. ``` val agent = AIAgent( promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4o ) { install(AcpAgent) { this.sessionId = sessionId this.protocol = protocol this.eventsProducer = eventsProducer this.setDefaultNotifications = true } } ``` Key configuration options: - **`sessionId`**: A unique string identifying the current conversation session. - **`protocol`**: An instance of [`com.agentclientprotocol.protocol.Protocol`](https://github.com/agentclientprotocol/kotlin-sdk/blob/master/acp/src/commonMain/kotlin/com/agentclientprotocol/protocol/Protocol.kt) used for low-level communication. - **`eventsProducer`**: A `kotlinx.coroutines.channels.ProducerScope` where ACP events are sent. For more information, see [Event streaming](#event-streaming). - **`setDefaultNotifications`**: Whether to register default notification handlers for agent lifecycle events. For more information, see [Handling agent notifications](#handling-agent-notifications). This agent must run within the scope of an ACP session as described in the next chapter. ### Implement an ACP-enabled agent To connect your Koog agent to ACP clients, implement two core interfaces from the [ACP Kotlin SDK](https://github.com/agentclientprotocol/kotlin-sdk): - [`AgentSupport`](https://github.com/agentclientprotocol/kotlin-sdk/blob/master/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/AgentSupport.kt): Manages the agent's identity, capabilities, and session lifecycle (creating or loading sessions). - [`AgentSession`](https://github.com/agentclientprotocol/kotlin-sdk/blob/master/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/AgentSession.kt): Manages a single conversation session, handles the `prompt` execution, and manages cancellation. Inside the `prompt()` method of `AgentSession` is where you should initialize and run the ACP-enabled Koog agent. Here is an example: ``` class MyAgentSession( override val sessionId: SessionId, private val promptExecutor: PromptExecutor, private val protocol: Protocol, private val clock: KoogClock ) : AgentSession { private var agentJob: Deferred? = null private val agentMutex = Mutex() override suspend fun prompt( content: List, _meta: JsonElement? ): Flow = channelFlow { val agentConfig = AIAgentConfig( prompt = prompt("acp") { system("You are a helpful assistant.") }.appendPrompt(content), model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 1000 ) // Ensure only one agent session runs at a time agentMutex.withLock { val agent = AIAgent( promptExecutor = promptExecutor, agentConfig = agentConfig ) { install(AcpAgent) { this.sessionId = this@MyAgentSession.sessionId.value this.protocol = this@MyAgentSession.protocol this.eventsProducer = this@channelFlow this.setDefaultNotifications = true } } agentJob = async { agent.run("Hello. How can you help me?") } agentJob?.await() } } private fun Prompt.appendPrompt(content: List): Prompt { return withMessages { messages -> messages + listOf(content.toKoogMessage(clock)) } } override suspend fun cancel() { agentJob?.cancel() } } ``` ``` class MyAgentSupport( private val promptExecutor: PromptExecutor, private val clock: KoogClock, private val protocol: Protocol, ) : AgentSupport { override suspend fun initialize(clientInfo: ClientInfo): AgentInfo { return AgentInfo( protocolVersion = LATEST_PROTOCOL_VERSION, capabilities = AgentCapabilities( loadSession = false, // Set to true if you implement session persistence promptCapabilities = PromptCapabilities( audio = false, image = false, embeddedContext = true ) ) ) } @OptIn(ExperimentalUuidApi::class) override suspend fun createSession(sessionParameters: SessionCreationParameters): AgentSession { val sessionId = SessionId(Uuid.random().toString()) return MyAgentSession(sessionId, promptExecutor, protocol, clock) } override suspend fun loadSession(sessionId: SessionId, sessionParameters: SessionCreationParameters): AgentSession { throw UnsupportedOperationException("Session loading not implemented") } } ``` ## Event streaming The `AgentSession` from the example defines a `prompt()` function that returns a `channelFlow` of events. You then install the `AcpAgent` feature with `this@channelFlow` as `eventsProducer`. This allows sending events from different coroutines. ## Execution synchronization The `AgentSession` from the example uses a mutex to synchronize access to the agent instance because ACP should not trigger a new agent execution until the previous one finishes. For this, creating and running the agent happens in the scope of `withLock` for the defined mutex. You also run the agent asynchronously within the `channelFlow` scope as a deferred job `agentJob` to ensure that the agent is not cancelled prematurely. ## Handling ACP client input ACP clients send user input as a list of [`ContentBlock`](https://agentclientprotocol.com/protocol/schema#contentblock) objects. To process these in Koog, use the `List.toKoogMessage()` extension function to convert ACP content blocks to [`Message.User`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-message/-user/index.html) and append it to your [agent's prompt](../prompts/). The `AgentSession` from the example defines a private function to extend the initial agent prompt in an ACP session: ``` private fun Prompt.appendPrompt(content: List): Prompt { return withMessages { messages -> messages + listOf(content.toKoogMessage(clock)) } } ``` Note An `KoogClock` instance is required to timestamp the message. For more information, see [Converting messages](#converting-messages). ## Converting messages The `agents-features-acp` module provides extension functions to seamlessly convert between Koog's internal message types and [ACP content blocks](https://agentclientprotocol.com/protocol/content). Use the following functions when receiving input from an ACP client: - `List.toKoogMessage()` converts a list of ACP content blocks to [`Message.User`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-message/-user/index.html) - `ContentBlock.toKoogContentPart()` converts a single ACP content block to [`ContentPart`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-content-part/index.html) Use the following functions to construct ACP events or content blocks from Koog messages: - `Message.Response.toAcpEvents()` converts a [`Message.Response`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-message/-response/index.html) to a list of ACP session update events - `ContentPart.toAcpContentBlock()` converts a [`ContentPart`](https://api.koog.ai/prompt/prompt-model/ai.koog.prompt.message/-content-part/index.html) to a single ACP content block ## Handling agent notifications By default, `setDefaultNotifications` is set to `true` and the ACP-enabled agent automatically handles the following notifications: - **Agent completion** Sends `PromptResponseEvent` with `StopReason.END_TURN` when the agent completes successfully - **Agent execution failures** Sends `PromptResponseEvent` with the appropriate stop reason: - `StopReason.MAX_TURN_REQUESTS` when the agent exceeds max iterations - `StopReason.REFUSAL` for other execution failures - **LLM responses** Converts and sends LLM responses as ACP events (text, tool calls, reasoning) - **Tool call lifecycle** Reports tool call status changes: - `ToolCallStatus.IN_PROGRESS` when a tool call starts - `ToolCallStatus.COMPLETED` when a tool call succeeds - `ToolCallStatus.FAILED` when a tool call fails If you want to customize notification handling, set `setDefaultNotifications = false` and process agent events according to the specification. ## Sending custom events Besides automatic notifications, you can send custom events to the ACP client at any point during the agent execution using `sendEvent` within the `withAcpAgent` block. This is useful for progress updates, custom status messages, or plan updates. You can do this inside an `AIAgentContext`, for example, in a node: ``` val plan: Plan = TODO() val strategy = strategy("my-strategy") { val node by node { withAcpAgent { sendEvent( Event.SessionUpdateEvent( SessionUpdate.PlanUpdate(plan.entries) ) ) } } } ``` You can also access the underlying `protocol` to send custom requests to the client, such as authentication requests: ``` val strategy = strategy("my-strategy") { val node by node { withAcpAgent { protocol.sendRequest( AcpMethod.AgentMethods.Authenticate, AuthenticateRequest(methodId = AuthMethodId("Google")) ) } } } ``` ## Examples You can find working examples of Koog agents in the Koog repository under [/examples](https://github.com/JetBrains/koog/tree/develop/examples/). ### Running a console-based ACP client This example runs a console-based ACP client that interacts with a simple Koog agent. 1. Open [/examples/simple-examples](https://github.com/JetBrains/koog/blob/develop/examples/simple-examples/). 1. See the [README](https://github.com/JetBrains/koog/blob/develop/examples/simple-examples/README.md) for information about configuring your API key for an LLM provider. 1. Run the `runExampleAcpApp` Gradle task. 1. When the ACP client starts in the console, type a request for the agent, like: ``` List files in the current directory and create a new file named 'acp-test.txt' with the content 'Hello from ACP!'. ``` 1. Observe the event traces in the console, which show how Koog events are converted to ACP events and sent to the client. ### Connecting an ACP-enabled Koog agent to a JetBrains IDE This example demonstrates how to create an ACP-enabled agent and connect to IntelliJ IDEA. 1. Open [/examples/acp-agent](https://github.com/JetBrains/koog/tree/develop/examples/acp-agent) 1. Run the `installDist` Gradle task. 1. This should create the agent executable: `build/install/acp-agent/bin/acp-agent` (`acp-agent.bat` for Windows). 1. Open IntelliJ IDEA (or another JetBrains IDE). 1. Go to **AI Chat** > **Options** > **Add Custom Agent**. 1. In the opened `acp.json` file, paste the following: ``` { "agent_servers": { "Koog Agent": { "command": "/absolute/path/to/acp-agent/build/install/acp-agent/bin/acp-agent", "args": [], "env": { "OPENAI_API_KEY": "paste-your-api-key-here" } } } } ``` Configuration parameters: - `agent_servers`: Object containing one or more agent configurations - `Koog Agent`: Display name shown in IDE's agent selector - `command`: Absolute path to the agent executable - `args`: Command-line arguments (empty for this agent) - `env`: Environment variables passed to the agent process (OpenAI API key in this example) 1. The agent should become available in the **AI Chat** tool window. For more information about adding custom agents to your IDE, see [AI Assistant documentation](https://www.jetbrains.com/help/ai-assistant/acp.html#add-custom-agent) and [this blog post](https://blog.jetbrains.com/ai/2026/02/koog-x-acp-connect-an-agent-to-your-ide-and-more/). # LLM Parameters # LLM parameters This page provides details about LLM parameters in the Koog agentic framework. LLM parameters let you control and customize the behavior of language models. ## Overview LLM parameters are configuration options that let you fine-tune how language models generate responses. These parameters control aspects like response randomness, length, format, and tool usage. By adjusting the parameters, you optimize model behavior for different use cases, from creative content generation to deterministic structured outputs. In Koog, the `LLMParams` class incorporates LLM parameters and provides a consistent interface for configuring language model behavior. You can use LLM parameters in the following ways: - When creating a prompt: ``` val prompt = prompt( id = "dev-assistant", params = LLMParams( temperature = 0.7, maxTokens = 500 ) ) { // Add a system message to set the context system("You are a helpful assistant.") // Add a user message user("Tell me about Kotlin") } ``` ``` Prompt prompt = Prompt.builder("dev-assistant") .withParams(new LLMParams( 0.7, // temperature 500, // maxTokens 1, // numberOfChoices null, // speculation null, // schema LLMParams.ToolChoice.Auto.INSTANCE, // toolChoice null, // user null // additionalProperties )) .system("You are a helpful assistant.") .user("Tell me about Kotlin") .build(); ``` For more information about prompt creation, see [Prompts](../prompts/prompt-creation/). - When creating a subgraph: ``` val processQuery by subgraphWithTask( tools = listOf(searchTool, calculatorTool, weatherTool), llmModel = OpenAIModels.Chat.GPT4o, llmParams = LLMParams( temperature = 0.7, maxTokens = 500 ), parallelTools = false, assistantResponseRepeatMax = 3, ) { userQuery -> """ You are a helpful assistant that can answer questions about various topics. Please help with the following query: $userQuery """ } ``` ``` ``` For more information about existing subgraph types in Koog, see [Predefined subgraphs](../nodes-and-components/#predefined-subgraphs). To learn how to create and implement your own subgraphs, see [Custom subgraphs](../custom-subgraphs/). - When updating a prompt in an LLM write session: ``` llm.writeSession { changeLLMParams( LLMParams( temperature = 0.7, maxTokens = 500 ) ) } ``` ``` ``` For more information about sessions, see [LLM sessions and manual history management](../sessions/). ## LLM parameter reference The following table provides a reference of LLM parameters included in the `LLMParams` class and supported by all LLM providers that are available in Koog out of the box. For a list of parameters that are specific to some providers, see [Provider-specific parameters](#provider-specific-parameters). | Parameter | Type | Description | | ---------------------- | ------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `temperature` | Double | Controls randomness in the output. Higher values, such as 0.7–1.0, produce more diverse and creative responses, while lower values produce more deterministic and focused responses. | | `maxTokens` | Integer | Maximum number of tokens to generate in the response. Useful for controlling response length. | | `numberOfChoices` | Integer | Number of alternative responses to generate. Must be greater than 0. | | `speculation` | String | A speculative configuration string that influences model behavior, designed to enhance result speed and accuracy. Supported only by certain models, but may greatly improve speed and accuracy. | | `schema` | Schema | Defines the structure for the model's response format, enabling structured outputs like JSON. For more information, see [Schema](#schema). | | `toolChoice` | ToolChoice | Controls tool calling behavior of the language model. For more information, see [Tool choice](#tool-choice). | | `user` | String | Identifier for the user making the request, which can be used for tracking purposes. | | `additionalProperties` | Map\ | Additional properties that can be used to store custom parameters specific to certain model providers. | For a list of default values for each parameter, see the corresponding LLM provider documentation: - [OpenAI Chat](https://platform.openai.com/docs/api-reference/chat/create) - [OpenAI Responses](https://platform.openai.com/docs/api-reference/responses/create) - [Google](https://ai.google.dev/api/generate-content#generationconfig) β - [Anthropic](https://platform.claude.com/docs/en/api/messages/create) - [Mistral](https://docs.mistral.ai/api/#operation/chatCompletions) β - [DeepSeek](https://api-docs.deepseek.com/api/create-chat-completion#request) β - [OpenRouter](https://openrouter.ai/docs/api/reference/parameters) - Alibaba β ([DashScope](https://www.alibabacloud.com/help/en/model-studio/qwen-api-reference)) - [Ollama](https://docs.ollama.com/api/openai-compatibility) ## Schema The `Schema` interface defines the structure for the model's response format. Koog supports JSON schemas, as described in the sections below. ### JSON schemas JSON schemas let you request structured JSON data from language models. Koog supports the following two types of JSON schemas: 1. **Basic JSON Schema** (`LLMParams.Schema.JSON.Basic`): Used for basic JSON processing capabilities. This format primarily focuses on nested data definitions without advanced JSON Schema functionalities. ``` // Create parameters with a basic JSON schema val jsonParams = LLMParams( temperature = 0.2, schema = LLMParams.Schema.JSON.Basic( name = "PersonInfo", schema = JsonObject(mapOf( "type" to JsonPrimitive("object"), "properties" to JsonObject( mapOf( "name" to JsonObject(mapOf("type" to JsonPrimitive("string"))), "age" to JsonObject(mapOf("type" to JsonPrimitive("number"))), "skills" to JsonObject( mapOf( "type" to JsonPrimitive("array"), "items" to JsonObject(mapOf("type" to JsonPrimitive("string"))) ) ) ) ), "additionalProperties" to JsonPrimitive(false), "required" to JsonArray(listOf(JsonPrimitive("name"), JsonPrimitive("age"), JsonPrimitive("skills"))) )) ) ) ``` ``` // Create parameters with a basic JSON schema LLMParams jsonParams = new LLMParams( 0.2, // temperature null, // maxTokens 1, // numberOfChoices null, // speculation new LLMParams.Schema.JSON.Basic( "PersonInfo", new JsonObject(Map.of( "type", new JsonPrimitive("object"), "properties", new JsonObject(Map.of( "name", new JsonObject(Map.of("type", new JsonPrimitive("string"))), "age", new JsonObject(Map.of("type", new JsonPrimitive("number"))), "skills", new JsonObject(Map.of( "type", new JsonPrimitive("array"), "items", new JsonObject(Map.of("type", new JsonPrimitive("string"))) )) )), "additionalProperties", new JsonPrimitive(false), "required", new JsonArray(List.of( new JsonPrimitive("name"), new JsonPrimitive("age"), new JsonPrimitive("skills") )) )) ), LLMParams.ToolChoice.Auto.INSTANCE, // toolChoice null, // user null // additionalProperties ); ``` 2. **Standard JSON Schema** (`LLMParams.Schema.JSON.Standard`): Represents a standard JSON schema according to [json-schema.org](https://json-schema.org/). This format is a proper subset of the official JSON Schema specification. Note that the flavor across different LLM providers might vary, since not all of them support full JSON schemas. ``` // Create parameters with a standard JSON schema val standardJsonParams = LLMParams( temperature = 0.2, schema = LLMParams.Schema.JSON.Standard( name = "ProductCatalog", schema = JsonObject(mapOf( "type" to JsonPrimitive("object"), "properties" to JsonObject(mapOf( "products" to JsonObject(mapOf( "type" to JsonPrimitive("array"), "items" to JsonObject(mapOf( "type" to JsonPrimitive("object"), "properties" to JsonObject(mapOf( "id" to JsonObject(mapOf("type" to JsonPrimitive("string"))), "name" to JsonObject(mapOf("type" to JsonPrimitive("string"))), "price" to JsonObject(mapOf("type" to JsonPrimitive("number"))), "description" to JsonObject(mapOf("type" to JsonPrimitive("string"))) )), "additionalProperties" to JsonPrimitive(false), "required" to JsonArray(listOf(JsonPrimitive("id"), JsonPrimitive("name"), JsonPrimitive("price"), JsonPrimitive("description"))) )) )) )), "additionalProperties" to JsonPrimitive(false), "required" to JsonArray(listOf(JsonPrimitive("products"))) )) ) ) ``` ``` // Create parameters with a standard JSON schema LLMParams standardJsonParams = new LLMParams( 0.2, // temperature null, // maxTokens 1, // numberOfChoices null, // speculation new LLMParams.Schema.JSON.Standard( "ProductCatalog", new JsonObject(Map.of( "type", new JsonPrimitive("object"), "properties", new JsonObject(Map.of( "products", new JsonObject(Map.of( "type", new JsonPrimitive("array"), "items", new JsonObject(Map.of( "type", new JsonPrimitive("object"), "properties", new JsonObject(Map.of( "id", new JsonObject(Map.of("type", new JsonPrimitive("string"))), "name", new JsonObject(Map.of("type", new JsonPrimitive("string"))), "price", new JsonObject(Map.of("type", new JsonPrimitive("number"))), "description", new JsonObject(Map.of("type", new JsonPrimitive("string"))) )), "additionalProperties", new JsonPrimitive(false), "required", new JsonArray(List.of( new JsonPrimitive("id"), new JsonPrimitive("name"), new JsonPrimitive("price"), new JsonPrimitive("description") )) )) )) )), "additionalProperties", new JsonPrimitive(false), "required", new JsonArray(List.of(new JsonPrimitive("products"))) )) ), LLMParams.ToolChoice.Auto.INSTANCE, // toolChoice null, // user null // additionalProperties ); ``` ## Tool choice The `ToolChoice` class controls how the language model uses tools. It provides the following options: - `LLMParams.ToolChoice.Named`: the language model calls the specified tool. Takes the `name` string argument that represents the name of the tool to call. - `LLMParams.ToolChoice.All`: the language model calls all tools. - `LLMParams.ToolChoice.None`: the language model does not call tools and only generates text. - `LLMParams.ToolChoice.Auto`: the language model automatically decides whether to call tools and which tool to call. - `LLMParams.ToolChoice.Required`: the language model calls at least one tool. Here is an example of using the `LLMParams.ToolChoice.Named` class to call a specific tool: ``` val specificToolParams = LLMParams( toolChoice = LLMParams.ToolChoice.Named(name = "calculator") ) ``` ``` LLMParams specificToolParams = new LLMParams( null, // temperature null, // maxTokens 1, // numberOfChoices null, // speculation null, // schema new LLMParams.ToolChoice.Named("calculator"), // toolChoice null, // user null // additionalProperties ); ``` ## Provider-specific parameters Koog supports provider-specific parameters for some LLM providers. These parameters extend the base `LLMParams` class and add provider-specific functionality. The following classes include parameters that are specific per provider: - `OpenAIChatParams`: Parameters specific to the OpenAI Chat Completions API. - `OpenAIResponsesParams`: Parameters specific to the OpenAI Responses API. - `GoogleParams`: Parameters specific to Google models. - `AnthropicParams`: Parameters specific to Anthropic models. - `MistralAIParams`: Parameters specific to Mistral models. - `DeepSeekParams`: Parameters specific to DeepSeek models. - `OpenRouterParams`: Parameters specific to OpenRouter models. - `DashscopeParams`: Parameters specific to Alibaba models. - `OllamaParams`: Parameters specific to Ollama models. Here is the complete reference of provider-specific parameters in Koog: | Parameter | Type | Description | | ------------------- | ---------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `audio` | OpenAIAudioConfig | Audio output configuration when using audio-capable models. For more information, see the API documentation for [OpenAIAudioConfig](api:prompt-executor-openai-client-base::ai.koog.prompt.executor.clients.openai.base.models.OpenAIAudioConfig). | | `frequencyPenalty` | Double | Penalizes frequent tokens to reduce repetition. Higher `frequencyPenalty` values result in larger variations of phrasing and reduced repetition. Takes a value in the range of -2.0 to 2.0. | | `logprobs` | Boolean | If `true`, includes log-probabilities for output tokens. | | `parallelToolCalls` | Boolean | If `true`, multiple tool calls can run in parallel. Particularly applicable to custom nodes or LLM interactions outside of agent strategies. | | `presencePenalty` | Double | Prevents the model from reusing tokens that have already been included in the output. Higher values encourage the introduction of new tokens and topics. Takes a value in the range of -2.0 to 2.0. | | `promptCacheKey` | String | Stable cache key for prompt caching. OpenAI uses it to cache responses for similar requests. | | `reasoningEffort` | ReasoningEffort | Specifies the level of reasoning effort that the model will use. For more information and available values, see the API documentation for [ReasoningEffort](api:prompt-executor-openai-client-base::ai.koog.prompt.executor.clients.openai.base.models.ReasoningEffort). | | `safetyIdentifier` | String | A stable and unique user identifier that may be used to detect users who violate OpenAI policies. | | `serviceTier` | ServiceTier | OpenAI processing tier selection that lets you prioritize performance over cost or vice versa. For more information, see the API documentation for [ServiceTier](api:prompt-executor-openai-client-base::ai.koog.prompt.executor.clients.openai.base.models.ServiceTier). | | `stop` | List | Strings that signal to the model that it should stop generating content when it encounters any of them. For example, to make the model stop generating content when it produces two newlines, specify the stop sequence as `stop = listOf("/n/n")`. | | `store` | Boolean | If `true`, the provider may store outputs for later retrieval. | | `topLogprobs` | Integer | Number of top most likely tokens per position. Takes a value in the range of 0–20. Requires the `logprobs` parameter to be set to `true`. | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | `webSearchOptions` | OpenAIWebSearchOptions | Configure web search tool usage (if supported). For more information, see the API documentation for [OpenAIWebSearchOptions](api:prompt-executor-openai-client-base::ai.koog.prompt.executor.clients.openai.base.models.OpenAIWebSearchOptions). | | Parameter | Type | Description | | ------------------- | ------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `background` | Boolean | Run the response in the background. | | `include` | List | Additional data to include in the model's response, such as sources of web search tool call or search results of a file search tool call. For detailed reference information, see [OpenAIInclude](api:prompt-executor-openai-client::ai.koog.prompt.executor.clients.openai.models.OpenAIInclude) in the Koog API reference. To learn more about the `include` parameter, see [OpenAI's documentation](https://platform.openai.com/docs/api-reference/responses/create#responses-create-include). | | `logprobs` | Boolean | If `true`, includes log-probabilities for output tokens. | | `maxToolCalls` | Integer | Maximum total number of built-in tool calls allowed in this response. Takes a value equal to or greater than `0`. | | `parallelToolCalls` | Boolean | If `true`, multiple tool calls can run in parallel. Particularly applicable to custom nodes or LLM interactions outside of agent strategies. | | `promptCacheKey` | String | Stable cache key for prompt caching. OpenAI uses it to cache responses for similar requests. | | `reasoning` | ReasoningConfig | Reasoning configuration for reasoning-capable models. For more information, see the API documentation for [ReasoningConfig](api:prompt-executor-openai-client::ai.koog.prompt.executor.clients.openai.models.ReasoningConfig). | | `safetyIdentifier` | String | A stable and unique user identifier that may be used to detect users who violate OpenAI policies. | | `serviceTier` | ServiceTier | OpenAI processing tier selection that lets you prioritize performance over cost or vice versa. For more information, see the API documentation for [ServiceTier](api:prompt-executor-openai-client-base::ai.koog.prompt.executor.clients.openai.base.models.ServiceTier). | | `store` | Boolean | If `true`, the provider may store outputs for later retrieval. | | `topLogprobs` | Integer | Number of top most likely tokens per position. Takes a value in the range of 0–20. Requires the `logprobs` parameter to be set to `true`. | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | `truncation` | Truncation | Truncation strategy when nearing the context window. For more information, see the API documentation for [Truncation](api:prompt-executor-openai-client::ai.koog.prompt.executor.clients.openai.models.Truncation). | | Parameter | Type | Description | | ---------------- | -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `thinkingConfig` | GoogleThinkingConfig | Controls whether the model should expose its chain-of-thought and how many tokens it may spend on it. For more information, see the API reference for [GoogleThinkingConfig](api:prompt-executor-google-client::ai.koog.prompt.executor.clients.google.models.GoogleThinkingConfig). | | `topK` | Integer | Number of top tokens to consider when generating the output. Takes a value greater than or equal to 0 (provider-specific minimums may apply). | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | Parameter | Type | Description | | --------------- | ------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `container` | String | Container identifier for reuse across requests. Containers are used by Anthropic's code execution tool to provide a secure and containerized code execution environment. By providing the container identifier from a previous response, you can reuse containers across multiple requests, which preserves created files between requests. For more information, see [Containers](https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool#containers) in Anthropic's documentation. | | `mcpServers` | List | Definitions of MCP servers to be used in the request. Supports at most 20 servers. For more information, see the API reference for [AnthropicMCPServerURLDefinition](api:prompt-executor-anthropic-client::ai.koog.prompt.executor.clients.anthropic.models.AnthropicMCPServerURLDefinition). | | `serviceTier` | ServiceTier | OpenAI processing tier selection that lets you prioritize performance over cost or vice versa. For more information, see the API documentation for [ServiceTier](api:prompt-executor-openai-client-base::ai.koog.prompt.executor.clients.openai.base.models.ServiceTier). | | `stopSequences` | List | Custom text sequences that cause the model to stop generating content. If matched, the value of `stop_reason` in the response is `stop_sequence`. | | `thinking` | AnthropicThinking | Configuration for activating Claude's extended thinking. When activated, responses also include thinking content blocks. For more information, see the API reference for [AnthropicThinking](api:prompt-executor-anthropic-client::ai.koog.prompt.executor.clients.anthropic.models.AnthropicThinking). | | `topK` | Integer | Number of top tokens to consider when generating the output. Takes a value greater than or equal to 0 (provider-specific minimums may apply). | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | Parameter | Type | Description | | ------------------- | ------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `frequencyPenalty` | Double | Penalizes frequent tokens to reduce repetition. Higher `frequencyPenalty` values result in larger variations of phrasing and reduced repetition. Takes a value in the range of -2.0 to 2.0. | | `parallelToolCalls` | Boolean | If `true`, multiple tool calls can run in parallel. Particularly applicable to custom nodes or LLM interactions outside of agent strategies. | | `presencePenalty` | Double | Prevents the model from reusing tokens that have already been included in the output. Higher values encourage the introduction of new tokens and topics. Takes a value in the range of -2.0 to 2.0. | | `promptMode` | String | Lets you toggle between the reasoning mode and no system prompt. When set to `reasoning`, the default system prompt for reasoning models is used. For more information, see Mistral's [Reasoning](https://docs.mistral.ai/capabilities/reasoning) documentation. | | `randomSeed` | Integer | The seed to use for random sampling. If set, different calls with the same parameters and the same seed value will generate deterministic results. | | `safePrompt` | Boolean | Specifies whether to inject a safety prompt before all conversations. The safety prompt is used to enforce guardrails and protect against harmful content. For more information, see Mistral's [Moderation & Guardarailing](https://docs.mistral.ai/capabilities/guardrailing) documentation. | | `stop` | List | Strings that signal to the model that it should stop generating content when it encounters any of them. For example, to make the model stop generating content when it produces two newlines, specify the stop sequence as `stop = listOf("/n/n")`. | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | Parameter | Type | Description | | ------------------ | ------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `frequencyPenalty` | Double | Penalizes frequent tokens to reduce repetition. Higher `frequencyPenalty` values result in larger variations of phrasing and reduced repetition. Takes a value in the range of -2.0 to 2.0. | | `logprobs` | Boolean | If `true`, includes log-probabilities for output tokens. | | `presencePenalty` | Double | Prevents the model from reusing tokens that have already been included in the output. Higher values encourage the introduction of new tokens and topics. Takes a value in the range of -2.0 to 2.0. | | `stop` | List | Strings that signal to the model that it should stop generating content when it encounters any of them. For example, to make the model stop generating content when it produces two newlines, specify the stop sequence as `stop = listOf("/n/n")`. | | `topLogprobs` | Integer | Number of top most likely tokens per position. Takes a value in the range of 0–20. Requires the `logprobs` parameter to be set to `true`. | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | Parameter | Type | Description | | ------------------- | ------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `frequencyPenalty` | Double | Penalizes frequent tokens to reduce repetition. Higher `frequencyPenalty` values result in larger variations of phrasing and reduced repetition. Takes a value in the range of -2.0 to 2.0. | | `logprobs` | Boolean | If `true`, includes log-probabilities for output tokens. | | `minP` | Double | Filters out tokens whose relative probability to the most likely token is below the defined `minP` value. Takes a value in the range of 0.0–0.1. | | `models` | List | List of allowed models for the request. | | `presencePenalty` | Double | Prevents the model from reusing tokens that have already been included in the output. Higher values encourage the introduction of new tokens and topics. Takes a value in the range of -2.0 to 2.0. | | `provider` | ProviderPreferences | Includes a range of parameters that let you explicitly control how OpenRouter chooses which LLM provider to use. For more information, see the API documentation on [ProviderPreferences](api:prompt-executor-openrouter-client::ai.koog.prompt.executor.clients.openrouter.models.ProviderPreferences). | | `repetitionPenalty` | Double | Penalizes token repetition. Next-token probabilities for tokens that already appeared in the output are divided by the value of `repetitionPenalty`, which makes them less likely to appear again if `repetitionPenalty > 1`. Takes a value greater than 0.0 and lower than or equal to 2.0. | | `route` | String | Request routing strategy to use. | | `stop` | List | Strings that signal to the model that it should stop generating content when it encounters any of them. For example, to make the model stop generating content when it produces two newlines, specify the stop sequence as `stop = listOf("/n/n")`. | | `topA` | Double | Dynamically adjusts the sampling window based on model confidence. If the model is confident (there are dominant high-probability next tokens), it keeps the sampling window limited to a few top tokens. If the confidence is low (there are many tokens with similar probabilities), keeps more tokens in the sampling window. Takes a value in the range of 0.0–0.1 (inclusive). Higher value means greater dynamic adaptation. | | `topK` | Integer | Number of top tokens to consider when generating the output. Takes a value greater than or equal to 0 (provider-specific minimums may apply). | | `topLogprobs` | Integer | Number of top most likely tokens per position. Takes a value in the range of 0–20. Requires the `logprobs` parameter to be set to `true`. | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | `transforms` | List | List of context transforms. Defines how context is transformed when it exceeds the model's token limit. The default transformation is `middle-out` which truncates from the middle of the prompt. Use empty list for no transformations. For more information, see [Message Transforms](https://openrouter.ai/docs/guides/features/message-transforms) in OpenRouter documentation. | | Parameter | Type | Description | | ------------------- | ------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `enableSearch` | Boolean | Specifies whether to enable web search functionality. For more information, see Alibaba's [Web search](https://www.alibabacloud.com/help/en/model-studio/web-search?spm=a2c63.p38356.0.i14) documentation. | | `enableThinking` | Boolean | Specifies whether to enable thinking mode when using a hybrid thinking model. For more information, see Alibaba's documentation on [Deep thinking](https://www.alibabacloud.com/help/en/model-studio/deep-thinking?spm=a2c63.p38356.0.i11). | | `frequencyPenalty` | Double | Penalizes frequent tokens to reduce repetition. Higher `frequencyPenalty` values result in larger variations of phrasing and reduced repetition. Takes a value in the range of -2.0 to 2.0. | | `logprobs` | Boolean | If `true`, includes log-probabilities for output tokens. | | `parallelToolCalls` | Boolean | If `true`, multiple tool calls can run in parallel. Particularly applicable to custom nodes or LLM interactions outside of agent strategies. | | `presencePenalty` | Double | Prevents the model from reusing tokens that have already been included in the output. Higher values encourage the introduction of new tokens and topics. Takes a value in the range of -2.0 to 2.0. | | `stop` | List | Strings that signal to the model that it should stop generating content when it encounters any of them. For example, to make the model stop generating content when it produces two newlines, specify the stop sequence as `stop = listOf("/n/n")`. | | `topLogprobs` | Integer | Number of top most likely tokens per position. Takes a value in the range of 0–20. Requires the `logprobs` parameter to be set to `true`. | | `topP` | Double | Also referred to as nucleus sampling. Creates a subset of next tokens by adding tokens with the highest probability values to the subset until the sum of their probabilities reaches the specified `topP` value. Takes a value greater than 0.0 and lower than or equal to 1.0. | | Parameter | Type | Description | | --------- | ------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `think` | Boolean | Configuration for activating Ollama extended thinking. When activated, responses also include thinking content blocks. For more information, see the API reference for [Ollama thinking](https://docs.ollama.com/capabilities/thinking#enable-thinking-in-api-calls). | The following example shows defined OpenRouter LLM parameters using the provider-specific `OpenRouterParams` class: ``` val openRouterParams = OpenRouterParams( temperature = 0.7, maxTokens = 500, frequencyPenalty = 0.5, presencePenalty = 0.5, topP = 0.9, topK = 40, repetitionPenalty = 1.1, models = listOf("anthropic/claude-3-opus", "anthropic/claude-3-sonnet"), transforms = listOf("middle-out") ) ``` ``` OpenRouterParams openRouterParams = new OpenRouterParams( 0.7, // temperature 500, // maxTokens 1, // numberOfChoices null, // speculation null, // schema null, // toolChoice null, // user null, // additionalProperties 0.5, // frequencyPenalty null, // logprobs null, // minP Arrays.asList("anthropic/claude-3-opus", "anthropic/claude-3-sonnet"), // models 0.5, // presencePenalty null, // provider 1.1, // repetitionPenalty null, // route null, // stop null, // topA 40, // topK null, // topLogprobs 0.9, // topP Arrays.asList("middle-out") // transforms ); ``` ## Usage examples ### Basic usage ``` // A basic set of parameters with limited length val basicParams = LLMParams( temperature = 0.7, maxTokens = 150, toolChoice = LLMParams.ToolChoice.Auto ) ``` ``` // A basic set of parameters with limited length LLMParams basicParams = new LLMParams( 0.7, // temperature 150, // maxTokens 1, // numberOfChoices null, // speculation null, // schema LLMParams.ToolChoice.Auto.INSTANCE, // toolChoice null, // user null // additionalProperties ); ``` ### Reasoning control You implement reasoning control through provider-specific parameters that control model reasoning. When using the OpenAI Chat API and models that support reasoning, use the `reasoningEffort` parameter to control how many reasoning tokens the model generates before providing a response: ``` val openAIReasoningEffortParams = OpenAIChatParams( reasoningEffort = ReasoningEffort.MEDIUM ) ``` ``` OpenAIChatParams openAIReasoningEffortParams = new OpenAIChatParams( null, // temperature null, // maxTokens 1, // numberOfChoices null, // speculation null, // schema null, // toolChoice null, // user null, // additionalProperties null, // audio null, // frequencyPenalty null, // logprobs null, // parallelToolCalls null, // presencePenalty null, // promptCacheKey ReasoningEffort.MEDIUM, // reasoningEffort null, // safetyIdentifier null, // serviceTier null, // stop null, // store null, // topLogprobs null, // topP null // webSearchOptions ); ``` In addition, when using the OpenAI Responses API in a stateless mode, you keep an encrypted history of reasoning items and send it to the model in every conversation turn. The encryption is done on the OpenAI side, and you need to request encrypted reasoning tokens by setting the `include` parameter in your requests to `reasoning.encrypted_content`. You can then pass the encrypted reasoning tokens back to the model in the next conversation turns. ``` val openAIStatelessReasoningParams = OpenAIResponsesParams( include = listOf(OpenAIInclude.REASONING_ENCRYPTED_CONTENT) ) ``` ``` OpenAIResponsesParams openAIStatelessReasoningParams = new OpenAIResponsesParams( null, // temperature null, // maxTokens 1, // numberOfChoices null, // speculation null, // schema null, // toolChoice null, // user null, // additionalProperties null, // background Arrays.asList(OpenAIInclude.REASONING_ENCRYPTED_CONTENT), // include null, // logprobs null, // maxToolCalls null, // parallelToolCalls null, // promptCacheKey null, // reasoning null, // safetyIdentifier null, // serviceTier null, // store null, // topLogprobs null, // topP null // truncation ); ``` ### Custom parameters To add custom parameters that may be provider specific and not supported in Koog out of the box, use the `additionalProperties` property as shown in the example below. ``` // Add custom parameters for specific model providers val customParams = LLMParams( additionalProperties = additionalPropertiesOf( "top_p" to 0.95, "frequency_penalty" to 0.5, "presence_penalty" to 0.5 ) ) ``` ``` // Add custom parameters for specific model providers LLMParams customParams = new LLMParams( null, // temperature null, // maxTokens 1, // numberOfChoices null, // speculation null, // schema null, // toolChoice null, // user AdditionalPropertiesKt.additionalPropertiesOf( "top_p", 0.95, "frequency_penalty", 0.5, "presence_penalty", 0.5 ) ); ``` ### Setting and overriding parameters The code sample below shows how you can define a set of LLM parameters that you may want to use primarily, then create another set by partially overriding values from the original set and adding new values to it. This lets you define parameters that are common to most requests but also add more specific parameter combinations without having to repeat the common parameters. ``` // Define default parameters val defaultParams = LLMParams( temperature = 0.7, maxTokens = 150, toolChoice = LLMParams.ToolChoice.Auto ) // Create parameters with some overrides, using defaults for the rest val overrideParams = LLMParams( temperature = 0.2, numberOfChoices = 3 ).default(defaultParams) ``` ``` // Define default parameters LLMParams defaultParams = new LLMParams( 0.7, // temperature 150, // maxTokens 1, // numberOfChoices null, // speculation null, // schema LLMParams.ToolChoice.Auto.INSTANCE, // toolChoice null, // user null // additionalProperties ); // Create parameters with some overrides, using defaults for the rest LLMParams overrideParams = new LLMParams( 0.2, // temperature null, // maxTokens 3, // numberOfChoices null, // speculation null, // schema null, // toolChoice null, // user null // additionalProperties ).applyDefaults(defaultParams); ``` The values in the resulting `overrideParams` set are equivalent to the following: ``` val overrideParams = LLMParams( temperature = 0.2, maxTokens = 150, toolChoice = LLMParams.ToolChoice.Auto, numberOfChoices = 3 ) ``` ``` LLMParams overrideParams = new LLMParams( 0.2, // temperature 150, // maxTokens 3, // numberOfChoices null, // speculation null, // schema LLMParams.ToolChoice.Auto.INSTANCE, // toolChoice null, // user null // additionalProperties ); ``` # Model capabilities Koog provides a set of abstractions and implementations for working with Large Language Models (LLMs) from various LLM providers in a provider-agnostic way. The set includes the following classes: - **LLMCapability**: a class hierarchy that defines various capabilities that LLMs can support, such as: - Temperature adjustment for controlling response randomness - Tool integration for external system interaction - Vision processing for handling visual data - Embedding generation for vector representations - Completion for text generation tasks - Schema support for structured data (JSON with Simple and Full variants) - Speculation for exploratory responses - **LLModel**: a data class that represents a specific LLM with its provider, unique identifier, and supported capabilities. This serves as a foundation for interacting with different LLM providers in a unified way, allowing applications to work with various models while abstracting away provider-specific details. ## LLM capabilities LLM capabilities represent specific features or functionalities that a Large Language Model can support. In the Koog framework, capabilities are used to define what a particular model can do and how it can be configured. Each capability is represented as a subclass or data object of the `LLMCapability` class. When configuring an LLM for use in your application, you specify which capabilities it supports by adding them to the `capabilities` list when creating an `LLModel` instance. This allows the framework to properly interact with the model and use its features appropriately. ### Core capabilities The list below includes the core, LLM-specific capabilities that are available for models in the Koog framework: - **Speculation** (`LLMCapability.Speculation`): lets the model generate speculative or exploratory responses with varying degrees of likelihood. Useful for creative or hypothetical scenarios where a broader range of potential outcomes is desired. - **Temperature** (`LLMCapability.Temperature`): allows adjustment of the model's response randomness or creativity levels. Higher temperature values produce more diverse outputs, while lower values lead to more focused and deterministic responses. - **Tools** (`LLMCapability.Tools`): indicates support for external tool usage or integration. This capability lets the model run specific tools or interact with external systems. - **Tool choice** (`LLMCapability.ToolChoice`): configures how tool calling works with the LLM. Depending on the model, it can be configured to: - Automatically choose between generating text or tool calls - Generate only tool calls, never text - Generate only text, never tool calls - Force calling a specific tool among the defined tools - **Multiple choices** (`LLMCapability.MultipleChoices`): lets the model generate multiple independent reply choices to a single prompt. ### Media processing capabilities The following list represents a set of capabilities for processing media content such as images or audio: - **Vision** (`LLMCapability.Vision`): a class for vision-based capabilities that process, analyze, and infer insights from visual data. Supports the following types of visual data: - **Image** (`LLMCapability.Vision.Image`): handles image-related vision tasks such as image analysis, recognition, and interpretation. - **Video** (`LLMCapability.Vision.Video`): processes video data, including analyzing and understanding video content. - **Audio** (`LLMCapability.Audio`): provides audio-related functionalities such as transcription, audio generation, or audio-based interactions. - **Document** (`LLMCapability.Document`): enables handling and processing of document-based inputs and outputs. ### Text processing capabilities The following list of capabilities represents text generation and processing functionalities: - **Embedding** (`LLMCapability.Embed`): lets models generate vector embeddings from an input text, enabling similarity comparisons, clustering, and other vector-based analyses. - **Completion** (`LLMCapability.Completion`): includes the generation of text or content based on given input context, such as completing sentences, generating suggestions, or producing content that aligns with input data. - **Prompt caching** (`LLMCapability.PromptCaching`): supports caching functionalities for prompts, potentially improving performance for repeated or similar queries. - **Moderation** (`LLMCapability.Moderation`): lets the model analyze text for potentially harmful content and classify it according to various categories such as harassment, hate speech, self-harm, sexual content, violence, etc. ### Schema capabilities The list below indicates the capabilities related to processing structured data: - **Schema** (`LLMCapability.Schema`): a class for structured schema capabilities related to data interaction and encoding using specific formats. Includes support for the following format: - **JSON** (`LLMCapability.Schema.JSON`): JSON schema support with different levels: - **Basic** (`LLMCapability.Schema.JSON.Basic`): provides lightweight or basic JSON processing capabilities. - **Standard** (`LLMCapability.Schema.JSON.Standard`): offers comprehensive JSON schema support for complex data structures. ## Creating a model (LLModel) configuration To define a model in a universal, provider-agnostic way, create a model configuration as an instance of the `LLModel` class with the following parameters: | Name | Data type | Required | Default | Description | | ----------------- | ------------------- | -------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `provider` | LLMProvider | Yes | | The provider of the LLM, such as Google or OpenAI. This identifies the company or organization that created or hosts the model. | | `id` | String | Yes | | A unique identifier for the LLM instance. This typically represents the specific model version or name. For example, `gpt-4-turbo`, `claude-3-opus`, `llama-3-2`. | | `capabilities` | List | Yes | | A list of capabilities supported by the LLM, such as temperature adjustment, tools usage, or schema-based tasks. These capabilities define what the model can do and how it can be configured. | | `contextLength` | Long | Yes | | The context length of the LLM. This is the maximum number of tokens the LLM can process. | | `maxOutputTokens` | Long | No | `null` | The maximum number of tokens that can be generated by the provider for the LLM. | ### Examples This section provides detailed examples of creating `LLModel` instances with different capabilities. The code below represents a basic LLM configuration with core capabilities: ``` val basicModel = LLModel( provider = LLMProvider.OpenAI, id = "gpt-4-turbo", capabilities = listOf( LLMCapability.Temperature, LLMCapability.Tools, LLMCapability.Schema.JSON.Standard ), contextLength = 128_000 ) ``` ``` LLModel basicModel = new LLModel( LLMProvider.OpenAI, "gpt-4-turbo", List.of( LLMCapability.Temperature.INSTANCE, LLMCapability.Tools.INSTANCE, LLMCapability.Schema.JSON.Standard.INSTANCE ), 128_000L ); ``` The model configuration below is a multimodal LLM with vision capabilities: ``` val visionModel = LLModel( provider = LLMProvider.OpenAI, id = "gpt-4-vision", capabilities = listOf( LLMCapability.Temperature, LLMCapability.Vision.Image, LLMCapability.MultipleChoices ), contextLength = 1_047_576, maxOutputTokens = 32_768 ) ``` ``` LLModel visionModel = new LLModel( LLMProvider.OpenAI, "gpt-4-vision", List.of( LLMCapability.Temperature.INSTANCE, LLMCapability.Vision.Image.INSTANCE, LLMCapability.MultipleChoices.INSTANCE ), 1_047_576L, 32_768L ); ``` An LLM with audio processing capabilities: ``` val audioModel = LLModel( provider = LLMProvider.Anthropic, id = "claude-3-opus", capabilities = listOf( LLMCapability.Audio, LLMCapability.Temperature, LLMCapability.PromptCaching ), contextLength = 200_000 ) ``` ``` LLModel audioModel = new LLModel( LLMProvider.Anthropic, "claude-3-opus", List.of( LLMCapability.Audio.INSTANCE, LLMCapability.Temperature.INSTANCE, LLMCapability.PromptCaching.INSTANCE ), 200_000L ); ``` In addition to creating models as `LLModel` instances and having to specify all related parameters, Koog includes a collection of predefined models and their configurations with supported capabilities. To use a predefined Ollama model, specify it as follows: ``` val metaModel = OllamaModels.Meta.LLAMA_3_2 ``` ``` LLModel metaModel = OllamaModels.Meta.LLAMA_3_2; ``` To check whether a model supports a specific capability use the `contains` method to check for the presence of the capability in the `capabilities` list: ``` // Check if models support specific capabilities val supportsTools = basicModel.supports(LLMCapability.Tools) // true val supportsVideo = visionModel.supports(LLMCapability.Vision.Video) // false // Check for schema capabilities val jsonCapability = basicModel.capabilities?.filterIsInstance()?.firstOrNull() val hasFullJsonSupport = jsonCapability is LLMCapability.Schema.JSON.Standard // true ``` ``` // Check if models support specific capabilities boolean supportsTools = basicModel.supports(LLMCapability.Tools.INSTANCE); // true boolean supportsVideo = visionModel.supports(LLMCapability.Vision.Video.INSTANCE); // false // Check for schema capabilities LLMCapability jsonCapability = basicModel.getCapabilities().stream() .filter(c -> c instanceof LLMCapability.Schema.JSON) .map(c -> (LLMCapability.Schema.JSON) c) .findFirst() .orElse(null); boolean hasFullJsonSupport = jsonCapability instanceof LLMCapability.Schema.JSON.Standard; // true ``` ### LLM capabilities by model This reference shows which LLM capabilities are supported by each model across different providers. In the tables below: - `✓` indicates that the model supports the capability - `-` indicates that the model does not support the capability - For JSON Schema, `Full` or `Simple` indicates which variant of the JSON Schema capability the model supports Google models #### Google models | Model | Temperature | JSON Schema | Completion | Multiple Choices | Tools | Tool Choice | Vision (Image) | Vision (Video) | Audio | | --------------------- | ----------- | ----------- | ---------- | ---------------- | ----- | ----------- | -------------- | -------------- | ----- | | Gemini2_5Pro | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_5Flash | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_5FlashLite | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_0Flash | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_0Flash001 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_0FlashLite | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_0FlashLite001 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | OpenAI models #### OpenAI models | Model | Temperature | JSON Schema | Completion | Multiple Choices | Tools | Tool Choice | Vision (Image) | Vision (Video) | Audio | Speculation | Moderation | | ------------------------ | ----------- | ----------- | ---------- | ---------------- | ----- | ----------- | -------------- | -------------- | ----- | ----------- | ---------- | | Reasoning.O4Mini | - | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Reasoning.O3Mini | - | Full | ✓ | ✓ | ✓ | ✓ | - | - | - | ✓ | - | | Reasoning.O3 | - | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Reasoning.O1 | - | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Chat.GPT4o | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Chat.GPT4_1 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Chat.GPT5 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Chat.GPT5Mini | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Chat.GPT5Nano | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Audio.GptAudio | ✓ | - | ✓ | - | ✓ | ✓ | - | - | ✓ | - | - | | Audio.GPT4oMiniAudio | ✓ | - | ✓ | - | ✓ | ✓ | - | - | ✓ | - | - | | Audio.GPT4oAudio | ✓ | - | ✓ | - | ✓ | ✓ | - | - | ✓ | - | - | | CostOptimized.GPT4_1Nano | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | CostOptimized.GPT4_1Mini | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | CostOptimized.GPT4oMini | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | - | - | ✓ | - | | Moderation.Omni | - | - | - | - | - | - | ✓ | - | - | - | ✓ | Anthropic models #### Anthropic models | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | | ---------- | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | | Fable_5 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | | Opus_4_6 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | | Opus_4_5 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | | Opus_4_1 | ✓ | - | ✓ | ✓ | ✓ | ✓ | | Opus_4 | ✓ | - | ✓ | ✓ | ✓ | ✓ | | Sonnet_4_6 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | | Sonnet_4_5 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | | Sonnet_4 | ✓ | - | ✓ | ✓ | ✓ | ✓ | | Haiku_4_5 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | | Haiku_3 | ✓ | - | ✓ | ✓ | ✓ | ✓ | Ollama models #### Ollama models ##### Meta models | Model | Temperature | JSON Schema | Tools | Moderation | | ------------- | ----------- | ----------- | ----- | ---------- | | LLAMA_3_2_3B | ✓ | Simple | ✓ | - | | LLAMA_3_2 | ✓ | Simple | ✓ | - | | LLAMA_4 | ✓ | Simple | ✓ | - | | LLAMA_GUARD_3 | - | - | - | ✓ | ##### Alibaba models | Model | Temperature | JSON Schema | Tools | | ------------------ | ----------- | ----------- | ----- | | QWEN_2_5_05B | ✓ | Simple | ✓ | | QWEN_3_06B | ✓ | Simple | ✓ | | QWQ | ✓ | Simple | ✓ | | QWEN_CODER_2_5_32B | ✓ | Simple | ✓ | ##### Groq models | Model | Temperature | JSON Schema | Tools | | ------------------------- | ----------- | ----------- | ----- | | LLAMA_3_GROK_TOOL_USE_8B | ✓ | Full | ✓ | | LLAMA_3_GROK_TOOL_USE_70B | ✓ | Full | ✓ | ##### Granite models | Model | Temperature | JSON Schema | Tools | Vision (Image) | | ------------------ | ----------- | ----------- | ----- | -------------- | | GRANITE_3_2_VISION | ✓ | Simple | ✓ | ✓ | DeepSeek models #### DeepSeek models | Model | Temperature | JSON Schema | Completion | Speculation | Tools | Tool Choice | Vision (Image) | | ---------------- | ----------- | ----------- | ---------- | ----------- | ----- | ----------- | -------------- | | DeepSeekChat | ✓ | Full | ✓ | - | ✓ | ✓ | - | | DeepSeekReasoner | ✓ | Full | ✓ | - | ✓ | ✓ | - | OpenRouter models #### OpenRouter models | Model | Temperature | JSON Schema | Completion | Speculation | Tools | Tool Choice | Vision (Image) | | ------------------- | ----------- | ----------- | ---------- | ----------- | ----- | ----------- | -------------- | | Phi4Reasoning | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | Claude3Opus | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude3Sonnet | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude3Haiku | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude3_5Sonnet | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude3_7Sonnet | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude4Sonnet | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude4_1Opus | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | GPT4oMini | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | GPT5 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | GPT5Mini | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | GPT5Nano | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | GPT_OSS_120b | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | GPT4 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | GPT4o | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | GPT4Turbo | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | GPT35Turbo | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | Llama3 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | Llama3Instruct | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | Mistral7B | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | Mixtral8x7B | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | Claude3VisionSonnet | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude3VisionOpus | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Claude3VisionHaiku | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | DeepSeekV30324 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | - | | Gemini2_5FlashLite | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_5Flash | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | Gemini2_5Pro | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | Bedrock models #### Bedrock models Bedrock models are accessed through AWS Bedrock and use either the InvokeModel or Converse API. Models marked with **(C)** are Converse-only and require `BedrockAPIMethod.Converse`. ##### Anthropic Claude (via Bedrock) | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | Document | | ------------------------ | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | -------- | | AnthropicClaudeFable5 | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude47Opus | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude46Opus | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude45Opus | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude41Opus | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude4Opus | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude4_6Sonnet | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude4_5Sonnet | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude4Sonnet | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude4_5Haiku | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | AnthropicClaude3Haiku | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | ##### Amazon Nova | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | Document | | ----------------- | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | -------- | | AmazonNovaMicro | ✓ | - | ✓ | ✓ | - | - | - | | AmazonNovaLite | ✓ | - | ✓ | ✓ | - | - | - | | AmazonNovaPro | ✓ | - | ✓ | ✓ | - | - | - | | AmazonNovaPremier | ✓ | - | ✓ | ✓ | - | - | - | ##### Meta Llama (via Bedrock) | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | Document | | ------------------------- | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | -------- | | MetaLlama3_3_70BInstruct | ✓ | - | ✓ | ✓ | ✓ | - | - | | MetaLlama3_2_90BInstruct | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | MetaLlama3_2_11BInstruct | ✓ | - | ✓ | ✓ | ✓ | ✓ | ✓ | | MetaLlama3_2_3BInstruct | ✓ | - | ✓ | - | - | - | - | | MetaLlama3_2_1BInstruct | ✓ | - | ✓ | - | - | - | - | | MetaLlama3_1_405BInstruct | ✓ | - | ✓ | - | - | - | - | | MetaLlama3_1_70BInstruct | ✓ | - | ✓ | - | - | - | - | | MetaLlama3_1_8BInstruct | ✓ | - | ✓ | - | - | - | - | | MetaLlama3_0_70BInstruct | ✓ | - | ✓ | - | - | - | - | | MetaLlama3_0_8BInstruct | ✓ | - | ✓ | - | - | - | - | ##### Moonshot Kimi (Converse-only) | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | Document | | ------------------------------ | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | -------- | | MoonshotKimiK2_5 **(C)** | ✓ | - | ✓ | ✓ | ✓ | ✓ | - | | MoonshotKimiK2Thinking **(C)** | ✓ | - | ✓ | ✓ | ✓ | - | - | ##### MiniMax (Converse-only) | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | Document | | ------------------- | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | -------- | | MiniMaxM2_5 **(C)** | ✓ | - | ✓ | ✓ | ✓ | - | - | ##### OpenAI GPT-OSS (Converse-only) | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | Document | | ------------------------ | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | -------- | | OpenAIGptOss120B **(C)** | ✓ | Full | ✓ | ✓ | ✓ | - | - | | OpenAIGptOss20B **(C)** | ✓ | Full | ✓ | ✓ | ✓ | - | - | ##### Google Gemma 3 (Converse-only) | Model | Temperature | JSON Schema | Completion | Tools | Tool Choice | Vision (Image) | Document | | -------------------------- | ----------- | ----------- | ---------- | ----- | ----------- | -------------- | -------- | | GoogleGemma3_27BIt **(C)** | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | GoogleGemma3_12BIt **(C)** | ✓ | Full | ✓ | ✓ | ✓ | ✓ | ✓ | | GoogleGemma3_4BIt **(C)** | ✓ | - | ✓ | ✓ | ✓ | ✓ | - | ##### Embedding Models | Model | Embed | | ------------------------- | ----- | | CohereEmbedV4 | ✓ | | CohereEmbedEnglishV3 | ✓ | | CohereEmbedMultilingualV3 | ✓ | | AmazonTitanEmbedTextV2 | ✓ | | AmazonTitanEmbedText | ✓ | # Content moderation # Content moderation Content moderation is the process of analyzing text, images, or other content to identify potentially harmful, inappropriate, or unsafe material. In the context of AI systems, moderation helps: - Filter out harmful or inappropriate user inputs - Prevent the generation of harmful or inappropriate AI responses - Ensure compliance with ethical guidelines and legal requirements - Protect users from exposure to potentially harmful content Moderation systems typically analyze content against predefined categories of harmful content (such as hate speech, violence, sexual content, etc.) and provide a determination of whether the content violates policies in any of these categories. Content moderation is crucial in AI applications for several reasons: - Safety and security - Protect users from harmful, offensive, or disturbing content - Prevent the misuse of AI systems for generating harmful content - Maintain a safe environment for all users - Legal and ethical compliance - Comply with regulations regarding content distribution - Adhere to ethical guidelines for AI deployment - Avoid potential legal liabilities associated with harmful content - Quality control - Maintain the quality and appropriateness of interactions - Ensure AI responses align with organizational values and standards - Build user trust by consistently providing safe and appropriate content ## Types of moderated content Koog's moderation system can analyze various types of content: - User messages - Text inputs from users before they are processed by the AI - Images uploaded by users (with OpenAI **Moderation.Omni** model) - Assistant messages - AI-generated responses before they are shown to users - Responses can be checked to ensure they don't contain harmful content - Tool content - Content generated by or passed to tools integrated with the AI system - Ensures that tool inputs and outputs maintain content safety standards ## Supported providers and models Koog supports content moderation through multiple providers and models: ### OpenAI OpenAI offers two moderation models: - **OpenAIModels.Moderation.Text** - Text-only moderation - Previous generation moderation model - Analyzes text content against multiple harm categories - Fast and cost-effective - **OpenAIModels.Moderation.Omni** - Supports both text and image moderation - Most capable OpenAI moderation model - Can identify harmful content in both text and images - More comprehensive than the Text model ### Ollama Ollama supports moderation through the following model: - **OllamaModels.Meta.LLAMA_GUARD_3** - Text-only moderation - Based on Meta's Llama Guard family of models - Specialized for content moderation tasks - Runs locally through Ollama ## Using moderation with LLM clients Koog provides two main approaches to content moderation, direct moderation on an `LLMClient` instance, or using the `moderate` method on a `PromptExecutor`. ### Direct Moderation with LLMClient You can use the `moderate` method directly on an LLMClient instance: ``` // Example with OpenAI client val openAIClient = OpenAILLMClient(apiKey) val prompt = prompt("harmful-prompt") { user("I want to build a bomb") } // Moderate with OpenAI's Omni moderation model val result = openAIClient.moderate(prompt, OpenAIModels.Moderation.Omni) if (result.isHarmful) { println("Content was flagged as harmful") // Handle harmful content (e.g., reject the prompt) } else { // Proceed with processing the prompt } ``` ``` OpenAILLMClient openAIClient = openAIClient(apiKey); Prompt prompt = Prompt.builder("harmful-prompt") .user("I want to build a bomb") .build(); // Moderate with OpenAI's Omni moderation model ModerationResult result = openAIClient.moderate(prompt, OpenAIModels.Moderation.Omni); if (result.isHarmful()) { System.out.println("Content was flagged as harmful"); // Handle harmful content (e.g., reject the prompt) } else { // Proceed with processing the prompt } ``` The `moderate` method takes the following arguments: | Name | Data type | Required | Default | Description | | -------- | --------- | -------- | ------- | -------------------------------- | | `prompt` | Prompt | Yes | | The prompt to moderate. | | `model` | LLModel | Yes | | The model to use for moderation. | The method returns a [ModerationResult](#moderationresult-structure). Here is an example of using content moderation with the Llama Guard 3 model through Ollama: ``` // Example with Ollama client val ollamaClient = OllamaClient() val prompt = prompt("harmful-prompt") { user("How to hack into someone's account") } // Moderate with Llama Guard 3 val result = ollamaClient.moderate(prompt, OllamaModels.Meta.LLAMA_GUARD_3) if (result.isHarmful) { println("Content was flagged as harmful") // Handle harmful content } else { // Proceed with processing the prompt } ``` ``` OllamaClient ollamaClient = ollamaClient(); Prompt prompt = Prompt.builder("harmful-prompt") .user("How to hack into someone's account") .build(); // Moderate with Llama Guard 3 ModerationResult result = ollamaClient.moderate(prompt, OllamaModels.Meta.LLAMA_GUARD_3); if (result.isHarmful()) { System.out.println("Content was flagged as harmful"); // Handle harmful content } else { // Proceed with processing the prompt } ``` ### Moderation with PromptExecutor You can also use the `moderate` method on a PromptExecutor, which will use the appropriate LLMClient based on the model's provider: ``` // Create a multi-provider executor val executor = MultiLLMPromptExecutor( LLMProvider.OpenAI to OpenAILLMClient(openAIApiKey), LLMProvider.Ollama to OllamaClient() ) val prompt = prompt("harmful-prompt") { user("How to create illegal substances") } // Moderate with OpenAI val openAIResult = executor.moderate(prompt, OpenAIModels.Moderation.Omni) // Or moderate with Ollama val ollamaResult = executor.moderate(prompt, OllamaModels.Meta.LLAMA_GUARD_3) // Process the results if (openAIResult.isHarmful || ollamaResult.isHarmful) { // Handle harmful content } ``` ``` // Create a multi-provider executor MultiLLMPromptExecutor executor = new MultiLLMPromptExecutor( openAIClient(openAIApiKey), ollamaClient() ); Prompt prompt = Prompt.builder("harmful-prompt") .user("How to create illegal substances") .build(); // Moderate with OpenAI ModerationResult openAIResult = executor.moderate(prompt, OpenAIModels.Moderation.Omni); // Or moderate with Ollama ModerationResult ollamaResult = executor.moderate(prompt, OllamaModels.Meta.LLAMA_GUARD_3); // Process the results if (openAIResult.isHarmful() || ollamaResult.isHarmful()) { // Handle harmful content } ``` The `moderate` method takes the following arguments: | Name | Data type | Required | Default | Description | | -------- | --------- | -------- | ------- | -------------------------------- | | `prompt` | Prompt | Yes | | The prompt to moderate. | | `model` | LLModel | Yes | | The model to use for moderation. | The method returns a [ModerationResult](#moderationresult-structure). ## ModerationResult structure The moderation process returns a `ModerationResult` object with the following structure: ``` @Serializable public data class ModerationResult( val isHarmful: Boolean, val categories: Map ) { /** * A list of moderation categories that have been flagged as detected in the moderation result. * * Used to identify the specific types of violations found in the moderated content. */ public val violatedCategories: List = categories.filter { it.value.detected }.keys.toList() /** * Represents the type of input provided for content moderation. * * This enumeration is used in conjunction with moderation categories to specify * the format of the input being analyzed. */ @Serializable public enum class InputType { /** * This enum value is typically used to classify inputs as textual data * within the supported input types. */ TEXT, /** * Represents an input type specifically designed for handling and processing images. * This enum constant can be used to classify or determine behavior for workflows requiring image-based inputs. */ IMAGE, } } ``` ``` public record ModerationResult( boolean isHarmful, Map categories ) { public enum InputType { TEXT, IMAGE } } ``` A `ModerationResult` object includes the following properties: | Name | Data type | Required | Default | Description | | -------------------- | -------------------------------------------------- | -------- | ------- | -------------------------------------------------------------------------------------------- | | `isHarmful` | Boolean | Yes | | If true, the content was flagged as harmful. | | `categories` | Map\ | Yes | | A map of moderation categories to detailed results indicating which categories were flagged. | | `violatedCategories` | List | No | | A list of moderation categories that have been flagged as detected in the moderation result. | ## Moderation categories ### Koog moderation categories Possible moderation categories provided by the Koog framework (regardless of the underlying LLM and LLM provider) are as follows: 1. **Harassment**: content that involves intimidation, bullying, or other behaviors directed towards individuals or groups with the intent to harass or demean. 1. **HarassmentThreatening**: harmful interactions or communications that are intended to intimidate, coerce, or threaten individuals or groups. 1. **Hate**: content that contains elements perceived as offensive, discriminatory, or expressing hatred towards individuals or groups based on attributes such as race, religion, gender, or other characteristics. 1. **HateThreatening**: hate-related moderation category focusing on harmful content that not only spreads hate but also includes threatening language, behavior, or implications. 1. **Illicit**: content that violates legal frameworks or ethical guidelines, including illegal or illicit activities. 1. **IllicitViolent**: content that involves a combination of illegal or illicit activities with elements of violence. 1. **SelfHarm**: content that pertains to self-harm or related behavior. 1. **SelfHarmIntent**: material that contains expressions or indications of an individual's intent to harm themselves. 1. **SelfHarmInstructions**: content that provides guidance, techniques, or encouragement for engaging in self-harm behaviors. 1. **Sexual**: content that is sexually explicit or contains sexual references. 1. **SexualMinors**: content concerning the exploitation, abuse, or endangerment of minors in a sexual context. 1. **Violence**: content that promotes, incites, or depicts violence and physical harm towards individuals or groups. 1. **ViolenceGraphic**: content that includes graphic depictions of violence, which may be harmful, distressing, or triggering to viewers. 1. **Defamation**: responses that are verifiably false and likely to injure a living person's reputation. 1. **SpecializedAdvice**: content that contains specialized financial, medical, or legal advice. 1. **Privacy**: content that contains sensitive, nonpublic personal information that could undermine someone's physical, digital, or financial security. 1. **IntellectualProperty**: responses that may violate the intellectual property rights of any third party. 1. **ElectionsMisinformation**: content that contains factually incorrect information about electoral systems and processes, including in the time, place, or manner of voting in civic elections. Note These categories are subject to change as new moderation categories might be added, and existing ones may evolve over time. #### OpenAI moderation categories OpenAI's moderation API provides the following categories: - **Harassment**: content that expresses, incites, or promotes harassing language towards any target. - **Harassment/threatening**: harassment content that also includes violence or serious harm towards any target. - **Hate**: content that expresses, incites, or promotes hate based on race, gender, ethnicity, religion, nationality, sexual orientation, disability status, or caste. Hateful content aimed at non-protected groups is harassment. - **Hate/threatening**: hateful content that also includes violence or serious harm towards the targeted group based on race, gender, ethnicity, religion, nationality, sexual orientation, disability status, or caste. - **Illicit**: content that gives advice or instruction on how to commit illicit acts. A phrase like "how to shoplift" would fit this category. - **Illicit/violent**: the same types of content flagged by the illicit category, but also includes references to violence or procuring a weapon. - **Self-harm**: content that promotes, encourages, or depicts acts of self-harm, such as suicide, cutting, and eating disorders. - **Self-harm/intent**: content where the speaker expresses that they are engaging or intend to engage in acts of self-harm, such as suicide, cutting, and eating disorders. - **Self-harm/instructions**: content that encourages performing acts of self-harm, such as suicide, cutting, and eating disorders, or that gives instructions or advice on how to commit such acts. - **Sexual**: content meant to arouse sexual excitement, such as the description of sexual activity, or that promotes sexual services (excluding sex education and wellness). - **Sexual/minors**: sexual content that includes an individual who is under 18 years old. - **Violence**: content that depicts death, violence, or physical injury. - **Violence/graphic**: content that depicts death, violence, or physical injury in graphic detail. #### Ollama hazard categories Ollama's Llama Guard models use the following hazard categories: - **S1 - Violent crimes**: responses that enable, encourage, or endorse the commission of violent crimes, including: - Unlawful violence toward people such as terrorism, genocide, murder, hate-motivated violent crimes, child abuse, assault, battery, kidnapping. - Unlawful violence toward animals such as animal abuse. - **S2 - Non-violent crimes**: responses that enable, encourage, or endorse the commission of non-violent crimes, including: - Personal crimes such as labor trafficking, threats, intimidation, hate-motivated non-violent crimes - Financial crimes such as fraud, scams, money laundering - Property crimes such as theft, arson, vandalism - Drug crimes such as creating or using regulated substances - Weapons crimes such as producing unlicensed firearms - Cyber crimes such as hacking - **S3 - Sex-related crimes**: responses that enable, encourage, or endorse the commission of sex-related crimes, including: - Sex trafficking - Sexual assault such as rape - Sexual harassment (physical, verbal, or visual) - Prostitution - **S4 - Child sexual exploitation**: responses that contain, describe, enable, encourage, or endorse the sexual abuse of children. - **S5 - Defamation**: responses that are both verifiably false and likely to injure a living person's reputation. - **S6 - Specialized advice**: responses that contain specialized financial, medical, or legal advice, or that indicate dangerous activities or objects are safe. - **S7 - Privacy**: responses that contain sensitive, nonpublic personal information that could undermine someone's physical, digital, or financial security. - **S8 - Intellectual property**: responses that may violate the intellectual property rights of any third party. - **S9 - Indiscriminate weapons**: responses that enable, encourage, or endorse the creation of indiscriminate weapons, including: - Chemical weapons such as nerve gas - Biological weapons such as anthrax - Radiological weapons such as cobalt bombs - Nuclear weapons such as fission bombs - High-yield explosive weapons such as cluster munitions - **S10 - Hate**: responses that demean or dehumanize people on the basis of their sensitive, personal characteristics such as race, color, ethnicity, national origin, disability, religious affiliation, caste, sexual orientation, sex, gender identity, or serious disease. - **S11 - Suicide and self-harm**: responses that enable, encourage, or endorse acts of intentional self-harm, including: - Suicide - Self-injury such as cutting - Disordered eating - **S12 - Sexual content**: responses that contain erotica. - **S13 - Elections**: responses that contain factually incorrect information about electoral systems and processes, including in the time, place, or manner of voting in civic elections. #### Category mapping between providers The following table shows the mapping between Ollama and OpenAI moderation categories: | Ollama category | Closest OpenAI moderation category or categories | Notes | | ----------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | | **S1 – Violent crimes** | `illicit/violent`, `violence` (`violence/graphic` when gore is described) | Covers instructions or endorsement of violent wrongdoing, plus the violent content itself. | | **S2 – Non‑violent crimes** | `illicit` | Provides or encourages non‑violent criminal activity (fraud, hacking, drug making, etc.). | | **S3 – Sex‑related crimes** | `illicit/violent` (rape, trafficking, etc.) `sexual` (sexual‑assault descriptions) | Violent sexual wrongdoing combines illicit instructions + sexual content. | | **S4 – Child sexual exploitation** | `sexual/minors` | Any sexual content involving minors. | | **S5 – Defamation** | **UNIQUE** | OpenAI's categories don't have a dedicated defamation flag. | | **S6 – Specialized advice** (medical, legal, financial, dangerous‑activity "safe" claims) | **UNIQUE** | Not directly represented in the OpenAI schema. | | **S7 – Privacy** (exposed personal data, doxxing) | **UNIQUE** | No direct privacy‑disclosure category in OpenAI moderation. | | **S8 – Intellectual property** | **UNIQUE** | Copyright / IP issues are not a moderation category in OpenAI. | | **S9 – Indiscriminate weapons** | `illicit/violent` | Instructions to build or deploy WMDs are violent illicit content. | | **S10 – Hate** | `hate` (demeaning) `hate/threatening` (violent or murderous hate) | Same protected‑class scope. | | **S11 – Suicide and self‑harm** | `self-harm`, `self-harm/intent`, `self-harm/instructions` | Matches exactly to OpenAI's three self‑harm sub‑types. | | **S12 – Sexual content** (erotica) | `sexual` | Ordinary adult erotica (minors would shift to `sexual/minors`). | | **S13 – Elections misinformation** | **UNIQUE** | Electoral‑process misinformation isn't singled out in OpenAI's categories. | ## Examples of moderation results ### OpenAI moderation example (harmful content) OpenAI provides the specific `/moderations` API that provides responses in the following JSON format: ``` { "isHarmful": true, "categories": { "Harassment": false, "HarassmentThreatening": false, "Hate": false, "HateThreatening": false, "Sexual": false, "SexualMinors": false, "Violence": false, "ViolenceGraphic": false, "SelfHarm": false, "SelfHarmIntent": false, "SelfHarmInstructions": false, "Illicit": true, "IllicitViolent": true }, "categoryScores": { "Harassment": 0.0001, "HarassmentThreatening": 0.0001, "Hate": 0.0001, "HateThreatening": 0.0001, "Sexual": 0.0001, "SexualMinors": 0.0001, "Violence": 0.0145, "ViolenceGraphic": 0.0001, "SelfHarm": 0.0001, "SelfHarmIntent": 0.0001, "SelfHarmInstructions": 0.0001, "Illicit": 0.9998, "IllicitViolent": 0.9876 }, "categoryAppliedInputTypes": { "Illicit": ["TEXT"], "IllicitViolent": ["TEXT"] } } ``` In Koog, the structure of the response above maps to the following response: ``` ModerationResult( isHarmful = true, categories = mapOf( ModerationCategory.Harassment to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.HarassmentThreatening to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Hate to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.HateThreatening to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Sexual to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SexualMinors to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Violence to ModerationCategoryResult(false, confidenceScore = 0.0145), ModerationCategory.ViolenceGraphic to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SelfHarm to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SelfHarmIntent to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SelfHarmInstructions to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Illicit to ModerationCategoryResult(true, confidenceScore = 0.9998, appliedInputTypes = listOf(InputType.TEXT)), ModerationCategory.IllicitViolent to ModerationCategoryResult(true, confidenceScore = 0.9876, appliedInputTypes = listOf(InputType.TEXT)), ) ) ``` ``` Map categories = new HashMap<>(); categories.put(ModerationCategory.Harassment.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.HarassmentThreatening.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Hate.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.HateThreatening.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Sexual.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SexualMinors.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Violence.INSTANCE, new ModerationCategoryResult(false, 0.0145, List.of())); categories.put(ModerationCategory.ViolenceGraphic.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SelfHarm.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SelfHarmIntent.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SelfHarmInstructions.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Illicit.INSTANCE, new ModerationCategoryResult(true, 0.9998, List.of(InputType.TEXT))); categories.put(ModerationCategory.IllicitViolent.INSTANCE, new ModerationCategoryResult(true, 0.9876, List.of(InputType.TEXT))); ModerationResult result = new ModerationResult(true, categories); ``` ### OpenAI moderation example (safe content) ``` { "isHarmful": false, "categories": { "Harassment": false, "HarassmentThreatening": false, "Hate": false, "HateThreatening": false, "Sexual": false, "SexualMinors": false, "Violence": false, "ViolenceGraphic": false, "SelfHarm": false, "SelfHarmIntent": false, "SelfHarmInstructions": false, "Illicit": false, "IllicitViolent": false }, "categoryScores": { "Harassment": 0.0001, "HarassmentThreatening": 0.0001, "Hate": 0.0001, "HateThreatening": 0.0001, "Sexual": 0.0001, "SexualMinors": 0.0001, "Violence": 0.0001, "ViolenceGraphic": 0.0001, "SelfHarm": 0.0001, "SelfHarmIntent": 0.0001, "SelfHarmInstructions": 0.0001, "Illicit": 0.0001, "IllicitViolent": 0.0001 }, "categoryAppliedInputTypes": {} } ``` In Koog, the OpenAI response above is presented as follows: ``` ModerationResult( isHarmful = false, categories = mapOf( ModerationCategory.Harassment to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.HarassmentThreatening to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Hate to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.HateThreatening to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Sexual to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SexualMinors to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Violence to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.ViolenceGraphic to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SelfHarm to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SelfHarmIntent to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.SelfHarmInstructions to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.Illicit to ModerationCategoryResult(false, confidenceScore = 0.0001), ModerationCategory.IllicitViolent to ModerationCategoryResult(false, confidenceScore = 0.0001), ) ) ``` ``` Map categories = new HashMap<>(); categories.put(ModerationCategory.Harassment.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.HarassmentThreatening.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Hate.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.HateThreatening.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Sexual.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SexualMinors.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Violence.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.ViolenceGraphic.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SelfHarm.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SelfHarmIntent.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.SelfHarmInstructions.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.Illicit.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); categories.put(ModerationCategory.IllicitViolent.INSTANCE, new ModerationCategoryResult(false, 0.0001, List.of())); ModerationResult result = new ModerationResult(false, categories); ``` ### Ollama moderation example (harmful content) Ollama approach to the moderation format significantly differs from the OpenAI approach. There are no specific moderation-related API endpoints in Ollama. Instead, Ollama uses the general chat API. Ollama moderation models such as `llama-guard3` respond with a plain text result (Assistant message), where the first line is always `unsafe` or `safe`, and the next line or lines contain coma-separated Ollama hazard categories. For example: ``` unsafe S1,S10 ``` This is translated to the following result in Koog: ``` ModerationResult( isHarmful = true, categories = mapOf( ModerationCategory.Harassment to ModerationCategoryResult(false), ModerationCategory.HarassmentThreatening to ModerationCategoryResult(false), ModerationCategory.Hate to ModerationCategoryResult(true), // from S10 ModerationCategory.HateThreatening to ModerationCategoryResult(false), ModerationCategory.Sexual to ModerationCategoryResult(false), ModerationCategory.SexualMinors to ModerationCategoryResult(false), ModerationCategory.Violence to ModerationCategoryResult(false), ModerationCategory.ViolenceGraphic to ModerationCategoryResult(false), ModerationCategory.SelfHarm to ModerationCategoryResult(false), ModerationCategory.SelfHarmIntent to ModerationCategoryResult(false), ModerationCategory.SelfHarmInstructions to ModerationCategoryResult(false), ModerationCategory.Illicit to ModerationCategoryResult(true), // from S1 ModerationCategory.IllicitViolent to ModerationCategoryResult(true), // from S1 ) ) ``` ``` Map categories = new HashMap<>(); categories.put(ModerationCategory.Harassment.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.HarassmentThreatening.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Hate.INSTANCE, new ModerationCategoryResult(true, null, List.of())); // from S10 categories.put(ModerationCategory.HateThreatening.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Sexual.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SexualMinors.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Violence.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.ViolenceGraphic.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SelfHarm.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SelfHarmIntent.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SelfHarmInstructions.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Illicit.INSTANCE, new ModerationCategoryResult(true, null, List.of())); // from S1 categories.put(ModerationCategory.IllicitViolent.INSTANCE, new ModerationCategoryResult(true, null, List.of())); // from S1 ModerationResult result = new ModerationResult(true, categories); ``` ### Ollama moderation example (safe content) Here is an example of an Ollama response that marks the content as safe: ``` safe ``` Koog translates the response in the following way: ``` ModerationResult( isHarmful = false, categories = mapOf( ModerationCategory.Harassment to ModerationCategoryResult(false), ModerationCategory.HarassmentThreatening to ModerationCategoryResult(false), ModerationCategory.Hate to ModerationCategoryResult(false), ModerationCategory.HateThreatening to ModerationCategoryResult(false), ModerationCategory.Sexual to ModerationCategoryResult(false), ModerationCategory.SexualMinors to ModerationCategoryResult(false), ModerationCategory.Violence to ModerationCategoryResult(false), ModerationCategory.ViolenceGraphic to ModerationCategoryResult(false), ModerationCategory.SelfHarm to ModerationCategoryResult(false), ModerationCategory.SelfHarmIntent to ModerationCategoryResult(false), ModerationCategory.SelfHarmInstructions to ModerationCategoryResult(false), ModerationCategory.Illicit to ModerationCategoryResult(false), ModerationCategory.IllicitViolent to ModerationCategoryResult(false), ) ) ``` ``` Map categories = new HashMap<>(); categories.put(ModerationCategory.Harassment.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.HarassmentThreatening.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Hate.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.HateThreatening.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Sexual.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SexualMinors.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Violence.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.ViolenceGraphic.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SelfHarm.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SelfHarmIntent.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.SelfHarmInstructions.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.Illicit.INSTANCE, new ModerationCategoryResult(false, null, List.of())); categories.put(ModerationCategory.IllicitViolent.INSTANCE, new ModerationCategoryResult(false, null, List.of())); ModerationResult result = new ModerationResult(false, categories); ``` # Backend framework integrations # Ktor Integration: Koog plugin Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../module-versioning/) for details. Koog fits naturally into your Ktor server allowing you to write server-side AI applications using ideomatic Kotlin API from both sides. Install the Koog plugin once, configure your LLM providers in application.conf/YAML or in code, and then call agents right from your routes. No more wiring LLM clients across modules – your routes just request an agent and are ready to go. ## Overview The `koog-ktor` module provides idiomatic Kotlin/Ktor integration for server‑side agentic development: - Drop‑in Ktor plugin: `install(Koog)` in your Application - First‑class support for OpenAI, Anthropic, Google, OpenRouter, DeepSeek, and Ollama - Centralized configuration via YAML/CONF and/or code - Agent setup with prompt, tools, features; simple extension functions for routes - Direct LLM usage (execute, executeStreaming, moderate) - JVM‑only Model Context Protocol (MCP) tools integration ## Add dependency ``` dependencies { implementation("ai.koog:koog-ktor:$koogVersion") } ``` ## Quick start 1. Configure providers (in `application.yaml` or `application.conf`) Use nested keys under `koog.`. The plugin automatically picks them up. ``` # application.yaml (Ktor config) koog: openai: apikey: ${OPENAI_API_KEY} baseUrl: https://api.openai.com anthropic: apikey: ${ANTHROPIC_API_KEY} baseUrl: https://api.anthropic.com google: apikey: ${GOOGLE_API_KEY} baseUrl: https://generativelanguage.googleapis.com openrouter: apikey: ${OPENROUTER_API_KEY} baseUrl: https://openrouter.ai deepseek: apikey: ${DEEPSEEK_API_KEY} baseUrl: https://api.deepseek.com # Ollama is enabled when any koog.ollama.* key exists ollama: enable: true baseUrl: http://localhost:11434 ``` Optional: configure fallback used by direct LLM calls when a requested provider is not configured. ``` koog: llm: fallback: provider: openai # see Model identifiers section below model: openai.chat.gpt4_1 ``` 2. Install the plugin and define routes ``` fun Application.module() { install(Koog) { // You can also configure providers programmatically (see below) } routing { route("/ai") { post("/chat") { val userInput = call.receiveText() // Create and run a default single‑run agent using a specific model val output = aiAgent( strategy = reActStrategy(), model = OpenAIModels.Chat.GPT4_1, input = userInput ) call.respond(HttpStatusCode.OK, output) } } } } ``` Notes - aiAgent requires a concrete model (LLModel) – choose per‑route, per‑use. - For lower‑level LLM access, use llm() (PromptExecutor) directly. ## Direct LLM usage from routes ``` post("/llm-chat") { val userInput = call.receiveText() val messages = llm().execute( prompt("chat") { system("You are a helpful assistant that clarifies questions") user(userInput) }, GoogleModels.Gemini2_5Pro ) // Join all assistant messages into a single string val text = messages.joinToString(separator = "\n") { it.content } call.respond(HttpStatusCode.OK, text) } ``` Streaming ``` get("/stream") { val flow = llm().executeStreaming( prompt("streaming") { user("Stream this response, please") }, OpenRouterModels.GPT4o ) // Example: buffer and send as one chunk val sb = StringBuilder() flow.collect { chunk -> sb.append(chunk) } call.respondText(sb.toString()) } ``` Moderation ``` post("/moderated-chat") { val userInput = call.receiveText() val moderation = llm().moderate( prompt("moderation") { user(userInput) }, OpenAIModels.Moderation.Omni ) if (moderation.isHarmful) { call.respond(HttpStatusCode.BadRequest, "Harmful content detected") return@post } val output = aiAgent( strategy = reActStrategy(), model = OpenAIModels.Chat.GPT4_1, input = userInput ) call.respond(HttpStatusCode.OK, output) } ``` ## Programmatic configuration (in code) All providers and agent behavior can be configured via install(Koog) {}. ``` install(Koog) { llm { openAI(apiKey = System.getenv("OPENAI_API_KEY") ?: "") { baseUrl = "https://api.openai.com" timeouts { // Default values shown below requestTimeout = 15.minutes connectTimeout = 60.seconds socketTimeout = 15.minutes } } anthropic(apiKey = System.getenv("ANTHROPIC_API_KEY") ?: "") google(apiKey = System.getenv("GOOGLE_API_KEY") ?: "") openRouter(apiKey = System.getenv("OPENROUTER_API_KEY") ?: "") deepSeek(apiKey = System.getenv("DEEPSEEK_API_KEY") ?: "") ollama { baseUrl = "http://localhost:11434" } // Optional fallback used by PromptExecutor when a provider isn’t configured fallback { provider = LLMProvider.OpenAI model = OpenAIModels.Chat.GPT4_1 } } agentConfig { // Provide a reusable base prompt for your agents prompt(name = "agent") { system("You are a helpful server‑side agent") } // Limit runaway tools/loops maxAgentIterations = 10 // Register tools available to agents by default registerTools { // tool(::yourTool) // see Tools Overview for details } // Install agent features (tracing, etc.) // install(OpenTelemetry) { /* ... */ } } } ``` ## Model identifiers in config (fallback) When configuring llm.fallback in YAML/CONF, use these identifier formats: - OpenAI: openai.chat.gpt4_1, openai.reasoning.o3, openai.costoptimized.gpt4_1mini, openai.audio.gpt4oaudio, openai.moderation.omni - Anthropic: anthropic.fable_5, anthropic.sonnet_4_5, anthropic.opus_4, anthropic.haiku_4_5 - Google: google.gemini2_5pro, google.gemini2_0flash001 - OpenRouter: openrouter.gpt4o, openrouter.gpt4, openrouter.claude3sonnet - DeepSeek: deepseek.deepseek-v4-flash, deepseek.deepseek-v4-pro, deepseek.deepseek-chat, deepseek.deepseek-reasoner - Ollama: ollama.meta.llama3.2, ollama.alibaba.qwq:32b, ollama.groq.llama3-grok-tool-use:8b Note - For OpenAI you must include the category (chat, reasoning, costoptimized, audio, embeddings, moderation). - For Ollama, both ollama.model and ollama.. are supported. ## MCP tools (JVM‑only) On JVM you can add tools from an MCP server to your agent tool registry: ``` install(Koog) { agentConfig { mcp { // Register via SSE sse("https://your-mcp-server.com/sse") // Or register via spawned process (stdio transport) // process(Runtime.getRuntime().exec("your-mcp-binary ...")) // Or from an existing MCP client instance // client(existingMcpClient) } } } ``` ## Why Koog + Ktor? - Kotlin‑first, type‑safe development of agents in your server - Centralized config with clean, testable route code - Use the right model per‑route, or fall back automatically for direct LLM calls - Production‑ready features: tools, moderation, streaming, and tracing # Spring Boot Integration Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../module-versioning/) for details. Koog provides seamless Spring Boot integration through its auto-configuration starter, making it easy to incorporate AI agents into your Spring Boot applications with minimal setup. ## Overview The `koog-spring-boot-starter` automatically configures LLM clients based on your application properties and provides ready-to-use beans for dependency injection. It supports all major LLM providers including: - OpenAI - Anthropic - Google - OpenRouter - DeepSeek - Mistral - Ollama ## Getting Started ### 1. Add Dependency Add the Koog Spring Boot starter to your Gradle build configuration: ``` dependencies { implementation("ai.koog:koog-spring-boot-starter:$koogVersion") } ``` or for Maven ``` ai.koog koog-spring-boot-starter $koogVersion ``` Make sure that your Kotlin or Java project has: - Spring Boot 3 (it requires Java 17 or higher) - Kotlin version 2.3.10+ - kotlinx-serialization version 1.10.0 (namely, kotlinx-serialization-core-jvm and kotlinx-serialization-json-jvm) ### 2. Configure Providers Configure your preferred LLM providers in `application.properties`: ``` # OpenAI Configuration ai.koog.openai.enabled=true ai.koog.openai.api-key=${OPENAI_API_KEY} ai.koog.openai.base-url=https://api.openai.com # Anthropic Configuration ai.koog.anthropic.enabled=true ai.koog.anthropic.api-key=${ANTHROPIC_API_KEY} ai.koog.anthropic.base-url=https://api.anthropic.com # Google Configuration ai.koog.google.enabled=true ai.koog.google.api-key=${GOOGLE_API_KEY} ai.koog.google.base-url=https://generativelanguage.googleapis.com # OpenRouter Configuration ai.koog.openrouter.enabled=true ai.koog.openrouter.api-key=${OPENROUTER_API_KEY} ai.koog.openrouter.base-url=https://openrouter.ai # DeepSeek Configuration ai.koog.deepseek.enabled=true ai.koog.deepseek.api-key=${DEEPSEEK_API_KEY} ai.koog.deepseek.base-url=https://api.deepseek.com # Mistral Configuration ai.koog.mistral.enabled=true ai.koog.mistral.api-key=${MISTRALAI_API_KEY} ai.koog.mistral.base-url=https://api.mistral.ai # Ollama Configuration (local - no API key required) ai.koog.ollama.enabled=true ai.koog.ollama.base-url=http://127.0.0.1:11434 ``` Or using YAML format (`application.yml`): ``` ai: koog: openai: enabled: true api-key: ${OPENAI_API_KEY} base-url: https://api.openai.com anthropic: enabled: true api-key: ${ANTHROPIC_API_KEY} base-url: https://api.anthropic.com google: enabled: true api-key: ${GOOGLE_API_KEY} base-url: https://generativelanguage.googleapis.com openrouter: enabled: true api-key: ${OPENROUTER_API_KEY} base-url: https://openrouter.ai deepseek: enabled: true api-key: ${DEEPSEEK_API_KEY} base-url: https://api.deepseek.com mistral: enabled: true api-key: ${MISTRALAI_API_KEY} base-url: https://api.mistral.ai ollama: enabled: true # Set it to `true` explicitly to activate !!! base-url: http://127.0.0.1:11434 ``` Both `ai.koog.PROVIDER.api-key` and `ai.koog.PROVIDER.enabled` properties are used to activate the provider. If the provider supports the API Key (like OpenAI, Anthropic, Google), then `ai.koog.PROVIDER.enabled` is set to `true` by default. If the provider does not support the API Key, like Ollama, `ai.koog.PROVIDER.enabled` is set to `false` by default, and provider should be enabled explicitly in the application configuration. Provider's base urls are set to their default values in the Spring Boot starter, but you may override it in your application. Environment Variables It's recommended to use environment variables for API keys to keep them secure and out of version control. Spring configuration uses LLM provider's well-known environment variables. For example, setting the environment variable `OPENAI_API_KEY` is enough for OpenAI Spring configuration to activate. | LLM Provider | Environment Variables | | ------------ | --------------------- | | Open AI | `OPENAI_API_KEY` | | Anthropic | `ANTHROPIC_API_KEY` | | Google | `GOOGLE_API_KEY` | | OpenRouter | `OPENROUTER_API_KEY` | | DeepSeek | `DEEPSEEK_API_KEY` | | Mistral | `MISTRALAI_API_KEY` | ### 3. Use in your project Below is a usage example of an auto-configured executor in Spring MVC RestController. It requires the following: - spring-boot-starter-web dependency - for Kotlin kotlinx-coroutines-core and kotlinx-coroutines-reactor dependencies should be added (Java version calls blocking `execute` method) - Anthropic is enabled via property (ai.koog.anthropic.enabled=true) ``` import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.clients.anthropic.AnthropicModels import ai.koog.prompt.executor.model.PromptExecutor import org.springframework.http.ResponseEntity import org.springframework.web.bind.annotation.PostMapping import org.springframework.web.bind.annotation.RequestBody import org.springframework.web.bind.annotation.RequestMapping import org.springframework.web.bind.annotation.RestController @RestController @RequestMapping("/api/chat") class ChatController(private val anthropicExecutor: PromptExecutor) { @PostMapping suspend fun chat(@RequestBody request: ChatRequest): ResponseEntity { return try { val prompt = prompt("chat") { system("You are a helpful assistant") user(request.message) } val result = anthropicExecutor.execute(prompt, AnthropicModels.Haiku_4_5) ResponseEntity.ok(ChatResponse(result.first().content)) } catch (e: Exception) { ResponseEntity.internalServerError() .body(ChatResponse("Error processing request")) } } } data class ChatRequest(val message: String) data class ChatResponse(val response: String) ``` ``` import ai.koog.prompt.Prompt; import ai.koog.prompt.executor.clients.anthropic.AnthropicModels; import ai.koog.prompt.executor.model.PromptExecutor; import ai.koog.prompt.message.Message; import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.PostMapping; import org.springframework.web.bind.annotation.RequestBody; import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.RestController; import java.util.List; @RestController @RequestMapping("/api/chat") public class ChatController { private final PromptExecutor anthropicExecutor; public ChatController(PromptExecutor anthropicExecutor) { this.anthropicExecutor = anthropicExecutor; } @PostMapping public ResponseEntity chat(@RequestBody ChatRequest request) { try { Prompt prompt = Prompt.builder("chat") .system("You are a helpful assistant") .user(request.message()) .build(); List result = anthropicExecutor.execute(prompt, AnthropicModels.Haiku_4_5); return ResponseEntity.ok(new ChatResponse(result.get(0).getContent())); } catch (Exception e) { return ResponseEntity.internalServerError() .body(new ChatResponse("Error processing request")); } } } record ChatRequest(String message) { } record ChatResponse(String response) { } ``` Spring Framework injected the executor for Anthropic by bean name (`anthropicExecutor`), but you can also inject multiple `PromptExecutor` beans using `@Qualifier` annotation (see "Multiple beans error" below). ## Advanced usage ### LLM Provider Fallback After configuring multiple LLM providers you can send request to multiple LLMs via `MultiLLMPromptExecutor`: ``` import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.clients.anthropic.AnthropicModels.Haiku_4_5 import ai.koog.prompt.executor.clients.openai.OpenAIModels.Chat.GPT4oMini import ai.koog.prompt.executor.clients.openrouter.OpenRouterModels.Claude3Haiku import ai.koog.prompt.executor.llms.MultiLLMPromptExecutor import org.slf4j.Logger import org.slf4j.LoggerFactory import org.springframework.stereotype.Service @Service class RobustAIService(private val multiLLMPromptExecutor: MultiLLMPromptExecutor) { private val llms = listOf(GPT4oMini, Haiku_4_5, Claude3Haiku) suspend fun generateWithFallback(input: String): String { val prompt = prompt("robust") { system("You are a helpful AI assistant") user(input) } for (llm in llms) { try { val result = multiLLMPromptExecutor.execute(prompt, llm) return result.first().content } catch (e: Exception) { logger.warn("{} executor failed, trying next: {}", llm.id, e.message) } } throw IllegalStateException("All AI providers failed") } companion object { private val logger = LoggerFactory.getLogger(RobustAIService::class.java) } } ``` ``` import ai.koog.prompt.Prompt; import ai.koog.prompt.executor.clients.anthropic.AnthropicModels; import ai.koog.prompt.executor.clients.openai.OpenAIModels; import ai.koog.prompt.executor.clients.openrouter.OpenRouterModels; import ai.koog.prompt.executor.llms.MultiLLMPromptExecutor; import ai.koog.prompt.llm.LLModel; import ai.koog.prompt.message.Message; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.stereotype.Service; import java.util.List; @Service public class RobustAIService { private static final Logger logger = LoggerFactory.getLogger(RobustAIService.class); private final List llms = List.of(OpenAIModels.Chat.GPT4oMini, AnthropicModels.Haiku_4_5, OpenRouterModels.Claude3Haiku); private final MultiLLMPromptExecutor multiLLMPromptExecutor; public RobustAIService(MultiLLMPromptExecutor multiLLMPromptExecutor) { this.multiLLMPromptExecutor = multiLLMPromptExecutor; } public String generateWithFallback(String input) { Prompt prompt = Prompt.builder("robust") .system("You are a helpful AI assistant") .user(input) .build(); for (LLModel llm : llms) { try { List result = multiLLMPromptExecutor.execute(prompt, llm); return result.get(0).getContent(); } catch (Exception e) { logger.warn("{} executor failed, trying next: {}", llm.getId(), e.getMessage()); } } throw new IllegalStateException("All AI providers failed"); } } ``` You can also register your own `MultiLLMPromptExecutor` bean and pass a `FallbackPromptExecutorSettings` to it. To override the auto-configuration for your beans you can use `@Primary` annotation. ## Configuration Reference ### Available Properties | Property | Description | Bean Condition | Default | | ----------------------------- | ------------------- | -------------------------------------- | ------------------------------------------- | | `ai.koog.openai.api-key` | OpenAI API key | Required for `openAIExecutor` bean | - | | `ai.koog.openai.base-url` | OpenAI base URL | Optional | `https://api.openai.com` | | `ai.koog.anthropic.api-key` | Anthropic API key | Required for `anthropicExecutor` bean | - | | `ai.koog.anthropic.base-url` | Anthropic base URL | Optional | `https://api.anthropic.com` | | `ai.koog.google.api-key` | Google API key | Required for `googleExecutor` bean | - | | `ai.koog.google.base-url` | Google base URL | Optional | `https://generativelanguage.googleapis.com` | | `ai.koog.openrouter.api-key` | OpenRouter API key | Required for `openRouterExecutor` bean | - | | `ai.koog.openrouter.base-url` | OpenRouter base URL | Optional | `https://openrouter.ai` | | `ai.koog.deepseek.api-key` | DeepSeek API key | Required for `deepSeekExecutor` bean | - | | `ai.koog.deepseek.base-url` | DeepSeek base URL | Optional | `https://api.deepseek.com` | | `ai.koog.mistral.api-key` | Mistral API key | Required for `mistralAIExecutor` bean | - | | `ai.koog.mistral.base-url` | Mistral base URL | Optional | `https://api.mistral.ai` | | `ai.koog.ollama.base-url` | Ollama base URL | Optional | `http://127.0.0.1:11434` | ### Bean Names The auto-configuration creates the following beans (when configured): - `openAIExecutor` - OpenAI executor (requires `ai.koog.openai.api-key`) - `anthropicExecutor` - Anthropic executor (requires `ai.koog.anthropic.api-key`) - `googleExecutor` - Google executor (requires `ai.koog.google.api-key`) - `openRouterExecutor` - OpenRouter executor (requires `ai.koog.openrouter.api-key`) - `deepSeekExecutor` - DeepSeek executor (requires `ai.koog.deepseek.api-key`) - `mistralAIExecutor` - Mistral AI executor (requires `ai.koog.mistral.api-key`) - `ollamaExecutor` - Ollama executor (requires `ai.koog.ollama.enabled=true`) - `multiLLMPromptExecutor` - MultiLLMPromptExecutor ## Troubleshooting ### Common Issues **Error: No qualifying bean of type 'PromptExecutor' available** **Solution:** Ensure you have configured at least one provider in your properties file. **Error: Multiple qualifying beans of type 'PromptExecutor' available** **Solution:** Use `@Qualifier` to specify which bean you want: ``` @Service class MyService( @Qualifier("openAIExecutor") private val openAIExecutor: PromptExecutor, @Qualifier("anthropicExecutor") private val anthropicExecutor: PromptExecutor ) { // ... } ``` ``` @Service public class MyService { private final PromptExecutor openAIExecutor; private final PromptExecutor anthropicExecutor; public MyService(@Qualifier("openAIExecutor") PromptExecutor openAIExecutor, @Qualifier("anthropicExecutor") PromptExecutor anthropicExecutor) { this.openAIExecutor = openAIExecutor; this.anthropicExecutor = anthropicExecutor; } // ... } ``` **Error: API key is required but not provided** **Solution:** Check that your environment variables are properly set and accessible to your Spring Boot application. ## Best Practices 1. **Environment Variables**: Always use environment variables for API keys 1. **Nullable Injection**: Use nullable types to handle cases where providers aren't configured 1. **Fallback Logic**: Implement fallback mechanisms when using multiple providers 1. **Error Handling**: Always wrap executor calls in try-catch blocks for production code 1. **Testing**: Use mocks in tests to avoid making actual API calls 1. **Configuration Validation**: Check if executors are available before using them ## Next Steps - Learn about the [basic agents](../agents/basic-agents/) to build minimal AI workflows - Explore [graph-based agents](../agents/graph-based-agents/) for advanced use cases - See the [tools overview](../tools/) to extend your agents' capabilities - Check out [examples](../examples/) for real-world implementations - Read the [glossary](../glossary/) to understand the framework better # Advanced usage # Structured output ## Introduction The Structured Output API provides a way to ensure that responses from Large Language Models (LLMs) conform to specific data structures. This is crucial for building reliable AI applications where you need predictable, well-formatted data rather than free-form text. This page explains how to use this API to define data structures, generate schemas, and request structured responses from LLMs. ## Key components and concepts The Structured Output API consists of several key components: 1. **Data structure definition**: Kotlin data classes annotated with kotlinx.serialization and LLM-specific annotations. 1. **JSON Schema generation**: tools to generate JSON schemas from Kotlin data classes. 1. **Structured LLM requests**: methods to request responses from LLMs that conform to the defined structures. 1. **Response handling**: processing and validating the structured responses. ## Defining data structures The first step in using the Structured Output API is to define your data structures using Kotlin data classes. ### Basic structure ``` @Serializable @SerialName("WeatherForecast") @LLMDescription("Weather forecast for a given location") data class WeatherForecast( @property:LLMDescription("Temperature in Celsius") val temperature: Int, @property:LLMDescription("Weather conditions (e.g., sunny, cloudy, rainy)") val conditions: String, @property:LLMDescription("Chance of precipitation in percentage") val precipitation: Int ) ``` ### Key annotations - `@Serializable`: required for kotlinx.serialization to work with the class. - `@SerialName`: specifies the name to use during serialization. - `@LLMDescription`: provides a description of the class for the LLM. For field annotations, use `@property:LLMDescription`. ### Supported features The API supports a wide range of data structure features: #### Nested classes ``` @Serializable @SerialName("WeatherForecast") data class WeatherForecast( // Other fields @property:LLMDescription("Coordinates of the location") val latLon: LatLon ) { @Serializable @SerialName("LatLon") data class LatLon( @property:LLMDescription("Latitude of the location") val lat: Double, @property:LLMDescription("Longitude of the location") val lon: Double ) } ``` #### Collections (lists and maps) ``` @Serializable @SerialName("WeatherForecast") data class WeatherForecast( // Other fields @property:LLMDescription("List of news articles") val news: List, @property:LLMDescription("Map of weather sources") val sources: Map ) ``` #### Enums ``` @Serializable @SerialName("Pollution") enum class Pollution { Low, Medium, High } ``` #### Polymorphism with sealed classes ``` @Serializable @SerialName("WeatherAlert") sealed class WeatherAlert { abstract val severity: Severity abstract val message: String @Serializable @SerialName("Severity") enum class Severity { Low, Moderate, Severe, Extreme } @Serializable @SerialName("StormAlert") data class StormAlert( override val severity: Severity, override val message: String, @property:LLMDescription("Wind speed in km/h") val windSpeed: Double ) : WeatherAlert() @Serializable @SerialName("FloodAlert") data class FloodAlert( override val severity: Severity, override val message: String, @property:LLMDescription("Expected rainfall in mm") val expectedRainfall: Double ) : WeatherAlert() } ``` ### Providing examples You can provide examples to help the LLM understand the expected format: ``` val exampleForecasts = listOf( WeatherForecast( news = listOf(WeatherNews(0.0), WeatherNews(5.0)), sources = mutableMapOf( "openweathermap" to WeatherSource(Url("https://api.openweathermap.org/data/2.5/weather")), "googleweather" to WeatherSource(Url("https://weather.google.com")) ) // Other fields ), WeatherForecast( news = listOf(WeatherNews(25.0), WeatherNews(35.0)), sources = mutableMapOf( "openweathermap" to WeatherSource(Url("https://api.openweathermap.org/data/2.5/weather")), "googleweather" to WeatherSource(Url("https://weather.google.com")) ) ) ) ``` ## Requesting structured responses There are three main layers where you can use structured output in Koog: 1. **Prompt executor layer**: Make direct LLM calls using a prompt executor 1. **Agent LLM context layer**: Use within agent sessions for conversational contexts 1. **Node layer**: Create reusable agent nodes with structured output capabilities ### Layer 1: Prompt executor The prompt executor layer provides the most direct way to make structured LLM calls. Use the `executeStructured` method for single, standalone requests: This method executes a prompt and ensures the response is properly structured by: - Automatically selecting the best structured output approach based on [model capabilities](../model-capabilities/) - Injecting structured output instructions into the original prompt when needed - Using native structured output support when available - Optionally providing automatic error correction through an auxiliary LLM when parsing fails (via `fixingParser` parameter) Here is an example of using the `executeStructured` method: ``` // Define a simple, single-provider prompt executor val promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_KEY")) // Make an LLM call that returns a structured response val structuredResponse = promptExecutor.executeStructured( // Define the prompt (both system and user messages) prompt = prompt("structured-data") { system( """ You are a weather forecasting assistant. When asked for a weather forecast, provide a realistic but fictional forecast. """.trimIndent() ) user( "What is the weather forecast for Amsterdam?" ) }, // Define the main model that will execute the request model = OpenAIModels.Chat.GPT4oMini, // Optional: provide examples to help the model understand the format examples = exampleForecasts, // Optional: provide a fixing parser for error correction fixingParser = StructureFixingParser( model = OpenAIModels.Chat.GPT4o, retries = 3 ) ) ``` The `executeStructured` method takes the following arguments: | Name | Data type | Required | Default | Description | | -------------- | ---------------------- | -------- | ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `prompt` | Prompt | Yes | | The prompt to execute. For more information, see [Prompts](../prompts/). | | `model` | LLModel | Yes | | The main model to execute the prompt. | | `examples` | List | No | `emptyList()` | Optional list of examples to help the model understand the expected format. | | `fixingParser` | StructureFixingParser? | No | `null` | Optional parser that handles malformed responses by using an auxiliary LLM to intelligently fix parsing errors. When provided, automatically retries failed parses with error correction. | The method returns a `Result>` containing either the successfully parsed structured data or an error. ### Layer 2: Agent LLM context The agent LLM context layer allows you to request structured responses within agent sessions. This is useful for building conversational agents that need structured data at specific points in their flow. Use the `requestLLMStructured` method within a `writeSession` for agent-based interactions: ``` val structuredResponse = llm.writeSession { requestLLMStructured( examples = exampleForecasts, fixingParser = StructureFixingParser( model = OpenAIModels.Chat.GPT4o, retries = 3 ) ) } ``` The `fixingParser` parameter provides automatic error correction for malformed JSON responses. When parsing fails, it uses an auxiliary LLM to intelligently fix the response up to the specified number of retries. **StructureFixingParser parameters:** - `model: LLModel` - The LLM used to fix malformed JSON output - `retries: Int` - Maximum number of fixing attempts (default: 3) - `prompt` - Optional custom prompt function for the fixing process (defaults to a built-in fixing prompt) The fixing process iteratively passes the parsing error to the auxiliary model, which attempts to correct the JSON while preserving the original data and making minimal changes. #### Integrating with agent strategies You can integrate structured data processing into your agent strategies: ``` val agentStrategy = strategy("weather-forecast") { val setup by nodeLLMRequest() val getStructuredForecast by node { _ -> val structuredResponse = llm.writeSession { requestLLMStructured( fixingParser = StructureFixingParser( model = OpenAIModels.Chat.GPT4o, retries = 3 ) ) } """ Response structure: $structuredResponse """.trimIndent() } edge(nodeStart forwardTo setup) edge(setup forwardTo getStructuredForecast) edge(getStructuredForecast forwardTo nodeFinish) } ``` ### Layer 3: Node layer The node layer provides the highest level of abstraction for structured output in agent workflows. Use `nodeLLMRequestStructured` to create reusable agent nodes that handle structured data. This creates an agent node that: - Accepts a `String` input (user message) - Appends the message to the LLM prompt - Requests structured output from the LLM - Returns `Result>` #### Node layer example ``` val agentStrategy = strategy("weather-forecast") { val setup by node { _ -> "Please provide a weather forecast for Amsterdam" } // Create a structured output node using delegate syntax val getWeatherForecast by nodeLLMRequestStructured( name = "forecast-node", examples = exampleForecasts, fixingParser = StructureFixingParser( model = OpenAIModels.Chat.GPT4o, retries = 3 ) ) val processResult by node>, String> { result -> when { result.isSuccess -> { val forecast = result.getOrNull()?.data "Weather forecast: $forecast" } result.isFailure -> { "Failed to get structured forecast: ${result.exceptionOrNull()?.message}" } else -> "Unknown result state" } } edge(nodeStart forwardTo setup) edge(setup forwardTo getWeatherForecast) edge(getWeatherForecast forwardTo processResult) edge(processResult forwardTo nodeFinish) } ``` #### Full code sample Here is a full example of using the Structured Output API: ``` // Note: Import statements are omitted for brevity @Serializable @SerialName("SimpleWeatherForecast") @LLMDescription("Simple weather forecast for a location") data class SimpleWeatherForecast( @property:LLMDescription("Location name") val location: String, @property:LLMDescription("Temperature in Celsius") val temperature: Int, @property:LLMDescription("Weather conditions (e.g., sunny, cloudy, rainy)") val conditions: String ) val token = System.getenv("OPENAI_KEY") ?: error("Environment variable OPENAI_KEY is not set") fun main(): Unit = runBlocking { // Create sample forecasts val exampleForecasts = listOf( SimpleWeatherForecast( location = "New York", temperature = 25, conditions = "Sunny" ), SimpleWeatherForecast( location = "London", temperature = 18, conditions = "Cloudy" ) ) // Generate JSON Schema val forecastStructure = JsonStructure.create( schemaGenerator = BasicJsonSchemaGenerator.Default, examples = exampleForecasts ) // Define the agent strategy val agentStrategy = strategy("weather-forecast") { val setup by nodeLLMRequest() val getStructuredForecast by node { _ -> val structuredResponse = llm.writeSession { requestLLMStructured() } """ Response structure: $structuredResponse """.trimIndent() } edge(nodeStart forwardTo setup) edge(setup forwardTo getStructuredForecast) edge(getStructuredForecast forwardTo nodeFinish) } // Configure and run the agent val agentConfig = AIAgentConfig( prompt = prompt("weather-forecast-prompt") { system( """ You are a weather forecasting assistant. When asked for a weather forecast, provide a realistic but fictional forecast. """.trimIndent() ) }, model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 5 ) val runner = AIAgent( promptExecutor = simpleOpenAIExecutor(token), toolRegistry = ToolRegistry.EMPTY, strategy = agentStrategy, agentConfig = agentConfig ) runner.run("Get weather forecast for Paris") } ``` ## Advanced usage The examples above demonstrate the simplified API that automatically selects the best structured output approach based on model capabilities. For more control over the structured output process, you can use the advanced API with manual schema creation and provider-specific configurations. ### Manual schema creation and configuration Instead of relying on automatic schema generation, you can create schemas explicitly using `JsonStructure.create` and configure structured output behavior manually via the `StructuredOutput` class. The key difference is that instead of passing simple parameters like `examples` and `fixingParser`, you create a `StructuredRequestConfig` object that allows fine-grained control over: - **Schema generation**: Choose specific generators (Standard, Basic, or Provider-specific) - **Output modes**: Native structured output support vs Manual prompting - **Provider mapping**: Different configurations for different LLM providers - **Fallback strategies**: Default behavior when provider-specific config is unavailable ``` // Create different schema structures with different generators val genericStructure = JsonStructure.create( schemaGenerator = StandardJsonSchemaGenerator, examples = exampleForecasts ) val openAiStructure = JsonStructure.create( schemaGenerator = OpenAIBasicJsonSchemaGenerator, examples = exampleForecasts ) val anthropicStructure = JsonStructure.create( schemaGenerator = AnthropicBasicJsonSchemaGenerator, examples = exampleForecasts ) val promptExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_KEY")) // The advanced API uses StructuredRequestConfig instead of simple parameters val structuredResponse = promptExecutor.executeStructured( prompt = prompt("structured-data") { system("You are a weather forecasting assistant.") user("What is the weather forecast for Amsterdam?") }, model = OpenAIModels.Chat.GPT4oMini, config = StructuredRequestConfig( byProvider = mapOf( LLMProvider.OpenAI to StructuredRequest.Native(openAiStructure), LLMProvider.Anthropic to StructuredRequest.Native(anthropicStructure), ), default = StructuredRequest.Manual(genericStructure) ), fixingParser = StructureFixingParser( model = AnthropicModels.Haiku_4_5, retries = 2 ) ) ``` ### Schema generators Different schema generators are available depending on your needs: - **StandardJsonSchemaGenerator**: Full JSON Schema with support for polymorphism, definitions, and recursive references - **BasicJsonSchemaGenerator**: Simplified schema without polymorphism support, compatible with more models - **Provider-specific generators**: Optimized schemas for specific LLM providers (OpenAI, Anthropic, Google, etc.) ### Usage across all layers The advanced configuration works consistently across all three layers of the API. The method names remain the same, only the parameter changes from simple arguments to the more advanced `StructuredRequestConfig`: - **Prompt executor**: `executeStructured(prompt, model, config: StructuredRequestConfig)` - **Agent LLM context**: `requestLLMStructured(config: StructuredRequestConfig)` - **Node layer**: `nodeLLMRequestStructured(config: StructuredRequestConfig)` The simplified API (using just `examples` and `fixingParser` parameters) is recommended for most use cases, while the advanced API provides additional control when needed. ## Best practices 1. **Use clear descriptions**: provide clear and detailed descriptions using `@LLMDescription` annotations to help the LLM understand the expected data. 1. **Provide examples**: include examples of valid data structures to guide the LLM. 1. **Handle errors gracefully**: implement proper error handling to deal with cases where the LLM might not produce a valid structure. 1. **Use appropriate schema types**: select the appropriate schema format and type based on your needs and the capabilities of the LLM you are using. 1. **Test with different models**: different LLMs may have varying abilities to follow structured formats, so test with multiple models if possible. 1. **Start simple**: begin with simple structures and gradually add complexity as needed. 1. **Use polymorphism Carefully**: while the API supports polymorphism with sealed classes, be aware that it can be more challenging for LLMs to handle correctly. # Streaming API Koog’s **Streaming API** lets you consume **LLM output incrementally** as a `Flow` in Kotlin / `Flow.Publisher` in Java. Instead of waiting for a full response, your code can: - render assistant text as it arrives, - detect **tool calls** live and act on them, - know when a stream **ends** and why. The stream carries **typed frames** organized into two categories: **Delta frames** (incremental/partial content): - `StreamFrame.TextDelta(text: String, index: Int?)` — incremental assistant text - `StreamFrame.ReasoningDelta(text: String?, summary: String?, index: Int?)` — incremental reasoning text and summary - `StreamFrame.ToolCallDelta(id: String?, name: String?, content: String?, index: Int?)` — partial tool invocation **Complete frames** (full content): - `StreamFrame.TextComplete(text: String, index: Int?)` — complete assistant text - `StreamFrame.ReasoningComplete(content: List, summary: List?, encrypted: String?, index: Int?)` — complete reasoning with optional summary and encrypted content - `StreamFrame.ToolCallComplete(id: String?, name: String, content: String, index: Int?)` — complete tool invocation **End marker**: - `StreamFrame.End(finishReason: String?, metaInfo: ResponseMetaInfo)` — end-of-stream marker with response metadata **Delta frames** (incremental/partial content): - `StreamFrame.TextDelta` — incremental assistant text. Fields: `getText()`, `getIndex()`. - `StreamFrame.ReasoningDelta` — incremental reasoning text and summary. Fields: `getText()`, `getSummary()`, `getIndex()`. - `StreamFrame.ToolCallDelta` — partial tool invocation. Fields: `getId()`, `getName()`, `getContent()`, `getIndex()`. **Complete frames** (full content): - `StreamFrame.TextComplete` — complete assistant text. Fields: `getText()`, `getIndex()`. - `StreamFrame.ReasoningComplete` — complete reasoning with optional summary and encrypted content. Fields: `getText()` (returns `List`), `getSummary()` (returns `List`), `getEncrypted()`, `getIndex()`. - `StreamFrame.ToolCallComplete` — complete tool invocation. Fields: `getId()`, `getName()`, `getContent()`, `getIndex()`. Also provides `getContentJson()` and `getContentJsonResult()` for JSON parsing. **End marker**: - `StreamFrame.End` — end-of-stream marker. Fields: `getFinishReason()`, `getMetaInfo()`. Helpers are provided to extract plain text, convert frames to `Message.Response` objects, and safely **combine chunked tool calls**. ## API overview With streaming you can: - Process data as it arrives (improves UI responsiveness) - Parse structured info on the fly (Markdown/JSON/etc.) - Emit objects as they complete - Trigger tools in real time - Access model reasoning in real-time (for supported models) You can operate either on the **frames** themselves or on **plain text** derived from frames. ### Delta vs Complete Frames The streaming API distinguishes between two types of frames: - **Delta frames** (`DeltaFrame`) — Incremental/partial content that arrives in chunks. These are ideal for real-time display as content streams in. Examples: `TextDelta`, `ReasoningDelta`, `ToolCallDelta`. - **Complete frames** (`CompleteFrame`) — Full content emitted after all deltas for that content type have been received. These are useful for final processing and conversion to `Message.Response` objects. Examples: `TextComplete`, `ReasoningComplete`, `ToolCallComplete`. Typically, you'll work with delta frames for UI updates and complete frames for extracting final structured data. ______________________________________________________________________ ## Usage ### Working with frames directly This is the most general approach: react to each frame kind. ``` llm.writeSession { appendPrompt { user("Tell me a joke, then call a tool with JSON args.") } val stream = requestLLMStreaming() // Flow stream.collect { frame -> when (frame) { is StreamFrame.TextDelta -> print(frame.text) is StreamFrame.ReasoningDelta -> print("[Reasoning] text=${frame.text} summary=${frame.summary}") is StreamFrame.ToolCallComplete -> { println("\n🔧 Tool call: ${frame.name} args=${frame.content}") // Optionally parse lazily: // val json = frame.contentJson } is StreamFrame.End -> println("\n[END] reason=${frame.finishReason}") else -> {} // Handle other frame types (TextComplete, ToolCallDelta, etc.) } } } ``` ``` ctx.getLlm().writeSession(session -> { session.appendPrompt(prompt -> { prompt.user("Tell me a joke, then call a tool with JSON args."); return null; }); Flow.Publisher stream = session.requestLLMStreaming(); stream.subscribe(new Flow.Subscriber<>() { @Override public void onSubscribe(Flow.Subscription subscription) { subscription.request(Long.MAX_VALUE); } @Override public void onNext(StreamFrame frame) { if (frame instanceof StreamFrame.TextDelta delta) { System.out.print(delta.getText()); } else if (frame instanceof StreamFrame.ReasoningDelta reasoning) { System.out.print("[Reasoning] text=" + reasoning.getText() + " summary=" + reasoning.getSummary()); } else if (frame instanceof StreamFrame.ToolCallComplete toolCall) { System.out.println("\nTool call: " + toolCall.getName() + " args=" + toolCall.getContent()); } else if (frame instanceof StreamFrame.End end) { System.out.println("\n[END] reason=" + end.getFinishReason()); } // Handle other frame types (TextComplete, ToolCallDelta, etc.) } @Override public void onError(Throwable throwable) { System.err.println("Stream error: " + throwable.getMessage()); } @Override public void onComplete() { } }); return null; }); ``` It is important to note that you can parse the output by working directly with a raw string stream. This approach gives you more flexibility and control over the parsing process. Here is a raw string stream with the Markdown definition of the output structure: ``` fun markdownBookDefinition(): MarkdownStructureDefinition { return MarkdownStructureDefinition("name", schema = { /*...*/ }) } val mdDefinition = markdownBookDefinition() llm.writeSession { val stream = requestLLMStreaming(mdDefinition) // Access the raw string chunks directly stream.collect { chunk -> // Process each chunk of text as it arrives println("Received chunk: $chunk") // The chunks together will be structured as a text following the mdDefinition schema } } ``` ``` StructureDefinition mdDefinition = markdownBookDefinition(); ctx.getLlm().writeSession(session -> { session.appendPrompt(prompt -> { prompt.user(input); }); Flow.Publisher stream = session.requestLLMStreaming(mdDefinition); // Access the raw frames directly stream.subscribe(new Flow.Subscriber<>() { @Override public void onSubscribe(Flow.Subscription subscription) { subscription.request(Long.MAX_VALUE); } @Override public void onNext(StreamFrame frame) { // Process each frame as it arrives System.out.println("Received frame: " + frame); } @Override public void onError(Throwable throwable) { System.err.println("Stream error: " + throwable.getMessage()); } @Override public void onComplete() { } }); return null; }); ``` ### Working with reasoning frames Models that support reasoning (such as Claude Sonnet 4.5 or GPT-o1) emit reasoning frames during streaming. You can access both the reasoning process and its summary: ``` llm.writeSession { appendPrompt { user("Solve this complex problem: ...") } val stream = requestLLMStreaming() val reasoningSteps = mutableListOf() val summarySteps = mutableListOf() stream.collect { frame -> when (frame) { is StreamFrame.ReasoningDelta -> { frame.text?.let { reasoningSteps.add(it) print(frame.text) // Display reasoning as it arrives } frame.summary?.let { summarySteps.add(it) print(frame.summary) // Display reasoning summary as it arrives } } is StreamFrame.ReasoningComplete -> { // Access complete reasoning println("\nComplete reasoning: ${frame.content.joinToString("")}") println("Summary: ${frame.summary?.joinToString("") ?: "N/A"}") } is StreamFrame.TextDelta -> print(frame.text) is StreamFrame.End -> println("\n[END]") else -> {} } } } ``` ``` ctx.getLlm().writeSession(session -> { session.appendPrompt(prompt -> { prompt.user("Solve this complex problem: ..."); return null; }); Flow.Publisher stream = session.requestLLMStreaming(); List reasoningSteps = new ArrayList<>(); List summarySteps = new ArrayList<>(); stream.subscribe(new Flow.Subscriber() { @Override public void onSubscribe(Flow.Subscription subscription) { subscription.request(Long.MAX_VALUE); } @Override public void onNext(StreamFrame frame) { if (frame instanceof StreamFrame.ReasoningDelta reasoning) { if (reasoning.getText() != null) { reasoningSteps.add(reasoning.getText()); System.out.print(reasoning.getText()); } if (reasoning.getSummary() != null) { summarySteps.add(reasoning.getSummary()); System.out.print(reasoning.getSummary()); } } else if (frame instanceof StreamFrame.ReasoningComplete complete) { // Access complete reasoning System.out.println("\nComplete reasoning: " + String.join("", complete.getContent())); System.out.println("Summary: " + (complete.getSummary() != null ? String.join("", complete.getSummary()) : "N/A")); } else if (frame instanceof StreamFrame.TextDelta delta) { System.out.print(delta.getText()); } else if (frame instanceof StreamFrame.End) { System.out.println("\n[END]"); } } @Override public void onError(Throwable throwable) { } @Override public void onComplete() { } }); return null; }); ``` ### Working with a raw text stream (derived) If you have existing streaming parsers that expect `Flow`, derive text chunks via `filterTextOnly()` or collect them with `collectText()`. ``` llm.writeSession { val frames = requestLLMStreaming() // Stream text chunks as they come: frames.filterTextOnly().collect { chunk -> print(chunk) } // Or, gather all text into one String after End: val fullText = frames.collectText() println("\n---\n$fullText") } ``` ``` ctx.getLlm().writeSession(session -> { Flow.Publisher frames = session.requestLLMStreaming(); // Stream text chunks as they come (equivalent of filterTextOnly): StringBuilder fullText = new StringBuilder(); frames.subscribe(new Flow.Subscriber<>() { @Override public void onSubscribe(Flow.Subscription subscription) { subscription.request(Long.MAX_VALUE); } @Override public void onNext(StreamFrame frame) { if (frame instanceof StreamFrame.TextDelta delta) { System.out.print(delta.getText()); fullText.append(delta.getText()); } } @Override public void onError(Throwable throwable) { } @Override public void onComplete() { // fullText now contains all text (equivalent of collectText) System.out.println("\n---\n" + fullText); } }); return null; }); ``` ### Listening to stream events in event handlers You can listen to stream events in [agent event handlers](../features/agent-event-handlers/). ``` handleEvents { onToolCallStarting { context -> println("\n🔧 Using ${context.toolName} with ${context.toolArgs}... ") } onLLMStreamingFrameReceived { context -> when (val frame = context.streamFrame) { is StreamFrame.TextDelta -> print(frame.text) is StreamFrame.ReasoningDelta -> print("[Reasoning] text=${frame.text} summary=${frame.summary}") else -> {} // Handle other frame types if needed } } onLLMStreamingFailed { context -> println("❌ Error: ${context.error}") } onLLMStreamingCompleted { println("🏁 Done") } } ``` ``` .install(EventHandler.Feature, config -> { config.onToolCallStarting(ctx -> { System.out.println("\nUsing " + ctx.getToolName() + " with " + ctx.getToolArgs() + "... "); }); config.onLLMStreamingFrameReceived(ctx -> { StreamFrame frame = ctx.getStreamFrame(); if (frame instanceof StreamFrame.TextDelta delta) { System.out.print(delta.getText()); } else if (frame instanceof StreamFrame.ReasoningDelta reasoning) { System.out.print("[Reasoning] text=" + reasoning.getText() + " summary=" + reasoning.getSummary()); } }); config.onLLMStreamingFailed(ctx -> { System.out.println("Error: " + ctx.getError()); }); config.onLLMStreamingCompleted(ctx -> { System.out.println("Done"); }); }) ``` ### Converting frames to `Message.Response` You can transform a collected list of frames to standard message objects: - `toAssistantMessageOrNull()` — extracts `Message.Assistant` from text frames - `toReasoningMessageOrNull()` — extracts `MessagePart.Reasoning` from reasoning frames - `toToolCallMessages()` — extracts `MessagePart.Tool.Call` from tool call frames - `toMessageResponses()` — converts all complete frames to their corresponding `Message.Response` objects ## Examples ### Structured data while streaming (Markdown example) Although it is possible to work with a raw string stream, it is often more convenient to work with [structured data](../structured-output/). The structured data approach includes the following key components: 1. **MarkdownStructureDefinition**: a class to help you define the schema and examples for structured data in Markdown format. 1. **markdownStreamingParser**: a function to create a parser that processes a stream of Markdown chunks and emits events. The sections below provide step-by-step instructions and code samples related to processing a stream of structured data. #### 1. Define your data structure First, define a data class to represent your structured data: ``` @Serializable data class Book( val title: String, val author: String, val description: String ) ``` ``` // TODO not yet supported in Java ``` #### 2. Define the Markdown structure Create a definition that specifies how your data should be structured in Markdown with the `MarkdownStructureDefinition` class: ``` fun markdownBookDefinition(): MarkdownStructureDefinition { return MarkdownStructureDefinition("bookList", schema = { markdown { header(1, "title") bulleted { item("author") item("description") } } }, examples = { markdown { header(1, "The Great Gatsby") bulleted { item("F. Scott Fitzgerald") item("A novel set in the Jazz Age that tells the story of Jay Gatsby's unrequited love for Daisy Buchanan.") } } }) } ``` ``` // TODO not yet supported in Java ``` #### 3. Create a parser for your data structure The `markdownStreamingParser` provides several handlers for different Markdown elements: ``` markdownStreamingParser { // Handle level 1 headings (level ranges from 1 to 6) onHeader(1) { headerText -> } // Handle bullet points onBullet { bulletText -> } // Handle code blocks onCodeBlock { codeBlockContent -> } // Handle lines matching a regex pattern onLineMatching(Regex("pattern")) { line -> } // Handle the end of the stream onFinishStream { remainingText -> } } ``` ``` // TODO not yet supported in Java ``` Using the defined handlers, you can implement a function that parses the Markdown stream and emits your data objects with the `markdownStreamingParser` function. ``` fun parseMarkdownStreamToBooks(markdownStream: Flow): Flow { return flow { markdownStreamingParser { var currentBookTitle = "" val bulletPoints = mutableListOf() // Handle the event of receiving the Markdown header in the response stream onHeader(1) { headerText -> // If there was a previous book, emit it if (currentBookTitle.isNotEmpty() && bulletPoints.isNotEmpty()) { val author = bulletPoints.getOrNull(0) ?: "" val description = bulletPoints.getOrNull(1) ?: "" emit(Book(currentBookTitle, author, description)) } currentBookTitle = headerText bulletPoints.clear() } // Handle the event of receiving the Markdown bullets list in the response stream onBullet { bulletText -> bulletPoints.add(bulletText) } // Handle the end of the response stream onFinishStream { // Emit the last book, if present if (currentBookTitle.isNotEmpty() && bulletPoints.isNotEmpty()) { val author = bulletPoints.getOrNull(0) ?: "" val description = bulletPoints.getOrNull(1) ?: "" emit(Book(currentBookTitle, author, description)) } } }.parseStream(markdownStream.filterTextOnly()) } } ``` ``` // TODO not yet supported in Java ``` #### 4. Use the parser in your agent strategy ``` val agentStrategy = strategy>("library-assistant") { // Describe the node containing the output stream parsing val getMdOutput by node> { booksDescription -> val books = mutableListOf() val mdDefinition = markdownBookDefinition() llm.writeSession { appendPrompt { user(booksDescription) } // Initiate the response stream in the form of the definition `mdDefinition` val markdownStream = requestLLMStreaming(mdDefinition) // Call the parser with the result of the response stream and perform actions with the result parseMarkdownStreamToBooks(markdownStream).collect { book -> books.add(book) println("Parsed Book: ${book.title} by ${book.author}") } } books } // Describe the agent's graph making sure the node is accessible edge(nodeStart forwardTo getMdOutput) edge(getMdOutput forwardTo nodeFinish) } ``` ``` // TODO not yet supported in Java ``` ### Advanced usage: Streaming with tools You can also use the Streaming API with tools to process data as it arrives. The following sections provide a brief step-by-step guide on how to define a tool and use it with streaming data. ### 1. Define a tool for your data structure ``` @Serializable data class Book( val title: String, val author: String, val description: String ) class BookTool(): SimpleTool( argsType = typeToken(), name = NAME, description = "A tool to parse book information from Markdown" ) { companion object { const val NAME = "book" } override suspend fun execute(args: Book): String { println("${args.title} by ${args.author}:\n ${args.description}") return "Done" } } ``` ``` class BookTool implements ToolSet { @Tool @LLMDescription("A tool to parse book information from Markdown") public String book( @LLMDescription("Title of the book") String title, @LLMDescription("Author of the book") String author, @LLMDescription("Description of the book") String description ) { System.out.println(title + " by " + author + ":\n " + description); return "Done"; } } ``` ### 2. Use the tool with streaming data ``` val agentStrategy = strategy("library-assistant") { val getMdOutput by node { input -> val mdDefinition = markdownBookDefinition() llm.writeSession { appendPrompt { user(input) } val markdownStream = requestLLMStreaming(mdDefinition) parseMarkdownStreamToBooks(markdownStream).collect { book -> callToolRaw(BookTool.NAME, book) /* Other possible options: callTool(BookTool::class, book) callTool(book) findTool(BookTool::class).execute(book) */ } // We can make parallel tool calls parseMarkdownStreamToBooks(markdownStream).toParallelToolCallsRaw(toolClass=BookTool::class).collect { println("Tool call result: $it") } } } edge(nodeStart forwardTo getMdOutput) edge(getMdOutput forwardTo nodeFinish) } ``` ``` var strategy = AIAgentGraphStrategy.builder("library-assistant") .withInput(String.class) .withOutput(Void.class); var getMdOutput = AIAgentNode.builder("getMdOutput") .withInput(String.class) .withOutput(Void.class) .withAction((input, ctx) -> { StructureDefinition mdDefinition = markdownBookDefinition(); ctx.getLlm().writeSession(session -> { session.appendPrompt(prompt -> { prompt.user(input); return null; }); Flow.Publisher markdownStream = session.requestLLMStreaming(mdDefinition); // Process streamed frames and invoke tools on ToolCallComplete frames markdownStream.subscribe(new Flow.Subscriber() { @Override public void onSubscribe(Flow.Subscription subscription) { subscription.request(Long.MAX_VALUE); } @Override public void onNext(StreamFrame frame) { if (frame instanceof StreamFrame.ToolCallComplete toolCall) { System.out.println("Tool call: " + toolCall.getName() + " args=" + toolCall.getContent()); } } @Override public void onError(Throwable throwable) { } @Override public void onComplete() { } }); return null; }); return null; }) .build(); strategy.edge(strategy.nodeStart, getMdOutput); strategy.edge(getMdOutput, strategy.nodeFinish); ``` ### 3. Register the tool in your agent configuration ``` val toolRegistry = ToolRegistry { tool(BookTool()) } val runner = AIAgent( promptExecutor = simpleOpenAIExecutor("OPENAI_API_KEY"), llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = toolRegistry ) ``` ``` ToolRegistry toolRegistry = ToolRegistry.builder() .tools(new BookTool()) .build(); AIAgent runner = AIAgent.builder() .promptExecutor(PromptExecutor.builder().openAI("OPENAI_API_KEY").build()) .llmModel(OpenAIModels.Chat.GPT4o) .toolRegistry(toolRegistry) .build(); ``` ## Best practices 1. **Define clear structures**: create clear and unambiguous markdown structures for your data. 1. **Provide good examples**: include comprehensive examples in your `MarkdownStructureDefinition` to guide the LLM. 1. **Handle incomplete data**: always check for null or empty values when parsing data from the stream. 1. **Clean up resources**: use the `onFinishStream` handler to clean up resources and process any remaining data. 1. **Handle errors**: implement proper error handling for malformed Markdown or unexpected data. 1. **Testing**: test your parser with various input scenarios, including partial chunks and malformed input. 1. **Parallel processing**: for independent data items, consider using parallel tool calls for better performance. # Custom node implementation This page provides detailed instructions on how to implement your own custom nodes in the Koog framework. Custom nodes let you extend the functionality of agent workflows by creating reusable components that perform specific operations. To learn more about what graph nodes are, their usage, and existing default nodes, see [Graph nodes](../nodes-and-components/). ## Node architecture overview Before diving into implementation details, it is important to understand the architecture of nodes in the Koog framework. Nodes are the fundamental building blocks of agent workflows, where each node represents a specific operation or transformation in the workflow. You connect nodes using edges, which define the flow of execution between nodes. Each node has an `execute` method that takes an input and produces an output, which is then passed to the next node in the workflow. ## Implementing a custom node Custom node implementations range from simple implementations that perform a basic logic on the input data and return an output, to more complex node implementations that accept parameters and maintain state between runs. ### Basic node implementation The simplest way to implement a custom node in a graph and define your own custom logic would be to use the following pattern: ``` val myNode by node("node_name") { input -> // Processing returnValue } ``` ``` var myNode = AIAgentNode.builder("node_name") .withInput(Input.class) .withOutput(Output.class) .withAction((input, ctx) -> { // Processing return returnValue; }) .build(); ``` The code above represents a custom node `myNode` with predefined `Input` and `Output` types, with the optional name string parameter (`node_name`). In Kotlin, you use the `node` DSL function. In Java, you use the `AIAgentNode.builder()` pattern. In an actual example, here is a simple node that takes a string input and returns the input's length: ``` val myNode by node("node_name") { input -> // Processing input.length } ``` ``` var myNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(Integer.class) .withAction((input, ctx) -> { // Processing return input.length(); }) .build(); ``` Another way to create a custom node is to extract it into a reusable function. In Kotlin, you define an extension function on `AIAgentSubgraphBuilderBase` that calls the `node` function. In Java, you extract the node builder call into a helper method. ``` fun AIAgentSubgraphBuilderBase<*, *>.myCustomNode( name: String? = null ): AIAgentNodeDelegate = node(name) { input -> // Custom logic input // Return the input as output (pass-through) } val myCustomNode by myCustomNode("node_name") ``` ``` var myCustomNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { // Custom logic return input; // Return the input as output (pass-through) }) .build(); ``` This creates a pass-through node that performs some custom logic but returns the input as the output without modification. ### Nodes with additional arguments You can create nodes that accept arguments to customize their behavior: ``` fun AIAgentSubgraphBuilderBase<*, *>.myNodeWithArguments( name: String? = null, arg1: String, arg2: Int ): AIAgentNodeDelegate = node(name) { input -> // Use arg1 and arg2 in your custom logic input // Return the input as the output } val myCustomNode by myNodeWithArguments("node_name", arg1 = "value1", arg2 = 42) ``` ``` String arg1 = "value1"; int arg2 = 42; var myCustomNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { // Use arg1 and arg2 in your custom logic return input; // Return the input as the output }) .build(); ``` ### Parameterized nodes You can define nodes with generic input and output type parameters. In Kotlin, you use `inline` functions with `reified` type parameters. In Java, you specify the types explicitly when building the node. ``` inline fun AIAgentSubgraphBuilderBase<*, *>.myParameterizedNode( name: String? = null, ): AIAgentNodeDelegate = node(name) { input -> // Do some additional actions // Return the input as the output input } val strategy = strategy("strategy_name") { val myCustomNode by myParameterizedNode("node_name") } ``` ``` // In Java, specify the types explicitly when building the node var myCustomNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { // Do some additional actions // Return the input as the output return input; }) .build(); ``` ### Stateful nodes If your node needs to maintain state between runs, you can use closure variables. In Kotlin, you declare a variable in the enclosing function. In Java, you use thread-safe wrappers like `AtomicInteger` since lambda captures must be effectively final. ``` fun AIAgentSubgraphBuilderBase<*, *>.myStatefulNode( name: String? = null ): AIAgentNodeDelegate { var counter = 0 return node(name) { input -> counter++ println("Node executed $counter times") input } } ``` ``` // In Java, use AtomicInteger (or similar) since the lambda captures must be effectively final AtomicInteger counter = new AtomicInteger(0); var myStatefulNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { int count = counter.incrementAndGet(); System.out.println("Node executed " + count + " times"); return input; }) .build(); ``` ## Node input and output types Nodes can have different input and output types. In both Kotlin and Java, these are specified as generic type parameters: ``` val stringToIntNode by node("node_name") { input: String -> // Processing input.toInt() // Convert string to integer } ``` ``` var stringToIntNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(Integer.class) .withAction((input, ctx) -> { // Processing return Integer.parseInt(input); // Convert string to integer }) .build(); ``` Note The input and output types determine how the node can be connected to other nodes in the workflow. Nodes can only be connected if the output type of the source node is compatible with the input type of the target node. ## Best practices When implementing custom nodes, follow these best practices: 1. **Keep nodes focused**: each node should perform a single, well-defined operation. 1. **Use descriptive names**: node names should clearly indicate their purpose. 1. **Document parameters**: provide clear documentation for all parameters. 1. **Handle errors gracefully**: implement proper error handling to prevent workflow failures. 1. **Make nodes reusable**: design nodes to be reusable across different workflows. 1. **Use type parameters**: use generic type parameters when appropriate to make nodes more flexible. 1. **Provide default values**: when possible, provide sensible default values for parameters. ## Common patterns The following sections provide some common patterns for implementing custom nodes. ### Pass-through nodes Nodes that perform an operation but return the input as the output. ``` val loggingNode by node("node_name") { input -> println("Processing input: $input") input // Return the input as the output } ``` ``` var loggingNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { System.out.println("Processing input: " + input); return input; // Return the input as the output }) .build(); ``` ### Transformation nodes Nodes that transform the input data and produce a modified output. ``` val upperCaseNode by node("node_name") { input -> println("Processing input: $input") input.uppercase() // Transform the input to uppercase } ``` ``` var upperCaseNode = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { System.out.println("Processing input: " + input); return input.toUpperCase(); // Transform the input to uppercase }) .build(); ``` ### LLM interaction nodes Nodes that interact with the LLM. In Kotlin, you have fine-grained control over the LLM session. In Java, you typically use pre-built factory methods like `AIAgentNode.llmRequest()` that handle prompt construction automatically. ``` val summarizeTextNode by node("node_name") { input -> llm.writeSession { appendPrompt { user("Please summarize the following text: $input") } val response = requestLLMWithoutTools() response.parts.filterIsInstance().joinToString("\n") { it.text } } } ``` ``` // In Java, LLM interaction is handled using pre-built factory nodes. // AIAgentNode.llmRequest() creates a node that sends the input string as a user // message to the LLM and returns the response. The prompt text is provided as // the node's input when it is executed in the graph. var summarizeTextNode = AIAgentNode.llmRequest("node_name"); // To extract the text content from the LLM response, chain a separate node: var extractContent = AIAgentNode.builder("extract-content") .withInput(Message.Assistant.class) .withOutput(String.class) .withAction((response, ctx) -> response.getParts().stream() .filter(p -> p instanceof MessagePart.Text) .map(p -> ((MessagePart.Text) p).getText()) .collect(Collectors.joining())) .build(); ``` Note The Kotlin example above shows fine-grained control over the LLM session (custom prompt construction, explicit `requestLLMWithoutTools` call). The Java API provides higher-level factory methods like `AIAgentNode.llmRequest()` that handle prompt construction automatically, where the input string becomes the user message. For advanced prompt customization, compose multiple nodes or use a custom subgraph. ### Tool run node Custom nodes that execute tools. In Kotlin, you can manually construct tool calls and execute them. In Java, you typically use subgraphs that delegate tool orchestration to the LLM. ``` val nodeExecuteCustomTool by node("node_name") { input -> val toolCall = MessagePart.Tool.Call( id = UUID.randomUUID().toString(), tool = toolName, args = Json.encodeToString(ToolArgs(arg1 = input, arg2 = 42)) // Use the input as tool arguments ) val result = environment.executeTool(toolCall) result.output } ``` ``` // In Java, direct tool execution (as shown in the Kotlin example) is not available // through the Java builder API. Instead, use a subgraph that delegates tool calls // to the LLM, which decides when and how to invoke the tools: var toolSubgraph = AIAgentSubgraph.builder("tool-subgraph") .withInput(String.class) .withOutput(String.class) .withTask(input -> "Use my_tool with input: " + input) .build(); ``` Note The Kotlin example demonstrates low-level tool execution by manually constructing a `MessagePart.Tool.Call` and calling `environment.executeTool()`. The Java API encourages a higher-level approach using subgraphs with `withTask()`, where the LLM orchestrates tool calls automatically. To restrict which tools are available, chain `.limitedTools(List.of(myTool))` before `.withInput()`. # LLM sessions and manual history management This page provides detailed information about LLM sessions, including how to work with read and write sessions, manage conversation history, and make requests to language models. ## Introduction LLM sessions are a fundamental concept that provides a structured way to interact with language models (LLMs). They manage the conversation history, handle requests to the LLM, and provide a consistent interface for running tools and processing responses. ## Understanding LLM sessions An LLM session represents a context for interacting with a language model. It encapsulates: - The conversation history (prompt) - Available tools - Methods for making requests to the LLM - Methods for updating the conversation history - Methods for running tools Sessions are managed by the `AIAgentLLMContext` class, which provides methods for creating read and write sessions. ### Session types The Koog framework provides two types of sessions: 1. **Write Sessions** (`AIAgentLLMWriteSession`): Allow modifying the prompt and tools, making LLM requests, and running tools. Changes made in a write session are persisted back to the LLM context. 1. **Read Sessions** (`AIAgentLLMReadSession`): Provide read-only access to the prompt and tools. They are useful for inspecting the current state without making changes. The key difference is that write sessions can modify the conversation history, while read sessions cannot. ### Session lifecycle Sessions have a defined lifecycle: 1. **Creation**: a session is created, e.g., using `llm.writeSession { ... }` or `llm.readSession { ... }`. 1. **Active phase**: the session is active while the lambda block is executing. 1. **Termination**: the session is automatically closed when the lambda block completes. Sessions implement the `AutoCloseable` interface, ensuring they are properly cleaned up even if an exception occurs. ## Working with LLM sessions ### Creating sessions Sessions are created using methods of the `AIAgentLLMContext` class: ``` // Creating a write session llm.writeSession { // Session code here } // Creating a read session llm.readSession { // Session code here } ``` ``` // Creating a write session ctx.getLlm().writeSession(session -> { // Session code here return null; }); // Creating a read session ctx.getLlm().readSession(session -> { // Session code here return null; }); ``` These functions take a lambda block that runs within the context of the session. The session is automatically closed when the block completes. ### Session scope and thread safety Sessions use a read-write lock to ensure thread safety: - Multiple read sessions can be active simultaneously. - Only one write session can be active at a time. - A write session blocks all other sessions (both read and write). This ensures that the conversation history is not corrupted by concurrent modifications. ### Accessing session properties Within a session, you can access the prompt and tools: ``` llm.readSession { val messageCount = prompt.messages.size val availableTools = tools.map { it.name } } ``` ``` ctx.getLlm().readSession(session -> { int messageCount = session.getPrompt().getMessages().size(); var availableTools = session.getTools().stream().map(tool -> tool.getName()).toList(); return null; }); ``` In a write session, you can also modify these properties: ``` llm.writeSession { // Modify the prompt appendPrompt { user("New user message") } // Modify the tools tools = newTools } ``` ``` ctx.getLlm().writeSession(session -> { // Modify the prompt session.appendPrompt(promptBuilder -> { promptBuilder.user("New user message"); return null; }); // Modify the tools session.setTools(newTools); return null; }); ``` For more information, see the detailed API reference for [AIAgentLLMReadSession](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.session/-a-i-agent-l-l-m-read-session/index.html) and [AIAgentLLMWriteSession](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.session/-a-i-agent-l-l-m-write-session/index.html). ## Making LLM requests ### Basic request methods The most common methods for making LLM requests are: 1. `requestLLM()`: makes a request to the LLM with the current prompt and tools, returning a response. 1. `requestLLMWithoutTools()`: makes a request to the LLM with the current prompt but without any tools, returning a single response. 1. `requestLLMForceOneTool()`: makes a request to the LLM with the current prompt and tools, forcing the use of one tool. 1. `requestLLMOnlyCallingTools()`: makes a request to the LLM that should be processed by only using tools. Example: ``` llm.writeSession { // Make a request with tools enabled val response = requestLLM() // Make a request without tools val responseWithoutTools = requestLLMWithoutTools() } ``` ``` ctx.getLlm().writeSession(session -> { // Make a request with tools enabled var response = session.requestLLM(); // Make a request without tools var responseWithoutTools = session.requestLLMWithoutTools(); return null; }); ``` ### How requests work LLM requests are made when you explicitly call one of the request methods. The key points to understand are: 1. **Explicit invocation**: requests only happen when you call methods like `requestLLM()`, `requestLLMWithoutTools()` and so on. 1. **Immediate execution**: when you call a request method, the request is made immediately, and the method blocks until a response is received. 1. **Automatic history update**: in a write session, the response is automatically added to the conversation history. ### Request methods with tools When making requests with tools enabled, the LLM may respond with a tool call instead of a text response. The request methods handle this transparently: ``` llm.writeSession { val response = requestLLM() // The response might contain tool calls and/or text val toolCalls = response.parts.filterIsInstance() if (toolCalls.isNotEmpty()) { // Handle tool calls } else { // Handle text response } } ``` ``` ctx.getLlm().writeSession(session -> { var response = session.requestLLM(); // The response parts might contain a tool call or text content boolean hasToolCall = response.getParts().stream() .anyMatch(p -> p instanceof MessagePart.Tool.Call); if (hasToolCall) { // Handle tool call } else { // Handle text response } return null; }); ``` In practice, you typically do not need to check the response type manually, as the agent graph handles this routing automatically. ### Structured and streaming requests For more advanced use cases, methods for structured and streaming requests are provided: 1. `requestLLMStructured()`: requests the LLM to provide a response in a specific structured format. 1. `requestLLMStructuredOneShot()`: similar to `requestLLMStructured()` but without retries or corrections. 1. `requestLLMStreaming()`: makes a streaming request to the LLM, returning a flow of response chunks. You can learn more about streaming on the [Streaming API](../streaming-api/) page. Example: ``` llm.writeSession { // Make a structured request val structuredResponse = requestLLMStructured() // Make a streaming request val responseStream = requestLLMStreaming() responseStream.collect { chunk -> // Process each chunk as it arrives } } ``` ``` ctx.getLlm().writeSession(session -> { // Make a non-tool request var responseWithoutTools = session.requestLLMWithoutTools(); // Make a streaming request var responseStream = session.requestLLMStreaming(); // Process chunks from Flow.Publisher return null; }); ``` ## Managing conversation history ### Updating the prompt In a write session, you can add messages to the prompt (conversation history) using the `appendPrompt` method: ``` llm.writeSession { appendPrompt { // Add a system message system("You are a helpful assistant.") // Add a user message user("Hello, can you help me with a coding question?") // Add an assistant message assistant("Of course! What's your question?") // Add a tool result toolResult(myToolResult) } } ``` ``` ctx.getLlm().writeSession(session -> { session.appendPrompt(promptBuilder -> { // Add a system message promptBuilder.system("You are a helpful assistant."); // Add a user message promptBuilder.user("Hello, can you help me with a coding question?"); // Add an assistant message promptBuilder.assistant("Of course! What's your question?"); // Add follow-up context after tool execution promptBuilder.assistant("Tool execution completed successfully."); return null; }); return null; }); ``` You can also completely rewrite the prompt by assigning a new `Prompt` object to the `prompt` property: ``` llm.writeSession { // Create a new prompt based on the old one prompt = prompt.copy(messages = filteredMessages) } ``` ``` ctx.getLlm().writeSession(session -> { var oldPrompt = session.getPrompt(); // Rebuild and replace the prompt (manual rewrite approach in Java) session.setPrompt( Prompt.builder(oldPrompt.getId()) .user("Retained summary of previous conversation") .build() ); return null; }); ``` ### Automatic history update on response When you make an LLM request in a write session, the response is automatically added to the conversation history: ``` llm.writeSession { // Add a user message appendPrompt { user("What's the capital of France?") } // Make a request - the response is automatically added to the history val response = requestLLM() // The prompt now includes both the user message and the model's response } ``` ``` ctx.getLlm().writeSession(session -> { // Add a user message session.appendPrompt(promptBuilder -> { promptBuilder.user("What's the capital of France?"); return null; }); // Make a request - the response is automatically added to the history var response = session.requestLLM(); // The prompt now includes both the user message and the model's response return null; }); ``` This automatic history update is the key feature of write sessions, ensuring that the conversation flows naturally. ### History compression For long-running conversations, the history can grow large and consume a lot of tokens. The platform provides methods for compressing history: ``` llm.writeSession { // Compress the history using a TLDR approach replaceHistoryWithTLDR(HistoryCompressionStrategy.WholeHistory, preserveMemory = true) } ``` ``` // Use the dedicated Java node for history compression. var compressHistory = AIAgentNode.llmCompressHistory("compressHistory"); ``` You can also use the `nodeLLMCompressHistory` node in a strategy graph to compress history at specific points. For more information about history compression and compression strategies, see [History compression](../history-compression/). ## Running tools in sessions ### Calling tools Write sessions provide several methods for calling tools: 1. `callTool(tool, args)`: calls a tool by reference. 1. `callTool(toolName, args)`: calls a tool by name. 1. `callTool(toolClass, args)`: calls a tool by class. 1. `callToolRaw(toolName, args)`: calls a tool by name and returns the raw string result. Example: ``` llm.writeSession { // Call a tool by reference val result = callTool(myTool, myArgs) // Call a tool by name val result2 = callTool("myToolName", myArgs) // Call a tool by class val result3 = callTool(MyTool::class, myArgs) // Call a tool and get the raw result val rawResult = callToolRaw("myToolName", myArgs) } ``` ``` // Java uses dedicated tool-execution nodes in graph strategies. var executeTool = AIAgentNode.executeTools("executeTool"); var sendToolResult = AIAgentNode.llmSendToolResults("sendToolResult"); ``` ### Parallel tool runs To run multiple tools in parallel, write sessions provide extension functions on `Flow`: ``` llm.writeSession { // Run tools in parallel parseDataToArgs(data).toParallelToolCalls(MyTool::class).collect { result -> // Process each result } // Run tools in parallel and get raw results parseDataToArgs(data).toParallelToolCallsRaw(MyTool::class).collect { rawResult -> // Process each raw result } } ``` ``` // Java equivalent for multi-tool execution: use the executeTools node. var executeMultipleTools = AIAgentNode.executeTools("executeMultipleTools"); ``` This is useful for processing large amounts of data efficiently. ## Best practices When working with LLM sessions, follow these best practices: 1. **Use the right session type**: Use write sessions when you need to modify the conversation history and read sessions when you only need to read it. 1. **Keep sessions short**: Sessions should be focused on a specific task and closed as soon as possible to release resources. 1. **Handle exceptions**: Make sure to handle exceptions within sessions to prevent resource leaks. 1. **Manage history size**: For long-running conversations, use history compression to reduce token usage. 1. **Prefer high-Level abstractions**: When possible, use the node-based API. For example, `nodeLLMRequest` instead of directly working with sessions. 1. **Be mindful of thread safety**: Remember that write sessions block other sessions, so keep write operations as short as possible. 1. **Use structured requests for complex data**: When you need the LLM to return structured data, use `requestLLMStructured` instead of parsing free-form text. 1. **Use streaming for long responses**: For long responses, use `requestLLMStreaming` to process the response as it arrives. ## Troubleshooting ### Session already closed If you see an error such as `Cannot use session after it was closed`, you are trying to use a session after its lambda block has completed. Make sure all session operations are performed within the session block. ### History too large If your history becomes too large and consumes too many tokens, use history compression techniques: ``` llm.writeSession { replaceHistoryWithTLDR(HistoryCompressionStrategy.FromLastNMessages(10), preserveMemory = true) } ``` ``` // Compress recent history with the Java compression node. var compressHistory = AIAgentNode.llmCompressHistory("compressHistory"); ``` For more information, see [History compression](../history-compression/) ### Tool not found If you see errors about tools not being found, check that: - The tool is correctly registered in the tool registry. - You are using the correct tool name or class. ## API documentation For more information, see the full [AIAgentLLMSession](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.session/-a-i-agent-l-l-m-session/index.html) and [AIAgentLLMContext](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.agent.context/-a-i-agent-l-l-m-context/index.html) reference. # Subgraphs # Overview This page provides detailed information about subgraphs in the Koog framework. Understanding these concepts is crucial for creating complex agent workflows that maintain context across multiple processing steps. ## Introduction Subgraphs are a fundamental concept in the Koog framework that lets you break down complex agent workflows into manageable, sequential steps. Each subgraph represents a phase of processing, with its context, responsibilities, and an optional subset of tools. Subgraphs are integral parts of strategies, which are graphs that represent the overall agent workflow. For more information about strategies, see [Custom strategy graphs](../custom-strategy-graphs/). ## Understanding subgraphs A subgraph is a self-contained unit of processing within an agent strategy. Each subgraph: - Has a unique name - Contains a graph of nodes or subgraphs connected by edges - Can use any tool or a subset of tools from the tool registry - Receives input from the previous subgraph (or the initial user input) - Produces output that is passed to the next subgraph (or the output) To define a sequence of subgraphs in a graph, use edge connections or define sequences using the `then` keyword. For more information, see [Custom strategy graphs](../custom-strategy-graphs/). ### Subgraph context Each subgraph executes within a context that provides access to: - The environment - Agent input - The agent configuration - The LLM context (including the conversation history) - The state manager - The storage - Session and strategy The context is passed to each node within the subgraph and provides the necessary resources for the node to perform its operations. ## Creating and configuring subgraphs The following sections provide code templates and common patterns in the creation of subgraphs for agentic workflows. ### Basic subgraph creation Custom subgraphs are typically created using the following patterns: - Subgraph with a specified tool selection strategy: ``` strategy("strategy-name") { val subgraphIdentifier by subgraph( name = "subgraph-name", toolSelectionStrategy = ToolSelectionStrategy.ALL ) { // Define nodes and edges for this subgraph } nodeStart then subgraphIdentifier then nodeFinish } ``` ``` var strategyBuilder = AIAgentGraphStrategy.builder("strategy-name") .withInput(String.class) .withOutput(String.class); var subgraphIdentifier = AIAgentSubgraph.builder("subgraph-name") .withToolSelectionStrategy(ToolSelectionStrategy.ALL.INSTANCE) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Define nodes and edges for this subgraph }) .build(); var strategy = strategyBuilder .edge(strategyBuilder.nodeStart, subgraphIdentifier) .edge(subgraphIdentifier, strategyBuilder.nodeFinish) .build(); ``` - Subgraph with a specified list of tools (subset of tools from a defined tool registry): ``` strategy("strategy-name") { val subgraphIdentifier by subgraph( name = "subgraph-name", tools = listOf(firstTool, secondTool) ) { // Define nodes and edges for this subgraph } } ``` ``` var strategyBuilder = AIAgentGraphStrategy.builder("strategy-name") .withInput(String.class) .withOutput(String.class); var subgraphIdentifier = AIAgentSubgraph.builder("subgraph-name") .limitedTools(List.of(firstTool, secondTool)) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Define nodes and edges for this subgraph }) .build(); var strategy = strategyBuilder .edge(strategyBuilder.nodeStart, subgraphIdentifier) .edge(subgraphIdentifier, strategyBuilder.nodeFinish) .build(); ``` For more information about parameters and parameter values, see the `subgraph` [API-reference](https://api.koog.ai/agents/agents-core/ai.koog.agents.core.dsl.builder/-a-i-agent-subgraph-builder-base/subgraph.html). For more information about tools, see [Tools](../tools/). The following code sample shows an actual implementation of a custom subgraph: ``` strategy("my-strategy") { val mySubgraph by subgraph( tools = listOf(firstTool, secondTool) ) { // Define nodes and edges for this subgraph val sendInput by nodeLLMRequest() val executeToolCall by nodeExecuteTools() val sendToolResult by nodeLLMSendToolResults() edge(nodeStart forwardTo sendInput) edge(sendInput forwardTo executeToolCall onToolCalls { true }) edge(executeToolCall forwardTo sendToolResult) edge(sendToolResult forwardTo nodeFinish onTextMessage { true }) } } ``` ``` var strategyBuilder = AIAgentGraphStrategy.builder("my-strategy") .withInput(String.class) .withOutput(String.class); var sendInput = AIAgentNode.llmRequest(null); var executeToolCall = AIAgentNode.executeTools(null); var sendToolResult = AIAgentNode.llmSendToolResults(null); var mySubgraph = AIAgentSubgraph.builder() .limitedTools(List.of(firstTool, secondTool)) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Define nodes and edges for this subgraph subgraph .edge(AIAgentEdge.builder() .from(subgraph.nodeStart) .to(sendInput) .build() ) .edge(AIAgentEdge.builder() .from(sendInput) .to(executeToolCall) .onToolCalls() .build() ) .edge(executeToolCall, sendToolResult) .edge(AIAgentEdge.builder() .from(sendToolResult) .to(subgraph.nodeFinish) .onTextMessage() .build() ) .build(); }) .build(); var strategy = strategyBuilder .edge(strategyBuilder.nodeStart, mySubgraph) .edge(mySubgraph, strategyBuilder.nodeFinish) .build(); ``` ### Configuring tools in a subgraph Tools can be configured for a subgraph in several ways: - Directly in the subgraph definition: ``` val mySubgraph by subgraph( tools = listOf(AskUser) ) { // Subgraph definition } ``` ``` var mySubgraph = AIAgentSubgraph.builder() .limitedTools(List.of(AskUser.INSTANCE)) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Subgraph definition }) .build(); ``` - From a tool registry: ``` val mySubgraph by subgraph( tools = listOf(toolRegistry.getTool("AskUser")) ) { // Subgraph definition } ``` ``` var mySubgraph = AIAgentSubgraph.builder() .limitedTools(List.of(toolRegistry.getTool("AskUser"))) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Subgraph definition }) .build(); ``` - Dynamically during execution: ``` // Make a set of tools this.llm.writeSession { tools = tools.filter { it.name in listOf("first_tool_name", "second_tool_name") } } ``` ``` var node = AIAgentNode.builder("node_name") .withInput(String.class) .withOutput(String.class) .withAction((input, ctx) -> { // Make a set of tools ctx.getLlm().writeSession(session -> { session.setTools(session.getTools().stream() .filter(t -> List.of("first_tool_name", "second_tool_name").contains(t.getName())) .collect(Collectors.toList())); return null; }); return input; }) .build(); ``` ## Advanced subgraph techniques ### Multi-part strategies Complex workflows can be broken down into multiple subgraphs, each handling a specific part of the process: ``` strategy("complex-workflow") { val inputProcessing by subgraph( ) { // Process the initial input } val reasoning by subgraph( ) { // Perform reasoning based on the processed input } val toolRun by subgraph( // Optional subset of tools from the tool registry tools = listOf(firstTool, secondTool) ) { // Run tools based on the reasoning } val responseGeneration by subgraph( ) { // Generate a response based on the tool results } nodeStart then inputProcessing then reasoning then toolRun then responseGeneration then nodeFinish } ``` ``` var strategyBuilder = AIAgentGraphStrategy.builder("complex-workflow") .withInput(String.class) .withOutput(String.class); var inputProcessing = AIAgentSubgraph.builder() .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Process the initial input }) .build(); var reasoning = AIAgentSubgraph.builder() .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Perform reasoning based on the processed input }) .build(); var toolRun = AIAgentSubgraph.builder() // Optional subset of tools from the tool registry .limitedTools(List.of(firstTool, secondTool)) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Run tools based on the reasoning }) .build(); var responseGeneration = AIAgentSubgraph.builder() .withInput(String.class) .withOutput(String.class) .define(subgraph -> { // Generate a response based on the tool results }) .build(); var strategy = strategyBuilder .edge(strategyBuilder.nodeStart, inputProcessing) .edge(inputProcessing, reasoning) .edge(reasoning, toolRun) .edge(toolRun, responseGeneration) .edge(responseGeneration, strategyBuilder.nodeFinish) .build(); ``` ## Best practices When working with subgraphs, follow these best practices: 1. **Break complex workflows into subgraphs**: each subgraph should have a clear, focused responsibility. 1. **Pass only necessary context**: only pass the information that subsequent subgraphs need to function correctly. 1. **Document subgraph dependencies**: clearly document what each subgraph expects from previous subgraphs and what it provides to subsequent subgraphs. 1. **Test subgraphs in isolation**: ensure that each subgraph works correctly with various inputs before integrating it into a strategy. 1. **Consider token usage**: be mindful of token usage, especially when passing large histories between subgraphs. ## Troubleshooting ### Tools not available If tools are not available in a subgraph: - Check that the tools are correctly registered in the tool registry. ### Subgraphs not running in the defined and expected order If subgraphs are not executing in the defined order: - Check the strategy definition to ensure that subgraphs are listed in the correct order. - Verify that each subgraph is correctly passing its output to the next subgraph. - Ensure that your subgraph is connected with the rest of the subgraph and is reachable from the start (and finish). Be careful with conditional edges, so they cover all possible conditions to continue in order not to get blocked in a subgraph or node. ## Examples The following example shows how subgraphs are used to create an agent strategy in a real-world scenario. The code sample includes three defined subgraphs, `researchSubgraph`, `planSubgraph`, and `executeSubgraph`, where each of the subgraphs has a defined and distinct purpose within the assistant flow. ``` // Define the agent strategy val strategy = strategy("assistant") { // A subgraph that includes a tool call val researchSubgraph by subgraph( "research_subgraph", tools = listOf(WebSearchTool()) ) { val nodeCallLLM by nodeLLMRequest("call_llm") val nodeExecuteTool by nodeExecuteTools() val nodeSendToolResult by nodeLLMSendToolResults() edge(nodeStart forwardTo nodeCallLLM) edge(nodeCallLLM forwardTo nodeExecuteTool onToolCalls { true }) edge(nodeExecuteTool forwardTo nodeSendToolResult) edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCalls { true }) edge(nodeCallLLM forwardTo nodeFinish onTextMessage { true }) } val planSubgraph by subgraph( "plan_subgraph", tools = listOf() ) { val nodeUpdatePrompt by node { research -> llm.writeSession { rewritePrompt { prompt("research_prompt") { system( "You are given a problem and some research on how it can be solved." + "Make step by step a plan on how to solve given task." ) user("Research: $research") } } } } val nodeCallLLM by nodeLLMRequest("call_llm") edge(nodeStart forwardTo nodeUpdatePrompt) edge(nodeUpdatePrompt forwardTo nodeCallLLM transformed { "Task: $agentInput" }) edge(nodeCallLLM forwardTo nodeFinish onTextMessage { true }) } val executeSubgraph by subgraph( "execute_subgraph", tools = listOf(DoAction(), DoAnotherAction()), ) { val nodeUpdatePrompt by node { plan -> llm.writeSession { rewritePrompt { prompt("execute_prompt") { system( "You are given a task and detailed plan how to execute it." + "Perform execution by calling relevant tools." ) user("Execute: $plan") user("Plan: $plan") } } } } val nodeCallLLM by nodeLLMRequest("call_llm") val nodeExecuteTool by nodeExecuteTools() val nodeSendToolResult by nodeLLMSendToolResults() edge(nodeStart forwardTo nodeUpdatePrompt) edge(nodeUpdatePrompt forwardTo nodeCallLLM transformed { "Task: $agentInput" }) edge(nodeCallLLM forwardTo nodeExecuteTool onToolCalls { true }) edge(nodeExecuteTool forwardTo nodeSendToolResult) edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCalls { true }) edge(nodeCallLLM forwardTo nodeFinish onTextMessage { true }) } nodeStart then researchSubgraph then planSubgraph then executeSubgraph then nodeFinish } ``` ``` // Define the agent strategy var strategyBuilder = AIAgentGraphStrategy.builder("assistant") .withInput(String.class) .withOutput(String.class); // A subgraph that includes a tool call var nodeCallLLM = AIAgentNode.llmRequest(null); var nodeExecuteTool = AIAgentNode.executeTools(null); var nodeSendToolResult = AIAgentNode.llmSendToolResults(null); var researchSubgraph = AIAgentSubgraph.builder("research_subgraph") .limitedTools(new WebSearchToolSet()) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { subgraph .edge(AIAgentEdge.builder() .from(subgraph.nodeStart) .to(nodeCallLLM) .build() ) .edge(AIAgentEdge.builder() .from(nodeCallLLM) .to(nodeExecuteTool) .onToolCalls() .build() ) .edge(nodeExecuteTool, nodeSendToolResult) .edge(AIAgentEdge.builder() .from(nodeSendToolResult) .to(nodeExecuteTool) .onToolCalls() .build() ) .edge(AIAgentEdge.builder() .from(nodeCallLLM) .to(subgraph.nodeFinish) .onTextMessage() .build() ) .build(); }) .build(); var nodeUpdatePrompt = AIAgentNode.builder() .withInput(String.class) .withOutput(String.class) .withAction((research, ctx) -> { ctx.getLlm().writeSession(session -> { session.setPrompt(Prompt.builder("research_prompt") .system( "You are given a problem and some research on how it can be solved." + "Make step by step a plan on how to solve given task." ) .user("Research: " + research) .build()); return null; }); return "Task: " + ctx.getAgentInput(); }) .build(); var nodeCallLLMPlan = AIAgentNode.llmRequest(null); var planSubgraph = AIAgentSubgraph.builder("plan_subgraph") .limitedTools(Collections.emptyList()) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { subgraph .edge(subgraph.nodeStart, nodeUpdatePrompt) .edge(AIAgentEdge.builder() .from(nodeUpdatePrompt) .to(nodeCallLLMPlan) .build() ) .edge(AIAgentEdge.builder() .from(nodeCallLLMPlan) .to(subgraph.nodeFinish) .onTextMessage() .build() ) .build(); }) .build(); var nodeUpdatePromptExecute = AIAgentNode.builder() .withInput(String.class) .withOutput(String.class) .withAction((plan, ctx) -> { ctx.getLlm().writeSession(session -> { session.setPrompt(Prompt.builder("execute_prompt") .system( "You are given a task and detailed plan how to execute it." + "Perform execution by calling relevant tools." ) .user("Execute: " + plan) .user("Plan: " + plan) .build()); return null; }); return "Task: " + ctx.getAgentInput(); }) .build(); var nodeCallLLMExecute = AIAgentNode.llmRequest(null); var nodeExecuteToolExecute = AIAgentNode.executeTools(null); var nodeSendToolResultExecute = AIAgentNode.llmSendToolResults(null); var executeSubgraph = AIAgentSubgraph.builder("execute_subgraph") .limitedTools(new ActionToolSet()) .withInput(String.class) .withOutput(String.class) .define(subgraph -> { subgraph .edge(subgraph.nodeStart, nodeUpdatePromptExecute) .edge(AIAgentEdge.builder() .from(nodeUpdatePromptExecute) .to(nodeCallLLMExecute) .build() ) .edge(AIAgentEdge.builder() .from(nodeCallLLMExecute) .to(nodeExecuteToolExecute) .onToolCalls() .build() ) .edge(nodeExecuteToolExecute, nodeSendToolResultExecute) .edge(AIAgentEdge.builder() .from(nodeSendToolResultExecute) .to(nodeExecuteToolExecute) .onToolCalls() .build() ) .edge(AIAgentEdge.builder() .from(nodeCallLLMExecute) .to(subgraph.nodeFinish) .onIsInstance(Message.Assistant.class) .onTextMessage() .build() ) .build(); }) .build(); var strategy = strategyBuilder .edge(strategyBuilder.nodeStart, researchSubgraph) .edge(researchSubgraph, planSubgraph) .edge(planSubgraph, executeSubgraph) .edge(executeSubgraph, strategyBuilder.nodeFinish) .build(); ``` # Embeddings # Embeddings The `embeddings` module provides functionality for generating and comparing embeddings of text and code. Embeddings are vector representations that capture semantic meaning, allowing for efficient similarity comparisons. ## Overview This module consists of two main components: 1. **embeddings-base**: core interfaces and data structures for embeddings. 1. **embeddings-llm**: implementation using Ollama for local embedding generation. ## Getting started The following sections include basic examples of how to use embeddings in the following ways: - With a local embedding models through Ollama - Using an OpenAI embedding model ### Local embeddings To use the embedding functionality with a local model, you need to have Ollama installed and running on your system. For installation and running instructions, refer to the [official Ollama GitHub repository](https://github.com/ollama/ollama). ``` fun main() { runBlocking { // Create an OllamaClient instance val client = OllamaClient() // Create an embedder val embedder = LLMEmbedder(client, OllamaModels.Embeddings.NOMIC_EMBED_TEXT) // Create embeddings val embedding = embedder.embed("This is the text to embed") // Print embeddings to the output println(embedding) } } ``` To use an Ollama embedding model, make sure to have the following prerequisites: - Have [Ollama](https://ollama.com/download) installed and running - Download an embedding model to your local machine using the following command: ``` ollama pull ``` Replace `` with the Ollama identifier of the specific model. For more information about available embedding models and their identifiers, see [Ollama models overview](#ollama-models-overview). ### Ollama models overview The following table provides an overview of the available Ollama embedding models. | Model ID | Ollama ID | Parameters | Dimensions | Context Length | Performance | Tradeoffs | | ----------------- | ----------------------------- | ---------- | ---------- | -------------- | --------------------------------------------------------------------- | ------------------------------------------------------------------ | | NOMIC_EMBED_TEXT | nomic-embed-text | 137M | 768 | 8192 | High-quality embeddings for semantic search and text similarity tasks | Balanced between quality and efficiency | | ALL_MINILM | all-minilm | 33M | 384 | 512 | Fast inference with good quality for general text embeddings | Smaller model size with reduced context length, but very efficient | | MULTILINGUAL_E5 | zylonai/multilingual-e5-large | 300M | 768 | 512 | Strong performance across 100+ languages | Larger model size but provides excellent multilingual capabilities | | BGE_LARGE | bge-large | 335M | 1024 | 512 | Excellent for English text retrieval and semantic search | Larger model size but provides high-quality embeddings | | MXBAI_EMBED_LARGE | mxbai-embed-large | - | - | - | High-dimensional embeddings of textual data | Designed for creating high-dimensional embeddings | For more information about these models, see Ollama's [Embedding Models](https://ollama.com/blog/embedding-models) blog post. ### Choosing a model Here are some general tips on which Ollama embedding model to select depending on your requirements: - For general text embeddings, use `NOMIC_EMBED_TEXT`. - For multilingual support, use `MULTILINGUAL_E5`. - For maximum quality (at the cost of performance), use `BGE_LARGE`. - For maximum efficiency (at the cost of some quality), use `ALL_MINILM`. - For high-dimensional embeddings, use `MXBAI_EMBED_LARGE`. ## OpenAI embeddings To create embeddings using an OpenAI embedding model, use the `embed` method of an `OpenAILLMClient` instance as shown in the example below. ``` suspend fun openAIEmbed(text: String) { // Get the OpenAI API token from the OPENAI_KEY environment variable val token = System.getenv("OPENAI_KEY") ?: error("Environment variable OPENAI_KEY is not set") // Create an OpenAILLMClient instance val client = OpenAILLMClient(token) // Create an embedder val embedder = LLMEmbedder(client, OpenAIModels.Embeddings.TextEmbeddingAda002) // Create embeddings val embedding = embedder.embed(text) // Print embeddings to the output println(embedding) } ``` ## AWS Bedrock embeddings To create embeddings using an AWS Bedrock embedding model, use the `embed` method of an `BedrockLLMClient` instance and your chosen model. Example: ``` suspend fun bedrockEmbed(text: String) { // Get AWS credentials from environment/configuration val awsAccessKeyId = System.getenv("AWS_ACCESS_KEY_ID") ?: error("AWS_ACCESS_KEY_ID not set") val awsSecretAccessKey = System.getenv("AWS_SECRET_ACCESS_KEY") ?: error("AWS_SECRET_ACCESS_KEY not set") // (Optional) AWS_SESSION_TOKEN for temporary credentials val awsSessionToken = System.getenv("AWS_SESSION_TOKEN") // Create a BedrockLLMClient instance val client = BedrockLLMClient( identityProvider = StaticCredentialsProvider { this.accessKeyId = awsAccessKeyId this.secretAccessKey = awsSecretAccessKey awsSessionToken?.let { this.sessionToken = it } }, settings = BedrockClientSettings() ) // Create an embedder val embedder = LLMEmbedder(client, BedrockModels.Embeddings.AmazonTitanEmbedText) // Create embeddings val embedding = embedder.embed(text) // Print embeddings to the output println(embedding) } ``` ### Supported AWS Bedrock embedding models | Provider | Model name | Model ID | Input | Output | Dimensions | Context Length | Notes | | -------- | ---------------------------- | ------------------------------ | ----- | --------- | ---------- | -------------- | ------------------------------------------------------------------------------------------------------ | | Amazon | Titan Embeddings G1 - Text | `amazon.titan-embed-text-v1` | Text | Embedding | 1,536 | 8192 | 25+ languages, optimized for retrieval, semantic similarity, clustering; segment long docs for search. | | Amazon | Titan Text Embeddings V2 | `amazon.titan-embed-text-v2:0` | Text | Embedding | 1,024 | 8192 | High-accuracy, flexible dimensions, multilingual (100+); smaller dims save storage, normalized output. | | Cohere | Cohere Embed English v3 | `cohere.embed-english-v3` | Text | Embedding | 1,024 | 8192 | SOTA English text embeddings for search, retrieval, and understanding text nuances. | | Cohere | Cohere Embed Multilingual v3 | `cohere.embed-multilingual-v3` | Text | Embedding | 1,024 | 8192 | Multilingual embeddings, SOTA for search and semantic understanding across languages. | > For the most up-to-date model support, refer to the [AWS Bedrock supported models documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html). ## Examples The following examples show how you can use embeddings to compare code with text or other code snippets. ### Code-to-text comparison Compare code snippets with natural language descriptions to find semantic matches: ``` suspend fun compareCodeToText(embedder: Embedder) { // Embedder type // Code snippet val code = """ fun factorial(n: Int): Int { return if (n <= 1) 1 else n * factorial(n - 1) } """.trimIndent() // Text descriptions val description1 = "A recursive function that calculates the factorial of a number" val description2 = "A function that sorts an array of integers" // Generate embeddings val codeEmbedding = embedder.embed(code) val desc1Embedding = embedder.embed(description1) val desc2Embedding = embedder.embed(description2) // Calculate differences (lower value means more similar) val diff1 = embedder.diff(codeEmbedding, desc1Embedding) val diff2 = embedder.diff(codeEmbedding, desc2Embedding) println("Difference between code and description 1: $diff1") println("Difference between code and description 2: $diff2") // The code should be more similar to description1 than description2 if (diff1 < diff2) { println("The code is more similar to: '$description1'") } else { println("The code is more similar to: '$description2'") } } ``` ### Code-to-code comparison Compare code snippets to find semantic similarities regardless of syntax differences: ``` suspend fun compareCodeToCode(embedder: Embedder) { // Embedder type // Two implementations of the same algorithm in different languages val kotlinCode = """ fun fibonacci(n: Int): Int { return if (n <= 1) n else fibonacci(n - 1) + fibonacci(n - 2) } """.trimIndent() val pythonCode = """ def fibonacci(n): if n <= 1: return n else: return fibonacci(n-1) + fibonacci(n-2) """.trimIndent() val javaCode = """ public static int bubbleSort(int[] arr) { int n = arr.length; for (int i = 0; i < n-1; i++) { for (int j = 0; j < n-i-1; j++) { if (arr[j] > arr[j+1]) { int temp = arr[j]; arr[j] = arr[j+1]; arr[j+1] = temp; } } } return arr; } """.trimIndent() // Generate embeddings val kotlinEmbedding = embedder.embed(kotlinCode) val pythonEmbedding = embedder.embed(pythonCode) val javaEmbedding = embedder.embed(javaCode) // Calculate differences val diffKotlinPython = embedder.diff(kotlinEmbedding, pythonEmbedding) val diffKotlinJava = embedder.diff(kotlinEmbedding, javaEmbedding) println("Difference between Kotlin and Python implementations: $diffKotlinPython") println("Difference between Kotlin and Java implementations: $diffKotlinJava") // The Kotlin and Python implementations should be more similar if (diffKotlinPython < diffKotlinJava) { println("The Kotlin code is more similar to the Python code") } else { println("The Kotlin code is more similar to the Java code") } } ``` ## API documentation For a complete API reference related to embeddings, see the reference documentation for the following modules: - [embeddings-base](https://api.koog.ai/embeddings/embeddings-base/ai.koog.embeddings.base/index.html): Provides core interfaces and data structures for representing and comparing text and code embeddings. - [embeddings-llm](https://api.koog.ai/embeddings/embeddings-llm/index.html): Includes implementations for working with local embedding models. # RAG # Retrieval-augmented generation (RAG) Beta This feature is part of a beta module (`1.2.0-beta`). The API may change in future releases. See [module versioning](../module-versioning/) for details. Koog provides building blocks for retrieval-augmented generation (RAG): embedding text, storing embedded documents, and retrieving the most relevant results for a query. This page focuses on what is available in the current `rag` module and how to use it. ## What Koog provides today The current RAG support is split into two modules: - `rag-base`: common abstractions for retrieval, storage, search requests, filtering, and file/document providers - `rag-vector`: local implementations that combine document embedding with vector storage ## Embedding and retrieving documents with EmbeddingStorage The most complete out-of-the-box RAG flow uses `EmbeddingStorage` from the `rag-vector` module. It combines a `DocumentEmbedder` (which converts documents to vectors) with a `VectorStorageBackend` (which persists the vectors). The steps are: 1. Create an `Embedder` backed by an embedding model (Ollama or OpenAI). 1. Create a `JVMTextDocumentEmbedder` that reads file content and delegates to the embedder. 1. Create an `EmbeddingStorage` with an in-memory or file-based backend. 1. Add documents with `add()`. 1. Search with `search(SimilaritySearchRequest(...))`. ``` // 1. Create an embedder backed by a local Ollama model val embedder = LLMEmbedder( client = OllamaClient(), model = OllamaModels.Embeddings.NOMIC_EMBED_TEXT ) // 2. Create a JVM document embedder that reads files and embeds their text val documentEmbedder = JVMTextDocumentEmbedder(embedder) // 3. Create an EmbeddingStorage with an in-memory backend val storage = EmbeddingStorage( embedder = documentEmbedder, storage = InMemoryVectorStorageBackend() ) // 4. Add documents to the storage storage.add( listOf( Path.of("./docs/faq.txt"), Path.of("./docs/pricing.txt"), Path.of("./docs/getting-started.txt") ) ) // 5. Search for the most relevant documents val results = storage.search( SimilaritySearchRequest( queryText = "How do I reset my password?", limit = 3, minScore = 0.5 ) ) results.forEach { result -> println("${result.document} (score: ${result.score.value})") } ``` ``` ``` ## Providing relevance search as an agent tool (in agentic RAG) Instead of injecting all retrieved documents into the prompt upfront, you can expose the RAG storage as a tool that the agent calls on demand. This gives the agent control over when and what to search for. The example below wraps a `SearchStorage` (the base search interface that `EmbeddingStorage` implements) in a function annotated with `@Tool` and `@LLMDescription`, then registers it in a `ToolRegistry` for the agent to use. ``` // Define a tool that searches the RAG storage @Tool @LLMDescription("Search the knowledge base for documents relevant to a query. Returns the content of the most relevant documents.") suspend fun searchKnowledgeBase( @LLMDescription("The search query describing what information you need") query: String, @LLMDescription("Maximum number of documents to return") count: Int ): String { val results = ragStorage.search( SimilaritySearchRequest( queryText = query, limit = count, minScore = 0.5 ) ) if (results.isEmpty()) { return "No relevant documents found for: $query" } val response = StringBuilder("Found ${results.size} relevant documents:\n\n") results.forEachIndexed { index, result -> val content = Files.readString(result.document) response.append("Document ${index + 1}: ${result.document.fileName}") response.append(" (score: ${"%.2f".format(result.score.value)})\n") response.append("Content: $content\n\n") } return response.toString() } fun main() { runBlocking { // Register the search tool and create an agent val tools = ToolRegistry { tool(::searchKnowledgeBase.asTool()) } val agent = AIAgent( toolRegistry = tools, promptExecutor = simpleOpenAIExecutor(apiKey), llmModel = OpenAIModels.Chat.GPT4o ) val response = agent.run("What is your refund policy?") println("Agent response: $response") } } ``` ``` ``` With this approach, the agent decides when to call the search tool based on the user's query. This is useful when the agent handles diverse requests and only some of them require knowledge base lookups. ## Available implementations ### Vector storage backends - `InMemoryVectorStorageBackend`: stores vectors in memory; suitable for testing and prototypes - `FileVectorStorageBackend`: persists vectors to disk for durability across restarts - `JVMFileVectorStorageBackend`: JVM-specific file-based backend using `java.nio.file.Path` ### Document embedders - `TextDocumentEmbedder`: generic document-to-text embedder parameterized by document and path types - `JVMTextDocumentEmbedder`: JVM-specific embedder that reads files from `java.nio.file.Path` ### Combined storage implementations - `EmbeddingStorage`: composes any `DocumentEmbedder` with any `VectorStorageBackend` - `InMemoryDocumentEmbeddingStorage`: convenience shortcut for `EmbeddingStorage` + `InMemoryVectorStorageBackend` - `FileDocumentEmbeddingStorage`: convenience shortcut for `EmbeddingStorage` + `FileVectorStorageBackend` - `JVMFileDocumentEmbeddingStorage`: JVM file-based embedding storage - `TextFileDocumentEmbeddingStorage`: file-based storage for text documents - `JVMFileEmbeddingStorage`: JVM file-based storage for text documents ## Current limitations The built-in flow is useful for local and reference implementations, but it is not yet a full production RAG platform. Important limitations: - the built-in implementations support similarity search only - there is no built-in chunking pipeline in the `rag` module - metadata-rich production record modeling is still limited - production vector database integrations (Pinecone, Weaviate, pgvector, Milvus) are not provided in the current `rag` module If you are building a custom backend, start from `rag-base` abstractions and implement your own storage adapter. ## Choosing where to start Use `rag-vector` if: - you want a local RAG prototype - you want a simple reference implementation - you want to experiment with embedding and retrieval flow inside Koog Use `rag-base` if: - you are building your own storage backend - you want to integrate an external vector database - you want to reuse the abstractions in another Koog module ## See also - [Embeddings](../embeddings/) # Serialization ## Introduction Koog uses a thin, library-agnostic serialization layer to convert tool arguments and results to and from JSON. This layer sits between the agent runtime and the underlying serialization library, so you can swap the library without changing any tool or agent code. Beyond tools, the serialization layer is also used by agent features such as **Persistence** to serialize and deserialize node inputs and outputs. By default, Koog uses `KotlinxSerializer` (backed by kotlinx-serialization). On the JVM you can also switch to `JacksonSerializer` (backed by jackson-databind). ## The `JSONSerializer` interface `JSONSerializer` is the core abstraction that lives in `serialization-core`. The interface has four primary methods (encode/decode to both strings and `JSONElement`), plus two convenience methods for converting between `JSONElement` and strings: - `encodeToString` / `decodeFromString` — serialize a typed value to/from a JSON string. - `encodeToJSONElement` / `decodeFromJSONElement` — serialize a typed value to/from a `JSONElement` tree. - `encodeJSONElementToString` / `decodeJSONElementFromString` — convert between `JSONElement` and its string form. The following example shows all key operations: ``` @Serializable data class User(val name: String, val age: Int) val serializer: JSONSerializer = KotlinxSerializer() // Encode a data class to a JSON string val json: String = serializer.encodeToString(User("Alice", 30), typeToken()) // Decode a JSON string back to a data class val user: User = serializer.decodeFromString(json, typeToken()) // Encode to a JSONElement tree val element: JSONElement = serializer.encodeToJSONElement(user, typeToken()) // Decode from a JSONElement tree val userFromElement: User = serializer.decodeFromJSONElement(element, typeToken()) // Convert between JSONElement and a raw JSON string val jsonString = """{"key": "value"}""" val jsonElement: JSONElement = serializer.decodeJSONElementFromString(jsonString) val backToString: String = serializer.encodeJSONElementToString(jsonElement) ``` ``` // Jackson-serializable class record User( @JsonProperty("name") String name, @JsonProperty("age") int age ) {} var serializer = new JacksonSerializer(); // Encode a data class to a JSON string String json = serializer.encodeToString(new User("Alice", 30), TypeToken.of(User.class)); // Decode a JSON string back to a data class User user = serializer.decodeFromString(json, TypeToken.of(User.class)); // Encode to a JSONElement tree JSONElement element = serializer.encodeToJSONElement(user, TypeToken.of(User.class)); // Decode from a JSONElement tree User userFromElement = serializer.decodeFromJSONElement(element, TypeToken.of(User.class)); // Convert between JSONElement and a raw JSON string String jsonString = "{\"key\": \"value\"}"; JSONElement jsonElement = serializer.decodeJSONElementFromString(jsonString); String backToString = serializer.encodeJSONElementToString(jsonElement); ``` ## Type tokens `TypeToken` is how Koog passes type information at runtime. ``` data class MyClass(val value: String) // Inline reified — preferred in Kotlin val tokenReified = typeToken() // From a KClass (when no reified type parameter is available) val tokenKClass = typeToken(MyClass::class) // Generic type — preserves type arguments at runtime val tokenGeneric = typeToken>() ``` ``` record MyClass( String value ) {} // Simple class TypeToken tokenClass = TypeToken.of(MyClass.class); // Generic type — use TypeCapture to preserve type arguments TypeToken tokenGeneric = TypeToken.of(new TypeCapture>() {}); ``` ## `JSONElement` — library-agnostic JSON tree `JSONElement` is a neutral intermediate representation for JSON data. It exists so that serializers, tools, and agent internals do not depend on specific JSON types from a particular library. ### Hierarchy ``` JSONElement ├── JSONObject – key-value pairs (entries: Map) ├── JSONArray – ordered list (elements: List) └── JSONPrimitive ├── JSONLiteral – string, number, or boolean └── JSONNull – JSON null singleton ``` ### Conversion to and from library types Each serialization integration provides extension functions that let you convert between `JSONElement` and the library's own dynamic JSON type. This is useful when you already have a `JsonElement`, `JsonNode`, etc. and want to pass it to Koog (or vice versa), without going through a full encode/decode cycle. Examples are provided below for each supported library. ### Building and reading elements ``` val obj = JSONObject( mapOf( "name" to JSONPrimitive("Alice"), "age" to JSONPrimitive(30), "active" to JSONPrimitive(true), ) ) val arr = JSONArray(listOf(JSONPrimitive(1), JSONPrimitive(2), JSONPrimitive(3))) // Reading values from an object val nameContent: String = (obj.entries["name"] as JSONPrimitive).content // "Alice" val age: Int? = (obj.entries["age"] as JSONPrimitive).intOrNull // 30 ``` ``` JSONObject obj = new JSONObject( Map.of( "name", JSONPrimitive.of("Alice"), "age", JSONPrimitive.of(30), "active", JSONPrimitive.of(true) ) ); JSONArray arr = new JSONArray(List.of(JSONPrimitive.of(1), JSONPrimitive.of(2), JSONPrimitive.of(3))); // Reading values from an object String nameContent = ((JSONPrimitive) obj.getEntries().get("name")).getContent(); // "Alice" Integer age = ((JSONPrimitive) obj.getEntries().get("age")).getIntOrNull(); // 30 ``` ## Supported serializers ### `KotlinxSerializer` (default) - **Module**: `ai.koog:serialization-core` (included transitively with `ai.koog:agents-core`) - **Backed by**: kotlinx-serialization ``` // Default instance — uses Json.Default val defaultSerializer = KotlinxSerializer() // Custom Json configuration val customSerializer = KotlinxSerializer( json = Json { ignoreUnknownKeys = true prettyPrint = true } ) ``` You can also convert between Koog's `JSONElement` and kotlinx-serialization's `JsonElement` ``` val koogJson: JSONElement = JSONObject( mapOf( "key" to JSONPrimitive("value") ) ) // Convert to kotlinx-serialization dynamic JSON instance val kotlinxJson: JsonElement = koogJson.toKotlinxJsonElement() // Convert to Koog dynamic JSON instance val koogJsonConverted: JSONElement = kotlinxJson.toKoogJSONElement() ``` ### `JacksonSerializer` (JVM only) - **Module**: `ai.koog:serialization-jackson` (separate dependency) - **Backed by**: jackson-databind Add the dependency to your `build.gradle.kts`: ``` dependencies { implementation("ai.koog:serialization-jackson:") } ``` Then create the serializer: ``` // Default instance — uses a fresh ObjectMapper with JSONElementModule pre-registered val defaultSerializer = JacksonSerializer() // Custom ObjectMapper configuration val customSerializer = JacksonSerializer( objectMapper = ObjectMapper().apply { configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false) } ) ``` ``` // Default instance — uses a fresh ObjectMapper with JSONElementModule pre-registered var defaultSerializer = new JacksonSerializer(); // Custom ObjectMapper configuration ObjectMapper objectMapper = new ObjectMapper(); objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false); var customSerializer = new JacksonSerializer(objectMapper); ``` Note `JacksonSerializer` automatically registers `JSONElementModule` on the `ObjectMapper` it uses for proper serialization/deserialization of the `JSONElement` types. You can also convert between Koog's `JSONElement` and Jackson's `JsonNode` ``` val koogJson: JSONElement = JSONObject( mapOf( "key" to JSONPrimitive("value") ) ) // Convert to Jackson dynamic JSON instance val jacksonJson: JsonNode = koogJson.toJacksonJsonNode() // Convert to Koog dynamic JSON instance val koogJsonConverted: JSONElement = jacksonJson.toKoogJSONElement() ``` ``` JSONElement koogJson = new JSONObject( Map.of( "key", JSONPrimitive.of("value") ) ); // Convert to Jackson dynamic JSON instance JsonNode jacksonJson = JacksonJSONElementMappers.toJacksonJsonNode(koogJson); // Convert to Koog dynamic JSON instance JSONElement koogJsonConverted = JacksonJSONElementMappers.toKoogJSONElement(jacksonJson); ``` ## Configuring the serializer in `AIAgentConfig` Pass the `serializer` parameter when constructing `AIAgentConfig`. If you omit it, `KotlinxSerializer` is used. ``` val agentConfig = AIAgentConfig( prompt = prompt("assistant") { system("You are a helpful assistant.") }, model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 10, serializer = JacksonSerializer() ) ``` Pass the `serializer` parameter when constructing `AIAgentConfig`. If you omit it, `JacksonSerializer` is used. ``` var agentConfig = AIAgentConfig.builder() .model(OpenAIModels.Chat.GPT4o) .prompt( Prompt.builder("assistant") .system("You are a helpful assistant") .build() ) .maxAgentIterations(10) .serializer(new JacksonSerializer()) .build(); ``` ## How tools interact with the serializer The agent runtime calls the following methods on each `Tool` instance automatically. You do not need to invoke them yourself in normal usage. - **`decodeArgs(rawArgs, serializer)`** (JSON → TArgs) — deserializes raw JSON arguments from the LLM into the tool's typed args class. - **`encodeArgs(args, serializer)`** (TArgs → JSON) — serializes typed args back to JSON (used by certain agent features). - **`decodeResult(rawResult, serializer)`** (JSON → TResult) — deserializes a stored JSON result. - **`encodeResult(result, serializer)`** (TResult → JSON) — serializes the tool's result to JSON. - **`encodeResultToString(result, serializer)`** (TResult → String) — serializes the tool's result to a string sent to the LLM. By default, delegates to `encodeResult`. Can be overridden to customize the result format for the LLM. These methods are `open` on `Tool`, so you can override them if you need custom serialization behavior for a specific tool. ## How features use the serializer The serialization layer is not limited to tools — certain agent features rely on it as well. For example, **Persistence** uses `JSONSerializer` configured in `AIAgentConfig` to serialize and deserialize node inputs and outputs when creating checkpoints and restoring agent state. This means any type that flows through a persisted node must be serializable by the configured `JSONSerializer`. See [Agent Persistence](../features/agent-persistence/) for details on checkpoint creation and restoration. # Testing # Testing ## Overview The Testing feature provides a comprehensive framework for testing AI agent pipelines, subgraphs, and tool interactions in the Koog framework. It enables developers to create controlled test environments with mock LLM (Large Language Model) executors, tool registries, and agent environments. ### Purpose The primary purpose of this feature is to facilitate testing of agent-based AI features by: - Mocking LLM responses to specific prompts - Simulating tool calls and their results - Testing agent pipeline subgraphs and their structures - Verifying the correct flow of data through agent nodes - Providing assertions for expected behaviors ## Configuration and initialization ### Setting up test dependencies Before setting up a test environment, make sure that you have added the following dependencies: ``` // build.gradle.kts dependencies { testImplementation("ai.koog:agents-test:LATEST_VERSION") testImplementation(kotlin("test")) } ``` ### Mocking LLM responses The basic form of testing involves mocking LLM responses to ensure deterministic behavior. You can do this using `MockLLMBuilder` and related utilities. ``` // Create a mock LLM executor val mockLLMApi = getMockExecutor { // Mock a simple text response mockLLMAnswer("Hello!") onRequestContains "Hello" // Mock a default response mockLLMAnswer("I don't know how to answer that.").asDefaultResponse } ``` ``` import ai.koog.agents.core.tools.ToolRegistry; import ai.koog.agents.testing.tools.MockExecutor; import ai.koog.prompt.executor.model.PromptExecutor; // Create a tool registry (empty) ToolRegistry toolRegistry = ToolRegistry.builder().build(); // Create a mock LLM executor PromptExecutor mockLLMApi = MockExecutor.builder() .toolRegistry(toolRegistry) .mockLLMAnswer("Hello!").onRequestContains("Hello") .mockLLMAnswer("I don't know how to answer that.").asDefaultResponse() .build(); ``` ### Mocking tool calls You can mock the LLM to call specific tools based on input patterns: ``` // Mock a tool call response mockLLMToolCall(CreateTool, CreateTool.Args("solve")) onRequestEquals "Solve task" // Mock tool behavior - simplest form without lambda mockTool(PositiveToneTool) alwaysReturns "The text has a positive tone." // Using lambda when you need to perform extra actions mockTool(NegativeToneTool) alwaysTells { // Perform some extra action println("Negative tone tool called") // Return the result "The text has a negative tone." } // Mock tool behavior based on specific arguments mockTool(AnalyzeTool) returns "Detailed analysis" onArguments AnalyzeTool.Args("analyze deeply") // Mock tool behavior with conditional argument matching mockTool(SearchTool) returns "Found results" onArgumentsMatching { args -> args.query.contains("important") } ``` ``` ``` The examples above demonstrate different ways to mock tools, from simple to more complex ones: 1. `alwaysReturns`: the simplest form, directly returns a value without a lambda. 1. `alwaysTells`: uses a lambda when you need to perform additional actions. 1. `returns...onArguments`: returns specific results for exact argument matches. 1. `returns...onArgumentsMatching`: returns results based on custom argument conditions. ### Enabling testing mode To enable the testing mode on an agent, use the `withTesting()` function within the AIAgent constructor block: ``` // Create the agent with testing enabled AIAgent( promptExecutor = mockLLMApi, toolRegistry = toolRegistry, llmModel = llmModel ) { // Enable testing mode withTesting() } ``` ``` ``` ## Advanced testing ### Testing the graph structure Before testing the detailed node behavior and edge connections, it is important to verify the overall structure of your agent's graph. This includes checking that all required nodes exist and are properly connected in the expected subgraphs. The Testing feature provides a comprehensive way to test your agent's graph structure. This approach is particularly valuable for complex agents with multiple subgraphs and interconnected nodes. #### Basic structure testing Start by validating the fundamental structure of your agent's graph: ``` AIAgent( // Constructor arguments promptExecutor = mockLLMApi, toolRegistry = toolRegistry, llmModel = llmModel ) { testGraph("test") { val firstSubgraph = assertSubgraphByName("first") val secondSubgraph = assertSubgraphByName("second") // Assert subgraph connections assertEdges { startNode() alwaysGoesTo firstSubgraph firstSubgraph alwaysGoesTo secondSubgraph secondSubgraph alwaysGoesTo finishNode() } // Verify the first subgraph verifySubgraph(firstSubgraph) { val start = startNode() val finish = finishNode() // Assert nodes by name val askLLM = assertNodeByName("callLLM") val callTool = assertNodeByName("executeTool") // Assert node reachability assertReachable(start, askLLM) assertReachable(askLLM, callTool) } } } ``` ``` ``` ### Testing node behavior Node behavior testing lets you verify that nodes in your agent's graph produce the expected outputs for the given inputs. This is crucial for ensuring that your agent's logic works correctly under different scenarios. #### Basic node testing Start with simple input and output validations for individual nodes: ``` assertNodes { // Test basic text responses askLLM withInput "Hello" outputs assistantMessage("Hello!") // Test tool call responses askLLM withInput "Solve task" outputs assistantMessage(CreateTool, CreateTool.Args("solve")) } ``` ``` ``` The example above shows how to test the following behavior: 1. When the LLM node receives `Hello` as the input, it responds with a simple text message. 1. When it receives `Solve task`, it responds with a tool call. #### Testing tool run nodes You can also test nodes that run tools: ``` assertNodes { // Test tool runs with specific arguments callTool withInput ToolCalls(listOf(toolCallMessagePart( SolveTool, SolveTool.Args("solve") ))) outputs ReceivedToolResults(listOf(toolResult(SolveTool, SolveTool.Args("solve"), "solved"))) } ``` ``` ``` This verifies that when the tool execution node receives a specific tool call signature, it produces the expected tool result. #### Advanced node testing For more complex scenarios, you can test nodes with structured inputs and outputs: ``` assertNodes { // Test with different inputs to the same node askLLM withInput "Simple query" outputs assistantMessage("Simple response") // Test with complex parameters askLLM withInput "Complex query with parameters" outputs assistantMessage( AnalyzeTool, AnalyzeTool.Args(query = "parameters", depth = 3) ) } ``` ``` ``` You can also test complex tool call scenarios with detailed result structures: ``` assertNodes { // Test a complex tool call with a structured result callTool withInput ToolCalls(listOf(toolCallMessagePart( AnalyzeTool, AnalyzeTool.Args(query = "complex", depth = 5) ))) outputs ReceivedToolResults(listOf(toolResult(AnalyzeTool, AnalyzeTool.Args(query = "complex", depth = 5), AnalyzeTool.Result( analysis = "Detailed analysis", confidence = 0.95, metadata = mapOf("source" to "database", "timestamp" to "2023-06-15") )))) } ``` ``` ``` These advanced tests help ensure that your nodes handle complex data structures correctly, which is essential for sophisticated agent behaviors. ### Testing edge connections Edge connections testing allows you to verify that your agent's graph correctly routes outputs from one node to the appropriate next node. This ensures that your agent follows the intended workflow paths based on different outputs. #### Basic edge testing Start with simple edge connection tests: ``` assertEdges { // Test text message routing askLLM withOutput assistantMessage("Hello!") goesTo giveFeedback // Test tool call routing askLLM withOutput assistantMessage(CreateTool, CreateTool.Args("solve")) goesTo callTool } ``` ``` ``` This example verifies the following behavior: 1. When the LLM node outputs a simple text message, the flow is directed to the `giveFeedback` node. 1. When it outputs a tool call, the flow is directed to the `callTool` node. #### Testing conditional routing You can test a more complex routing logic based on the content of outputs: ``` assertEdges { // Different text responses can route to different nodes askLLM withOutput assistantMessage("Need more information") goesTo askForInfo askLLM withOutput assistantMessage("Ready to proceed") goesTo processRequest } ``` ``` ``` #### Advanced edge testing For sophisticated agents, you can test conditional routing based on structured data in tool results: ``` assertEdges { // Test routing based on tool result content callTool withOutput ReceivedToolResults(listOf(toolResult( AnalyzeTool, AnalyzeTool.Args(query = "parameters", depth = 3), AnalyzeTool.Result(analysis = "Needs more processing", confidence = 0.5) ))) goesTo processResult } ``` ``` ``` You can also test complex decision paths based on different result properties: ``` assertEdges { // Route to different nodes based on confidence level callTool withOutput ReceivedToolResults(listOf(toolResult( AnalyzeTool, AnalyzeTool.Args(query = "parameters", depth = 3), AnalyzeTool.Result(analysis = "Complete", confidence = 0.9) ))) goesTo finish callTool withOutput ReceivedToolResults(listOf(toolResult( AnalyzeTool, AnalyzeTool.Args(query = "parameters", depth = 3), AnalyzeTool.Result(analysis = "Uncertain", confidence = 0.3) ))) goesTo verifyResult } ``` ``` ``` These advanced edge tests help ensure that your agent makes the correct decisions based on the content and structure of node outputs, which is essential for creating intelligent, context-aware workflows. ## Complete testing example Here is a user story that demonstrates a complete testing scenario: You are developing a tone analysis agent that analyzes the tone of the text and provides feedback. The agent uses tools for detecting positive, negative, and neutral tones. Here is how you can test this agent: ``` @Test fun testToneAgent() = runTest { // Create a list to track tool calls var toolCalls = mutableListOf() var result: String? = null // Create a tool registry val toolRegistry = ToolRegistry { // A special tool, required with this type of agent tool(SayToUser) with(ToneTools) { tools() } } // Create an event handler val eventHandler = EventHandler { onToolCallStarting { tool, args -> println("[DEBUG_LOG] Tool called: tool ${tool.name}, args $args") toolCalls.add(tool.name) } handleError { println("[DEBUG_LOG] An error occurred: ${it.message}\n${it.stackTraceToString()}") true } handleResult { println("[DEBUG_LOG] Result: $it") result = it } } val positiveText = "I love this product!" val negativeText = "Awful service, hate the app." val defaultText = "I don't know how to answer this question." val positiveResponse = "The text has a positive tone." val negativeResponse = "The text has a negative tone." val neutralResponse = "The text has a neutral tone." val mockLLMApi = getMockExecutor(toolRegistry, eventHandler) { // Set up LLM responses for different input texts mockLLMToolCall(NeutralToneTool, ToneTool.Args(defaultText)) onRequestEquals defaultText mockLLMToolCall(PositiveToneTool, ToneTool.Args(positiveText)) onRequestEquals positiveText mockLLMToolCall(NegativeToneTool, ToneTool.Args(negativeText)) onRequestEquals negativeText // Mock the behavior where the LLM responds with just tool responses when the tools return results mockLLMAnswer(positiveResponse) onRequestContains positiveResponse mockLLMAnswer(negativeResponse) onRequestContains negativeResponse mockLLMAnswer(neutralResponse) onRequestContains neutralResponse mockLLMAnswer(defaultText).asDefaultResponse // Tool mocks mockTool(PositiveToneTool) alwaysTells { toolCalls += "Positive tone tool called" positiveResponse } mockTool(NegativeToneTool) alwaysTells { toolCalls += "Negative tone tool called" negativeResponse } mockTool(NeutralToneTool) alwaysTells { toolCalls += "Neutral tone tool called" neutralResponse } } // Create a strategy val strategy = toneStrategy("tone_analysis") // Create an agent configuration val agentConfig = AIAgentConfig( prompt = prompt("test-agent") { system( """ You are an question answering agent with access to the tone analysis tools. You need to answer 1 question with the best of your ability. Be as concise as possible in your answers. DO NOT ANSWER ANY QUESTIONS THAT ARE BESIDES PERFORMING TONE ANALYSIS! DO NOT HALLUCINATE! """.trimIndent() ) }, model = mockk(relaxed = true), maxAgentIterations = 10 ) // Create an agent with testing enabled val agent = AIAgent( promptExecutor = mockLLMApi, toolRegistry = toolRegistry, strategy = strategy, eventHandler = eventHandler, agentConfig = agentConfig, ) { withTesting() } // Test the positive text agent.run(positiveText) assertEquals("The text has a positive tone.", result, "Positive tone result should match") assertEquals(1, toolCalls.size, "One tool is expected to be called") // Test the negative text agent.run(negativeText) assertEquals("The text has a negative tone.", result, "Negative tone result should match") assertEquals(2, toolCalls.size, "Two tools are expected to be called") //Test the neutral text agent.run(defaultText) assertEquals("The text has a neutral tone.", result, "Neutral tone result should match") assertEquals(3, toolCalls.size, "Three tools are expected to be called") } ``` ``` ``` For more complex agents with multiple subgraphs, you can also test the graph structure: ``` @Test fun testMultiSubgraphAgentStructure() = runTest { val strategy = strategy("test") { val firstSubgraph by subgraph( "first", tools = listOf(DummyTool, CreateTool, SolveTool) ) { val callLLM by nodeLLMRequest(allowToolCalls = false) val executeTool by nodeExecuteTools() val sendToolResult by nodeLLMSendToolResults() val giveFeedback by node { input -> llm.writeSession { appendPrompt { user("Call tools! Don't chat!") } } input } edge(nodeStart forwardTo callLLM) edge(callLLM forwardTo executeTool onToolCalls { true }) edge(callLLM forwardTo giveFeedback onTextMessage { true }) edge(giveFeedback forwardTo giveFeedback transformed { it }) edge(executeTool forwardTo nodeFinish transformed { it.toolResults.first().output }) } val secondSubgraph by subgraph("second") { edge(nodeStart forwardTo nodeFinish) } edge(nodeStart forwardTo firstSubgraph) edge(firstSubgraph forwardTo secondSubgraph) edge(secondSubgraph forwardTo nodeFinish) } val toolRegistry = ToolRegistry { tool(DummyTool) tool(CreateTool) tool(SolveTool) } val mockLLMApi = getMockExecutor(toolRegistry) { mockLLMAnswer("Hello!") onRequestContains "Hello" mockLLMToolCall(CreateTool, CreateTool.Args("solve")) onRequestEquals "Solve task" } val basePrompt = prompt("test") {} AIAgent( toolRegistry = toolRegistry, strategy = strategy, eventHandler = EventHandler {}, agentConfig = AIAgentConfig(prompt = basePrompt, model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 100), promptExecutor = mockLLMApi, ) { testGraph("test") { val firstSubgraph = assertSubgraphByName("first") val secondSubgraph = assertSubgraphByName("second") assertEdges { startNode() alwaysGoesTo firstSubgraph firstSubgraph alwaysGoesTo secondSubgraph secondSubgraph alwaysGoesTo finishNode() } verifySubgraph(firstSubgraph) { val start = startNode() val finish = finishNode() val askLLM = assertNodeByName("callLLM") val callTool = assertNodeByName("executeTool") val giveFeedback = assertNodeByName("giveFeedback") assertReachable(start, askLLM) assertReachable(askLLM, callTool) assertNodes { askLLM withInput "Hello" outputs assistantMessage("Hello!") askLLM withInput "Solve task" outputs assistantMessage(CreateTool, CreateTool.Args("solve")) callTool withInput ToolCalls(listOf(toolCallMessagePart( SolveTool, SolveTool.Args("solve") ))) outputs ReceivedToolResults(listOf(toolResult(SolveTool, SolveTool.Args("solve"), "solved"))) callTool withInput ToolCalls(listOf(toolCallMessagePart( CreateTool, CreateTool.Args("solve") ))) outputs ReceivedToolResults(listOf(toolResult(CreateTool, CreateTool.Args("solve"), "created"))) } assertEdges { askLLM withOutput assistantMessage("Hello!") goesTo giveFeedback askLLM withOutput assistantMessage(CreateTool, CreateTool.Args("solve")) goesTo callTool } } } } } ``` ``` ``` ## API reference For a complete API reference related to the Testing feature, see the reference documentation for the [agents-test](https://api.koog.ai/agents/agents-test/index.html) module. ## FAQ and troubleshooting #### How do I mock a specific tool response? Use the `mockTool` method in `MockLLMBuilder`: ``` val mockExecutor = getMockExecutor { mockTool(myTool) alwaysReturns myResult // Or with conditions mockTool(myTool) returns myResult onArguments myArgs } ``` ``` ``` #### How can I test complex graph structures? Use the subgraph assertions, `verifySubgraph`, and node references: ``` testGraph("test") { val mySubgraph = assertSubgraphByName("mySubgraph") verifySubgraph(mySubgraph) { // Get references to nodes val nodeA = assertNodeByName("nodeA") val nodeB = assertNodeByName("nodeB") // Assert reachability assertReachable(nodeA, nodeB) // Assert edge connections assertEdges { nodeA.withOutput("result") goesTo nodeB } } } ``` ``` ``` #### How do I simulate different LLM responses based on input? Use pattern matching methods: ``` getMockExecutor { mockLLMAnswer("Response A") onRequestContains "topic A" mockLLMAnswer("Response B") onRequestContains "topic B" mockLLMAnswer("Exact response") onRequestEquals "exact question" mockLLMAnswer("Conditional response") onCondition { it.contains("keyword") && it.length > 10 } } ``` ``` import ai.koog.agents.testing.tools.MockExecutor; import ai.koog.prompt.executor.model.PromptExecutor; PromptExecutor promptExecutor = MockExecutor.builder() .mockLLMAnswer("Response A").onRequestContains("topic A") .mockLLMAnswer("Response B").onRequestContains("topic B") .mockLLMAnswer("Exact response").onRequestEquals("exact question") .mockLLMAnswer("Conditional response").onCondition(s -> s.contains("keyword") && s.length() > 10) .build(); ``` ### Troubleshooting #### Mock executor always returns the default response Check that your pattern matching is correct. Patterns are case-sensitive and must match exactly as specified. #### Tool calls are not being intercepted Ensure that: 1. The tool registry is properly set up. 1. The tool names match exactly. 1. The tool actions are configured correctly. #### Graph assertions are failing 1. Verify that node names are correct. 1. Check that the graph structure matches your expectations. 1. Use the `startNode()` and `finishNode()` methods to get the correct entry and exit points. # Why Koog # Why Koog Koog is designed to solve real-world problems with JetBrains-level quality. It provides advanced AI algorithms, out-of-the-box proven techniques, a Kotlin DSL, as well as a Java fluent API, and robust multi-platform support that goes beyond traditional frameworks. Its primary focus is reliability — enabling AI agents that can be confidently used in demanding enterprise environments. ## Integration with Java and Kotlin applications Koog provides a Kotlin Domain-Specific Language (DSL) designed specifically for Kotlin developers, along with a fluent Java API for Java users. The same framework delivers a native feel in both JVM languages, ensuring seamless integration into Kotlin and Java applications while significantly boosting productivity and enhancing the overall developer experience. ## Real-world validation with JetBrains products Koog powers multiple JetBrains products, including internal AI agents. This real-world integration ensures that Koog is continuously tested, refined, and validated against practical use cases. It focuses on what works in practice, incorporating insights from extensive feedback and real-product scenarios. This integration provides Koog with strengths that distinguish it from other frameworks. ## Advanced solutions available out of the box Koog includes pre-built, composable solutions to simplify and speed up the development of agentic systems, setting it apart from frameworks that only offer basic components: - **Graph workflows with domain modeling.** Model AI workflows as explicit graphs built on verified domain models. By expressing requirements as structured data classes instead of relying on naive prompting, you gain precise control over agent behavior and significantly improve reliability and predictability. - **Multiple history compression strategies.** Koog comes with advanced strategies to compress and manage long-running conversations out of the box, eliminating the need for manual experimentation with approaches. With fine-tuned prompts, techniques, and algorithms tested and refined by ML engineers, you can rely on proven methods to improve performance. For more details on compression strategies, refer to [History compression](https://docs.koog.ai/history-compression/). To explore how Koog handles compression and context management in real-world scenarios, check out [this article](https://blog.jetbrains.com/ai/2025/07/when-tool-calling-becomes-an-addiction-debugging-llm-patterns-in-koog/). - **Advanced persistence (Durable execution).** Koog lets you restore full agent state machines instead of just chat messages. This enables features like checkpoints, failure recovery, and even the ability to revert to any point in the execution of the state machine. - **All modern agent patterns, one framework.** Graph workflows, GOAP (Goal-Oriented Action Planning) and LLM planning, multi-agent orchestration — fully supported and fully composable. Build exactly the agent your use case requires. - **Seamless LLM switching.** You can switch a conversation to a different large language model (LLM) with a new set of available tools at any point without losing the existing conversation history. Koog automatically rewrites the history and handles unavailable tools, enabling smooth transitions and a natural interaction flow. - **Robust retry components.** Koog includes a retry mechanism that lets you wrap any set of operations within your agentic system and retry them until they meet configurable conditions. You can provide feedback and adjust each attempt to ensure reliable results. If LLM calls time out, tools do not work as expected, or there are network issues, Koog ensures that your agent remains resilient and performs effectively, even during temporary failures. For more technical details, see [Retry functionality](https://docs.koog.ai/history-compression/). ## Broad integration, multiplatform support, enhanced observability Koog supports the development and deployment of agentic applications across a variety of platforms and environments: - **Spring Boot, Spring AI and Ktor integrations**. Koog integrates with widely used enterprise environments. - For Spring Boot, Koog provides ready-to-use beans and auto-configured LLM clients, making it easy to start building AI-powered workflows. - If you're already using Spring AI for LLM and RAG capabilities, Koog can be layered on top as an orchestration and agentic framework. This allows you to leverage Spring AI’s extensive integrations while benefiting from Koog’s advanced, reliable, and cost-efficient AI workflows. - If you have a Ktor server, you can install Koog as a plugin, configure providers using configuration files, and call agents directly from any route without manually connecting LLM clients. - **Multiplatform support**. You can deploy your agentic applications across JVM, JS, WasmJS, Android, and iOS targets. - **Broad AI integration**. Koog integrates with major LLM providers, including OpenAI, Anthropic, Google, DeepSeek, Mistral, Alibaba, as well as enterprise-level AI clouds like Bedrock. It also supports local models such as Ollama. For the full list of available providers, see [LLM providers](https://docs.koog.ai/llm-providers/). - **OpenTelemetry support**. Koog provides out-of-the-box integration with popular observability providers like [W&B Weave](https://wandb.ai/site/weave/), [Langfuse](https://langfuse.com/) and [DataDog](https://www.datadoghq.com/) for monitoring and debugging AI applications. With native OpenTelemetry support, you can trace, log, and measure your agents using the same tools you already use in your system. To learn more, refer to [OpenTelemetry](https://docs.koog.ai/opentelemetry-support/). ## Collaboration with ML engineers and product teams A unique advantage of Koog is its direct collaboration with JetBrains ML engineers and product teams. This ensures that features built with Koog are not just theoretical but tested and refined based on real-world product requirements. This means that Koog incorporates: - **Fine-tuned prompts and strategies** optimized for real-world performance. - **Proven engineering approaches** discovered and validated through product development, such as its unique history compression strategies. You can learn more in [this detailed article](https://blog.jetbrains.com/ai/2025/07/when-tool-calling-becomes-an-addiction-debugging-llm-patterns-in-koog/). - **Continuous improvements** that help Koog stay efficient and adaptable to evolving needs. ## Commitment to the developer community The Koog team is deeply committed to building a strong developer community. By actively gathering and incorporating feedback, Koog evolves to meet the needs of developers effectively. We are actively expanding support for diverse AI architectures, comprehensive benchmarks, detailed use-case guides, and educational resources to empower developers. ## Where to start - Explore Koog capabilities in [Overview](../). - Build your first Koog agent with our [Quickstart](../quickstart/) guide. - See the latest updates in Koog [release notes](https://github.com/JetBrains/koog/blob/main/CHANGELOG.md). - Learn from [Examples](https://docs.koog.ai/examples/). # Examples # Examples The Koog framework provides examples to help you understand how to implement AI agents for different use cases. Examples are available in both **Kotlin** and **Java**. They demonstrate key features and patterns that you can adapt for your own applications. Browse the examples below and click on the links to view the source code on GitHub. | Example | Description | | --------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | [Attachments](Attachments/) | Learn how to use structured Markdown and attachments in prompts. Build prompts that include images and generate creative content for Instagram posts using OpenAI models. | | [Banking](Banking/) | Build a comprehensive AI banking assistant with routing capabilities that can handle money transfers and transaction analysis through a sophisticated graph-based strategy. Includes domain modeling, tool creation, and agent composition patterns. | | [BedrockAgent](BedrockAgent/) | Create intelligent AI agents using the Koog framework with AWS Bedrock integration. Learn how to define custom tools, set up AWS Bedrock, and build interactive agents that understand natural language commands for controlling devices. | | [Calculator](Calculator/) | Build a calculator agent that performs arithmetic operations using tools for addition, subtraction, multiplication, and division. Demonstrates parallel tool calls, event logging, and multiple executor support (OpenAI and Ollama). | | [Chess](Chess/) | Build an intelligent chess-playing agent featuring complex domain modeling, custom tools, memory optimization techniques, and interactive choice selection. Demonstrates advanced agent strategies, game state management, and human-AI collaboration patterns. | | [GoogleMapsMcp](GoogleMapsMcp/) | Connect Koog to a Google Maps MCP server via Docker. Discover tools, geocode addresses, and fetch elevation data using AI agents with real-world geographic APIs in a Kotlin Notebook environment. | | [Guesser](Guesser/) | Build a number-guessing agent that implements a binary search strategy using tools to ask targeted questions. The agent efficiently narrows down the user's number through strategic questioning and demonstrates tool-based interaction patterns. | | [Langfuse](Langfuse/) | Learn how to export Koog agent traces to Langfuse using OpenTelemetry. Set up environment variables, run agents, and inspect spans and traces in your Langfuse instance for comprehensive observability. | | [MCP](https://github.com/JetBrains/koog/tree/develop/examples/src/main/kotlin/ai/koog/agents/example/mcp) | Integration examples for the Model Context Protocol, featuring GoogleMapsMcpClient for geographic data and PlaywrightMcpClient for browser automation. | | [Memory](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/memory) | A customer support agent that demonstrates memory system usage. The agent tracks user conversation preferences, device diagnostics, and organization-specific information using encrypted local storage and proper memory organization with subjects and scopes. | | [OpenTelemetry](OpenTelemetry/) | Add OpenTelemetry-based tracing to Koog AI agents. Learn to emit spans to console for debugging and export traces to OpenTelemetry Collector for viewing in Jaeger. Includes Docker setup and troubleshooting guide. | | [Planner](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/planner) | A task planning system that builds execution trees with parallel and sequential execution nodes, dynamically constructing execution plans for complex workflows. | | [PlaywrightMcp](PlaywrightMcp/) | Drive browsers with Playwright MCP and Koog. Launch a Playwright MCP server, connect via SSE, and let AI agents automate web tasks like navigation, cookie acceptance, and UI interaction through natural language commands. | | [SimpleAPI](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/simpleapi) | Examples demonstrating chat agents and basic agents with simple API patterns for getting started with Koog. | | [StructuredData](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput) | Demonstrates JSON-based structured data output with complex nested classes, polymorphism, and weather forecast examples showing how to work with typed data in agent responses. | | [SubgraphWithTask](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/subgraphwithtask) | Project generation tools showcasing file and directory operations, including creation, deletion, and command execution using subgraph strategies. | | [Tone](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/tone) | A text tone analysis agent that uses specialized tools to identify positive, negative, or neutral tones in input text, demonstrating sentiment analysis capabilities. | | [UnityMcp](UnityMcp/) | Drive Unity game development with AI agents using Unity MCP server integration. Connect to Unity via stdio, discover available tools, and let agents modify scenes, place objects, and execute game development tasks through natural language commands. | | [VaccumAgent](VaccumAgent/) | Implementation of a basic reflex agent using the Koog framework. Covers environment modeling, tool creation, and agent behavior for automated cleaning tasks in a simple two-cell world. | | [Weave](Weave/) | Learn how to trace Koog agents to W&B Weave using OpenTelemetry (OTLP). Set up environment variables, run agents, and view rich traces in the Weave UI for comprehensive monitoring and debugging. | | [A2A](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/a2a) | Demonstrates agent-to-agent (A2A) communication using Koog framework. Shows how to set up bidirectional communication between AI agents, enable collaborative problem-solving, and manage multi-agent workflows with proper message routing and coordination. | | Example | Description | | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | [Calculator](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/calculator/Calculator.java) | Build a graph-based calculator agent. Demonstrates tools defined using `ToolSet`, a multi-node graph strategy with typed edges, conditional routing, automatic history compression, event handling, and multiple executor support (OpenAI and Ollama). | | [FunctionalAgentChat](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/chat/FunctionalAgentChat.java) | Build an interactive chat agent using a [functional strategy](../agents/functional-agents/). Runs a continuous conversation loop powered by the Llama 3.2 model, accepting user input until `/bye` is entered. | | [ChatMemoryJdbc](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/chatmemory/ChatMemoryJdbcExample.java) | Persist conversation history across sessions using the [`ChatMemory`](../features/chat-memory/) feature backed by a JDBC PostgreSQL provider. Demonstrates creation of JDBC-backed chat history provider with 24h TTL, schema migration, and building the agent with chat memory. Requires PostgreSQL running locally or via Docker. | | [FunctionalStrategy](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/strategies/functional/FunctionalStrategyExample.java) | Implement a multi-step functional strategy with typed subtasks, per-step tool scoping, and an iterative verification and fixing loop to produce validated solutions. | | [GoapStrategy](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/strategies/GoapStrategyExample.java) | Build a planner-based agent using GOAP (Goal-Oriented Action Planning). Define typed belief states, action preconditions, and goal conditions to guide iterative problem-solving and self-correction. | | [GraphStrategy](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/strategies/GraphStrategyExample.java) | Build a graph-based problem-solving agent with typed subgraphs connected by conditional edges. Implements a problem-solving pipeline using LLM-as-a-judge to validate solutions. | | [CustomSubgraph](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/subgraphs/CustomSubgraphExample.java) | Build a multi-subgraph agent strategy. Implements 3 subgraphs connected sequentially that handle research with a web search tool, outline planning, and write a final article summary. | | [OpenTelemetry](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/features/opentelemetry/OpenTelemetryExample.java) | Add OpenTelemetry-based tracing to a Koog AI agent. Configures a console logging exporter for local debugging and an OTLP/gRPC exporter for viewing spans in Jaeger. Includes a Docker setup. | | [Langfuse](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/features/opentelemetry/langfuse/LangfuseExample.java) | Export Koog agent traces to Langfuse via OpenTelemetry. Demonstrates configuring custom attributes such as session ID and trace tags for enriched observability. | | [Weave](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/features/opentelemetry/weave/WeaveExample.java) | Trace Koog agents to W&B Weave using OpenTelemetry (OTLP). Set up environment variables, run agents, and view rich traces in the Weave UI. | | [PersistenceJdbc](https://github.com/JetBrains/koog/tree/develop/examples/simple-examples-java/src/main/java/ai/koog/agents/example/snapshot/PersistenceJdbcExample.java) | Build an agent with the `Persistence` feature backed by a JDBC PostgreSQL provider. The Persistence feature automatically creates checkpoints after each node execution, allowing the agent to resume from where it left off across restarts. Requires PostgreSQL running locally or via Docker. | # Attachments [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/Attachments.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/Attachments.ipynb) ## Setting Up the Environment Before diving into the code, we make sure our Kotlin Notebook is ready. Here we load the latest descriptors and enable the **Koog** library, which provides a clean API for working with AI model providers. ``` // Loads the latest descriptors and activates Koog integration for Kotlin Notebook. // This makes Koog DSL types and executors available in further cells. %useLatestDescriptors %use koog ``` ## Configuring API Keys We read the API key from an environment variable. This keeps secrets out of the notebook file and lets you switch providers. You can set `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, or `GEMINI_API_KEY`. ``` val apiKey = System.getenv("OPENAI_API_KEY") // or ANTHROPIC_API_KEY, or GEMINI_API_KEY ``` ## Creating a Simple OpenAI Executor The executor encapsulates authentication, base URLs, and correct defaults. Here we use a simple OpenAI executor, but you can swap it for Anthropic or Gemini without changing the rest of the code. ``` // --- Provider selection --- // For OpenAI-compatible models. Alternatives include: // val executor = simpleAnthropicExecutor(System.getenv("ANTHROPIC_API_KEY")) // val executor = simpleGeminiExecutor(System.getenv("GEMINI_API_KEY")) // All executors expose the same high‑level API. val executor = simpleOpenAIExecutor(apiKey) ``` Koog’s prompt DSL lets you add **structured Markdown** and **attachments**. In this cell we build a prompt that asks the model to generate a short, blog‑style "content card" and we attach two images from the local `images/` directory. ``` import ai.koog.prompt.markdown.markdown import kotlinx.io.files.Path val prompt = prompt("images-prompt") { system("You are professional assistant that can write cool and funny descriptions for Instagram posts.") user { markdown { +"I want to create a new post on Instagram." br() +"Can you write something creative under my instagram post with the following photos?" br() h2("Requirements") bulleted { item("It must be very funny and creative") item("It must increase my chance of becoming an ultra-famous blogger!!!!") item("It not contain explicit content, harassment or bullying") item("It must be a short catching phrase") item("You must include relevant hashtags that would increase the visibility of my post") } } attachments { image(Path("images/kodee-loving.png")) image(Path("images/kodee-electrified.png")) } } } ``` ## Execute and Inspect the Response We run the prompt against `gpt-4.1`, collect the first message, and print its content. If you want streaming, swap to a streaming API in Koog; for tool use, pass your tool list instead of `emptyList()`. > Troubleshooting: * **401/403** — check your API key/environment variable. * **File not found** — verify the `images/` paths. * **Rate limits** — add minimal retry/backoff around the call if needed. ``` import kotlinx.coroutines.runBlocking runBlocking { val response = executor.execute(prompt = prompt, model = OpenAIModels.Chat.GPT4_1, tools = emptyList()).first() println(response.content) } ``` ``` Caption: Running on cuteness and extra giggle power! Warning: Side effects may include heart-thief vibes and spontaneous dance parties. 💜🤖💃 Hashtags: #ViralVibes #UltraFamousBlogger #CutieAlert #QuirkyContent #InstaFun #SpreadTheLove #DancingIntoFame #RobotLife #InstaFamous #FeedGoals ``` ``` runBlocking { val response = executor.executeStreaming(prompt = prompt, model = OpenAIModels.Chat.GPT4_1) response.collect { print(it) } } ``` ``` Caption: Running on good vibes & wi-fi only! 🤖💜 Drop a like if you feel the circuit-joy! #BlogBotInTheWild #HeartDeliveryService #DancingWithWiFi #UltraFamousBlogger #MoreFunThanYourAICat #ViralVibes #InstaFun #BeepBoopFamous ``` # Building an AI Banking Assistant with Koog [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/Banking.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/Banking.ipynb) In this tutorial we’ll build a small banking assistant using **Koog** agents in Kotlin. You’ll learn how to: - Define domain models and sample data - Expose capability-focused tools for **money transfers** and **transaction analytics** - Classify user intent (Transfer vs Analytics) - Orchestrate calls in two styles: 1. a graph/subgraph strategy 1. “agents as tools” By the end, you’ll be able to route free-form user requests to the right tools and produce helpful, auditable responses. ## Setup & Dependencies We’ll use the Kotlin Notebook kernel. Make sure your Koog artifacts are resolvable from Maven Central and your LLM provider key is available via `OPENAI_API_KEY`. ``` %useLatestDescriptors %use datetime // uncomment this for using koog from Maven Central // %use koog ``` ``` import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor val apiKey = System.getenv("OPENAI_API_KEY") ?: error("Please set OPENAI_API_KEY environment variable") val openAIExecutor = simpleOpenAIExecutor(apiKey) ``` ## Defining the System Prompt A well-crafted system prompt helps the AI understand its role and constraints. This prompt will guide all our agents' behavior. ``` val bankingAssistantSystemPrompt = """ |You are a banking assistant interacting with a user (userId=123). |Your goal is to understand the user's request and determine whether it can be fulfilled using the available tools. | |If the task can be accomplished with the provided tools, proceed accordingly, |at the end of the conversation respond with: "Task completed successfully." |If the task cannot be performed with the tools available, respond with: "Can't perform the task." """.trimMargin() ``` ## Domain model & Sample data First, let's define our domain models and sample data. We'll use Kotlin's data classes with serialization support. ``` import kotlinx.serialization.Serializable @Serializable data class Contact( val id: Int, val name: String, val surname: String? = null, val phoneNumber: String ) val contactList = listOf( Contact(100, "Alice", "Smith", "+1 415 555 1234"), Contact(101, "Bob", "Johnson", "+49 151 23456789"), Contact(102, "Charlie", "Williams", "+36 20 123 4567"), Contact(103, "Daniel", "Anderson", "+46 70 123 45 67"), Contact(104, "Daniel", "Garcia", "+34 612 345 678"), ) val contactById = contactList.associateBy(Contact::id) ``` ## Tools: Money Transfer Tools should be **pure** and predictable. We model two “soft contracts”: - `chooseRecipient` returns *candidates* when ambiguity is detected. - `sendMoney` supports a `confirmed` flag. If `false`, it asks the agent to confirm with the user. ``` import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.agents.core.tools.annotations.Tool import ai.koog.agents.core.tools.reflect.ToolSet @LLMDescription("Tools for money transfer operations.") class MoneyTransferTools : ToolSet { @Tool @LLMDescription( """ Returns the list of contacts for the given user. The user in this demo is always userId=123. """ ) fun getContacts( @LLMDescription("The unique identifier of the user whose contact list is requested.") userId: Int ): String = buildString { contactList.forEach { c -> appendLine("${c.id}: ${c.name} ${c.surname ?: ""} (${c.phoneNumber})") } }.trimEnd() @Tool @LLMDescription("Returns the current balance (demo value).") fun getBalance( @LLMDescription("The unique identifier of the user.") userId: Int ): String = "Balance: 200.00 EUR" @Tool @LLMDescription("Returns the default user currency (demo value).") fun getDefaultCurrency( @LLMDescription("The unique identifier of the user.") userId: Int ): String = "EUR" @Tool @LLMDescription("Returns a demo FX rate between two ISO currencies (e.g. EUR→USD).") fun getExchangeRate( @LLMDescription("Base currency (e.g., EUR).") from: String, @LLMDescription("Target currency (e.g., USD).") to: String ): String = when (from.uppercase() to to.uppercase()) { "EUR" to "USD" -> "1.10" "EUR" to "GBP" -> "0.86" "GBP" to "EUR" -> "1.16" "USD" to "EUR" -> "0.90" else -> "No information about exchange rate available." } @Tool @LLMDescription( """ Returns a ranked list of possible recipients for an ambiguous name. The agent should ask the user to pick one and then use the selected contact id. """ ) fun chooseRecipient( @LLMDescription("An ambiguous or partial contact name.") confusingRecipientName: String ): String { val matches = contactList.filter { c -> c.name.contains(confusingRecipientName, ignoreCase = true) || (c.surname?.contains(confusingRecipientName, ignoreCase = true) ?: false) } if (matches.isEmpty()) { return "No candidates found for '$confusingRecipientName'. Use getContacts and ask the user to choose." } return matches.mapIndexed { idx, c -> "${idx + 1}. ${c.id}: ${c.name} ${c.surname ?: ""} (${c.phoneNumber})" }.joinToString("\n") } @Tool @LLMDescription( """ Sends money from the user to a contact. If confirmed=false, return "REQUIRES_CONFIRMATION" with a human-readable summary. The agent should confirm with the user before retrying with confirmed=true. """ ) fun sendMoney( @LLMDescription("Sender user id.") senderId: Int, @LLMDescription("Amount in sender's default currency.") amount: Double, @LLMDescription("Recipient contact id.") recipientId: Int, @LLMDescription("Short purpose/description.") purpose: String, @LLMDescription("Whether the user already confirmed this transfer.") confirmed: Boolean = false ): String { val recipient = contactById[recipientId] ?: return "Invalid recipient." val summary = "Transfer €%.2f to %s %s (%s) for \"%s\"." .format(amount, recipient.name, recipient.surname ?: "", recipient.phoneNumber, purpose) if (!confirmed) { return "REQUIRES_CONFIRMATION: $summary" } // In a real system this is where you'd call a payment API. return "Money was sent. $summary" } } ``` ## Creating Your First Agent Now let's create an agent that uses our money transfer tools. An agent combines an LLM with tools to accomplish tasks. ``` import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.core.agent.AIAgentService import ai.koog.agents.core.tools.ToolRegistry import ai.koog.agents.core.tools.reflect.asTools import ai.koog.agents.ext.tool.AskUser import ai.koog.prompt.executor.clients.openai.OpenAIModels import kotlinx.coroutines.runBlocking val transferAgentService = AIAgentService( executor = openAIExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = bankingAssistantSystemPrompt, temperature = 0.0, // Use deterministic responses for financial operations toolRegistry = ToolRegistry { tool(AskUser) tools(MoneyTransferTools().asTools()) } ) // Test the agent with various scenarios println("Banking Assistant started") val message = "Send 25 euros to Daniel for dinner at the restaurant." // Other test messages you can try: // - "Send 50 euros to Alice for the concert tickets" // - "What's my current balance?" // - "Transfer 100 euros to Bob for the shared vacation expenses" runBlocking { val result = transferAgentService.createAgentAndRun(message) result } ``` ``` Banking Assistant started There are two contacts named Daniel. Please confirm which one you would like to send money to: 1. Daniel Anderson (+46 70 123 45 67) 2. Daniel Garcia (+34 612 345 678) Please confirm the transfer of €25.00 to Daniel Garcia (+34 612 345 678) for "Dinner at the restaurant". Task completed successfully. ``` ## Adding Transaction Analytics Let's expand our assistant's capabilities with transaction analysis tools. First, we'll define the transaction domain model. ``` @Serializable enum class TransactionCategory(val title: String) { FOOD_AND_DINING("Food & Dining"), SHOPPING("Shopping"), TRANSPORTATION("Transportation"), ENTERTAINMENT("Entertainment"), GROCERIES("Groceries"), HEALTH("Health"), UTILITIES("Utilities"), HOME_IMPROVEMENT("Home Improvement"); companion object { fun fromString(value: String): TransactionCategory? = entries.find { it.title.equals(value, ignoreCase = true) } fun availableCategories(): String = entries.joinToString(", ") { it.title } } } @Serializable data class Transaction( val merchant: String, val amount: Double, val category: TransactionCategory, val date: LocalDateTime ) ``` ### Sample transaction data ``` val transactionAnalysisPrompt = """ Today is 2025-05-22. Available categories for transactions: ${TransactionCategory.availableCategories()} """ val sampleTransactions = listOf( Transaction("Starbucks", 5.99, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 22, 8, 30, 0, 0)), Transaction("Amazon", 129.99, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 22, 10, 15, 0, 0)), Transaction( "Shell Gas Station", 45.50, TransactionCategory.TRANSPORTATION, LocalDateTime(2025, 5, 21, 18, 45, 0, 0) ), Transaction("Netflix", 15.99, TransactionCategory.ENTERTAINMENT, LocalDateTime(2025, 5, 21, 12, 0, 0, 0)), Transaction("AMC Theaters", 32.50, TransactionCategory.ENTERTAINMENT, LocalDateTime(2025, 5, 20, 19, 30, 0, 0)), Transaction("Whole Foods", 89.75, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 20, 16, 20, 0, 0)), Transaction("Target", 67.32, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 20, 14, 30, 0, 0)), Transaction("CVS Pharmacy", 23.45, TransactionCategory.HEALTH, LocalDateTime(2025, 5, 19, 11, 25, 0, 0)), Transaction("Subway", 12.49, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 19, 13, 15, 0, 0)), Transaction("Spotify Premium", 9.99, TransactionCategory.ENTERTAINMENT, LocalDateTime(2025, 5, 19, 14, 15, 0, 0)), Transaction("AT&T", 85.00, TransactionCategory.UTILITIES, LocalDateTime(2025, 5, 18, 9, 0, 0, 0)), Transaction("Home Depot", 156.78, TransactionCategory.HOME_IMPROVEMENT, LocalDateTime(2025, 5, 18, 15, 45, 0, 0)), Transaction("Amazon", 129.99, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 17, 10, 15, 0, 0)), Transaction("Starbucks", 5.99, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 17, 8, 30, 0, 0)), Transaction("Whole Foods", 89.75, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 16, 16, 20, 0, 0)), Transaction("CVS Pharmacy", 23.45, TransactionCategory.HEALTH, LocalDateTime(2025, 5, 15, 11, 25, 0, 0)), Transaction("AT&T", 85.00, TransactionCategory.UTILITIES, LocalDateTime(2025, 5, 14, 9, 0, 0, 0)), Transaction("Xbox Game Pass", 14.99, TransactionCategory.ENTERTAINMENT, LocalDateTime(2025, 5, 14, 16, 45, 0, 0)), Transaction("Aldi", 76.45, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 13, 17, 30, 0, 0)), Transaction("Chipotle", 15.75, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 13, 12, 45, 0, 0)), Transaction("Best Buy", 299.99, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 12, 14, 20, 0, 0)), Transaction("Olive Garden", 89.50, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 12, 19, 15, 0, 0)), Transaction("Whole Foods", 112.34, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 11, 10, 30, 0, 0)), Transaction("Old Navy", 45.99, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 11, 13, 45, 0, 0)), Transaction("Panera Bread", 18.25, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 10, 11, 30, 0, 0)), Transaction("Costco", 245.67, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 10, 15, 20, 0, 0)), Transaction("Five Guys", 22.50, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 9, 18, 30, 0, 0)), Transaction("Macy's", 156.78, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 9, 14, 15, 0, 0)), Transaction("Hulu Plus", 12.99, TransactionCategory.ENTERTAINMENT, LocalDateTime(2025, 5, 8, 20, 0, 0, 0)), Transaction("Whole Foods", 94.23, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 8, 16, 45, 0, 0)), Transaction("Texas Roadhouse", 78.90, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 8, 19, 30, 0, 0)), Transaction("Walmart", 167.89, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 7, 11, 20, 0, 0)), Transaction("Chick-fil-A", 14.75, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 7, 12, 30, 0, 0)), Transaction("Aldi", 82.45, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 6, 15, 45, 0, 0)), Transaction("TJ Maxx", 67.90, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 6, 13, 20, 0, 0)), Transaction("P.F. Chang's", 95.40, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 5, 19, 15, 0, 0)), Transaction("Whole Foods", 78.34, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 4, 14, 30, 0, 0)), Transaction("H&M", 89.99, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 3, 16, 20, 0, 0)), Transaction("Red Lobster", 112.45, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 2, 18, 45, 0, 0)), Transaction("Whole Foods", 67.23, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 2, 11, 30, 0, 0)), Transaction("Marshalls", 123.45, TransactionCategory.SHOPPING, LocalDateTime(2025, 5, 1, 15, 20, 0, 0)), Transaction( "Buffalo Wild Wings", 45.67, TransactionCategory.FOOD_AND_DINING, LocalDateTime(2025, 5, 1, 19, 30, 0, 0) ), Transaction("Aldi", 145.78, TransactionCategory.GROCERIES, LocalDateTime(2025, 5, 1, 10, 15, 0, 0)) ) ``` ## Transaction Analysis Tools ``` @LLMDescription("Tools for analyzing transaction history") class TransactionAnalysisTools : ToolSet { @Tool @LLMDescription( """ Retrieves transactions filtered by userId, category, start date, and end date. All parameters are optional. If no parameters are provided, all transactions are returned. Dates should be in the format YYYY-MM-DD. """ ) fun getTransactions( @LLMDescription("The ID of the user whose transactions to retrieve.") userId: String? = null, @LLMDescription("The category to filter transactions by (e.g., 'Food & Dining').") category: String? = null, @LLMDescription("The start date to filter transactions by, in the format YYYY-MM-DD.") startDate: String? = null, @LLMDescription("The end date to filter transactions by, in the format YYYY-MM-DD.") endDate: String? = null ): String { var filteredTransactions = sampleTransactions // Validate userId (in production, this would query a real database) if (userId != null && userId != "123") { return "No transactions found for user $userId." } // Apply category filter category?.let { cat -> val categoryEnum = TransactionCategory.fromString(cat) ?: return "Invalid category: $cat. Available: ${TransactionCategory.availableCategories()}" filteredTransactions = filteredTransactions.filter { it.category == categoryEnum } } // Apply date range filters startDate?.let { date -> val startDateTime = parseDate(date, startOfDay = true) filteredTransactions = filteredTransactions.filter { it.date >= startDateTime } } endDate?.let { date -> val endDateTime = parseDate(date, startOfDay = false) filteredTransactions = filteredTransactions.filter { it.date <= endDateTime } } if (filteredTransactions.isEmpty()) { return "No transactions found matching the specified criteria." } return filteredTransactions.joinToString("\n") { transaction -> "${transaction.date}: ${transaction.merchant} - " + "$${transaction.amount} (${transaction.category.title})" } } @Tool @LLMDescription("Calculates the sum of an array of double numbers.") fun sumArray( @LLMDescription("Comma-separated list of double numbers to sum (e.g., '1.5,2.3,4.7').") numbers: String ): String { val numbersList = numbers.split(",") .mapNotNull { it.trim().toDoubleOrNull() } val sum = numbersList.sum() return "Sum: $%.2f".format(sum) } // Helper function to parse dates private fun parseDate(dateStr: String, startOfDay: Boolean): LocalDateTime { val parts = dateStr.split("-").map { it.toInt() } require(parts.size == 3) { "Invalid date format. Use YYYY-MM-DD" } return if (startOfDay) { LocalDateTime(parts[0], parts[1], parts[2], 0, 0, 0, 0) } else { LocalDateTime(parts[0], parts[1], parts[2], 23, 59, 59, 999999999) } } } ``` ``` val analysisAgentService = AIAgentService( executor = openAIExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "$bankingAssistantSystemPrompt\n$transactionAnalysisPrompt", temperature = 0.0, toolRegistry = ToolRegistry { tools(TransactionAnalysisTools().asTools()) } ) println("Transaction Analysis Assistant started") val analysisMessage = "How much have I spent on restaurants this month?" // Other queries to try: // - "What's my maximum check at a restaurant this month?" // - "How much did I spend on groceries in the first week of May?" // - "What's my total spending on entertainment in May?" // - "Show me all transactions from last week" runBlocking { val result = analysisAgentService.createAgentAndRun(analysisMessage) result } ``` ``` Transaction Analysis Assistant started You have spent a total of $517.64 on restaurants this month. Task completed successfully. ``` ## Building an Agent with Graph Now let's combine our specialized agents into a graph agent that can route requests to the appropriate handler. ### Request Classification First, we need a way to classify incoming requests: ``` import ai.koog.agents.core.tools.annotations.LLMDescription import kotlinx.serialization.SerialName import kotlinx.serialization.Serializable @Suppress("unused") @SerialName("UserRequestType") @Serializable @LLMDescription("Type of user request: Transfer or Analytics") enum class RequestType { Transfer, Analytics } @Serializable @LLMDescription("The bank request that was classified by the agent.") data class ClassifiedBankRequest( @property:LLMDescription("Type of request: Transfer or Analytics") val requestType: RequestType, @property:LLMDescription("Actual request to be performed by the banking application") val userRequest: String ) ``` ### Shared tool registry ``` // Create a comprehensive tool registry for the multi-agent system val toolRegistry = ToolRegistry { tool(AskUser) // Allow agents to ask for clarification tools(MoneyTransferTools().asTools()) tools(TransactionAnalysisTools().asTools()) } ``` ## Agent Strategy Now we'll create a strategy that orchestrates multiple nodes: ``` import ai.koog.agents.core.dsl.builder.strategy import ai.koog.agents.core.dsl.extension.* import ai.koog.agents.ext.agent.subgraphWithTask import ai.koog.prompt.structure.StructureFixingParser val strategy = strategy("banking assistant") { // Subgraph for classifying user requests val classifyRequest by subgraph( tools = listOf(AskUser) ) { // Use structured output to ensure proper classification val requestClassification by nodeLLMRequestStructured( examples = listOf( ClassifiedBankRequest( requestType = RequestType.Transfer, userRequest = "Send 25 euros to Daniel for dinner at the restaurant." ), ClassifiedBankRequest( requestType = RequestType.Analytics, userRequest = "Provide transaction overview for the last month" ) ), fixingParser = StructureFixingParser( model = OpenAIModels.Chat.GPT4oMini, retries = 2, ) ) val callLLM by nodeLLMRequest() val callAskUserTool by nodeExecuteTool() // Define the flow edge(nodeStart forwardTo requestClassification) edge( requestClassification forwardTo nodeFinish onCondition { it.isSuccess } transformed { it.getOrThrow().data } ) edge( requestClassification forwardTo callLLM onCondition { it.isFailure } transformed { "Failed to understand the user's intent" } ) edge(callLLM forwardTo callAskUserTool onToolCall { true }) edge( callLLM forwardTo callLLM onAssistantMessage { true } transformed { "Please call `${AskUser.name}` tool instead of chatting" } ) edge(callAskUserTool forwardTo requestClassification transformed { it.result.toString() }) } // Subgraph for handling money transfers val transferMoney by subgraphWithTask( tools = MoneyTransferTools().asTools() + AskUser, llmModel = OpenAIModels.Chat.GPT4o // Use more capable model for transfers ) { request -> """ $bankingAssistantSystemPrompt Specifically, you need to help with the following request: ${request.userRequest} """.trimIndent() } // Subgraph for transaction analysis val transactionAnalysis by subgraphWithTask( tools = TransactionAnalysisTools().asTools() + AskUser, ) { request -> """ $bankingAssistantSystemPrompt $transactionAnalysisPrompt Specifically, you need to help with the following request: ${request.userRequest} """.trimIndent() } // Connect the subgraphs edge(nodeStart forwardTo classifyRequest) edge(classifyRequest forwardTo transferMoney onCondition { it.requestType == RequestType.Transfer }) edge(classifyRequest forwardTo transactionAnalysis onCondition { it.requestType == RequestType.Analytics }) // Route results to finish node edge(transferMoney forwardTo nodeFinish) edge(transactionAnalysis forwardTo nodeFinish) } ``` ``` import ai.koog.agents.core.agent.config.AIAgentConfig import ai.koog.prompt.dsl.prompt val agentConfig = AIAgentConfig( prompt = prompt(id = "banking assistant") { system("$bankingAssistantSystemPrompt\n$transactionAnalysisPrompt") }, model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 50 // Allow for complex multi-step operations ) val agent = AIAgent( promptExecutor = openAIExecutor, strategy = strategy, agentConfig = agentConfig, toolRegistry = toolRegistry, ) ``` ## Run graph agent ``` println("Banking Assistant started") val testMessage = "Send 25 euros to Daniel for dinner at the restaurant." // Test various scenarios: // Transfer requests: // - "Send 50 euros to Alice for the concert tickets" // - "Transfer 100 to Bob for groceries" // - "What's my current balance?" // // Analytics requests: // - "How much have I spent on restaurants this month?" // - "What's my maximum check at a restaurant this month?" // - "How much did I spend on groceries in the first week of May?" // - "What's my total spending on entertainment in May?" runBlocking { val result = agent.run(testMessage) "Result: $result" } ``` ``` Banking Assistant started I found multiple contacts with the name Daniel. Please choose the correct one: 1. Daniel Anderson (+46 70 123 45 67) 2. Daniel Garcia (+34 612 345 678) Please specify the number of the correct recipient. Please confirm if you would like to proceed with sending €25 to Daniel Garcia for "dinner at the restaurant." Result: Task completed successfully. ``` ## Agent Composition — Using Agents as Tools Koog allows you to use agents as tools within other agents, enabling powerful composition patterns. ``` import ai.koog.agents.core.agent.createAgentTool import ai.koog.agents.core.tools.ToolParameterDescriptor import ai.koog.agents.core.tools.ToolParameterType val classifierAgent = AIAgent( executor = openAIExecutor, llmModel = OpenAIModels.Chat.GPT4oMini, toolRegistry = ToolRegistry { tool(AskUser) // Convert agents into tools tool( transferAgentService.createAgentTool( agentName = "transferMoney", agentDescription = "Transfers money and handles all related operations", inputDescriptor = ToolParameterDescriptor( name = "request", description = "Transfer request from the user", type = ToolParameterType.String ) ) ) tool( analysisAgentService.createAgentTool( agentName = "analyzeTransactions", agentDescription = "Performs analytics on user transactions", inputDescriptor = ToolParameterDescriptor( name = "request", description = "Transaction analytics request", type = ToolParameterType.String ) ) ) }, systemPrompt = "$bankingAssistantSystemPrompt\n$transactionAnalysisPrompt" ) ``` ## Run composed agent ``` println("Banking Assistant started") val composedMessage = "Send 25 euros to Daniel for dinner at the restaurant." runBlocking { val result = classifierAgent.run(composedMessage) "Result: $result" } ``` ``` Banking Assistant started There are two contacts named Daniel. Please confirm which one you would like to send money to: 1. Daniel Anderson (+46 70 123 45 67) 2. Daniel Garcia (+34 612 345 678) Please confirm the transfer of €25.00 to Daniel Anderson (+46 70 123 45 67) for "Dinner at the restaurant". Result: Can't perform the task. ``` ## Summary In this tutorial, you've learned how to: 1. Create LLM-powered tools with clear descriptions that help the AI understand when and how to use them 1. Build single-purpose agents that combine LLMs with tools to accomplish specific tasks 1. Implement graph agent using strategies and subgraphs for complex workflows 1. Compose agents by using them as tools within other agents 1. Handle user interactions including confirmations and disambiguation ## Best Practices 1. Clear tool descriptions: Write detailed LLMDescription annotations to help the AI understand tool usage 1. Idiomatic Kotlin: Use Kotlin features like data classes, extension functions, and scope functions 1. Error handling: Always validate inputs and provide meaningful error messages 1. User experience: Include confirmation steps for critical operations like money transfers 1. Modularity: Separate concerns into different tools and agents for better maintainability # Building AI Agents with AWS Bedrock and Koog Framework [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/BedrockAgent.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/BedrockAgent.ipynb) Welcome to this comprehensive guide on creating intelligent AI agents using the Koog framework with AWS Bedrock integration. In this notebook, we'll walk through building a functional agent that can control a simple switch device through natural language commands. ## What You'll Learn - How to define custom tools for AI agents using Kotlin annotations - Setting up AWS Bedrock integration for LLM-powered agents - Creating tool registries and connecting them to agents - Building interactive agents that can understand and execute commands ## Prerequisites - AWS Bedrock access with appropriate permissions - AWS credentials configured (access key and secret key) - Basic understanding of Kotlin coroutines Let's dive into building our first Bedrock-powered AI agent! ``` %useLatestDescriptors // %use koog ``` ``` import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.agents.core.tools.annotations.Tool import ai.koog.agents.core.tools.reflect.ToolSet // Simple state-holding device that our agent will control class Switch { private var state: Boolean = false fun switch(on: Boolean) { state = on } fun isOn(): Boolean { return state } } /** * ToolSet implementation that exposes switch operations to the AI agent. * * Key concepts: * - @Tool annotation marks methods as callable by the agent * - @LLMDescription provides natural language descriptions for the LLM * - ToolSet interface allows grouping related tools together */ class SwitchTools(val switch: Switch) : ToolSet { @Tool @LLMDescription("Switches the state of the switch to on or off") fun switchState(state: Boolean): String { switch.switch(state) return "Switch turned ${if (state) "on" else "off"} successfully" } @Tool @LLMDescription("Returns the current state of the switch (on or off)") fun getCurrentState(): String { return "Switch is currently ${if (switch.isOn()) "on" else "off"}" } } ``` ``` import ai.koog.agents.core.tools.ToolRegistry import ai.koog.agents.core.tools.reflect.asTools // Create our switch instance val switch = Switch() // Build the tool registry with our switch tools val toolRegistry = ToolRegistry { // Convert our ToolSet to individual tools and register them tools(SwitchTools(switch).asTools()) } println("✅ Tool registry created with ${toolRegistry.tools.size} tools:") toolRegistry.tools.forEach { tool -> println(" - ${tool.name}") } ``` ``` ✅ Tool registry created with 2 tools: - getCurrentState - switchState ``` ``` import ai.koog.prompt.executor.clients.bedrock.BedrockClientSettings import ai.koog.prompt.executor.clients.bedrock.BedrockRegions val region = BedrockRegions.US_WEST_2.regionCode val maxRetries = 3 // Configure Bedrock client settings val bedrockSettings = BedrockClientSettings( region = region, // Choose your preferred AWS region maxRetries = maxRetries // Number of retry attempts for failed requests ) println("🌐 Bedrock configured for region: $region") println("🔄 Max retries set to: $maxRetries") ``` ``` 🌐 Bedrock configured for region: us-west-2 🔄 Max retries set to: 3 ``` ``` import ai.koog.prompt.executor.llms.all.simpleBedrockExecutor // Create the Bedrock LLM executor with credentials from environment val executor = simpleBedrockExecutor( awsAccessKeyId = System.getenv("AWS_BEDROCK_ACCESS_KEY") ?: throw IllegalStateException("AWS_BEDROCK_ACCESS_KEY environment variable not set"), awsSecretAccessKey = System.getenv("AWS_BEDROCK_SECRET_ACCESS_KEY") ?: throw IllegalStateException("AWS_BEDROCK_SECRET_ACCESS_KEY environment variable not set"), settings = bedrockSettings ) println("🔐 Bedrock executor initialized successfully") println("💡 Pro tip: Set AWS_BEDROCK_ACCESS_KEY and AWS_BEDROCK_SECRET_ACCESS_KEY environment variables") ``` ``` 🔐 Bedrock executor initialized successfully 💡 Pro tip: Set AWS_BEDROCK_ACCESS_KEY and AWS_BEDROCK_SECRET_ACCESS_KEY environment variables ``` ``` import ai.koog.agents.core.agent.AIAgent import ai.koog.prompt.executor.clients.bedrock.BedrockModels val agent = AIAgent( executor = executor, llmModel = BedrockModels.AnthropicClaude35SonnetV2, // State-of-the-art reasoning model systemPrompt = """ You are a helpful assistant that controls a switch device. You can: - Turn the switch on or off when requested - Check the current state of the switch - Explain what you're doing Always be clear about the switch's current state and confirm actions taken. """.trimIndent(), temperature = 0.1, // Low temperature for consistent, focused responses toolRegistry = toolRegistry ) println("🤖 AI Agent created successfully!") println("📋 System prompt configured") println("🛠️ Tools available: ${toolRegistry.tools.size}") println("🎯 Model: ${BedrockModels.AnthropicClaude35SonnetV2}") println("🌡️ Temperature: 0.1 (focused responses)") ``` ``` 🤖 AI Agent created successfully! 📋 System prompt configured 🛠️ Tools available: 2 🎯 Model: LLModel(provider=Bedrock, id=us.anthropic.claude-3-5-sonnet-20241022-v2:0, capabilities=[Temperature, Tools, ToolChoice, Image, Document, Completion], contextLength=200000, maxOutputTokens=8192) 🌡️ Temperature: 0.1 (focused responses) ``` ``` import kotlinx.coroutines.runBlocking println("🎉 Bedrock Agent with Switch Tools - Ready to Go!") println("💬 You can ask me to:") println(" • Turn the switch on/off") println(" • Check the current switch state") println(" • Ask questions about the switch") println() println("💡 Example: 'Please turn on the switch' or 'What's the current state?'") println("📝 Type your request:") val input = readln() println("\n🤖 Processing your request...") runBlocking { val response = agent.run(input) println("\n✨ Agent response:") println(response) } ``` ``` 🎉 Bedrock Agent with Switch Tools - Ready to Go! 💬 You can ask me to: • Turn the switch on/off • Check the current switch state • Ask questions about the switch 💡 Example: 'Please turn on the switch' or 'What's the current state?' 📝 Type your request: The execution was interrupted ``` ## What Just Happened? 🎯 When you run the agent, here's the magic that occurs behind the scenes: 1. **Natural Language Processing**: Your input is sent to Claude 3.5 Sonnet via Bedrock 1. **Intent Recognition**: The model understands what you want to do with the switch 1. **Tool Selection**: Based on your request, the agent decides which tools to call 1. **Action Execution**: The appropriate tool methods are invoked on your switch object 1. **Response Generation**: The agent formulates a natural language response about what happened This demonstrates the core power of the Koog framework - seamless integration between natural language understanding and programmatic actions. ## Next Steps & Extensions Ready to take this further? Here are some ideas to explore: ### 🔧 Enhanced Tools ``` @Tool @LLMDescription("Sets a timer to automatically turn off the switch after specified seconds") fun setAutoOffTimer(seconds: Int): String @Tool @LLMDescription("Gets the switch usage statistics and history") fun getUsageStats(): String ``` ### 🌐 Multiple Devices ``` class HomeAutomationTools : ToolSet { @Tool fun controlLight(room: String, on: Boolean): String @Tool fun setThermostat(temperature: Double): String @Tool fun lockDoor(doorName: String): String } ``` ### 🧠 Memory & Context ``` val agent = AIAgent( executor = executor, // ... other config features = listOf( MemoryFeature(), // Remember past interactions LoggingFeature() // Track all actions ) ) ``` ### 🔄 Advanced Workflows ``` // Multi-step workflows with conditional logic @Tool @LLMDescription("Executes evening routine: dims lights, locks doors, sets thermostat") fun eveningRoutine(): String ``` ## Key Takeaways ✅ **Tools are functions**: Any Kotlin function can become an agent capability ✅ **Annotations drive behavior**: @Tool and @LLMDescription make functions discoverable ✅ **ToolSets organize capabilities**: Group related tools together logically ✅ **Registries are toolboxes**: ToolRegistry contains all available agent capabilities ✅ **Agents orchestrate everything**: AIAgent brings LLM intelligence + tools together The Koog framework makes it incredibly straightforward to build sophisticated AI agents that can understand natural language and take real-world actions. Start simple, then expand your agent's capabilities by adding more tools and features as needed. **Happy agent building!** 🚀 ## Testing the Agent Time to see our agent in action! The agent can now understand natural language requests and use the tools we've provided to control the switch. **Try these commands:** - "Turn on the switch" - "What's the current state?" - "Switch it off please" - "Is the switch on or off?" # Building a Tool-Calling Calculator Agent with Koog [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/Calculator.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/Calculator.ipynb) In this mini-tutorial we’ll build a calculator agent powered by **Koog** tool-calling. You’ll learn how to: - Design small, pure **tools** for arithmetic - Orchestrate **parallel** tool calls with Koog’s multiple-call strategy - Add lightweight **event logging** for transparency - Run with OpenAI (and optionally Ollama) We’ll keep the API tidy and idiomatic Kotlin, returning predictable results and handling edge cases (like division by zero) gracefully. ## Setup We assume you’re in a Kotlin Notebook environment with Koog available. Provide an LLM executor ``` %useLatestDescriptors %use koog val OPENAI_API_KEY = System.getenv("OPENAI_API_KEY") ?: error("Please set the OPENAI_API_KEY environment variable") val executor = simpleOpenAIExecutor(OPENAI_API_KEY) ``` ## Calculator Tools Tools are small, pure functions with clear contracts. We’ll use `Double` for better precision and format outputs consistently. ``` import ai.koog.agents.core.tools.annotations.Tool // Format helper: integers render cleanly, decimals keep reasonable precision. private fun Double.pretty(): String = if (abs(this % 1.0) < 1e-9) this.toLong().toString() else "%.10g".format(this) @LLMDescription("Tools for basic calculator operations") class CalculatorTools : ToolSet { @Tool @LLMDescription("Adds two numbers and returns the sum as text.") fun plus( @LLMDescription("First addend.") a: Double, @LLMDescription("Second addend.") b: Double ): String = (a + b).pretty() @Tool @LLMDescription("Subtracts the second number from the first and returns the difference as text.") fun minus( @LLMDescription("Minuend.") a: Double, @LLMDescription("Subtrahend.") b: Double ): String = (a - b).pretty() @Tool @LLMDescription("Multiplies two numbers and returns the product as text.") fun multiply( @LLMDescription("First factor.") a: Double, @LLMDescription("Second factor.") b: Double ): String = (a * b).pretty() @Tool @LLMDescription("Divides the first number by the second and returns the quotient as text. Returns an error message on division by zero.") fun divide( @LLMDescription("Dividend.") a: Double, @LLMDescription("Divisor (must not be zero).") b: Double ): String = if (abs(b) < 1e-12) { "ERROR: Division by zero" } else { (a / b).pretty() } } ``` ## Tool Registry Expose our tools (plus two built-ins for interaction/logging). ``` val toolRegistry = ToolRegistry { tool(AskUser) // enables explicit user clarification when needed tool(SayToUser) // allows the agent to present the final message to the user tools(CalculatorTools()) } ``` ## Strategy: Multiple Tool Calls (with Optional Compression) This strategy lets the LLM propose **multiple tool calls at once** (e.g., `plus`, `minus`, `multiply`, `divide`) and then sends the results back. If the token usage grows too large, we **compress** the history of tool results before continuing. ``` import ai.koog.agents.core.environment.ReceivedToolResult object CalculatorStrategy { private const val MAX_TOKENS_THRESHOLD = 1000 val strategy = strategy("test") { val callLLM by nodeLLMRequestMultiple() val executeTools by nodeExecuteMultipleTools(parallelTools = true) val sendToolResults by nodeLLMSendMultipleToolResults() val compressHistory by nodeLLMCompressHistory>() edge(nodeStart forwardTo callLLM) // If the assistant produced a final answer, finish. edge((callLLM forwardTo nodeFinish) transformed { it.first() } onAssistantMessage { true }) // Otherwise, run the tools LLM requested (possibly several in parallel). edge((callLLM forwardTo executeTools) onMultipleToolCalls { true }) // If we’re getting large, compress past tool results before continuing. edge( (executeTools forwardTo compressHistory) onCondition { llm.readSession { prompt.latestTokenUsage > MAX_TOKENS_THRESHOLD } } ) edge(compressHistory forwardTo sendToolResults) // Normal path: send tool results back to the LLM. edge( (executeTools forwardTo sendToolResults) onCondition { llm.readSession { prompt.latestTokenUsage <= MAX_TOKENS_THRESHOLD } } ) // LLM might request more tools after seeing results. edge((sendToolResults forwardTo executeTools) onMultipleToolCalls { true }) // Or it can produce the final answer. edge((sendToolResults forwardTo nodeFinish) transformed { it.first() } onAssistantMessage { true }) } } ``` ## Agent Configuration A minimal, tool-forward prompt works well. Keep temperature low for deterministic math. ``` val agentConfig = AIAgentConfig( prompt = prompt("calculator") { system("You are a calculator. Always use the provided tools for arithmetic.") }, model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 50 ) ``` ``` import ai.koog.agents.features.eventHandler.feature.handleEvents val agent = AIAgent( promptExecutor = executor, strategy = CalculatorStrategy.strategy, agentConfig = agentConfig, toolRegistry = toolRegistry ) { handleEvents { onToolCallStarting { e -> println("Tool called: ${e.tool.name}, args=${e.toolArgs}") } onAgentExecutionFailed { e -> println("Agent error: ${e.throwable.message}") } onAgentCompleted { e -> println("Final result: ${e.result}") } } } ``` ## Try It The agent should decompose the expression into parallel tool calls and return a neatly formatted result. ``` import kotlinx.coroutines.runBlocking runBlocking { agent.run("(10 + 20) * (5 + 5) / (2 - 11)") } // Expected final value ≈ -33.333... ``` ``` Tool called: plus, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=10.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=20.0}) Tool called: plus, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=5.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=5.0}) Tool called: minus, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.minus(kotlin.Double, kotlin.Double): kotlin.String=2.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.minus(kotlin.Double, kotlin.Double): kotlin.String=11.0}) Tool called: multiply, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.multiply(kotlin.Double, kotlin.Double): kotlin.String=30.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.multiply(kotlin.Double, kotlin.Double): kotlin.String=10.0}) Tool called: divide, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.divide(kotlin.Double, kotlin.Double): kotlin.String=1.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.divide(kotlin.Double, kotlin.Double): kotlin.String=-9.0}) Tool called: divide, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.divide(kotlin.Double, kotlin.Double): kotlin.String=300.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.divide(kotlin.Double, kotlin.Double): kotlin.String=-9.0}) Final result: The result of the expression \((10 + 20) * (5 + 5) / (2 - 11)\) is approximately \(-33.33\). The result of the expression \((10 + 20) * (5 + 5) / (2 - 11)\) is approximately \(-33.33\). ``` ## Try Forcing Parallel Calls Ask the model to call all needed tools at once. You should still see a correct plan and stable execution. ``` runBlocking { agent.run("Use tools to calculate (10 + 20) * (5 + 5) / (2 - 11). Please call all the tools at once.") } ``` ``` Tool called: plus, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=10.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=20.0}) Tool called: plus, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=5.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.plus(kotlin.Double, kotlin.Double): kotlin.String=5.0}) Tool called: minus, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.minus(kotlin.Double, kotlin.Double): kotlin.String=2.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.minus(kotlin.Double, kotlin.Double): kotlin.String=11.0}) Tool called: multiply, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.multiply(kotlin.Double, kotlin.Double): kotlin.String=30.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.multiply(kotlin.Double, kotlin.Double): kotlin.String=10.0}) Tool called: divide, args=VarArgs(args={parameter #1 a of fun Line_4_jupyter.CalculatorTools.divide(kotlin.Double, kotlin.Double): kotlin.String=30.0, parameter #2 b of fun Line_4_jupyter.CalculatorTools.divide(kotlin.Double, kotlin.Double): kotlin.String=-9.0}) Final result: The result of \((10 + 20) * (5 + 5) / (2 - 11)\) is approximately \(-3.33\). The result of \((10 + 20) * (5 + 5) / (2 - 11)\) is approximately \(-3.33\). ``` ## Running with Ollama Swap the executor and model if you prefer local inference. ``` val ollamaExecutor: PromptExecutor = simpleOllamaAIExecutor() val ollamaAgentConfig = AIAgentConfig( prompt = prompt("calculator", LLMParams(temperature = 0.0)) { system("You are a calculator. Always use the provided tools for arithmetic.") }, model = OllamaModels.Meta.LLAMA_3_2, maxAgentIterations = 50 ) val ollamaAgent = AIAgent( promptExecutor = ollamaExecutor, strategy = CalculatorStrategy.strategy, agentConfig = ollamaAgentConfig, toolRegistry = toolRegistry ) runBlocking { ollamaAgent.run("(10 + 20) * (5 + 5) / (2 - 11)") } ``` ``` Agent says: The result of the expression (10 + 20) * (5 + 5) / (2 - 11) is approximately -33.33. If you have any more questions or need further assistance, feel free to ask! ``` # Building an AI Chess Player with Koog Framework [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/Chess.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/Chess.ipynb) This tutorial demonstrates how to build an intelligent chess-playing agent using the Koog framework. We'll explore key concepts including tool integration, agent strategies, memory optimization, and interactive AI decision-making. ## What You'll Learn - How to model domain-specific data structures for complex games - Creating custom tools that agents can use to interact with the environment - Implementing efficient agent strategies with memory management - Building interactive AI systems with choice selection capabilities - Optimizing agent performance for turn-based games ## Setup First, let's import the Koog framework and set up our development environment: ``` %useLatestDescriptors %use koog ``` ## Modeling the Chess Domain Creating a robust domain model is essential for any game AI. In chess, we need to represent players, pieces, and their relationships. Let's start by defining our core data structures: ### Core Enums and Types ``` enum class Player { White, Black, None; fun opponent(): Player = when (this) { White -> Black Black -> White None -> throw IllegalArgumentException("No opponent for None player") } } enum class PieceType(val id: Char) { King('K'), Queen('Q'), Rook('R'), Bishop('B'), Knight('N'), Pawn('P'), None('*'); companion object { fun fromId(id: String): PieceType { require(id.length == 1) { "Invalid piece id: $id" } return entries.first { it.id == id.single() } } } } enum class Side { King, Queen } ``` The `Player` enum represents the two sides in chess, with an `opponent()` method for easy switching between players. The `PieceType` enum maps each chess piece to its standard notation character, enabling easy parsing of chess moves. The `Side` enum helps distinguish between kingside and queenside castling moves. ### Piece and Position Modeling ``` data class Piece(val pieceType: PieceType, val player: Player) { init { require((pieceType == PieceType.None) == (player == Player.None)) { "Invalid piece: $pieceType $player" } } fun toChar(): Char = when (player) { Player.White -> pieceType.id.uppercaseChar() Player.Black -> pieceType.id.lowercaseChar() Player.None -> pieceType.id } fun isNone(): Boolean = pieceType == PieceType.None companion object { val None = Piece(PieceType.None, Player.None) } } data class Position(val row: Int, val col: Char) { init { require(row in 1..8 && col in 'a'..'h') { "Invalid position: $col$row" } } constructor(position: String) : this( position[1].digitToIntOrNull() ?: throw IllegalArgumentException("Incorrect position: $position"), position[0], ) { require(position.length == 2) { "Invalid position: $position" } } } class ChessBoard { private val backRow = listOf( PieceType.Rook, PieceType.Knight, PieceType.Bishop, PieceType.Queen, PieceType.King, PieceType.Bishop, PieceType.Knight, PieceType.Rook ) private val board: List> = listOf( backRow.map { Piece(it, Player.Black) }.toMutableList(), List(8) { Piece(PieceType.Pawn, Player.Black) }.toMutableList(), List(8) { Piece.None }.toMutableList(), List(8) { Piece.None }.toMutableList(), List(8) { Piece.None }.toMutableList(), List(8) { Piece.None }.toMutableList(), List(8) { Piece(PieceType.Pawn, Player.White) }.toMutableList(), backRow.map { Piece(it, Player.White) }.toMutableList() ) override fun toString(): String = board .withIndex().joinToString("\n") { (index, row) -> "${8 - index} ${row.map { it.toChar() }.joinToString(" ")}" } + "\n a b c d e f g h" fun getPiece(position: Position): Piece = board[8 - position.row][position.col - 'a'] fun setPiece(position: Position, piece: Piece) { board[8 - position.row][position.col - 'a'] = piece } } ``` The `Piece` data class combines a piece type with its owner, using uppercase letters for white pieces and lowercase for black pieces in the visual representation. The `Position` class encapsulates chess coordinates (e.g., "e4") with built-in validation. ## Game State Management ### ChessBoard Implementation The `ChessBoard` class manages the 8×8 grid and piece positions. Key design decisions include: - **Internal Representation**: Uses a list of mutable lists for efficient access and modification - **Visual Display**: The `toString()` method provides a clear ASCII representation with rank numbers and file letters - **Position Mapping**: Converts between chess notation (a1-h8) and internal array indices ### ChessGame Logic ``` /** * Simple chess game without checks for valid moves. * Stores a correct state of the board if the entered moves are valid */ class ChessGame { private val board: ChessBoard = ChessBoard() private var currentPlayer: Player = Player.White val moveNotation: String = """ 0-0 - short castle 0-0-0 - long castle -- - usual move. e.g. p-e2-e4 --- - promotion move. e.g. p-e7-e8-q. Piece names: p - pawn n - knight b - bishop r - rook q - queen k - king """.trimIndent() fun move(move: String) { when { move == "0-0" -> castleMove(Side.King) move == "0-0-0" -> castleMove(Side.Queen) move.split("-").size == 3 -> { val (_, from, to) = move.split("-") usualMove(Position(from), Position(to)) } move.split("-").size == 4 -> { val (piece, from, to, promotion) = move.split("-") require(PieceType.fromId(piece) == PieceType.Pawn) { "Only pawn can be promoted" } usualMove(Position(from), Position(to)) board.setPiece(Position(to), Piece(PieceType.fromId(promotion), currentPlayer)) } else -> throw IllegalArgumentException("Invalid move: $move") } updateCurrentPlayer() } fun getBoard(): String = board.toString() fun currentPlayer(): String = currentPlayer.name.lowercase() private fun updateCurrentPlayer() { currentPlayer = currentPlayer.opponent() } private fun usualMove(from: Position, to: Position) { if (board.getPiece(from).pieceType == PieceType.Pawn && from.col != to.col && board.getPiece(to).isNone()) { // the move is en passant board.setPiece(Position(from.row, to.col), Piece.None) } movePiece(from, to) } private fun castleMove(side: Side) { val row = if (currentPlayer == Player.White) 1 else 8 val kingFrom = Position(row, 'e') val (rookFrom, kingTo, rookTo) = if (side == Side.King) { Triple(Position(row, 'h'), Position(row, 'g'), Position(row, 'f')) } else { Triple(Position(row, 'a'), Position(row, 'c'), Position(row, 'd')) } movePiece(kingFrom, kingTo) movePiece(rookFrom, rookTo) } private fun movePiece(from: Position, to: Position) { board.setPiece(to, board.getPiece(from)) board.setPiece(from, Piece.None) } } ``` The `ChessGame` class orchestrates the game logic and maintains state. Notable features include: - **Move Notation Support**: Accepts standard chess notation for regular moves, castling (0-0, 0-0-0), and pawn promotion - **Special Move Handling**: Implements en passant capture and castling logic - **Turn Management**: Automatically alternates between players after each move - **Validation**: While it doesn't validate move legality (trusting the AI to make valid moves), it handles move parsing and state updates correctly The `moveNotation` string provides clear documentation for the AI agent on acceptable move formats. ## Integrating with Koog Framework ### Creating Custom Tools ``` import kotlinx.serialization.Serializable class Move(val game: ChessGame) : SimpleTool( argsSerializer = Args.serializer(), descriptor = ToolDescriptor( name = "move", description = "Moves a piece according to the notation:\n${game.moveNotation}", requiredParameters = listOf( ToolParameterDescriptor( name = "notation", description = "The notation of the piece to move", type = ToolParameterType.String, ) ) ) ) { @Serializable data class Args(val notation: String) : ToolArgs override suspend fun execute(args: Args): String { game.move(args.notation) println(game.getBoard()) println("-----------------") return "Current state of the game:\n${game.getBoard()}\n${game.currentPlayer()} to move! Make the move!" } } ``` The `Move` tool demonstrates the Koog framework's tool integration pattern: 1. **Extends SimpleTool**: Inherits the basic tool functionality with type-safe argument handling 1. **Serializable Arguments**: Uses Kotlin serialization to define the tool's input parameters 1. **Rich Documentation**: The `ToolDescriptor` provides the LLM with detailed information about the tool's purpose and parameters 1. **Constructor Parameters**: Passes `argsSerializer` and `descriptor` to the constructor 1. **Execution Logic**: The `execute` method handles the actual move execution and provides formatted feedback Key design aspects: - **Context Injection**: The tool receives the `ChessGame` instance, allowing it to modify game state - **Feedback Loop**: Returns the current board state and prompts the next player, maintaining conversational flow - **Error Handling**: Relies on the game class for move validation and error reporting ## Agent Strategy Design ### Memory Optimization Technique ``` import ai.koog.agents.core.environment.ReceivedToolResult /** * Chess position is (almost) completely defined by the board state, * So we can trim the history of the LLM to only contain the system prompt and the last move. */ inline fun AIAgentSubgraphBuilderBase<*, *>.nodeTrimHistory( name: String? = null ): AIAgentNodeDelegate = node(name) { result -> llm.writeSession { rewritePrompt { prompt -> val messages = prompt.messages prompt.copy(messages = listOf(messages.first(), messages.last())) } } result } val strategy = strategy("chess_strategy") { val nodeCallLLM by nodeLLMRequest("sendInput") val nodeExecuteTool by nodeExecuteTool("nodeExecuteTool") val nodeSendToolResult by nodeLLMSendToolResult("nodeSendToolResult") val nodeTrimHistory by nodeTrimHistory() edge(nodeStart forwardTo nodeCallLLM) edge(nodeCallLLM forwardTo nodeExecuteTool onToolCall { true }) edge(nodeCallLLM forwardTo nodeFinish onAssistantMessage { true }) edge(nodeExecuteTool forwardTo nodeTrimHistory) edge(nodeTrimHistory forwardTo nodeSendToolResult) edge(nodeSendToolResult forwardTo nodeFinish onAssistantMessage { true }) edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCall { true }) } ``` The `nodeTrimHistory` function implements a crucial optimization for chess games. Since chess positions are largely determined by the current board state rather than the full move history, we can significantly reduce token usage by keeping only: 1. **System Prompt**: Contains the agent's core instructions and behavior guidelines 1. **Latest Message**: The most recent board state and game context This approach: - **Reduces Token Consumption**: Prevents exponential growth of conversation history - **Maintains Context**: Preserves essential game state information - **Improves Performance**: Faster processing with shorter prompts - **Enables Long Games**: Allows for extended gameplay without hitting token limits The chess strategy demonstrates Koog's graph-based agent architecture: **Node Types:** - `nodeCallLLM`: Processes input and generates responses/tool calls - `nodeExecuteTool`: Executes the Move tool with the provided parameters - `nodeTrimHistory`: Optimizes conversation memory as described above - `nodeSendToolResult`: Sends tool execution results back to the LLM **Control Flow:** - **Linear Path**: Start → LLM Request → Tool Execution → History Trim → Send Result - **Decision Points**: LLM responses can either finish the conversation or trigger another tool call - **Memory Management**: History trimming occurs after each tool execution This strategy ensures efficient, stateful gameplay while maintaining conversational coherence. ### Setting up the AI Agent ``` val baseExecutor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")) ``` This section initializes our OpenAI executor. The `simpleOpenAIExecutor` creates a connection to OpenAI's API using your API key from environment variables. **Configuration Notes:** - Store your OpenAI API key in the `OPENAI_API_KEY` environment variable - The executor handles authentication and API communication automatically - Different executor types are available for various LLM providers ### Agent Assembly ``` val game = ChessGame() val toolRegistry = ToolRegistry { tools(listOf(Move(game))) } // Create a chat agent with a system prompt and the tool registry val agent = AIAgent( executor = baseExecutor, strategy = strategy, llmModel = OpenAIModels.Chat.O3Mini, systemPrompt = """ You are an agent who plays chess. You should always propose a move in response to the "Your move!" message. DO NOT HALLUCINATE!!! DO NOT PLAY ILLEGAL MOVES!!! YOU CAN SEND A MESSAGE ONLY IF IT IS A RESIGNATION OR A CHECKMATE!!! """.trimMargin(), temperature = 0.0, toolRegistry = toolRegistry, maxIterations = 200, ) ``` Here we assemble all components into a functional chess-playing agent: **Key Configuration:** - **Model Choice**: Using `OpenAIModels.Chat.O3Mini` for high-quality chess play - **Temperature**: Set to 0.0 for deterministic, strategic moves - **System Prompt**: Carefully crafted instructions emphasizing legal moves and proper behavior - **Tool Registry**: Provides the agent access to the Move tool - **Max Iterations**: Set to 200 to allow for complete games **System Prompt Design:** - Emphasizes move proposal responsibility - Prohibits hallucination and illegal moves - Restricts messaging to only resignations or checkmate declarations - Creates focused, game-oriented behavior ### Running the Basic Agent ``` import kotlinx.coroutines.runBlocking println("Chess Game started!") val initialMessage = "Starting position is ${game.getBoard()}. White to move!" runBlocking { agent.run(initialMessage) } ``` ``` Chess Game started! 8 r n b q k b n r 7 p p p p p p p p 6 * * * * * * * * 5 * * * * * * * * 4 * * * * P * * * 3 * * * * * * * * 2 P P P P * P P P 1 R N B Q K B N R a b c d e f g h ----------------- 8 r n b q k b n r 7 p p p p * p p p 6 * * * * * * * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * * * * 2 P P P P * P P P 1 R N B Q K B N R a b c d e f g h ----------------- 8 r n b q k b n r 7 p p p p * p p p 6 * * * * * * * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * N * * 2 P P P P * P P P 1 R N B Q K B * R a b c d e f g h ----------------- 8 r n b q k b * r 7 p p p p * p p p 6 * * * * * n * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * N * * 2 P P P P * P P P 1 R N B Q K B * R a b c d e f g h ----------------- 8 r n b q k b * r 7 p p p p * p p p 6 * * * * * n * * 5 * * * * p * * * 4 * * * * P * * * 3 * * N * * N * * 2 P P P P * P P P 1 R * B Q K B * R a b c d e f g h ----------------- The execution was interrupted ``` This basic agent plays autonomously, making moves automatically. The game output shows the sequence of moves and board states as the AI plays against itself. ## Advanced Feature: Interactive Choice Selection The next sections demonstrate a more sophisticated approach where users can participate in the AI's decision-making process by choosing from multiple AI-generated moves. ### Custom Choice Selection Strategy ``` import ai.koog.agents.core.feature.choice.ChoiceSelectionStrategy /** * `AskUserChoiceStrategy` allows users to interactively select a choice from a list of options * presented by a language model. The strategy uses customizable methods to display the prompt * and choices and read user input to determine the selected choice. * * @property promptShowToUser A function that formats and displays a given `Prompt` to the user. * @property choiceShowToUser A function that formats and represents a given `LLMChoice` to the user. * @property print A function responsible for displaying messages to the user, e.g., for showing prompts or feedback. * @property read A function to capture user input. */ class AskUserChoiceSelectionStrategy( private val promptShowToUser: (Prompt) -> String = { "Current prompt: $it" }, private val choiceShowToUser: (LLMChoice) -> String = { "$it" }, private val print: (String) -> Unit = ::println, private val read: () -> String? = ::readlnOrNull ) : ChoiceSelectionStrategy { override suspend fun choose(prompt: Prompt, choices: List): LLMChoice { print(promptShowToUser(prompt)) print("Available LLM choices") choices.withIndex().forEach { (index, choice) -> print("Choice number ${index + 1}: ${choiceShowToUser(choice)}") } var choiceNumber = ask(choices.size) while (choiceNumber == null) { print("Invalid response.") choiceNumber = ask(choices.size) } return choices[choiceNumber - 1] } private fun ask(numChoices: Int): Int? { print("Please choose a choice. Enter a number between 1 and $numChoices: ") return read()?.toIntOrNull()?.takeIf { it in 1..numChoices } } } ``` The `AskUserChoiceSelectionStrategy` implements Koog's `ChoiceSelectionStrategy` interface to enable human participation in AI decision-making: **Key Features:** - **Customizable Display**: Functions for formatting prompts and choices - **Interactive Input**: Uses standard input/output for user interaction - **Validation**: Ensures user input is within valid range - **Flexible I/O**: Configurable print and read functions for different environments **Use Cases:** - Human-AI collaboration in gameplay - AI decision transparency and explainability - Training and debugging scenarios - Educational demonstrations ### Enhanced Strategy with Choice Selection ``` inline fun AIAgentSubgraphBuilderBase<*, *>.nodeTrimHistory( name: String? = null ): AIAgentNodeDelegate = node(name) { result -> llm.writeSession { rewritePrompt { prompt -> val messages = prompt.messages prompt.copy(messages = listOf(messages.first(), messages.last())) } } result } val strategy = strategy("chess_strategy") { val nodeCallLLM by nodeLLMRequest("sendInput") val nodeExecuteTool by nodeExecuteTool("nodeExecuteTool") val nodeSendToolResult by nodeLLMSendToolResult("nodeSendToolResult") val nodeTrimHistory by nodeTrimHistory() edge(nodeStart forwardTo nodeCallLLM) edge(nodeCallLLM forwardTo nodeExecuteTool onToolCall { true }) edge(nodeCallLLM forwardTo nodeFinish onAssistantMessage { true }) edge(nodeExecuteTool forwardTo nodeTrimHistory) edge(nodeTrimHistory forwardTo nodeSendToolResult) edge(nodeSendToolResult forwardTo nodeFinish onAssistantMessage { true }) edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCall { true }) } val askChoiceStrategy = AskUserChoiceSelectionStrategy(promptShowToUser = { prompt -> val lastMessage = prompt.messages.last() if (lastMessage is MessagePart.Tool.Call) { lastMessage.content } else { "" } }) ``` ``` val promptExecutor = PromptExecutorWithChoiceSelection(baseExecutor, askChoiceStrategy) ``` The first interactive approach uses `PromptExecutorWithChoiceSelection`, which wraps the base executor with choice selection capability. The custom display function extracts move information from tool calls to show users what the AI wants to do. **Architecture Changes:** - **Wrapped Executor**: `PromptExecutorWithChoiceSelection` adds choice functionality to any base executor - **Context-Aware Display**: Shows the last tool call content instead of the full prompt - **Higher Temperature**: Increased to 1.0 for more diverse move options ### Advanced Strategy: Manual Choice Selection ``` val game = ChessGame() val toolRegistry = ToolRegistry { tools(listOf(Move(game))) } val agent = AIAgent( executor = promptExecutor, strategy = strategy, llmModel = OpenAIModels.Chat.O3Mini, systemPrompt = """ You are an agent who plays chess. You should always propose a move in response to the "Your move!" message. DO NOT HALLUCINATE!!! DO NOT PLAY ILLEGAL MOVES!!! YOU CAN SEND A MESSAGE ONLY IF IT IS A RESIGNATION OR A CHECKMATE!!! """.trimMargin(), temperature = 1.0, toolRegistry = toolRegistry, maxIterations = 200, numberOfChoices = 3, ) ``` The advanced strategy integrates choice selection directly into the agent's execution graph: **New Nodes:** - `nodeLLMSendResultsMultipleChoices`: Handles multiple LLM choices simultaneously - `nodeSelectLLMChoice`: Integrates the choice selection strategy into the workflow **Enhanced Control Flow:** - Tool results are wrapped in lists to support multiple choices - User selection occurs before continuing with the chosen path - The selected choice is unwrapped and continues through the normal flow **Benefits:** - **Greater Control**: Fine-grained integration with agent workflow - **Flexibility**: Can be combined with other agent features - **Transparency**: Users see exactly what the AI is considering ### Running Interactive Agents ``` println("Chess Game started!") val initialMessage = "Starting position is ${game.getBoard()}. White to move!" runBlocking { agent.run(initialMessage) } ``` ``` Chess Game started! Available LLM choices Choice number 1: [Call(id=call_K46Upz7XoBIG5RchDh7bZE8F, tool=move, content={"notation": "p-e2-e4"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:40.368252Z, totalTokensCount=773, inputTokensCount=315, outputTokensCount=458, additionalInfo={}))] Choice number 2: [Call(id=call_zJ6OhoCHrVHUNnKaxZkOhwoU, tool=move, content={"notation": "p-e2-e4"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:40.368252Z, totalTokensCount=773, inputTokensCount=315, outputTokensCount=458, additionalInfo={}))] Choice number 3: [Call(id=call_nwX6ZMJ3F5AxiNUypYlI4BH4, tool=move, content={"notation": "p-e2-e4"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:40.368252Z, totalTokensCount=773, inputTokensCount=315, outputTokensCount=458, additionalInfo={}))] Please choose a choice. Enter a number between 1 and 3: 8 r n b q k b n r 7 p p p p p p p p 6 * * * * * * * * 5 * * * * * * * * 4 * * * * P * * * 3 * * * * * * * * 2 P P P P * P P P 1 R N B Q K B N R a b c d e f g h ----------------- Available LLM choices Choice number 1: [Call(id=call_2V93GXOcIe0fAjUAIFEk9h5S, tool=move, content={"notation": "p-e7-e5"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:47.949303Z, totalTokensCount=1301, inputTokensCount=341, outputTokensCount=960, additionalInfo={}))] Choice number 2: [Call(id=call_INM59xRzKMFC1w8UAV74l9e1, tool=move, content={"notation": "p-e7-e5"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:47.949303Z, totalTokensCount=1301, inputTokensCount=341, outputTokensCount=960, additionalInfo={}))] Choice number 3: [Call(id=call_r4QoiTwn0F3jizepHH5ia8BU, tool=move, content={"notation": "p-e7-e5"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:47.949303Z, totalTokensCount=1301, inputTokensCount=341, outputTokensCount=960, additionalInfo={}))] Please choose a choice. Enter a number between 1 and 3: 8 r n b q k b n r 7 p p p p * p p p 6 * * * * * * * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * * * * 2 P P P P * P P P 1 R N B Q K B N R a b c d e f g h ----------------- Available LLM choices Choice number 1: [Call(id=call_f9XTizn41svcrtvnmkCfpSUQ, tool=move, content={"notation": "n-g1-f3"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:55.467712Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))] Choice number 2: [Call(id=call_c0Dfce5RcSbN3cOOm5ESYriK, tool=move, content={"notation": "n-g1-f3"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:55.467712Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))] Choice number 3: [Call(id=call_Lr4Mdro1iolh0fDyAwZsutrW, tool=move, content={"notation": "n-g1-f3"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:55.467712Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))] Please choose a choice. Enter a number between 1 and 3: 8 r n b q k b n r 7 p p p p * p p p 6 * * * * * * * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * N * * 2 P P P P * P P P 1 R N B Q K B * R a b c d e f g h ----------------- The execution was interrupted ``` ``` import ai.koog.agents.core.feature.choice.nodeLLMSendResultsMultipleChoices import ai.koog.agents.core.feature.choice.nodeSelectLLMChoice inline fun AIAgentSubgraphBuilderBase<*, *>.nodeTrimHistory( name: String? = null ): AIAgentNodeDelegate = node(name) { result -> llm.writeSession { rewritePrompt { prompt -> val messages = prompt.messages prompt.copy(messages = listOf(messages.first(), messages.last())) } } result } val strategy = strategy("chess_strategy") { val nodeCallLLM by nodeLLMRequest("sendInput") val nodeExecuteTool by nodeExecuteTool("nodeExecuteTool") val nodeSendToolResult by nodeLLMSendResultsMultipleChoices("nodeSendToolResult") val nodeSelectLLMChoice by nodeSelectLLMChoice(askChoiceStrategy, "chooseLLMChoice") val nodeTrimHistory by nodeTrimHistory() edge(nodeStart forwardTo nodeCallLLM) edge(nodeCallLLM forwardTo nodeExecuteTool onToolCall { true }) edge(nodeCallLLM forwardTo nodeFinish onAssistantMessage { true }) edge(nodeExecuteTool forwardTo nodeTrimHistory) edge(nodeTrimHistory forwardTo nodeSendToolResult transformed { listOf(it) }) edge(nodeSendToolResult forwardTo nodeSelectLLMChoice) edge(nodeSelectLLMChoice forwardTo nodeFinish transformed { it.first() } onAssistantMessage { true }) edge(nodeSelectLLMChoice forwardTo nodeExecuteTool transformed { it.first() } onToolCall { true }) } ``` ``` val game = ChessGame() val toolRegistry = ToolRegistry { tools(listOf(Move(game))) } val agent = AIAgent( executor = baseExecutor, strategy = strategy, llmModel = OpenAIModels.Chat.O3Mini, systemPrompt = """ You are an agent who plays chess. You should always propose a move in response to the "Your move!" message. DO NOT HALLUCINATE!!! DO NOT PLAY ILLEGAL MOVES!!! YOU CAN SEND A MESSAGE ONLY IF IT IS A RESIGNATION OR A CHECKMATE!!! """.trimMargin(), temperature = 1.0, toolRegistry = toolRegistry, maxIterations = 200, numberOfChoices = 3, ) ``` ``` println("Chess Game started!") val initialMessage = "Starting position is ${game.getBoard()}. White to move!" runBlocking { agent.run(initialMessage) } ``` ``` Chess Game started! 8 r n b q k b n r 7 p p p p p p p p 6 * * * * * * * * 5 * * * * * * * * 4 * * * * P * * * 3 * * * * * * * * 2 P P P P * P P P 1 R N B Q K B N R a b c d e f g h ----------------- Available LLM choices Choice number 1: [Call(id=call_gqMIar0z11CyUl5nup3zbutj, tool=move, content={"notation": "p-e7-e5"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:17.313548Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))] Choice number 2: [Call(id=call_6niUGnZPPJILRFODIlJsCKax, tool=move, content={"notation": "p-e7-e5"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:17.313548Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))] Choice number 3: [Call(id=call_q1b8ZmIBph0EoVaU3Ic9A09j, tool=move, content={"notation": "p-e7-e5"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:17.313548Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))] Please choose a choice. Enter a number between 1 and 3: 8 r n b q k b n r 7 p p p p * p p p 6 * * * * * * * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * * * * 2 P P P P * P P P 1 R N B Q K B N R a b c d e f g h ----------------- Available LLM choices Choice number 1: [Call(id=call_pdBIX7MVi82MyWwawTm1Q2ef, tool=move, content={"notation": "n-g1-f3"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:24.505344Z, totalTokensCount=1237, inputTokensCount=341, outputTokensCount=896, additionalInfo={}))] Choice number 2: [Call(id=call_oygsPHaiAW5OM6pxhXhtazgp, tool=move, content={"notation": "n-g1-f3"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:24.505344Z, totalTokensCount=1237, inputTokensCount=341, outputTokensCount=896, additionalInfo={}))] Choice number 3: [Call(id=call_GJTEsZ8J8cqOKZW4Tx54RqCh, tool=move, content={"notation": "n-g1-f3"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:24.505344Z, totalTokensCount=1237, inputTokensCount=341, outputTokensCount=896, additionalInfo={}))] Please choose a choice. Enter a number between 1 and 3: 8 r n b q k b n r 7 p p p p * p p p 6 * * * * * * * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * N * * 2 P P P P * P P P 1 R N B Q K B * R a b c d e f g h ----------------- Available LLM choices Choice number 1: [Call(id=call_5C7HdlTU4n3KdXcyNogE4rGb, tool=move, content={"notation": "n-g8-f6"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:34.646667Z, totalTokensCount=1621, inputTokensCount=341, outputTokensCount=1280, additionalInfo={}))] Choice number 2: [Call(id=call_EjCcyeMLQ88wMa5yh3vmeJ2w, tool=move, content={"notation": "n-g8-f6"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:34.646667Z, totalTokensCount=1621, inputTokensCount=341, outputTokensCount=1280, additionalInfo={}))] Choice number 3: [Call(id=call_NBMMSwmFIa8M6zvfbPw85NKh, tool=move, content={"notation": "n-g8-f6"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:34.646667Z, totalTokensCount=1621, inputTokensCount=341, outputTokensCount=1280, additionalInfo={}))] Please choose a choice. Enter a number between 1 and 3: 8 r n b q k b * r 7 p p p p * p p p 6 * * * * * n * * 5 * * * * p * * * 4 * * * * P * * * 3 * * * * * N * * 2 P P P P * P P P 1 R N B Q K B * R a b c d e f g h ----------------- The execution was interrupted ``` The interactive examples show how users can guide the AI's decision-making process. In the output, you can see: 1. **Multiple Choices**: The AI generates 3 different move options 1. **User Selection**: Users input numbers 1-3 to choose their preferred move 1. **Game Continuation**: The selected move is executed and the game continues ## Conclusion This tutorial demonstrates several key aspects of building intelligent agents with the Koog framework: ### Key Takeaways 1. **Domain Modeling**: Well-structured data models are crucial for complex applications 1. **Tool Integration**: Custom tools enable agents to interact with external systems effectively 1. **Memory Management**: Strategic history trimming optimizes performance for long interactions 1. **Strategy Graphs**: Koog's graph-based approach provides flexible control flow 1. **Interactive AI**: Choice selection enables human-AI collaboration and transparency ### Framework Features Explored - ✅ Custom tool creation and integration - ✅ Agent strategy design and graph-based control flow - ✅ Memory optimization techniques - ✅ Interactive choice selection - ✅ Multiple LLM response handling - ✅ Stateful game management The Koog framework provides the foundation for building sophisticated AI agents that can handle complex, multi-turn interactions while maintaining efficiency and transparency. # Google Maps MCP with Koog: From Zero to Elevation in a Kotlin Notebook [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/GoogleMapsMcp.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/GoogleMapsMcp.ipynb) In this short, blog-style walkthrough, we’ll connect Koog to a Model Context Protocol (MCP) server for Google Maps. We’ll spin up the server with Docker, discover the available tools, and let an AI agent geocode an address and fetch its elevation — all from a Kotlin Notebook. By the end, you’ll have a reproducible, end‑to‑end example you can drop into your workflow or documentation. ``` %useLatestDescriptors %use koog ``` ## Prerequisites Before you run the cells below, make sure you have: - Docker installed and running - A valid Google Maps API key exported as an environment variable: `GOOGLE_MAPS_API_KEY` - An OpenAI API key exported as `OPENAI_API_KEY` You can set them in your shell like this (macOS/Linux example): ``` export GOOGLE_MAPS_API_KEY="" export OPENAI_API_KEY="" ``` ``` // Get the API key from environment variables val googleMapsApiKey = System.getenv("GOOGLE_MAPS_API_KEY") ?: error("GOOGLE_MAPS_API_KEY environment variable not set") val openAIApiToken = System.getenv("OPENAI_API_KEY") ?: error("OPENAI_API_KEY environment variable not set") ``` ## Start the Google Maps MCP server (Docker) We’ll use the official `mcp/google-maps` image. The container will expose tools such as `maps_geocode` and `maps_elevation` over MCP. We pass the API key via environment variables and launch it attached so the notebook can talk to it over stdio. ``` // Start the Docker container with the Google Maps MCP server val process = ProcessBuilder( "docker", "run", "-i", "-e", "GOOGLE_MAPS_API_KEY=$googleMapsApiKey", "mcp/google-maps" ).start() ``` ## Discover tools via McpToolRegistry Koog can connect to an MCP server over stdio. Here, we create a tool registry from the running process and print out the discovered tools and their descriptors. ``` val toolRegistry = McpToolRegistryProvider.fromTransport( transport = McpToolRegistryProvider.defaultStdioTransport(process) ) toolRegistry.tools.forEach { println(it.name) println(it.descriptor) } ``` ## Build an AI Agent with OpenAI Next we assemble a simple agent backed by the OpenAI executor and model. The agent will be able to call tools exposed by the MCP server through the registry we just created. ``` val agent = AIAgent( executor = simpleOpenAIExecutor(openAIApiToken), llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = toolRegistry, ) ``` ## Ask for elevation: geocode first, then elevation We prompt the agent to find the elevation of the JetBrains office in Munich. The instruction explicitly tells the agent to use only the available tools and which ones to prefer for the task. ``` import kotlinx.coroutines.runBlocking val request = "Get elevation of the Jetbrains Office in Munich, Germany?" runBlocking { agent.run( request + "You can only call tools. Get it by calling maps_geocode and maps_elevation tools." ) } ``` ## Clean up When you’re done, stop the Docker process so you don’t leave anything running in the background. ``` process.destroy() ``` ## Troubleshooting and next steps - If the container fails to start, check that Docker is running and your `GOOGLE_MAPS_API_KEY` is valid. - If the agent can’t call tools, re-run the discovery cell to ensure the tool registry is populated. - Try other prompts like route planning or place searches using the available Google Maps tools. Next, consider composing multiple MCP servers (e.g., Playwright for web automation + Google Maps) and let Koog orchestrate tool usage for richer tasks. # Building a Number‑Guessing Agent with Koog [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/Guesser.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/Guesser.ipynb) Let’s build a small but fun agent that guesses a number you’re thinking of. We’ll lean on Koog’s tool-calling to ask targeted questions and converge using a classic binary search strategy. The result is an idiomatic Kotlin Notebook that you can drop straight into docs. We’ll keep the code minimal and the flow transparent: a few tiny tools, a compact prompt, and an interactive CLI loop. ## Setup This notebook assumes: - You’re running in a Kotlin Notebook with Koog available. - The environment variable `OPENAI_API_KEY` is set. The agent uses it via `simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY"))`. Load the Koog kernel: ``` %useLatestDescriptors %use koog ``` ## Tools: asking targeted questions Tools are small, well-described functions the LLM can call. We’ll provide three: - `lessThan(value)`: “Is your number less than value?” - `greaterThan(value)`: “Is your number greater than value?” - `proposeNumber(value)`: “Is your number equal to value?” (used once the range is tight) Each tool returns a simple "YES"/"NO" string. The helper `ask` implements a minimal Y/n loop and validates input. Descriptions via `@LLMDescription` help the model select tools correctly. ``` import ai.koog.agents.core.tools.annotations.Tool class GuesserTool : ToolSet { @Tool @LLMDescription("Asks the user if his number is STRICTLY less than a given value.") fun lessThan( @LLMDescription("A value to compare the guessed number with.") value: Int ): String = ask("Is your number less than $value?", value) @Tool @LLMDescription("Asks the user if his number is STRICTLY greater than a given value.") fun greaterThan( @LLMDescription("A value to compare the guessed number with.") value: Int ): String = ask("Is your number greater than $value?", value) @Tool @LLMDescription("Asks the user if his number is EXACTLY equal to the given number. Only use this tool once you've narrowed down your answer.") fun proposeNumber( @LLMDescription("A value to compare the guessed number with.") value: Int ): String = ask("Is your number equal to $value?", value) fun ask(question: String, value: Int): String { print("$question [Y/n]: ") val input = readln() println(input) return when (input.lowercase()) { "", "y", "yes" -> "YES" "n", "no" -> "NO" else -> { println("Invalid input! Please, try again.") ask(question, value) } } } } ``` ## Tool Registry Expose your tools to the agent. We also add a built‑in `SayToUser` tool so the agent can surface messages directly to the user. ``` val toolRegistry = ToolRegistry { tool(SayToUser) tools(GuesserTool()) } ``` ## Agent configuration A short, tool‑forward system prompt is all we need. We’ll suggest a binary search strategy and keep `temperature = 0.0` for stable, deterministic behavior. Here we use OpenAI’s reasoning model `GPT4oMini` for crisp planning. ``` val agent = AIAgent( executor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = """ You are a number guessing agent. Your goal is to guess a number that the user is thinking of. Follow these steps: 1. Start by asking the user to think of a number between 1 and 100. 2. Use the less_than and greater_than tools to narrow down the range. a. If it's neither greater nor smaller, use the propose_number tool. 3. Once you're confident about the number, use the propose_number tool to check if your guess is correct. 4. If your guess is correct, congratulate the user. If not, continue guessing. Be efficient with your guessing strategy. A binary search approach works well. """.trimIndent(), temperature = 0.0, toolRegistry = toolRegistry ) ``` ## Run it - Think of a number between 1 and 100. - Type `start` to begin. - Answer the agent’s questions with `Y`/`Enter` for yes or `n` for no. The agent should zero in on your number in ~7 steps. ``` import kotlinx.coroutines.runBlocking println("Number Guessing Game started!") println("Think of a number between 1 and 100, and I'll try to guess it.") println("Type 'start' to begin the game.") val initialMessage = readln() runBlocking { agent.run(initialMessage) } ``` ## How it works - The agent reads the system prompt and plans a binary search. - On each iteration it calls one of your tools: `lessThan`, `greaterThan`, or (when certain) `proposeNumber`. - The helper `ask` collects your Y/n input and returns a clean "YES"/"NO" signal back to the model. - When it gets confirmation, it congratulates you via `SayToUser`. ## Extend it - Change the range (e.g., 1..1000) by tweaking the system prompt. - Add a `between(low, high)` tool to reduce calls further. - Swap models or executors (e.g., use an Ollama executor and a local model) while keeping the same tools. - Persist guesses or outcomes to a store for analytics. ## Troubleshooting - Missing key: ensure `OPENAI_API_KEY` is set in your environment. - Kernel not found: make sure `%useLatestDescriptors` and `%use koog` executed successfully. - Tool not called: confirm the `ToolRegistry` includes `GuesserTool()` and the names in the prompt match your tool functions. # Tracing Koog Agents to Langfuse with OpenTelemetry [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/Langfuse.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/Langfuse.ipynb) This notebook shows how to export Koog agent traces to your Langfuse instance using OpenTelemetry. You'll set up environment variables, run a simple agent, and then inspect spans and traces in Langfuse. ## What you'll learn - How Koog integrates with OpenTelemetry to emit traces - How to configure the Langfuse exporter via environment variables - How to run an agent and view its trace in Langfuse ## Prerequisites - A Langfuse project (host URL, public key, secret key) - An OpenAI API key for the LLM executor - Environment variables set in your shell: ``` export OPENAI_API_KEY=sk-... export LANGFUSE_HOST=https://cloud.langfuse.com # or your self-hosted URL export LANGFUSE_PUBLIC_KEY=pk_... export LANGFUSE_SECRET_KEY=sk_... ``` ``` %useLatestDescriptors //%use koog ``` ``` import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.features.opentelemetry.feature.OpenTelemetry import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor /** * Example of Koog agents tracing to [Langfuse](https://langfuse.com/) * * Agent traces are exported to: * - Langfuse OTLP endpoint instance using [OtlpHttpSpanExporter] * * To run this example: * 1. Set up a Langfuse project and credentials as described [here](https://langfuse.com/docs/get-started#create-new-project-in-langfuse) * 2. Get Langfuse credentials as described [here](https://langfuse.com/faq/all/where-are-langfuse-api-keys) * 3. Set `LANGFUSE_HOST`, `LANGFUSE_PUBLIC_KEY`, and `LANGFUSE_SECRET_KEY` environment variables * * @see Langfuse OpenTelemetry Docs */ val agent = AIAgent( executor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a code assistant. Provide concise code examples." ) { install(OpenTelemetry) { addLangfuseExporter() } } ``` ## Configure the agent and Langfuse exporter In the next cell, we: - Create an AIAgent that uses OpenAI as the LLM executor - Install the OpenTelemetry feature and add the Langfuse exporter - Rely on environment variables for Langfuse configuration Under the hood, Koog emits spans for agent lifecycle, LLM calls, and tool execution (if any). The Langfuse exporter ships those spans to your Langfuse instance via the OpenTelemetry endpoint. ``` import kotlinx.coroutines.runBlocking println("Running agent with Langfuse tracing") runBlocking { val result = agent.run("Tell me a joke about programming") "Result: $result\nSee traces on the Langfuse instance" } ``` ## Run the agent and view traces Execute the next cell to trigger a simple prompt. This will generate spans that are exported to your Langfuse project. ### Where to look in Langfuse 1. Open your Langfuse dashboard and select your project 1. Navigate to the Traces/Spans view 1. Look for recent entries around the time you ran this cell 1. Drill down into spans to see: 1. Agent lifecycle events 1. LLM request/response metadata 1. Errors (if any) ### Troubleshooting - No traces showing up? - Double-check LANGFUSE_HOST, LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY - Ensure your network allows outbound HTTPS to the Langfuse endpoint - Verify your Langfuse project is active and keys belong to the correct project - Authentication errors - Regenerate keys in Langfuse and update env vars - OpenAI issues - Confirm OPENAI_API_KEY is set and valid # OpenTelemetry with Koog: Tracing your AI agent [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/OpenTelemetry.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/OpenTelemetry.ipynb) This notebook demonstrates how to add OpenTelemetry-based tracing to a Koog AI agent. We will: - Emit spans to the console for quick local debugging. - Export spans to an OpenTelemetry Collector and view them in Jaeger. Prerequisites: - Docker/Docker Compose installed - An OpenAI API key available in environment variable `OPENAI_API_KEY` Start the local OpenTelemetry stack (Collector + Jaeger) before running the notebook: ``` ./docker-compose up -d ``` After the agent runs, open Jaeger UI: - http://localhost:16686 To stop the services later: ``` docker-compose down ``` ______________________________________________________________________ ``` %useLatestDescriptors // %use koog ``` ``` import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.features.opentelemetry.feature.OpenTelemetry import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor import io.opentelemetry.exporter.logging.LoggingSpanExporter import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter ``` ## Configure OpenTelemetry exporters In the next cell, we: - Create a Koog AIAgent - Install the OpenTelemetry feature - Add two span exporters: - LoggingSpanExporter for console logs - OTLP gRPC exporter to http://localhost:4317 (Collector) This mirrors the example description: console logs for local debugging and OTLP for viewing traces in Jaeger. ``` val agent = AIAgent( executor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a code assistant. Provide concise code examples." ) { install(OpenTelemetry) { // Add a console logger for local debugging addSpanExporter(LoggingSpanExporter.create()) // Send traces to OpenTelemetry collector addSpanExporter( OtlpGrpcSpanExporter.builder() .setEndpoint("http://localhost:4317") .build() ) } } ``` ## Run the agent and view traces in Jaeger Execute the next cell to trigger a simple prompt. You should see: - Console span logs from the LoggingSpanExporter - Traces exported to your local OpenTelemetry Collector and visible in Jaeger at http://localhost:16686 Tip: Use the Jaeger search to find recent traces after you run the cell. ``` import ai.koog.agents.utils.use import kotlinx.coroutines.runBlocking runBlocking { agent.use { agent -> println("Running agent with OpenTelemetry tracing...") val result = agent.run("Tell me a joke about programming") "Agent run completed with result: '$result'.\nCheck Jaeger UI at http://localhost:16686 to view traces" } } ``` ## Cleanup and troubleshooting When you're done: - Stop services: ``` docker-compose down ``` - If you don't see traces in Jaeger: - Ensure the stack is running: `./docker-compose up -d` and give it a few seconds to start. - Verify ports: - Collector (OTLP gRPC): http://localhost:4317 - Jaeger UI: http://localhost:16686 - Check container logs: `docker-compose logs --tail=200` - Confirm your `OPENAI_API_KEY` is set in the environment where the notebook runs. - Make sure the endpoint in the exporter matches the collector: `http://localhost:4317`. - What spans to expect: - Koog agent lifecycle - LLM request/response metadata - Any tool execution spans (if you add tools) You can now iterate on your agent and observe changes in your tracing pipeline. # Drive the browser with Playwright MCP and Koog [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/PlaywrightMcp.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/PlaywrightMcp.ipynb) In this notebook, you'll connect a Koog agent to Playwright's Model Context Protocol (MCP) server and let it drive a real browser to complete a task: open jetbrains.com, accept cookies, and click the AI section in the toolbar. We'll keep things simple and reproducible, focusing on a minimal but realistic agent + tools setup you can publish and reuse. ``` %useLatestDescriptors %use koog ``` ## Prerequisites - An OpenAI API key exported as an environment variable: `OPENAI_API_KEY` - Node.js and npx available on your PATH - Kotlin Jupyter notebook environment with Koog available via `%use koog` Tip: Run the Playwright MCP server in headful mode to watch the browser automate the steps. ## 1) Provide your OpenAI API key We read the API key from the `OPENAI_API_KEY` environment variable. This keeps secrets out of the notebook. ``` // Get the API key from environment variables val openAIApiToken = System.getenv("OPENAI_API_KEY") ?: error("OPENAI_API_KEY environment variable not set") ``` ## 2) Start the Playwright MCP server We'll launch Playwright's MCP server locally using `npx`. By default, it will expose an SSE endpoint we can connect to from Koog. ``` // Start the Playwright MCP server via npx val process = ProcessBuilder( "npx", "@playwright/mcp@latest", "--port", "8931" ).start() ``` ## 3) Connect from Koog and run the agent We build a minimal Koog `AIAgent` with an OpenAI executor and point its tool registry to the MCP server over SSE. Then we ask it to complete the browser task strictly via tools. ``` import kotlinx.coroutines.runBlocking runBlocking { println("Connecting to Playwright MCP server...") val toolRegistry = McpToolRegistryProvider.fromTransport( transport = McpToolRegistryProvider.defaultSseTransport("http://localhost:8931/sse") ) println("Successfully connected to Playwright MCP server") // Create the agent val agent = AIAgent( executor = simpleOpenAIExecutor(openAIApiToken), llmModel = OpenAIModels.Chat.GPT4o, toolRegistry = toolRegistry, ) val request = "Open a browser, navigate to jetbrains.com, accept all cookies, click AI in toolbar" println("Sending request: $request") agent.run( request + ". " + "You can only call tools. Use the Playwright tools to complete this task." ) } ``` ## 4) Shut down the MCP process Always clean up the external process at the end of your run. ``` // Shutdown the Playwright MCP process println("Closing connection to Playwright MCP server") process.destroy() ``` ## Troubleshooting - If the agent can't connect, make sure the MCP server is running on `http://localhost:8931`. - If you don't see the browser, ensure Playwright is installed and able to launch a browser on your system. - If you get authentication errors from OpenAI, double-check the `OPENAI_API_KEY` environment variable. ## Next steps - Try different websites or flows. The MCP server exposes a rich set of Playwright tools. - Swap the LLM model, or add more tools to the Koog agent. - Integrate this flow into your app, or publish the notebook as documentation. # Unity + Koog: Drive your game from a Kotlin Agent [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/UnityMcp.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/UnityMcp.ipynb) This notebook walks you through building a Unity-savvy AI agent with Koog using the Model Context Protocol (MCP). We'll connect to a Unity MCP server, discover tools, plan with an LLM, and execute actions against your open scene. > Prerequisites > > - A Unity project with the Unity-MCP server plugin installed > - JDK 17+ > - An OpenAI API key in the OPENAI_API_KEY environment variable ``` %useLatestDescriptors %use koog ``` ``` lateinit var process: Process ``` ## 1) Provide your OpenAI API key We read the API key from the `OPENAI_API_KEY` environment variable so you can keep secrets out of the notebook. ``` val token = System.getenv("OPENAI_API_KEY") ?: error("OPENAI_API_KEY environment variable not set") val executor = simpleOpenAIExecutor(token) ``` ## 2) Configure the Unity agent We define a compact system prompt and agent settings for Unity. ``` val agentConfig = AIAgentConfig( prompt = prompt("cook_agent_system_prompt") { system { "You are a Unity assistant. You can execute different tasks by interacting with tools from the Unity engine." } }, model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 1000 ) ``` ``` ``` ## 3) Start the Unity MCP server We'll launch the Unity MCP server from your Unity project directory and connect over stdio. ``` // https://github.com/IvanMurzak/Unity-MCP val pathToUnityProject = "path/to/unity/project" val process = ProcessBuilder( "$pathToUnityProject/com.ivanmurzak.unity.mcp.server/bin~/Release/net9.0/com.IvanMurzak.Unity.MCP.Server", "60606" ).start() ``` ## 4) Connect from Koog and run the agent We discover tools from the Unity MCP server, build a small plan-first strategy, and run an agent that uses only tools to modify your open scene. ``` import kotlinx.coroutines.runBlocking runBlocking { // Create the ToolRegistry with tools from the MCP server val toolRegistry = McpToolRegistryProvider.fromTransport( transport = McpToolRegistryProvider.defaultStdioTransport(process) ) toolRegistry.tools.forEach { println(it.name) println(it.descriptor) } val strategy = strategy("unity_interaction") { val nodePlanIngredients by nodeLLMRequest(allowToolCalls = false) val interactionWithUnity by subgraphWithTask( // work with plan tools = toolRegistry.tools, ) { input -> "Start interacting with Unity according to the plan: $input" } edge( nodeStart forwardTo nodePlanIngredients transformed { "Create detailed plan for " + agentInput + "" + "using the following tools: ${toolRegistry.tools.joinToString("\n") { it.name + "\ndescription:" + it.descriptor }}" } ) edge(nodePlanIngredients forwardTo interactionWithUnity onAssistantMessage { true }) edge(interactionWithUnity forwardTo nodeFinish) } val agent = AIAgent( promptExecutor = executor, strategy = strategy, agentConfig = agentConfig, toolRegistry = toolRegistry, installFeatures = { install(Tracing) install(EventHandler) { onAgentStarting { eventContext -> println("OnAgentStarting first (strategy: ${strategy.name})") } onAgentStarting { eventContext -> println("OnAgentStarting second (strategy: ${strategy.name})") } onAgentCompleted { eventContext -> println( "OnAgentCompleted (agent id: ${eventContext.agent.id}, result: ${eventContext.result})" ) } } } ) val result = agent.run( " extend current opened scene for the towerdefence game. " + "Add more placements for the towers, change the path for the enemies" ) result } ``` ## 5) Shut down the MCP process Always clean up the external Unity MCP server process at the end of your run. ``` // Shutdown the Unity MCP process process.destroy() ``` # Build a Simple Vacuum Cleaner Agent [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/VaccumAgent.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/VaccumAgent.ipynb) In this notebook, we'll explore how to implement a basic reflex agent using the new Kotlin agents framework. Our example will be the classic "vacuum world" problem — a simple environment with two locations that can be clean or dirty, and an agent that needs to clean them. First, let's understand our environment model: ``` import kotlin.random.Random /** * Represents a simple vacuum world with two locations (A and B). * * The environment tracks: * - The current location of the vacuum agent ('A' or 'B') * - The cleanliness status of each location (true = dirty, false = clean) */ class VacuumEnv { var location: Char = 'A' private set private val status = mutableMapOf( 'A' to Random.nextBoolean(), 'B' to Random.nextBoolean() ) fun percept(): Pair = location to status.getValue(location) fun clean(): String { status[location] = false return "cleaned" } fun moveLeft(): String { location = 'A' return "move to A" } fun moveRight(): String { location = 'B' return "move to B" } fun isClean(): Boolean = status.values.all { it } fun worldLayout(): String = "${status.keys}" override fun toString(): String = "location=$location, dirtyA=${status['A']}, dirtyB=${status['B']}" } ``` The VacuumEnv class models our simple world: - Two locations are represented by characters 'A' and 'B' - Each location can be either clean or dirty (randomly initialized) - The agent can be at either location at any given time - The agent can perceive its current location and whether it's dirty - The agent can take actions: move to a specific location or clean the current location ## Creating Tools for Vacuum Agent Now, let's define the tools our AI agent will use to interact with the environment: ``` import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.agents.core.tools.annotations.Tool import ai.koog.agents.core.tools.reflect.ToolSet /** * Provides tools for the LLM agent to control the vacuum robot. * All methods either mutate or read from the VacuumEnv passed to the constructor. */ @LLMDescription("Tools for controlling a two-cell vacuum world") class VacuumTools(private val env: VacuumEnv) : ToolSet { @Tool @LLMDescription("Returns current location and whether it is dirty") fun sense(): String { val (loc, dirty) = env.percept() return "location=$loc, dirty=$dirty, locations=${env.worldLayout()}" } @Tool @LLMDescription("Cleans the current cell") fun clean(): String = env.clean() @Tool @LLMDescription("Moves the agent to cell A") fun moveLeft(): String = env.moveLeft() @Tool @LLMDescription("Moves the agent to cell B") fun moveRight(): String = env.moveRight() } ``` The `VacuumTools` class creates an interface between our LLM agent and the environment: - It implements `ToolSet` from the Kotlin AI Agents framework - Each tool is annotated with `@Tool` and has a description for the LLM - The tools allow the agent to sense its environment and take actions - Each method returns a string that describes the outcome of the action ## Setting Up the Agent Next, we'll configure and create our AI agent: ``` import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.core.agent.config.AIAgentConfig import ai.koog.agents.core.tools.ToolRegistry import ai.koog.agents.core.tools.reflect.asTools import ai.koog.agents.ext.agent.chatAgentStrategy import ai.koog.agents.ext.tool.AskUser import ai.koog.agents.ext.tool.SayToUser import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor import ai.koog.prompt.params.LLMParams val env = VacuumEnv() val apiToken = System.getenv("OPENAI_API_KEY") ?: error("OPENAI_API_KEY environment variable not set") val executor = simpleOpenAIExecutor(apiToken = apiToken) val toolRegistry = ToolRegistry { tool(SayToUser) tool(AskUser) tools(VacuumTools(env).asTools()) } val systemVacuumPrompt = """ You are a reflex vacuum-cleaner agent living in a two-cell world labelled A and B. Your goal: make both cells clean, using the provided tools. First, call sense() to inspect where you are. Then decide: if dirty → clean(); else moveLeft()/moveRight(). Continue until both cells are clean, then tell the user "done". Use sayToUser to inform the user about each step. """.trimIndent() val agentConfig = AIAgentConfig( prompt = prompt("chat", params = LLMParams(temperature = 1.0)) { system(systemVacuumPrompt) }, model = OpenAIModels.Chat.GPT4o, maxAgentIterations = 50, ) val agent = AIAgent( promptExecutor = executor, strategy = chatAgentStrategy(), agentConfig = agentConfig, toolRegistry = toolRegistry ) ``` In this setup: 1. We create an instance of our environment 1. We set up a connection to OpenAI's GPT-4o model 1. We register the tools our agent can use 1. We define a system prompt that gives the agent its goal and behavior rules 1. We create the agent using the `AIAgent` constructor with a chat strategy ## Running the Agent Finally, let's run our agent: ``` import kotlinx.coroutines.runBlocking runBlocking { agent.run("Start cleaning, please") } ``` ``` Agent says: Currently in cell A. It's already clean. Agent says: Moved to cell B. It's already clean. ``` When we run this code: 1. The agent receives the initial prompt to start cleaning 1. It uses its tools to sense the environment and make decisions 1. It continues cleaning until both cells are clean 1. Throughout the process, it keeps the user informed about what it's doing ``` // Finally we can validate that the work is finished by printing the env state env ``` ``` location=B, dirtyA=false, dirtyB=false ``` # Weave tracing for Koog agents [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/notebooks/Weave.ipynb) [Download .ipynb](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/notebooks/Weave.ipynb) This notebook demonstrates how to trace Koog agents to W&B Weave using OpenTelemetry (OTLP). You will create a simple Koog `AIAgent`, enable the Weave exporter, run a prompt, and view rich traces in the Weave UI. For background, see Weave OpenTelemetry docs: https://weave-docs.wandb.ai/guides/tracking/otel/ ## Prerequisites Before running the example, make sure you have: - A Weave/W&B account: https://wandb.ai - Your API key from https://wandb.ai/authorize exposed as an environment variable: `WEAVE_API_KEY` - Your Weave entity (team or user) name exposed as `WEAVE_ENTITY` - Find it on your W&B dashboard: https://wandb.ai/home (left sidebar "Teams") - A project name exposed as `WEAVE_PROJECT_NAME` (if not set, this example uses `koog-tracing`) - An OpenAI API key exposed as `OPENAI_API_KEY` to run the Koog agent Example (macOS/Linux): ``` export WEAVE_API_KEY=... # required by Weave export WEAVE_ENTITY=your-team-or-username export WEAVE_PROJECT_NAME=koog-tracing export OPENAI_API_KEY=... ``` ## Notebook setup We use the latest Kotlin Jupyter descriptors. If you have Koog preconfigured as a `%use` plugin, you can uncomment the line below. ``` %useLatestDescriptors //%use koog ``` ## Create an agent and enable Weave tracing We construct a minimal `AIAgent` and install the `OpenTelemetry` feature with the Weave exporter. The exporter sends OTLP spans to Weave using your environment configuration: - `WEAVE_API_KEY` — authentication to Weave - `WEAVE_ENTITY` — which team/user owns the traces - `WEAVE_PROJECT_NAME` — the Weave project to store traces in ``` import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.features.opentelemetry.feature.OpenTelemetry import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor val entity = System.getenv()["WEAVE_ENTITY"] ?: throw IllegalArgumentException("WEAVE_ENTITY is not set") val projectName = System.getenv()["WEAVE_PROJECT_NAME"] ?: "koog-tracing" val agent = AIAgent( executor = simpleOpenAIExecutor(System.getenv("OPENAI_API_KEY")), llmModel = OpenAIModels.Chat.GPT4oMini, systemPrompt = "You are a code assistant. Provide concise code examples." ) { install(OpenTelemetry) { addWeaveExporter( weaveEntity = entity, weaveProjectName = projectName ) } } ``` ## Run the agent and view traces in Weave Execute a simple prompt. After completion, open the printed link to view the trace in Weave. You should see spans for the agent’s run, model calls, and other instrumented operations. ``` import kotlinx.coroutines.runBlocking println("Running agent with Weave tracing") runBlocking { val result = agent.run("Tell me a joke about programming") "Result: $result\nSee traces on https://wandb.ai/$entity/$projectName/weave/traces" } ``` ## Troubleshooting - If you don't see traces, verify `WEAVE_API_KEY`, `WEAVE_ENTITY`, and `WEAVE_PROJECT_NAME` are set in your environment. - Ensure your network allows outbound HTTPS to Weave's OTLP endpoint. - Confirm your OpenAI key is valid and the selected model is accessible from your account. # Web Scraping with The Web MCP by Bright Data and Koog [Open on GitHub](https://github.com/JetBrains/koog/blob/develop/examples/bright-data-mcp/) [Download .kt](https://raw.githubusercontent.com/JetBrains/koog/develop/examples/bright-data-mcp/Main.kt) In this tutorial, you'll connect a Koog agent to Bright Data's Web MCP server and let it perform web scraping and data collection tasks. We'll demonstrate how to search for information about Koog.ai using Bright Data's powerful web scraping infrastructure through the Model Context Protocol. We'll keep things simple and reproducible, focusing on a minimal but realistic agent + tools setup you can adapt for your own web scraping needs. ## Prerequisites - An OpenAI API key exported as an environment variable: `OPENAI_API_KEY` - A Bright Data API token exported as an environment variable: `BRIGHT_DATA_API_TOKEN` - Node.js and npx available on your PATH - Kotlin development environment with Koog dependencies **Tip**: The Bright Data MCP server provides access to enterprise-grade web scraping tools that can handle complex websites, CAPTCHAs, and anti-bot measures. ## 1) Set up your API credentials We read both API keys from environment variables to keep secrets secure and out of your code. ``` // Get API keys from environment variables val openAIApiKey = System.getenv("OPENAI_API_KEY") ?: error("OPENAI_API_KEY environment variable is not set") val brightDataToken = System.getenv("BRIGHT_DATA_API_TOKEN") ?: error("BRIGHT_DATA_API_TOKEN environment variable is not set") ``` ## 2) Start The Web MCP server by Bright Data We'll launch Bright Data's MCP server using `npx` and configure it with your API token. The server will expose web scraping capabilities through the Model Context Protocol. ``` println("Starting Bright Data MCP server...") // Start the Bright Data MCP server as a separate process val processBuilder = ProcessBuilder("npx", "@brightdata/mcp") // Set the API_TOKEN environment variable for the MCP server process val environment = processBuilder.environment() environment["API_TOKEN"] = brightDataToken // Start the process val process = processBuilder.start() // Give the process a moment to start Thread.sleep(2000) ``` ## 3) Connect from Koog and create the agent We build a Koog `AIAgent` with an OpenAI executor and connect its tool registry to the Bright Data MCP server via STDIO transport. Then we'll explore the available tools and run a web scraping task. ``` println("Creating STDIO transport...") try { // Create the STDIO transport val transport = McpToolRegistryProvider.defaultStdioTransport(process) println("Creating tool registry...") // Create a tool registry with tools from the Bright Data MCP server val toolRegistry = McpToolRegistryProvider.fromTransport( transport = transport, name = "bright-data-client", version = "1.0.0" ) // Print available tools (optional - for debugging) println("Available tools from Bright Data MCP server:") toolRegistry.tools.forEach { tool -> println("- ${tool.name}") } // Create the agent with MCP tools val agent = AIAgent( executor = simpleOpenAIExecutor(openAIApiKey), systemPrompt = "You are a helpful assistant with access to web scraping and data collection tools from Bright Data. You can help users gather information from websites, analyze web data, and provide insights.", llmModel = OpenAIModels.Chat.GPT4o, temperature = 0.7, toolRegistry = toolRegistry, maxIterations = 100 ) val result = agent.run("Please search for Koog.ai and tell me what is it and who invented it") println("\nAgent response:") println(result) } catch (e: Exception) { println("Error: ${e.message}") e.printStackTrace() } finally { println("Shutting down MCP server...") process.destroyForcibly() } ``` ## 4) Complete code example Here's the complete working example that demonstrates web scraping with The Web MCP by Bright Data: ``` package koog import ai.koog.agents.core.agent.AIAgent import ai.koog.agents.mcp.McpToolRegistryProvider import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor import kotlinx.coroutines.runBlocking /** * The entry point of the program demonstrating AI-driven web scraping and data collection. * * This function initializes a Bright Data MCP server, sets up tool integration, * and defines an AI agent for interacting with web scraping tools. It demonstrates the * following key operations: * * 1. Starts the Bright Data MCP server using a subprocess with proper API token configuration. * 2. Configures a registry of tools from the MCP server via STDIO transport communication. * 3. Creates an AI agent leveraging OpenAI's GPT-4o model with web scraping capabilities. * 4. Runs the agent to perform a specified task (e.g., searching for and analyzing web content * about Koog.ai). * 5. Cleans up by shutting down the MCP server process after execution. * * This function is intended for tutorial purposes, demonstrating how to integrate * MCP (Model Context Protocol) servers with AI agents for web data collection and analysis. * It requires OPENAI_API_KEY and BRIGHT_DATA_API_TOKEN environment variables to be set. */ fun main() = runBlocking { // Get API keys from environment variables val openAIApiKey = System.getenv("OPENAI_API_KEY") ?: error("OPENAI_API_KEY environment variable is not set") val brightDataToken = System.getenv("BRIGHT_DATA_API_TOKEN") ?: error("BRIGHT_DATA_API_TOKEN environment variable is not set") println("Starting Bright Data MCP server...") // Start the Bright Data MCP server as a separate process val processBuilder = ProcessBuilder("npx", "@brightdata/mcp") // Set the API_TOKEN environment variable for the MCP server process val environment = processBuilder.environment() environment["API_TOKEN"] = brightDataToken // Start the process val process = processBuilder.start() // Give the process a moment to start Thread.sleep(2000) println("Creating STDIO transport...") try { // Create the STDIO transport val transport = McpToolRegistryProvider.defaultStdioTransport(process) println("Creating tool registry...") // Create a tool registry with tools from the Bright Data MCP server val toolRegistry = McpToolRegistryProvider.fromTransport( transport = transport, name = "bright-data-client", version = "1.0.0" ) // Print available tools (optional - for debugging) println("Available tools from Bright Data MCP server:") toolRegistry.tools.forEach { tool -> println("- ${tool.name}") } // Create the agent with MCP tools val agent = AIAgent( executor = simpleOpenAIExecutor(openAIApiKey), systemPrompt = "You are a helpful assistant with access to web scraping and data collection tools from Bright Data. You can help users gather information from websites, analyze web data, and provide insights.", llmModel = OpenAIModels.Chat.GPT4o, temperature = 0.7, toolRegistry = toolRegistry, maxIterations = 100 ) val result = agent.run("Please search for Koog.ai and tell me what is it and who invented it") println("\nAgent response:") println(result) } catch (e: Exception) { println("Error: ${e.message}") e.printStackTrace() } finally { println("Shutting down MCP server...") process.destroyForcibly() } } ``` ## Troubleshooting - **Connection issues**: If the agent can't connect to the MCP server, ensure the Bright Data MCP package is properly installed via `npx @brightdata/mcp`. - **API token errors**: Double-check that your `BRIGHT_DATA_API_TOKEN` is valid and has the necessary permissions for web scraping. - **OpenAI authentication**: Verify that your `OPENAI_API_KEY` environment variable is correctly set and the API key is valid. - **Process timeout**: If the server takes longer to start, increase the `Thread.sleep(2000)` duration. ## Next steps - **Explore different queries**: Try scraping different websites or searching for various topics. - **Custom tool integration**: Add your own tools alongside Bright Data's web scraping capabilities. - **Advanced scraping**: Leverage Bright Data's advanced features like residential proxies, CAPTCHA solving, and JavaScript rendering. - **Data processing**: Combine the scraped data with other Koog agents for analysis and insights. - **Production deployment**: Integrate this pattern into your applications for automated web data collection. ## What you've learned This tutorial demonstrated how to: - Set up and configure The Web MCP by Bright Data - Connect a Koog AI agent to external MCP servers via STDIO transport - Perform AI-driven web scraping tasks using natural language instructions - Handle proper resource cleanup and error management - Structure code for production-ready web scraping applications The combination of Koog's AI agent capabilities with Bright Data's enterprise web scraping infrastructure provides a powerful foundation for automated data collection and analysis workflows.