Notes

Koog Graph Workflows: Building Complex AI Agent Behaviors

4 min read

Koog is JetBrains’ new Kotlin framework for building AI agents. One of its key features is graph workflows - a way to design complex agent behaviors using nodes and edges, similar to flowcharts.

Why Graph Workflows?

Traditional AI agents often use linear conversation flows or hardcoded if-else logic. Graph workflows solve this by:

  • Making complex logic visual and easier to understand
  • Breaking down operations into modular, reusable pieces
  • Allowing conditional routing based on context
  • Providing type safety through Kotlin’s type system

Basic Components

Nodes

Nodes are processing steps in your workflow. Each node has input and output types:

val processNode by node<InputType, OutputType> { input ->
    // Process the input and return output
    processInput(input)
}

Built-in node types:

  • nodeLLMRequest() - Sends requests to LLM
  • nodeExecuteTool() - Executes tools/APIs
  • nodeLLMSendToolResult() - Sends tool results to LLM
  • nodeStart and nodeFinish - Entry and exit points

Edges

Edges connect nodes and define data flow. Four types:

1. Basic Edge

edge(sourceNode forwardTo targetNode)

2. Conditional Edge

edge(sourceNode forwardTo targetNode onCondition { output ->
    output.contains("specific text")
})

3. Transformation Edge

edge(sourceNode forwardTo targetNode transformed { output ->
    "Modified: $output"
})

4. Combined Edge

edge(sourceNode forwardTo targetNode 
    onCondition { it.isNotEmpty() } 
    transformed { it.uppercase() })

Simple Example: Calculator Agent

Here’s a basic calculator agent that can use tools:

val agentStrategy = strategy("Simple calculator") {
    // Define nodes
    val nodeSendInput by nodeLLMRequest()
    val nodeExecuteTool by nodeExecuteTool()
    val nodeSendToolResult by nodeLLMSendToolResult()
    
    // Connect nodes with edges
    edge(nodeStart forwardTo nodeSendInput)
    
    // If LLM responds normally, finish
    edge(
        (nodeSendInput forwardTo nodeFinish) 
        transformed { it } 
        onAssistantMessage { true }
    )
    
    // If LLM wants to use a tool, execute it
    edge(
        (nodeSendInput forwardTo nodeExecuteTool) 
        onToolCall { true }
    )
    
    // Send tool result back to LLM
    edge(nodeExecuteTool forwardTo nodeSendToolResult)
    
    // After tool result, finish
    edge(
        (nodeSendToolResult forwardTo nodeFinish) 
        transformed { it } 
        onAssistantMessage { true }
    )
}

How it works:

  1. User input goes to LLM
  2. If LLM responds normally → finish
  3. If LLM wants to use a tool → execute tool → send result → finish

More Complex Example: Customer Service Bot

val customerServiceStrategy = strategy("Customer Service") {
    val classifyRequest by node<UserInput, RequestType> { input ->
        classifyUserRequest(input)
    }
    
    val handleComplaint by node<ComplaintData, Resolution> { complaint ->
        processComplaint(complaint)
    }
    
    val handleInquiry by node<InquiryData, Information> { inquiry ->
        processInquiry(inquiry)
    }
    
    val handleRefund by node<RefundRequest, RefundStatus> { request ->
        processRefund(request)
    }
    
    // Route based on request type
    edge(nodeStart forwardTo classifyRequest)
    
    edge(classifyRequest forwardTo handleComplaint 
        onCondition { it.type == RequestType.COMPLAINT })
    
    edge(classifyRequest forwardTo handleInquiry 
        onCondition { it.type == RequestType.INQUIRY })
    
    edge(classifyRequest forwardTo handleRefund 
        onCondition { it.type == RequestType.REFUND })
    
    // All paths end here
    edge(handleComplaint forwardTo nodeFinish)
    edge(handleInquiry forwardTo nodeFinish)
    edge(handleRefund forwardTo nodeFinish)
}

Subgraphs

For complex agents, you can create subgraphs - self-contained workflow units with their own tools and context. Information can be:

  • Kept within the subgraph
  • Shared between subgraphs using AgentMemory
  • Persisted across sessions

Integration with Koog Features

Event Handling

Monitor what happens in your workflow:

install(EventHandler) {
    onBeforeAgentStarted = { strategy, agent ->
        println("Starting: ${strategy.name}")
    }
    onAgentFinished = { strategyName, result ->
        println("Finished: $strategyName")
    }
}

Tracing

Debug your workflows:

install(Tracing) {
    traceLevel = TraceLevel.DETAILED
    includeNodeTransitions = true
    includeEdgeConditions = true
}

Key Benefits

Type Safety: Kotlin catches type mismatches at compile time

edge(stringNode forwardTo intNode) // ❌ Won't compile
edge(stringNode forwardTo stringProcessor) // ✅ Works

Modularity: Each node handles one responsibility Testability: Nodes can be tested independently Performance: Lazy evaluation, parallel execution where possible Debugging: Clear execution paths make issues easier to find

Best Practices

  1. Start simple - Begin with linear flows, add complexity gradually
  2. Use descriptive names - Make workflows self-documenting
  3. Leverage types - Let Kotlin catch errors early
  4. Test nodes separately - Each node should be testable on its own

When to Use Graph Workflows

Good for:

  • Customer service bots with multiple escalation paths
  • Content generation with review steps
  • Data processing with conditional validation
  • Multi-modal agents switching between different interaction types

Not necessary for:

  • Simple question-answer bots
  • Linear workflows without branching
  • Prototypes where you just want to test basic functionality

Getting Started

  1. Install Koog in your Kotlin project
  2. Start with Simple API to learn basics
  3. Move to AI Agent for full graph workflow features
  4. Build simple linear flows first
  5. Add branching and conditions as needed

Basic Setup

dependencies {
    implementation("ai.koog:koog-agents:VERSION")
}
fun main() = runBlocking {
    val apiKey = System.getenv("OPENAI_API_KEY")
    val promptExecutor = simpleOpenAIExecutor(apiKey)
    
    val agentConfig = AIAgentConfig.withSystemPrompt(
        prompt = "You are a helpful assistant."
    )
    
    val agent = AIAgent(
        config = agentConfig,
        strategy = yourStrategy,
        promptExecutor = promptExecutor
    )
    
    val result = agent.runAndGetResult("Hello!")
    println(result)
}

References

Graph View

Backlinks