Koog

Part of Java is Coming for Python’s Lunch - The Rise of JVM AI Agents. Slides 53–67.

Koog slide 53 ↗

the agent is a graph you can read

The same review becomes a graph of nodes and edges slide 54 ↗

tool call

tool call

text

text

accepted

rejected

start

review

run-tools

send-tool-results

structure-draft

judge

finish

Tools live in a `ToolRegistry` the agent is built with slide 55 ↗

integration

class KoogTravelTools(private val facts: TravelFacts) : ToolSet {
    @Tool
    @LLMDescription("Load a submitted trip with all bookings")
    fun getTrip(@LLMDescription("The trip id") tripId: Long): TripSnapshot =
        facts.trip(tripId)
}
val tools = ToolRegistry {
    tools(KoogTravelTools(facts))
}

AIAgent(promptExecutor = executor, strategy = strategy,
    agentConfig = config, toolRegistry = tools)

Every node is a typed value declared up front slide 56 ↗

type safety

strategy<String, ReviewOutcome>("trip-review") {
    var revisions = 0

    val review by nodeLLMRequest("review")
    val runTools by nodeExecuteTools("run-tools")
    val sendToolResults by nodeLLMSendToolResults("send-tool-results")
    val draft by nodeLLMRequestStructured<ReviewOutcome>("structure-draft")
    // edges on the next slide
}

Four edges are the whole tool loop slide 57 ↗

edge(nodeStart forwardTo review)
edge(review forwardTo runTools onToolCalls { true })
edge(runTools forwardTo sendToolResults)
edge(sendToolResults forwardTo runTools onToolCalls { true })

edge(review forwardTo draft onTextMessage { true } transformed { STRUCTURE })
edge(sendToolResults forwardTo draft onTextMessage { true } transformed { STRUCTURE })

Koog ships the whole tool loop as one `subgraphWithTask` slide 58 ↗

val review by nodeLLMRequest("review")
val runTools by nodeExecuteTools("run-tools")
val sendToolResults by nodeLLMSendToolResults("send-tool-results")
val draft by nodeLLMRequestStructured<ReviewOutcome>("structure-draft")

edge(review forwardTo runTools onToolCalls { true })
edge(runTools forwardTo sendToolResults)
edge(sendToolResults forwardTo runTools onToolCalls { true })
edge(review forwardTo draft onTextMessage { true } transformed { STRUCTURE })
edge(sendToolResults forwardTo draft onTextMessage { true } transformed { STRUCTURE })
val review by subgraphWithTask<String, ReviewOutcome>(name = "review") { task ->
    task
}

The judge is a `subgraphWithVerification` that never sees the reviewer's chat slide 59 ↗

type safety

val judge by subgraph<ReviewOutcome, CriticResult<ReviewOutcome>>(
    "judge", freshHistory = true,
) {
    val verify by subgraphWithVerification<ReviewOutcome>(tools = emptyList()) { draft ->
        llm.writeSession {
            clearHistory()
            appendPrompt { system(Prompts.JUDGE) }
        }
        Prompts.judge(facts.trip(tripId), facts.policies(), facts.costs(tripId), draft)
    }
    nodeStart then verify then nodeFinish
}

The revision loop stays a plain node and two edges slide 60 ↗

reliability

val revise by node<CriticResult<ReviewOutcome>, String>("revise") { critique ->
    revisions++
    Prompts.revise(tripId, critique.input, critique.feedback)
}

nodeStart then review then judge
edge(judge forwardTo nodeFinish
    onCondition { it.successful || revisions >= MAX_REVISIONS } transformed { it.input })
edge(judge forwardTo revise onCondition { !it.successful })
revise then review

Two library subgraphs and one node of mine are the whole agent slide 61 ↗

subgraphWithTask

subgraphWithVerification

tool calls

accepted

rejected

start

review

tools

judge

freshHistory = true

revise

finish

Koog library subgraph

my node and edges

Koog runs on the same Spring AI `ChatModel` slide 62 ↗

integration

@Bean
fun koogPromptExecutor(chatModel: ChatModel): PromptExecutor = MultiLLMPromptExecutor(
    SpringAiLLMClient.builder()
        .chatModel(chatModel)
        .provider(LLMProvider.OpenAI)
        .build()
)

AIAgent

SpringAiLLMClient

ChatModel

Spring AI traces the model calls, Koog adds the graph around them slide 63 ↗

observability

agent trip-review subgraph review gen_ai.client.operation tool getTrip gen_ai.client.operation subgraph judge gen_ai.client.operation

One `install` sends the whole graph to Langfuse slide 64 ↗

observability

AIAgent(
    promptExecutor = executor,
    strategy = reviewStrategy(tripId),
    agentConfig = agentConfig,
    toolRegistry = tools,
) {
    install(OpenTelemetry) {
        setServiceInfo("trip-review", "1.0.0")
        addLangfuseExporter()
    }
}

Chat memory is a feature I install, compression is a node I place slide 65 ↗

context

install(ChatMemory) {
    chatHistoryProvider = JdbcChatHistory(dataSource)
    windowSize(20)
}
val compress by nodeLLMCompressHistory<ReceivedToolResults>(
    strategy = HistoryCompressionStrategy.FromLastNMessages(5),
)

Every tool can register its undo, a rollback runs them backwards slide 66 ↗

reliability

install(Persistence) {
    storage = JVMFilePersistenceStorageProvider(checkpoints)
    enableAutomaticPersistence = true
    rollbackToolRegistry = RollbackToolRegistry {
        registerRollback(bookings::addHotel, bookings::removeHotel)
    }
}
context.persistence().rollbackToCheckpoint(checkpointId, context)

The graph is one strategy, Koog lets every agent pick its own slide 67 ↗

graph strategy functional loop ReAct and chat LLM planner GOAP planner subgraphs