Custom Graphs

Custom Strategy Graphs

Build advanced agent workflows as directed graphs using Koog's strategy builder. Custom graphs give you full control over node execution, edge conditions, and data flow.

Strategy Builder

Kotlin

import ai.koog.agents.core.dsl.builder.strategy
import ai.koog.agents.core.dsl.builder.forwardTo
import ai.koog.agents.core.dsl.builder.onAssistantMessage
import ai.koog.agents.core.dsl.builder.onToolCall

val myStrategy = strategy<String, String>("my-workflow") {
    // Define nodes
    val nodeProcessInput by nodeLLMRequest()
    val nodeExecuteTool by nodeExecuteTool()
    val nodeSendResult by nodeLLMSendToolResult()

    // Define edges
    edge(nodeStart forwardTo nodeProcessInput)
    edge(nodeProcessInput forwardTo nodeFinish onAssistantMessage { true })
    edge(nodeProcessInput forwardTo nodeExecuteTool onToolCall { true })
    edge(nodeExecuteTool forwardTo nodeSendResult)
    edge(nodeSendResult forwardTo nodeFinish onAssistantMessage { true })
    edge(nodeSendResult forwardTo nodeExecuteTool onToolCall { true })
}

Java

var graph = AIAgentGraphStrategy.builder("single_run")
    .withInput(String.class)
    .withOutput(String.class);
// ... define nodes and edges, then call graph.build()

Predefined Nodes

Custom Nodes

Create custom nodes using the node function:

// Simple transform node
val nodeTransform by node<String, String> { input ->
    input.uppercase()
}

// Async node with side effects
val nodeSaveResult by node<String, String> { input ->
    database.save(input)
    "Saved: $input"
}

// Node that calls the LLM with custom session
val nodeCustomLLM by node<String, String> { input ->
    llm.writeSession {
        model = OpenAIModels.Chat.GPT4o
        appendPrompt {
            system("You are a tone analyzer.")
            user(input)
        }
        requestLLM().content
    }
}

Edges and Conditions

Edges define transitions between nodes. Use the edge function with forwardTo:

Condition Types

Output Transformation

Transform data before passing to the next node:

edge(sourceNode forwardTo targetNode
    onCondition { input -> input.length > 10 }
    transformed { input -> input.uppercase() }
)

Conditional Branching

Route to different nodes based on input:

edge((nodeAnalyze forwardTo branchA) onCondition { it == "quick" })
edge((nodeAnalyze forwardTo branchB) onCondition { it == "deep" })

History Compression

Add conditional compression when the conversation grows too long:

val nodeCompressHistory by nodeLLMCompressHistory<ReceivedToolResult>()

edge(
    (nodeExecuteTool forwardTo nodeCompressHistory)
        onCondition { _ -> llm.readSession { prompt.messages.size > 100 } }
)
edge(
    (nodeExecuteTool forwardTo nodeSendToolResult)
        onCondition { _ -> llm.readSession { prompt.messages.size <= 100 } }
)
edge(nodeCompressHistory forwardTo nodeSendToolResult)

Subgraphs

Organize sections of a graph with their own tool subsets:

Kotlin

val firstSubgraph by subgraph<FirstInput, FirstOutput>(
    name = "first",
    tools = listOf(someTool)
) {
    val nodeProcess by nodeLLMRequest()
    val nodeExecute by nodeExecuteTool()
    val nodeSendResult by nodeLLMSendToolResult()

    edge(nodeStart forwardTo nodeProcess)
    edge(nodeProcess forwardTo nodeFinish onAssistantMessage { true })
    edge(nodeProcess forwardTo nodeExecute onToolCall { true })
    edge(nodeExecute forwardTo nodeSendResult)
    edge(nodeSendResult forwardTo nodeFinish onAssistantMessage { true })
}

Java

var subgraph = AIAgentSubgraph.builder("first")
    .withInput(FirstInput.class)
    .withOutput(FirstOutput.class)
    .limitedTools(someTool)
    .build();

Complete Tone Analysis Example

fun toneStrategy(name: String, toolRegistry: ToolRegistry): AIAgentGraphStrategy<String, String> {
    return strategy(name) {
        val nodeSendInput by nodeLLMRequest()
        val nodeExecuteTool by nodeExecuteTool()
        val nodeSendToolResult by nodeLLMSendToolResult()
        val nodeCompressHistory by nodeLLMCompressHistory<ReceivedToolResult>()

        edge(nodeStart forwardTo nodeSendInput)
        edge((nodeSendInput forwardTo nodeFinish) onAssistantMessage { true })
        edge((nodeSendInput forwardTo nodeExecuteTool) onToolCall { true })
        edge((nodeExecuteTool forwardTo nodeCompressHistory)
            onCondition { _ -> llm.readSession { prompt.messages.size > 100 } })
        edge(nodeCompressHistory forwardTo nodeSendToolResult)
        edge((nodeExecuteTool forwardTo nodeSendToolResult)
            onCondition { _ -> llm.readSession { prompt.messages.size <= 100 } })
        edge((nodeSendToolResult forwardTo nodeExecuteTool) onToolCall { true })
        edge((nodeSendToolResult forwardTo nodeFinish) onAssistantMessage { true })
    }
}

Visualization

On JVM, generate a Mermaid state diagram of your strategy:

val mermaidDiagram: String = myStrategy.asMermaidDiagram()

Java

String diagram = MermaidDiagramGenerator.INSTANCE.generate(myStrategy);

Using the Strategy

val agent = AIAgent(
    promptExecutor = simpleOpenAIExecutor(apiKey),
    llmModel = OpenAIModels.Chat.GPT4o,
    strategy = myStrategy,
    toolRegistry = toolRegistry,
    systemPrompt = "You are a research assistant."
)

val result = agent.run("Research the latest Kotlin coroutines features")

Best Practices

1. Start simple: Begin with a basic LLM ↔ Tool loop, then add complexity
2. Name your nodes: Use descriptive by delegation names for tracing
3. Keep nodes focused: Each node should do one thing well
4. Handle all paths: Ensure every node has a path to nodeFinish
5. Avoid infinite loops: Ensure conditions eventually lead to nodeFinish
6. Use subgraphs: Organize complex workflows into logical sections
7. Apply history compression: For long-running conversations
8. Test incrementally: Build graphs step by step, testing each addition

Troubleshooting