From 24df3b55b35b645393cdc82b28a48e8f2d6b60ca Mon Sep 17 00:00:00 2001 From: kami Date: Mon, 18 May 2026 21:39:27 +0400 Subject: [PATCH] feat(infra:llama_cpp): GBNF grammar converter, tool_calls in ChatMessage, tool definitions in request, L2 context handling --- .../inference/llama_cpp/build.gradle | 5 + .../llama/cpp/GbnfGrammarConverter.kt | 64 +++++++++++++ .../inference/llama/cpp/LlamaCppApiModels.kt | 11 ++- .../llama/cpp/LlamaCppInferenceProvider.kt | 34 +++++-- .../llama/cpp/GbnfGrammarConverterTest.kt | 93 +++++++++++++++++++ 5 files changed, 200 insertions(+), 7 deletions(-) create mode 100644 infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverter.kt create mode 100644 infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverterTest.kt diff --git a/infrastructure/inference/llama_cpp/build.gradle b/infrastructure/inference/llama_cpp/build.gradle index adf1af14..ec110dc8 100644 --- a/infrastructure/inference/llama_cpp/build.gradle +++ b/infrastructure/inference/llama_cpp/build.gradle @@ -13,10 +13,15 @@ ext { dependencies { implementation project(':core:inference') implementation project(':core:events') + implementation project(':core:artifacts') implementation project(':core:context') implementation project(':infrastructure:inference:commons') + + implementation("com.fasterxml.jackson.core:jackson-databind:2.17.0") + implementation("com.fasterxml.jackson.module:jackson-module-kotlin:2.17.0") implementation "io.ktor:ktor-client-core:$ktor_version" implementation "io.ktor:ktor-client-cio:$ktor_version" implementation "io.ktor:ktor-client-content-negotiation:$ktor_version" implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version" + testImplementation "org.junit.jupiter:junit-jupiter" } diff --git a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverter.kt b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverter.kt new file mode 100644 index 00000000..ad0595f3 --- /dev/null +++ b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverter.kt @@ -0,0 +1,64 @@ +package com.correx.infrastructure.inference.llama.cpp + +import com.correx.core.artifacts.kind.JsonSchema + +internal object GbnfGrammarConverter { + fun convert(schema: JsonSchema): String { + require(schema.type == "object") { "Only object schemas supported" } + + val required = schema.required.toSet() + val allKeys = schema.properties.keys.toList() + + val sb = StringBuilder() + sb.appendLine("""root ::= "{" ws members ws "}"""") + sb.appendLine("""ws ::= [ \t\n\r]*""") + val hexDigits = "[0-9a-fA-F]" + val unicodeEscape = "\"u\" $hexDigits $hexDigits $hexDigits $hexDigits" + val closingQuote = "\"\\\"\"" // GBNF string literal matching a single " + val charClass = """[^"\\\x7F\x00-\x1F]""" + val escapeSeq = """"\\" (["\\/bfnrt] | $unicodeEscape)""" + val stringRule = """string ::= "\"" ($charClass | $escapeSeq)* $closingQuote""" + sb.appendLine(stringRule) + sb.appendLine("""number ::= "-"? [0-9]+""") + + // Build members rule: + // - Required keys are always present, joined by commas. + // - Optional keys are wrapped in (...)?; they always carry a leading comma so they + // can be appended after whatever came before without complicating the grammar. + // - When there are no required keys, the first optional key cannot have a leading comma. + // We handle this by making the first key (required or optional) comma-free, and all + // subsequent keys comma-prefixed (optional ones wrapped in (...)? around the comma too). + val memberParts = mutableListOf() + + for ((index, key) in allKeys.withIndex()) { + val valueRule = valueRuleFor(schema, key) + val isRequired = key in required + val pair = "\"\\\"$key\\\"\" ws \":\" ws $valueRule" + if (index == 0) { + // First key: never a leading comma + if (isRequired) memberParts.add(pair) else memberParts.add("($pair)?") + } else { + // Subsequent keys always have a leading comma, gated by optional wrapper if needed + if (isRequired) { + memberParts.add("ws \",\" ws $pair") + } else { + memberParts.add("(ws \",\" ws $pair)?") + } + } + } + + sb.append("members ::= ") + if (memberParts.isEmpty()) { + sb.appendLine("ws") + } else { + sb.appendLine(memberParts.joinToString(" ")) + } + + return sb.toString() + } + + private fun valueRuleFor(schema: JsonSchema, key: String): String { + val prop = schema.properties[key] ?: return "string" + return if (prop.type == "integer") "number" else "string" + } +} diff --git a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppApiModels.kt b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppApiModels.kt index 6b116d4b..65c86897 100644 --- a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppApiModels.kt +++ b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppApiModels.kt @@ -1,5 +1,7 @@ package com.correx.infrastructure.inference.llama.cpp +import com.correx.core.inference.ToolCallRequest +import com.correx.core.inference.ToolDefinition import kotlinx.serialization.SerialName import kotlinx.serialization.Serializable @@ -13,10 +15,17 @@ data class ChatCompletionRequest( @SerialName("stop") val stopSequences: List = emptyList(), val seed: Long? = null, val stream: Boolean = false, + val grammar: String? = null, + val tools: List? = null, ) @Serializable -data class ChatMessage(val role: String, val content: String) +data class ChatMessage( + val role: String, + val content: String?, + @SerialName("tool_calls") val toolCalls: List = emptyList(), + @SerialName("tool_call_id") val toolCallId: String? = null, +) @Serializable data class ChatCompletionResponse( diff --git a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt index 42b70d20..5ff49ac5 100644 --- a/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt +++ b/infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt @@ -9,6 +9,7 @@ import com.correx.core.inference.InferenceProvider import com.correx.core.inference.InferenceRequest import com.correx.core.inference.InferenceResponse import com.correx.core.inference.ProviderHealth +import com.correx.core.inference.ResponseFormat import com.correx.core.inference.TokenUsage import com.correx.core.inference.Tokenizer import com.correx.infrastructure.inference.commons.ModelDescriptor @@ -61,6 +62,11 @@ class LlamaCppInferenceProvider( val messages = buildMessages(request.contextPack) + val grammar = when (val fmt = request.responseFormat) { + is ResponseFormat.Json -> GbnfGrammarConverter.convert(fmt.schema) + ResponseFormat.Text -> null + } + val body = ChatCompletionRequest( model = modelId, messages = messages, @@ -70,6 +76,8 @@ class LlamaCppInferenceProvider( stopSequences = request.generationConfig.stopSequences, seed = request.generationConfig.seed, stream = false, + grammar = grammar, + tools = request.tools.takeIf { it.isNotEmpty() }, ) val response = httpClient.post("$baseUrl/v1/chat/completions") { @@ -78,19 +86,23 @@ class LlamaCppInferenceProvider( setBody(body) }.body() - val finishReason = response.choices.first().finishReason.lowercase().let { - if (it == "length") FinishReason.Length else FinishReason.Stop + val message = response.choices.first().message + val finishReason = when (response.choices.first().finishReason.lowercase()) { + "length" -> FinishReason.Length + "tool_calls" -> FinishReason.ToolCall + else -> FinishReason.Stop } return InferenceResponse( requestId = request.requestId, - text = response.choices.first().message.content, + text = message.content ?: "", finishReason = finishReason, tokensUsed = TokenUsage( promptTokens = response.usage.promptTokens, completionTokens = response.usage.completionTokens, ), latencyMs = System.currentTimeMillis() - startTime, + toolCalls = message.toolCalls, ) } @@ -108,19 +120,29 @@ class LlamaCppInferenceProvider( override fun capabilities(): Set = descriptor.capabilities private fun buildMessages(contextPack: ContextPack): List { - val messages = contextPack.layers.flatMap { (layer, entries) -> + val sorted = contextPack.layers.entries.sortedBy { it.key.ordinal } + val systemContent = sorted + .filter { it.key == ContextLayer.L0 } + .flatMap { it.value } + .joinToString("\n\n") { it.content } + .takeIf { it.isNotBlank() } + val conversationMessages = sorted.filter { it.key != ContextLayer.L0 }.flatMap { (layer, entries) -> when (layer) { - ContextLayer.L0 -> entries.map { ChatMessage("system", it.content) } ContextLayer.L1 -> entries.map { ChatMessage("user", it.content) } ContextLayer.L2 -> entries.map { ChatMessage( role = if (it.sourceType == "assistant") "assistant" else "user", - content = it.content + content = it.content, ) } + else -> emptyList() } } + val messages = buildList { + systemContent?.let { add(ChatMessage("system", it)) } + addAll(conversationMessages) + } return messages.ifEmpty { listOf(ChatMessage("user", "")) } } } diff --git a/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverterTest.kt b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverterTest.kt new file mode 100644 index 00000000..5327c7c8 --- /dev/null +++ b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/GbnfGrammarConverterTest.kt @@ -0,0 +1,93 @@ +package com.correx.infrastructure.inference.llama.cpp + +import com.correx.core.artifacts.kind.JsonSchema +import com.correx.core.artifacts.kind.JsonSchemaProperty +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test + +class GbnfGrammarConverterTest { + + /** FileWrittenPayload-shaped schema: all required fields. */ + private val allRequiredSchema = JsonSchema( + type = "object", + properties = mapOf( + "path" to JsonSchemaProperty(type = "string", description = "relative file path"), + "content" to JsonSchemaProperty(type = "string", description = "file content"), + "mode" to JsonSchemaProperty(type = "string", description = "octal permission mode"), + ), + required = listOf("path", "content", "mode"), + ) + + /** Schema with some required and some optional fields. */ + private val mixedSchema = JsonSchema( + type = "object", + properties = mapOf( + "name" to JsonSchemaProperty(type = "string"), + "count" to JsonSchemaProperty(type = "integer"), + "description" to JsonSchemaProperty(type = "string"), + ), + required = listOf("name"), + ) + + /** Schema with all optional fields (no required entries). */ + private val allOptionalSchema = JsonSchema( + type = "object", + properties = mapOf( + "label" to JsonSchemaProperty(type = "string"), + "value" to JsonSchemaProperty(type = "integer"), + ), + required = emptyList(), + ) + + @Test + fun `string rule does not contain quoted character class`() { + val grammar = GbnfGrammarConverter.convert(allRequiredSchema) + assertFalse(grammar.contains("\"["), "string rule must not contain quoted character class (\"[)") + } + + @Test + fun `all-required schema produces non-empty valid GBNF`() { + val grammar = GbnfGrammarConverter.convert(allRequiredSchema) + + assertFalse(grammar.isBlank(), "grammar must not be blank") + assertTrue(grammar.contains("root"), "grammar must declare root rule") + assertTrue(grammar.contains("members"), "grammar must declare members rule") + assertTrue(grammar.contains("path"), "grammar must reference 'path' key") + assertTrue(grammar.contains("content"), "grammar must reference 'content' key") + assertTrue(grammar.contains("mode"), "grammar must reference 'mode' key") + // Optional wrapper must not appear for required keys + assertFalse(grammar.substringAfter("members").trimStart().startsWith("("), + "first member of all-required schema must not be wrapped in optional ()?") + } + + @Test + fun `mixed required and optional schema produces non-empty valid GBNF`() { + val grammar = GbnfGrammarConverter.convert(mixedSchema) + + assertFalse(grammar.isBlank(), "grammar must not be blank") + assertTrue(grammar.contains("root"), "grammar must declare root rule") + assertTrue(grammar.contains("members"), "grammar must declare members rule") + assertTrue(grammar.contains("name"), "grammar must reference required key 'name'") + assertTrue(grammar.contains("count"), "grammar must reference optional key 'count'") + assertTrue(grammar.contains("description"), "grammar must reference optional key 'description'") + // Optional keys should be wrapped in (...)? + val membersLine = grammar.lines().first { it.startsWith("members") } + assertTrue(membersLine.contains(")?"), "optional fields must use (...)? in members rule") + } + + @Test + fun `all-optional schema produces non-empty valid GBNF without leading comma on first member`() { + val grammar = GbnfGrammarConverter.convert(allOptionalSchema) + + assertFalse(grammar.isBlank(), "grammar must not be blank") + assertTrue(grammar.contains("root"), "grammar must declare root rule") + assertTrue(grammar.contains("members"), "grammar must declare members rule") + // The members rule line must not start the first optional key with a comma + val membersLine = grammar.lines().first { it.startsWith("members") } + // First token after "members ::= " must be "(" (the optional wrapper), not a comma + val afterAssign = membersLine.substringAfter("::=").trimStart() + assertFalse(afterAssign.startsWith(","), "first member must not begin with a comma") + assertTrue(afterAssign.startsWith("("), "first optional member must be wrapped in (...)?") + } +}