refactor: decouple tool calling from kllama, add ChatSession abstraction

michalharakal · claude · michalharakal · commit f42617aae307 · 2026-04-11T10:25:35.000+02:00
Phase 1 of the unified pipeline plan. Tool calling no longer requires
GGUFTokenizer — any Tokenizer implementation works.

- Extend Tokenizer interface with eosTokenId, bosTokenId, vocabSize
- Add ChatSession in llm-agent that bundles runtime + tokenizer + metadata
- Refactor ToolCallingDemo and AgentCli to accept Tokenizer, not GGUFTokenizer
- Remove GGUFTokenizer cast from kllama Main.kt chat/agent/demo dispatch
- Fix JavaAgentLoop instanceof hack with tokenizer.eosTokenId
- Update all Tokenizer implementations (GGUF, HF BPE, Tekken, BERT)

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/PLAN-unified-pipeline.md b/PLAN-unified-pipeline.md
@@ -0,0 +1,147 @@
+# Plan: Unified Model Pipeline with Decoupled Tool Calling
+
+## Context
+
+Currently SKaiNET-transformers has:
+- **5+ hand-coded runtimes** (LlamaRuntime, Qwen35Runtime, Gemma3nRuntime, ApertusRuntime, VoxtralRuntimes) — each reimplements the forward pass, weight loading, and layer execution
+- **Tool calling tightly coupled to kllama** — the AgentLoop, ToolCallingDemo, and chat modes only exist in the kllama runner. Other models (Gemma, Apertus) cannot use tool calling without duplicating code
+- **Two execution paths** — legacy hand-coded runtimes AND the newer `OptimizedLLMRuntime` with DSL/compute-graph/AOT. LlamaRuntime and ApertusRuntime are already marked deprecated
+
+The goal: converge on **one unified pipeline** where model definition, weight loading, tokenization, and tool calling are cleanly separated pipeline stages.
+
+## Architecture Overview
+
+```
+GGUF/SafeTensors File
+    |
+WeightLoader (parse metadata + tensors)
+    |
+DSL Network Definition (model-specific, declarative)
+    |
+ComputeGraph (DAG)
+    |
+Optimization Pipeline (TransposeElim -> WeightDedup -> LLMFusion -> DCE)
+    |
+ComputeGraphExecutor (fused kernels)
+    |
+InferenceRuntime (unified: forward + generate)
+    |
+TokenizationPipeline (encode/decode, special tokens, byte-level BPE)
+    |
+ChatPipeline (template formatting, tool calling, agent loop)
+```
+
+## Phase 1: Decouple Tool Calling from kllama (immediate value) -- DONE
+
+**What was done:**
+
+1. **Enhanced `Tokenizer` interface** with `eosTokenId`, `bosTokenId`, `vocabSize`
+   - Updated all implementations: `GGUFTokenizer`, `TokenizerImpl`, `HuggingFaceBPETokenizer`, `TekkenTokenizerAdapter`, `HuggingFaceTokenizer` (BERT)
+
+2. **Created `ChatSession` abstraction** in `llm-agent`
+   - File: `llm-agent/.../chat/ChatSession.kt`
+   - Bundles `InferenceRuntime` + `Tokenizer` + `ModelMetadata`
+   - Provides `createAgentLoop()` and `runSingleTurn()` for any runner
+
+3. **Refactored `ToolCallingDemo` and `AgentCli`** to use `Tokenizer` interface instead of `GGUFTokenizer`
+   - Both now accept any `Tokenizer`, not just `GGUFTokenizer`
+   - Both use `ChatSession` internally for agent loop creation
+
+4. **Removed `GGUFTokenizer` cast from kllama Main.kt** dispatch
+   - Chat/agent/demo modes now work with any `Tokenizer`
+
+5. **Fixed `JavaAgentLoop`** — replaced `GGUFTokenizer` instanceof hack with `tokenizer.eosTokenId`
+
+## Phase 2: Unified DSL-Based Model Definition (converge on OptimizedLLMRuntime)
+
+**Problem:** Each model has a hand-coded runtime. `OptimizedLLMRuntime` already supports DSL -> graph -> optimized execution, but only some models use it.
+
+**Changes:**
+
+1. **Define DSL networks for all model families:**
+   - `llamaNetwork(config)` — LLaMA/Mistral/Qwen2/3 (standard transformer)
+   - `qwen35Network(config)` — Qwen3.5 (hybrid DeltaNet + full attention)
+   - `gemmaNetwork(config)` — Gemma (GELU, MatFormer FFN, sliding window)
+   - `apertusNetwork(config)` — Apertus (xIELU, ungated MLP, QK-norm)
+   - Each is a pure function returning a `Network<T>` from the DSL
+
+2. **Unified model loading flow:**
+   ```
+   detectArchitecture(ggufMetadata) -> ModelFamily
+   ModelFamily.createNetwork(config) -> Network<T>
+   WeightLoader.loadAndMap(file, network) -> weights
+   OptimizedLLMRuntime(network, weights, mode=HYBRID) -> InferenceRuntime
+   ```
+
+3. **Remove deprecated hand-coded runtimes** once DSL equivalents are validated:
+   - `LlamaRuntime` -> `llamaNetwork()` + `OptimizedLLMRuntime`
+   - `ApertusRuntime` -> `apertusNetwork()` + `OptimizedLLMRuntime`
+
+**Critical files:**
+- `llm-core/.../OptimizedLLMRuntime.kt` — already exists, extend
+- `llm-core/.../dsl/TransformerDsl.kt` — already has embedding, MHA, SwiGLU, RMSNorm
+- `llm-core/.../weights/LLMWeightNameResolvers.kt` — already maps DSL paths -> GGUF names
+- New: per-model DSL network definitions
+
+## Phase 3: Tokenization as Pipeline Stage
+
+**Problem:** Tokenization is split between `GGUFTokenizer` (kllama module), `QwenByteLevelBPETokenizer` (llm-core), and model-specific code. The byte-level BPE fix we just made shows the fragility.
+
+**Changes:**
+
+1. **Enhance `Tokenizer` interface** (`llm-core`):
+   ```kotlin
+   interface Tokenizer {
+       fun encode(text: String): IntArray
+       fun decode(token: Int): String
+       fun decode(tokens: IntArray): String
+       val eosTokenId: Int
+       val bosTokenId: Int
+       val vocabSize: Int
+       val specialTokens: Set<String>
+   }
+   ```
+
+2. **Unified tokenizer factory:**
+   - `TokenizerFactory.fromGGUF(source)` — auto-detects BPE/SentencePiece/WordPiece
+   - `TokenizerFactory.fromTokenizerJson(json)` — HuggingFace format
+   - Returns the correct implementation (byte-level BPE for GPT-2/Qwen, SentencePiece for LLaMA, etc.)
+
+3. **Move `GGUFTokenizer` to `llm-core`** so all runners can use it without depending on kllama
+
+## Phase 4: Unified Runner (single CLI entry point)
+
+**Problem:** 6 separate CLI apps with duplicated argument parsing, model loading, and dispatch logic.
+
+**Changes:**
+
+1. **Single `skainet` CLI** that auto-detects model architecture from GGUF metadata:
+   ```bash
+   skainet -m model.gguf "prompt"                    # auto-detect, generate
+   skainet -m model.gguf --chat                      # auto-detect, chat mode
+   skainet -m model.gguf --demo "What is 2+2?"       # auto-detect, tool calling
+   ```
+
+2. **Architecture registry:**
+   ```kotlin
+   ModelRegistry.register("llama", ::llamaNetwork)
+   ModelRegistry.register("qwen3", ::qwenNetwork)
+   ModelRegistry.register("gemma", ::gemmaNetwork)
+   ```
+
+3. **Auto-detection from GGUF metadata** (already exists in `peekGgufMetadata()`)
+
+## Verification
+
+- All existing unit tests pass (`llm-agent`, `llm-runtime:kllama`, `llm-core`)
+- Smoke test suite passes (generation + tool calling)
+- Basic generation produces identical output for all model families
+- Tool calling works for any model that supports ChatML/Qwen/Llama3 templates
+- `OptimizedLLMRuntime` in HYBRID mode matches hand-coded runtime output
+
+## Suggested Implementation Order
+
+1. **Phase 1** first — immediately unblocks tool calling for all models
+2. **Phase 3** next — reduces fragility (the GGUFTokenizer byte-level BPE issue)
+3. **Phase 2** then — biggest refactor, needs per-model validation
+4. **Phase 4** last — depends on all other phases
diff --git a/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/chat/ChatSession.kt b/llm-agent/src/commonMain/kotlin/sk/ainet/apps/kllama/chat/ChatSession.kt
@@ -0,0 +1,113 @@
+package sk.ainet.apps.kllama.chat
+
+import sk.ainet.apps.llm.InferenceRuntime
+import sk.ainet.apps.llm.Tokenizer
+import sk.ainet.lang.types.DType
+
+/**
+ * Bundles an [InferenceRuntime] with a [Tokenizer] and [ModelMetadata] to provide
+ * chat, agent, and tool-calling capabilities for any model.
+ *
+ * This decouples tool calling from any specific runner — any runner that can
+ * produce an [InferenceRuntime] and a [Tokenizer] can create a [ChatSession]
+ * and get chat/agent/demo modes for free.
+ *
+ * Usage:
+ * ```kotlin
+ * val session = ChatSession(runtime, tokenizer, metadata)
+ * session.chat(maxTokens = 512, temperature = 0.7f)       // interactive chat
+ * session.agent(maxTokens = 512, temperature = 0.7f)      // interactive agent with tools
+ * session.demo(maxTokens = 256, temperature = 0.7f)       // interactive tool calling demo
+ * session.demoSingleShot("What is 2+2?", maxTokens = 256) // non-interactive single prompt
+ * ```
+ */
+public class ChatSession<T : DType>(
+    public val runtime: InferenceRuntime<T>,
+    public val tokenizer: Tokenizer,
+    public val metadata: ModelMetadata = ModelMetadata(),
+    templateName: String? = null
+) {
+    private val provider: ToolCallingSupport = ToolCallingSupportResolver.resolveOrFallback(metadata, templateName)
+    private val template: ChatTemplate = provider.createChatTemplate()
+
+    /**
+     * Run a single agent round with the given prompt and tools.
+     * Returns the final response text. Non-interactive — suitable for smoke tests.
+     *
+     * @param prompt The user prompt.
+     * @param tools Tools to register. If empty, uses default calculator + list_files.
+     * @param maxTokens Maximum tokens per generation round.
+     * @param temperature Sampling temperature.
+     * @param listener Optional listener for observing the agent loop.
+     * @return The final assistant response text.
+     */
+    public fun runSingleTurn(
+        prompt: String,
+        tools: List<Tool> = emptyList(),
+        maxTokens: Int = 256,
+        temperature: Float = 0.7f,
+        listener: AgentListener? = null
+    ): String {
+        val registry = ToolRegistry()
+        tools.forEach { registry.register(it) }
+
+        val agentLoop = AgentLoop(
+            runtime = runtime,
+            template = template,
+            toolRegistry = registry,
+            eosTokenId = tokenizer.eosTokenId,
+            config = AgentConfig(
+                maxToolRounds = 5,
+                maxTokensPerRound = maxTokens,
+                temperature = temperature
+            ),
+            decode = { tokenId -> tokenizer.decode(tokenId) }
+        )
+
+        val systemPrompt = "You are a helpful assistant with access to tools."
+        val messages = mutableListOf(
+            ChatMessage(role = ChatRole.SYSTEM, content = systemPrompt),
+            ChatMessage(role = ChatRole.USER, content = prompt)
+        )
+
+        return agentLoop.runWithEncoder(
+            messages = messages,
+            encode = { text -> tokenizer.encode(text) },
+            listener = listener
+        )
+    }
+
+    /**
+     * Create an [AgentLoop] configured for this session.
+     */
+    public fun createAgentLoop(
+        toolRegistry: ToolRegistry,
+        maxTokens: Int = 512,
+        temperature: Float = 0.7f
+    ): AgentLoop<T> {
+        return AgentLoop(
+            runtime = runtime,
+            template = template,
+            toolRegistry = toolRegistry,
+            eosTokenId = tokenizer.eosTokenId,
+            config = AgentConfig(
+                maxToolRounds = 5,
+                maxTokensPerRound = maxTokens,
+                temperature = temperature
+            ),
+            decode = { tokenId -> tokenizer.decode(tokenId) }
+        )
+    }
+
+    /** The resolved chat template for this session. */
+    public val chatTemplate: ChatTemplate get() = template
+
+    /** The resolved tool calling provider family name. */
+    public val providerFamily: String get() = provider.family
+
+    /** Encode text to token IDs using this session's tokenizer. */
+    public fun encode(text: String): IntArray = tokenizer.encode(text)
+
+    /** Decode a token ID to text using this session's tokenizer. */
+    public fun decode(tokenId: Int): String = tokenizer.decode(tokenId)
+}
diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/Tokenizer.kt b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/Tokenizer.kt
@@ -4,4 +4,13 @@ interface Tokenizer {
     fun encode(text: String): IntArray
     fun decode(tokens: IntArray): String
     fun decode(token: Int): String
+
+    /** End-of-sequence token ID. */
+    val eosTokenId: Int
+
+    /** Beginning-of-sequence token ID. */
+    val bosTokenId: Int
+
+    /** Total vocabulary size. */
+    val vocabSize: Int
 }
diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/tokenizer/HuggingFaceBPETokenizer.kt b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/tokenizer/HuggingFaceBPETokenizer.kt
@@ -15,14 +15,14 @@ public class HuggingFaceBPETokenizer internal constructor(
     private val vocab: List<String>,
     private val tokenToId: Map<String, Int>,
     private val scores: FloatArray,
-    private val bosTokenId: Int,
-    private val eosTokenId: Int,
+    override val bosTokenId: Int,
+    override val eosTokenId: Int,
     private val unkTokenId: Int,
     private val addBosToken: Boolean,
     private val addEosToken: Boolean
 ) : Tokenizer {
 
-    public val vocabSize: Int get() = vocab.size
+    override val vocabSize: Int get() = vocab.size
 
     override fun encode(text: String): IntArray {
         if (text.isEmpty()) return intArrayOf()
diff --git a/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/HuggingFaceTokenizer.kt b/llm-inference/bert/src/commonMain/kotlin/sk/ainet/models/bert/HuggingFaceTokenizer.kt
@@ -40,6 +40,9 @@ public class HuggingFaceTokenizer private constructor(
     private val maxLength: Int = 512
 ) : Tokenizer {
 
+    override val eosTokenId: Int get() = tokenToId[SEP_TOKEN] ?: 102
+    override val bosTokenId: Int get() = tokenToId[CLS_TOKEN] ?: 101
+
     public companion object {
         private const val CLS_TOKEN = "[CLS]"
         private const val SEP_TOKEN = "[SEP]"
@@ -112,7 +115,7 @@ public class HuggingFaceTokenizer private constructor(
     private val sepId: Int = tokenToId[SEP_TOKEN] ?: error("Vocab missing $SEP_TOKEN")
     private val unkId: Int = tokenToId[UNK_TOKEN] ?: error("Vocab missing $UNK_TOKEN")
 
-    public val vocabSize: Int get() = tokenToId.size
+    override val vocabSize: Int get() = tokenToId.size
 
     /**
      * Encode text into token IDs with [CLS] and [SEP] tokens.
diff --git a/llm-inference/voxtral/src/commonMain/kotlin/sk/ainet/models/voxtral/TekkenTokenizerAdapter.kt b/llm-inference/voxtral/src/commonMain/kotlin/sk/ainet/models/voxtral/TekkenTokenizerAdapter.kt
@@ -25,6 +25,10 @@ public class TekkenTokenizerAdapter(
 
     override fun decode(token: Int): String = tekken.decode(token)
 
+    override val eosTokenId: Int = 2
+    override val bosTokenId: Int = 1
+    override val vocabSize: Int get() = 32768 // Tekken default
+
     public companion object {
         /**
          * Parse a tekken.json string and return a [Tokenizer] instance.
diff --git a/llm-runtime/kllama/src/commonMain/kotlin/sk/ainet/apps/kllama/GGUFTokenizer.kt b/llm-runtime/kllama/src/commonMain/kotlin/sk/ainet/apps/kllama/GGUFTokenizer.kt
@@ -24,8 +24,8 @@ import sk.ainet.io.gguf.StreamingGGUFReader
 class GGUFTokenizer private constructor(
     private val vocab: List<String>,
     private val scores: FloatArray,
-    private val bosTokenId: Int,
-    private val eosTokenId: Int,
+    private val _bosTokenId: Int,
+    private val _eosTokenId: Int,
     private val unkTokenId: Int,
     private val strategy: TokenizerStrategy
 ) : Tokenizer {
@@ -516,16 +516,18 @@ class GGUFTokenizer private constructor(
         }
     }
 
-    val vocabSize: Int get() = vocab.size
-
     /** The detected tokenizer type/strategy in use */
     val tokenizerType: TokenizerType get() = strategy.type
 
-    /** The BOS (beginning of sentence) token ID */
-    val bosId: Int get() = bosTokenId
+    override val bosTokenId: Int get() = _bosTokenId
+    override val eosTokenId: Int get() = _eosTokenId
+    override val vocabSize: Int get() = vocab.size
+
+    @Deprecated("Use eosTokenId", replaceWith = ReplaceWith("eosTokenId"))
+    val eosId: Int get() = _eosTokenId
 
-    /** The EOS (end of sentence) token ID */
-    val eosId: Int get() = eosTokenId
+    @Deprecated("Use bosTokenId", replaceWith = ReplaceWith("bosTokenId"))
+    val bosId: Int get() = _bosTokenId
 
     // Build reverse lookup for encoding
     private val tokenToId: Map<String, Int> by lazy {
diff --git a/llm-runtime/kllama/src/commonMain/kotlin/sk/ainet/apps/kllama/TokenizerImpl.kt b/llm-runtime/kllama/src/commonMain/kotlin/sk/ainet/apps/kllama/TokenizerImpl.kt
@@ -5,9 +5,11 @@ import sk.ainet.apps.llm.Tokenizer
 class TokenizerImpl(
     private val vocab: Array<String?>,
     private val vocabScores: FloatArray,
+    override val bosTokenId: Int = 1,
+    override val eosTokenId: Int = 2,
 ) : Tokenizer {
 
-    private val vocabSize = vocab.size
+    override val vocabSize: Int = vocab.size
     // ----------------------------------------------------------------------------
     // byte pair encoding (BPE) tokenizer, encodes strings into tokens so we can prompt
     private fun strLookup(str: String, vocabSize: Int): Int {
diff --git a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/chat/java/JavaAgentLoop.kt b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/chat/java/JavaAgentLoop.kt
@@ -51,9 +51,7 @@ public class JavaAgentLoop private constructor(
             runtime = session.runtime,
             template = template,
             toolRegistry = toolRegistry,
-            eosTokenId = session.tokenizer.let {
-                if (it is sk.ainet.apps.kllama.GGUFTokenizer) it.eosId else 2
-            },
+            eosTokenId = session.tokenizer.eosTokenId,
             config = config,
             decode = { session.tokenizer.decode(it) }
         )
@@ -81,9 +79,7 @@ public class JavaAgentLoop private constructor(
             runtime = session.runtime,
             template = template,
             toolRegistry = toolRegistry,
-            eosTokenId = session.tokenizer.let {
-                if (it is sk.ainet.apps.kllama.GGUFTokenizer) it.eosId else 2
-            },
+            eosTokenId = session.tokenizer.eosTokenId,
             config = config,
             decode = { session.tokenizer.decode(it) }
         )
diff --git a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/cli/AgentMain.kt b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/cli/AgentMain.kt
diff --git a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/cli/Main.kt b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/cli/Main.kt
diff --git a/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/cli/ToolCallingDemo.kt b/llm-runtime/kllama/src/jvmMain/kotlin/sk/ainet/apps/kllama/cli/ToolCallingDemo.kt
