Add props, slots, tokenize and the cache-loss check to the llama client

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)

crates/loopd/src/llama/info.rs

@@ -0,0 +1,145 @@
+//! Small non-streaming server calls and the cache-loss check.
+//!
+//! The responses are the server's format, so unknown fields are ignored; we read only the fields we
+//! need and treat anything else as a shape error.
+
+use std::time::Duration;
+
+use super::{Client, InferError, Timings};
+use crate::http::{Connection, HttpError, Request, read_capped};
+
+pub const MAX_BODY: usize = 4 * 1024 * 1024; // bodies that are not streams
+pub const MAX_ERROR_BODY: usize = 4 * 1024; // how much of an error body is kept
+pub const CACHE_TOLERANCE: u64 = 64;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Props {
+    pub chat_template: String,
+    pub n_ctx: u64,
+    pub total_slots: u32,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize)]
+pub struct SlotInfo {
+    pub id: u32,
+    pub is_processing: bool,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CacheOutcome {
+    Hit,
+    Loss { expected: u64, got: u64 },
+}
+
+pub fn cache_outcome(previous: &Timings, current: &Timings) -> CacheOutcome {
+    let expected = previous
+        .cache_n
+        .saturating_add(previous.prompt_n)
+        .saturating_add(previous.predicted_n);
+    if current.cache_n.saturating_add(CACHE_TOLERANCE) >= expected {
+        CacheOutcome::Hit
+    } else {
+        CacheOutcome::Loss {
+            expected,
+            got: current.cache_n,
+        }
+    }
+}
+
+pub(crate) fn map_http(e: HttpError) -> InferError {
+    match e {
+        HttpError::Connect(e) => InferError::Connect(e),
+        HttpError::Timeout => InferError::Stalled,
+        HttpError::Closed => InferError::StreamClosedEarly,
+        other => InferError::Protocol(other.to_string()),
+    }
+}
+
+pub(crate) fn error_text(bytes: &[u8]) -> String {
+    let text = String::from_utf8_lossy(bytes).into_owned();
+    if text.len() <= MAX_ERROR_BODY {
+        return text;
+    }
+    let mut end = MAX_ERROR_BODY;
+    while end > 0 && !text.is_char_boundary(end) {
+        end -= 1;
+    }
+    text[..end].to_string()
+}
+
+impl Client {
+    pub(crate) fn call(
+        &self,
+        method: &str,
+        path: &str,
+        body: Option<&[u8]>,
+    ) -> Result<Vec<u8>, InferError> {
+        let mut conn = Connection::open(&self.cfg.infer.socket).map_err(map_http)?;
+        conn.set_read_timeout(Duration::from_millis(self.cfg.limits.liveness_ms))
+            .map_err(map_http)?;
+        conn.send(&Request { method, path, body })
+            .map_err(map_http)?;
+        let head = conn.read_head().map_err(map_http)?;
+        let mut body = conn.body(&head).map_err(map_http)?;
+        let bytes = read_capped(&mut body, MAX_BODY).map_err(map_http)?;
+        if head.status != 200 {
+            return Err(InferError::Http {
+                status: head.status,
+                body: error_text(&bytes),
+            });
+        }
+        Ok(bytes)
+    }
+
+    pub fn props(&self) -> Result<Props, InferError> {
+        let path = format!("/props?model={}", self.cfg.infer.model);
+        let bytes = self.call("GET", &path, None)?;
+        let value: serde_json::Value = serde_json::from_slice(&bytes)
+            .map_err(|_| InferError::Protocol("props response is not JSON".to_string()))?;
+        let chat_template = value
+            .get("chat_template")
+            .and_then(|v| v.as_str())
+            .ok_or_else(|| {
+                InferError::Protocol("props response has no chat_template".to_string())
+            })?;
+        let n_ctx = value
+            .get("default_generation_settings")
+            .and_then(|v| v.get("n_ctx"))
+            .and_then(|v| v.as_u64())
+            .ok_or_else(|| InferError::Protocol("props response has no n_ctx".to_string()))?;
+        let total_slots = value
+            .get("total_slots")
+            .and_then(|v| v.as_u64())
+            .ok_or_else(|| InferError::Protocol("props response has no total_slots".to_string()))?;
+        Ok(Props {
+            chat_template: chat_template.to_string(),
+            n_ctx,
+            total_slots: u32::try_from(total_slots)
+                .map_err(|_| InferError::Protocol("props total_slots is too large".to_string()))?,
+        })
+    }
+
+    pub fn slots(&self) -> Result<Vec<SlotInfo>, InferError> {
+        let path = format!("/slots?model={}", self.cfg.infer.model);
+        let bytes = self.call("GET", &path, None)?;
+        let slots: Vec<SlotInfo> = serde_json::from_slice(&bytes)
+            .map_err(|_| InferError::Protocol("slots response is not a list".to_string()))?;
+        Ok(slots)
+    }
+
+    pub fn tokenize(&self, text: &str) -> Result<usize, InferError> {
+        let body = serde_json::to_string(&serde_json::json!({
+            "model": self.cfg.infer.model,
+            "content": text,
+        }))
+        .map_err(|_| InferError::Protocol("could not encode tokenize request".to_string()))?;
+        let bytes = self.call("POST", "/tokenize", Some(body.as_bytes()))?;
+        let value: serde_json::Value = serde_json::from_slice(&bytes)
+            .map_err(|_| InferError::Protocol("tokenize response is not JSON".to_string()))?;
+        value
+            .get("tokens")
+            .and_then(|v| v.as_array())
+            .map(Vec::len)
+            .ok_or_else(|| InferError::Protocol("tokenize response has no tokens".to_string()))
+    }
+}

crates/loopd/src/llama/mod.rs

@@ -1,6 +1,7 @@
 //! The llama inference server's chat-completions types and the request builder.
 
 pub mod assemble;
+pub mod info;
 pub mod request;
 
 #[derive(Debug, Clone, PartialEq, Eq)]