Hey folks! I have been trying to move to processin...

# koog-agentic-framework

[Didier Villevalois]

Hey folks! I have been trying to move to processing multiple tool calls at once. However, it seems the approach of having two edges onMultipleToolCalls { true }) and transformed { it.first() } onAssistantMessage { true }) (copied from the only example using onMultipleToolCalls) is not very solid without adding a lot of prompt to force the LLM to either only make calls or only make chat.

What would be the best approach with the current graph/edge system to send calls on one edge and chat messages on another edge?

[Code in the thread]

[Didier Villevalois]

Here is my agent graph:

strategy("qa") {
    val qa by qaWithTools { calls -> vetToolCalls(calls, tools) }
    nodeStart then qa then nodeFinish
}

private fun AIAgentSubgraphBuilderBase<*, *>.qaWithTools(
    vetToolCalls: suspend (List<Message.Tool.Call>) -> List<Boolean>,
) = subgraph {
    val initialRequest by nodeLLMRequestMultiple()
    val processResponses by nodeDoNothing<List<Message.Response>>()
    val vetToolCalls by nodeVetToolCalls(vetToolCalls = vetToolCalls)
    val executeTools by nodeExecuteVettedToolCalls(parallelTools = true)
    val toolResultsRequest by nodeLLMSendMultipleToolResults()

    edge(nodeStart forwardTo initialRequest)
    edge(initialRequest forwardTo processResponses)

    edge(processResponses forwardTo vetToolCalls onMultipleToolCalls { true })
    edge(processResponses forwardTo nodeFinish transformed { it.first() } onAssistantMessage { true })

    edge(vetToolCalls forwardTo executeTools)
    edge(executeTools forwardTo toolResultsRequest)
    edge(toolResultsRequest forwardTo processResponses)
}

nodeVetToolCalls and nodeExecuteVettedToolCalls are custom nodes. nodeVetToolCalls does user-interaction to vet calls to certain tools and nodeExecuteVettedToolCalls only execute the accepted tool calls and rejects the others.

I am using llama4:16x17b (aka. Scout). And it does want to explain his reasoning before calling tools. I don't want to deprive it of that, because I think that reasoning is important content that have to be added to the context.

[Didier Villevalois]

I feel like the current agent graph design is "optimized" for the flow String -> 1 tool call -> String, etc. As soon as you want to do things a little more clever, like handling all the messages of the LLM, that design makes things really complicated.

I have experienced it while implementing my nodeVetToolCalls and nodeExecuteVettedToolCalls nodes, you have to manually partition messages and create intermediate data containers. It gets really cumbersome.

Or am I missing something obvious? Any help is greatly appreciated.

[Didier Villevalois]

For now, I fixed it by making a custom onAnyToolCalls :

infix fun <IncomingOutput, IntermediateOutput, OutgoingInput>
        AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput>.onAnyToolCalls(
    block: suspend (List<Message.Tool.Call>) -> Boolean
): AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Tool.Call>, OutgoingInput> {
    return onIsInstance(List::class)
        .transformed { it.filterIsInstance<Message.Tool.Call>() }
        .onCondition { toolCalls -> toolCalls.isNotEmpty() && block(toolCalls) }
}

Maybe, that should replace/complement onMultipleToolCalls (which checks that all messages are tool calls).

[Vadim Briliantov]

Hi Didier! Thanks for finding this out. Indeed, for Ollama because it doesn’t support toolChoice parameter (which can tell the LLM to ONLY call tools, for example, see ai.koog.prompt.params.LLMParams.ToolChoice) as bigger models (OpenAI, Anthropic, or in Gemini it’s called functionCallingConfig) do, you need to workaround this.
There are options:
a) adding a giveFeedbackToCallTools node by hand, something like:

val giveFeedbackToCallTools by node<String, Message.Response> { input ->
        llm.writeSession {
            updatePrompt {
                user("Don't chat with plain text! Call one of the available tools, instead: ${tools.joinToString(", ") { it.name }}")
            }

            requestLLM()
        }
}

b) You can actually write onCondition and check if tool calls are present in the response, and make your vetToolCalls node receive List<Message.Response> and do the filtering by itself.
c) Your approach with onAnyToolCalls or similar.

Probably, we should add this to our API and document this. Please feel free to make a PR. I would probably suggest a name like onToolCallsPresent . Maybe something like this:

infix fun <IncomingOutput, IntermediateOutput, OutgoingInput>
        AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput>.onToolCallsPresent(
    block: suspend (List<Message.Tool.Call>) -> Boolean
): AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Tool.Call>, OutgoingInput> =
    onToolCallsPresent(filterOutAssistantMessages = true) { 
        block(it as List<Message.Tool.Call>)
    }.transformed { it as List<Message.Tool.Call> }

fun <IncomingOutput, IntermediateOutput, OutgoingInput>
        AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput>.onToolCallsPresent(
    filterOutAssistantMessages: Boolean,
    block: suspend (List<Message.Response>) -> Boolean
): AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Response>, OutgoingInput> {
    return onIsInstance(List::class)
        .transformed { it.filterIsInstance<Message.Response>() }
        .transformed {
            if (filterOutAssistantMessages) it.filterIsInstance<Message.Tool.Call>() else it
        }
        .onCondition { toolCalls -> toolCalls.isNotEmpty() && block(toolCalls) }
}

[Vadim Briliantov]

And also probably we need to add onSingleAssistantMessage so that we can change

edge(processResponses forwardTo nodeFinish transformed { it.first() } onAssistantMessage { true })

to

edge(processResponses forwardTo nodeFinish onSingleAssistantMessage { true })

WDYT?

[Didier Villevalois]

Hi Vadim! Thanks a lot for your answer.

> Indeed, for Ollama because it doesn’t support toolChoice parameter
Dang, I don't have access to any proprietary LLMs, so I always forget about the tool choice you explained me about 3 weeks ago...

> Probably, we should add this to our API and document this. Please feel free to make a PR. I would probably suggest a name like onToolCallsPresent . Maybe something like this:
I like it. I will submit some tentative PR for it.

[Didier Villevalois]

And also probably we need to add onSingleAssistantMessage so that we can change

I really wonder what would we gain to mute the other assistant messages (and which to choose, the first, the last, etc.) and only keep one (even if all are added to the prompt by nodeLLMRequest*). Also, this seems very related to what you guys will decide about the thinking messages. So I guess I should wait on what you guys design about that first. For now, I just dump all the messages and it is fine for my use-case.