Hi, I cant really find the following use case in t...

# koog-agentic-framework

[Chrylo]

Hi, I cant really find the following use case in the documentation.

Flow:
prompt -> LLMRequest (Generate queries) -> tool call -> Model output

I could use the subgraphWithTask but it seems to be pretty slow. I think the tool result is returned to the LLM for further processing before finishing the node. However I want want the LLM to only call the tool once and immediately return the result to speed up node. (In Spring AI there was annotation for this)

How do I do that? Should I do a LLMRequest (Query generation) without tool calls and just execute the method manually in a custom node afterwards?

Thank you 🙂

@Tool
@LLMDescription(
    """
    Search with query.
    """,
)
fun search(
     @LLMDescription("Query")
     query: String,
): List<Model>

[Vadim Briliantov]

Hey! Indeed, subgraphWithTask is a full-fledged loop that would call multiple tools.


If you need finer-grained steps, you can use nodeLLMRequest() (also there’s nodeLLMRequestMultiple() with multi-message/multi-tool result to then execute N tools in parallel) and then nodeExecuteTool() (also there’s a -Multiple version of it).

val askLLM by nodeLLMRequest()
val callTool by nodeExecuteTool()

edge(askLLM forwardTo callTool onToolCall { true }) // onToolCall check is required in compile time, because `askLLM` returns Message.Response which can be either Message.Assistant or Message.Tool.Call
edge(askLLM forwardTo ... onAssistantMessage { true }) // handle the case if LLM didn't want to call a tool (ex: retry)
edge(callTool forwardTo nodeFinish transformed { it.result }) // would return JSONElement that you can then parse using tool.decodeResult()

If you know one specific tool you want to call (specifically, search) — then there’s nodeLLMRequestForceOneTool(toolDescriptor) :

val searchTool = ::search.asTool()

val askLLM by nodeLLMRequestForceOneTool(tool = searchTool.descriptor) // would work if your model supports toolChoice (ex.: Ollama doesn't)
val callTool by nodeExecuteTool()

edge(askLLM forwardTo callTool onToolCall { true })
edge(callTool forwardTo nodeFinish transformed transformed { searchTool.decodeResult(it.result!!, config.serializer) })

Please note also that it.result (on the last line) is nullable. And there’s also it.resultKind that could be Success / Failure(public val error: AIAgentError?) / ValidationError(public val error: AIAgentError) — so if you want you can actually handle errors from your tool then

[Vadim Briliantov]

@Chrylo if you want the most resilient way to do what you want, you’ll need something like this with retries:

sealed interface SearchToolResult {
    class Success(val result: List<Model>) : SearchToolResult
    class Failure(val error: AIAgentError?, val content: String) : SearchToolResult
}

...

val searchTool = ::search.asTool()

val guaranteedSearch by subgraphWithRetry(
    condition = { result ->
        when (result) {
            is SearchToolResult.Success -> ConditionResult.Approve
            is SearchToolResult.Failure -> ConditionResult.Reject("Fix the error: ${result.error?.message ?: result.content}")
        }
    },
    maxRetries = 3,
    name = "retry-search",
) {
    val askLLM by nodeLLMRequestForceOneTool(tool = searchTool.descriptor)
    val callTool by nodeExecuteTool()

    edge(nodeStart forwardTo askLLM)
    edge(askLLM forwardTo callTool onToolCall { true })
    edge(callTool forwardTo nodeFinish transformed {
        val resultKind = it.resultKind
        when (resultKind) {
            ToolResultKind.Success -> SearchToolResult.Success(
                searchTool.decodeResult(
                    it.result!!,
                    config.serializer
                )
            )

            is ToolResultKind.Failure -> SearchToolResult.Failure(resultKind.error, it.content)
            is ToolResultKind.ValidationError -> SearchToolResult.Failure(resultKind.error, it.content)
        }
    })
}

[Chrylo]

Doing it like that now, thanks it works.

val selectNode by nodeLLMRequestForceOneTool(name = "SelectNode", SelectToolSet::search.asTool())
val executeToolNode by nodeExecuteTool("ExecuteToolNode")

edge(
    selectNode forwardTo executeToolNode
        onToolCall { true },
)

edge(
    executeToolNode forwardTo sendSResultNode
        transformed { toolsResult ->
            val tool = XToolSet::search.asTool()
            val toolResult = toolsResult.toSafeResult(tool, config.serializer).asSuccessful().result
            Result(toolResult)
        },
)

[Vadim Briliantov]

Consider wrapping into subgraphWithRetry as I shown above to be on the safe side before deploying to prod 🙂

[Vadim Briliantov]

It wouldn’t add extra LLM calls unless an error happens