Llama 3 / 3.1 / 3.2 Tool Calling

dependencies {
    implementation(platform("sk.ainet.transformers:skainet-transformers-bom:0.31.0"))

    implementation("sk.ainet.transformers:skainet-transformers-runtime-kllama")
    implementation("sk.ainet.transformers:skainet-transformers-agent")
}

--enable-preview --add-modules jdk.incubator.vector

import sk.ainet.apps.kllama.java.KLlamaJava
import java.nio.file.Path

val session = KLlamaJava.loadGGUF(Path.of("models/Llama-3.2-1B-Instruct-Q8_0.gguf"))
// session.runtime  : InferenceRuntime<FP32>
// session.tokenizer: Tokenizer
// session is AutoCloseable — close it to release the Arena.

import kotlinx.serialization.json.*
import sk.ainet.apps.kllama.chat.Tool
import sk.ainet.apps.kllama.chat.ToolDefinition

class WeatherTool : Tool {
    override val definition = ToolDefinition(
        name = "get_weather",
        description = "Get the current weather for a city.",
        parameters = buildJsonObject {
            put("type", "object")
            putJsonObject("properties") {
                putJsonObject("city") {
                    put("type", "string")
                    put("description", "City name, e.g. 'Bratislava'.")
                }
            }
            putJsonArray("required") { add(JsonPrimitive("city")) }
        }
    )

    override fun execute(arguments: JsonObject): String {
        val city = arguments["city"]?.jsonPrimitive?.content
            ?: return "Error: missing 'city'"
        // Real call to your weather backend goes here.
        return """{"city":"$city","tempC":22,"condition":"sunny"}"""
    }
}

import sk.ainet.apps.kllama.chat.*

val chat = ChatSession(
    runtime = session.runtime,
    tokenizer = session.tokenizer,
    // family="llama" auto-resolves to Llama3ToolCallingSupport with the
    // bare-JSON format Llama 3.2 was fine-tuned on. Override only if you
    // know you need FUNCTION_TAG (see "The two formats" below).
    metadata = ModelMetadata(family = "llama", architecture = "llama"),
)

val tools = ToolRegistry().apply {
    register(WeatherTool())
}

val loop = chat.createAgentLoop(
    toolRegistry = tools,
    maxTokens    = 256,
    temperature  = 0.7f,
)

val messages = mutableListOf(
    ChatMessage(
        role = ChatRole.SYSTEM,
        content = "You are a helpful assistant with access to tools. " +
            "Always call get_weather when asked about weather — never guess."
    ),
    ChatMessage(role = ChatRole.USER, content = "What's the weather in Bratislava?"),
)

val finalAnswer = loop.runWithEncoder(
    messages = messages,
    encode   = { chat.encode(it) },
)
println(finalAnswer)

val listener = object : AgentListener {
    override fun onToken(token: String) { print(token) }
    override fun onAssistantMessage(text: String) {
        println("\n[raw assistant] $text")
    }
    override fun onToolCalls(calls: List<ToolCall>) {
        for (c in calls) println("[tool call] ${c.name}(${c.arguments})")
    }
    override fun onToolResult(call: ToolCall, result: String) {
        println("[tool result] ${call.name} -> $result")
    }
    override fun onToolCallValidationFailed(call: ToolCall, reason: String) {
        println("[tool call invalid] ${call.name}: $reason")
    }
    override fun onComplete(finalResponse: String) {}
}

loop.runWithEncoder(messages, encode = { chat.encode(it) }, listener = listener)

val rendered = chat.chatTemplate.apply(
    messages = messages,
    tools = tools.definitions(),
    addGenerationPrompt = true,
)
println("[prompt] (${rendered.length} chars)\n$rendered")

{"name": "list_files", "parameters": {"path": "/tmp"}}

<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a helpful assistant with tool calling capabilities.
When you receive a tool call response, use the output to format an answer to the original user question.

You have access to the following functions:

{"name":"list_files","description":"...","parameters":{...}}

If you choose to call a function, your reply MUST be a single JSON object on one line in the following format and nothing else:
{"name": <function-name>, "parameters": <arguments-object>}
Do not write the function definition. Do not include any prose. Do not use variables.<|eot_id|>
<|start_header_id|>user<|end_header_id|>

What files are in /tmp?<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>

<function=list_files>{"path": "/tmp"}</function>

...
If you choose to call a function, ONLY reply in the format below and nothing else:
<function=function_name>{"arg_name": "arg_value"}</function>
Function calls MUST be on a single line. Required parameters MUST be specified.
...