package net.kentonc.si.agent
import kotlinx.serialization.json.JsonElement
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.RequestBody.Companion.toRequestBody
import java.io.IOException
import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.TreeMap
import java.util.concurrent.TimeUnit
/**
* Minimal OpenAI-compatible chat client (llama-server / LiteLLM / vLLM), a port
* of the streaming half of src/llm.rs. One call posts the whole history and
* reads the SSE stream; tool-call deltas arrive as fragments spread across many
* frames and are reassembled here, keyed by their wire index.
*/
class LlmClient(
private val baseUrl: String,
private val apiKey: String,
private val model: String,
private val toolsSpec: JsonElement,
private val maxTokens: Int,
private val temperature: Float,
) {
/** The instruction to send and everything the caller collects from the stream. */
data class Reply(val text: String, val toolCalls: List<ToolCall>)
/** Stream one completion into `onItem`, returning the accumulated reply text
* and any fragmented tool calls that arrived complete. */
fun streamReply(
messages: List<WireMessage>,
onItem: (Item) -> Unit,
): Reply {
val body = json.encodeToString(
ChatRequest.serializer(),
ChatRequest(
model = model,
messages = messages,
tools = toolsSpec,
maxTokens = maxTokens,
temperature = temperature,
stream = true,
),
)
val url = "${baseUrl.trimEnd('/')}/v1/chat/completions"
val request = Request.Builder()
.url(url)
.header("Authorization", "Bearer $apiKey")
.post(body.toRequestBody("application/json".toMediaType()))
.build()
val response = client.newCall(request).execute()
response.use { resp ->
if (!resp.isSuccessful) {
val errorText = resp.body?.string()?.take(4000).orEmpty()
throw IOException("HTTP ${resp.code}: $errorText")
}
val source = resp.body?.byteStream()
?: throw IOException("empty response body")
val reader = BufferedReader(InputStreamReader(source, Charsets.UTF_8))
val reply = StringBuilder()
val accumulator = ToolCallAccumulator()
loop@ while (true) {
val line = reader.readLine() ?: break // connection ended
if (!line.startsWith("data:")) continue // comments, "event:" lines, blank lines
val data = line.removePrefix("data:").trim()
if (data == "[DONE]") break@loop
val chunk = try { json.decodeFromString(StreamChunk.serializer(), data) }
catch (e: Exception) { null } ?: continue // tolerate frames we don't understand
for (choice in chunk.choices) {
val delta = choice.delta ?: continue
val text = flattenContent(delta.content)
if (!text.isNullOrEmpty()) {
reply.append(text)
onItem(Item.Text(text))
}
accumulator.push(delta.toolCalls)
}
}
return Reply(reply.toString(), accumulator.finish())
}
}
/** Same shape as the Rust ChatRequest. */
@kotlinx.serialization.Serializable
private data class ChatRequest(
val model: String,
val messages: List<WireMessage>,
val tools: JsonElement?,
@kotlinx.serialization.SerialName("max_tokens") val maxTokens: Int,
val temperature: Float,
val stream: Boolean,
)
companion object {
/** One client for the whole app process: connection pool and threads reused. */
val client: OkHttpClient = OkHttpClient.Builder()
.connectTimeout(30, TimeUnit.SECONDS)
// Per-chunk idle timeout, not a total cap: long generations are fine,
// but a stalled connection fails instead of hanging forever.
.readTimeout(120, TimeUnit.SECONDS)
.build()
}
}
/** Accumulates fragmented tool calls, keyed by their wire index (the Rust
* ToolCallAccumulator): the name and id show up once, `arguments` is a JSON
* string built up a few characters at a time. */
class ToolCallAccumulator {
private val calls = TreeMap<Int, ToolCall>()
fun push(deltas: List<ToolCallDelta>) {
for (delta in deltas) {
var call = calls.getOrPut(delta.index) {
ToolCall(function = FunctionCall(name = "", arguments = ""))
}
delta.id?.let { call = call.copy(id = it) }
delta.function?.let { fn ->
if (fn.name != null) call = call.copy(function = call.function.copy(name = fn.name))
if (fn.arguments != null) {
call = call.copy(function = call.function.copy(arguments = call.function.arguments + fn.arguments))
}
}
calls[delta.index] = call
}
}
fun finish(): List<ToolCall> = calls.values.mapIndexed { n, call ->
// Some servers omit ids; synthesize one so the tool reply still matches.
if (call.id.isEmpty()) call.copy(id = "call_$n") else call
}
}