cuttlefish-host 0.8.2

//! Inference served by a local [Ollama](https://ollama.com) instance.
//!
//! Ollama is the first real backend because it is the shortest path to a real
//! model: it is already an HTTP server that streams, so this is a client rather
//! than an embedding of llama.cpp. Embedding llama.cpp directly remains
//! worthwhile later — it removes a process boundary and gives direct control of
//! the KV cache — but it is a much larger commitment, and nothing above
//! [`InferBackend`] can tell the difference.
//!
//! # The wire format
//!
//! `POST /api/generate` with `stream: true` responds with newline-delimited
//! JSON — one object per token, then a final object carrying counts:
//!
//! ```text
//! {"model":"llama3.2:1b","response":"Hello","done":false}
//! {"model":"llama3.2:1b","response":" there","done":false}
//! {"model":"llama3.2:1b","response":"","done":true,"done_reason":"stop",
//!  "prompt_eval_count":28,"eval_count":3, ...}
//! ```
//!
//! Token-at-a-time delivery is what makes a guest's early-stop verdict
//! meaningful: the host can act on it while generation is still running rather
//! than after the fact.

use crate::infer::{InferBackend, InferRequest, InferResult};
use async_trait::async_trait;
use futures_util::StreamExt;
use serde::Deserialize;

/// Where Ollama listens when nothing says otherwise.
pub const DEFAULT_HOST: &str = "http://localhost:11434";

/// One line of Ollama's streaming response.
///
/// Deliberately partial: Ollama sends more fields than this (timings, and a
/// `context` array that can run to thousands of integers). Naming only what is
/// used keeps the deserializer from being coupled to fields the project does
/// not care about, and avoids materializing that context array on every call.
#[derive(Debug, Deserialize)]
struct Chunk {
    #[serde(default)]
    response: String,
    #[serde(default)]
    done: bool,
    #[serde(default)]
    prompt_eval_count: Option<u32>,
    #[serde(default)]
    eval_count: Option<u32>,
    /// Present when Ollama itself reports a problem — an unknown model, say.
    #[serde(default)]
    error: Option<String>,
}

/// Serves inference from a local Ollama instance.
pub struct OllamaBackend {
    client: reqwest::Client,
    host: String,
    model: String,
}

impl OllamaBackend {
    /// Target `model` on the Ollama instance at `host`.
    ///
    /// `model` is the name as Ollama knows it, `:tag` included.
    pub fn new(host: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            client: reqwest::Client::new(),
            // A trailing slash here would produce `//api/generate`, which Ollama
            // rejects — easy to introduce via an environment variable.
            host: host.into().trim_end_matches('/').to_string(),
            model: model.into(),
        }
    }

    /// Read `OLLAMA_HOST` if set, otherwise [`DEFAULT_HOST`].
    ///
    /// Named for the variable Ollama's own tooling uses, so an operator who has
    /// already pointed their CLI at a non-default instance does not have to
    /// configure this separately.
    pub fn host_from_env() -> String {
        std::env::var("OLLAMA_HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string())
    }
}

#[async_trait]
impl InferBackend for OllamaBackend {
    async fn infer(
        &self,
        req: InferRequest<'_>,
        on_token: &mut (dyn for<'t> FnMut(&'t str) -> bool + Send),
    ) -> anyhow::Result<InferResult> {
        // Ollama takes images as base64 strings alongside the prompt. Whether
        // the *model* can use them is Ollama's business — asking a text-only
        // model produces a clear error from it, which is more useful than a
        // guess made here from a model name.
        let images: Vec<String> = req
            .images
            .iter()
            .map(|bytes| {
                use base64::Engine;
                base64::engine::general_purpose::STANDARD.encode(bytes)
            })
            .collect();

        let response = self
            .client
            .post(format!("{}/api/generate", self.host))
            .json(&serde_json::json!({
                "model": self.model,
                "prompt": req.prompt,
                "stream": true,
                "images": images,
                "options": { "num_predict": req.max_tokens },
            }))
            .send()
            .await
            .map_err(|e| {
                // The overwhelmingly likely cause is that Ollama is not running,
                // and reqwest's own message does not say so.
                anyhow::anyhow!(
                    "could not reach Ollama at {} ({e}). Is it running? \
                     Set OLLAMA_HOST to point elsewhere.",
                    self.host
                )
            })?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            anyhow::bail!("Ollama returned {status}: {body}");
        }

        let mut text = String::new();
        let mut tokens_in = 0;
        let mut tokens_out = 0;

        // Responses are newline-delimited JSON, and a chunk boundary need not
        // fall on a newline — a line can arrive split across two chunks, and one
        // chunk can carry several lines. Buffering and splitting on '\n' is what
        // makes this correct rather than usually-correct.
        let mut stream = response.bytes_stream();
        let mut buf = String::new();

        'outer: while let Some(chunk) = stream.next().await {
            buf.push_str(&String::from_utf8_lossy(&chunk?));

            while let Some(newline) = buf.find('\n') {
                // Take everything before the newline; leave the remainder in the
                // buffer as the start of the next line.
                let line: String = buf.drain(..=newline).collect();
                let line = line.trim();
                if line.is_empty() {
                    continue;
                }

                let chunk: Chunk = serde_json::from_str(line).map_err(|e| {
                    anyhow::anyhow!("malformed response from Ollama: {e} in {line}")
                })?;

                if let Some(error) = chunk.error {
                    anyhow::bail!("Ollama error: {error}");
                }

                if !chunk.response.is_empty() {
                    text.push_str(&chunk.response);
                    tokens_out += 1;
                    if !on_token(&chunk.response) {
                        // The guest asked to stop. Dropping the stream closes
                        // the connection, which is how Ollama learns to stop
                        // generating — there is no separate cancel call.
                        break 'outer;
                    }
                }

                if chunk.done {
                    // The final message carries authoritative counts; prefer
                    // them over the tokens counted above, which only sees
                    // non-empty responses.
                    tokens_in = chunk.prompt_eval_count.unwrap_or(0);
                    tokens_out = chunk.eval_count.unwrap_or(tokens_out);
                    break 'outer;
                }
            }
        }

        Ok(InferResult {
            text,
            tokens_in,
            tokens_out,
        })
    }

    fn model_name(&self) -> String {
        self.model.clone()
    }

    /// True, but the honest claim is narrower than the signature allows:
    /// whether images work depends on the *model*, not on this backend, and
    /// this type does not know which model it is pointed at until a request is
    /// made.
    ///
    /// Forwarding is nonetheless the better behaviour, because Ollama answers
    /// the question authoritatively and specifically — HTTP 400 with
    /// "Multimodal data provided, but model does not support multimodal
    /// requests" (verified against llama3.2:1b). Guessing here from a model
    /// name would produce a worse message and a new way to be wrong, since the
    /// set of vision models changes without this code changing.
    fn supports_images(&self) -> bool {
        true
    }

    async fn embed(&self, texts: &[String]) -> anyhow::Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }
        // `/api/embed` (not the older `/api/embeddings`) because it takes an
        // array and answers with one vector per input. The legacy endpoint
        // is one text per request, which on a corpus is the difference
        // between minutes and hours.
        let response = self
            .client
            .post(format!("{}/api/embed", self.host))
            .json(&serde_json::json!({ "model": self.model, "input": texts }))
            .send()
            .await
            .map_err(|e| anyhow::anyhow!("embedding via {}: {e}", self.host))?;

        let status = response.status();
        let body = response
            .text()
            .await
            .map_err(|e| anyhow::anyhow!("reading the embedding response: {e}"))?;
        if !status.is_success() {
            anyhow::bail!("embedding with `{}`: {status}: {body}", self.model);
        }

        let parsed: serde_json::Value = serde_json::from_str(&body)
            .map_err(|e| anyhow::anyhow!("embedding response was not JSON: {e}: {body}"))?;
        let rows = parsed
            .get("embeddings")
            .and_then(|v| v.as_array())
            .ok_or_else(|| {
                anyhow::anyhow!("embedding response had no `embeddings` array: {body}")
            })?;

        // One vector per input, in order. A mismatch means the rows can no
        // longer be paired with their texts, and pairing them wrongly is
        // worse than failing: every downstream similarity would be against
        // the wrong content.
        if rows.len() != texts.len() {
            anyhow::bail!(
                "asked `{}` for {} embeddings and got {} — refusing to pair vectors with \
                 texts by guesswork",
                self.model,
                texts.len(),
                rows.len()
            );
        }

        let mut out = Vec::with_capacity(rows.len());
        for (i, row) in rows.iter().enumerate() {
            let vector: Vec<f32> = row
                .as_array()
                .ok_or_else(|| anyhow::anyhow!("embedding {i} was not an array: {row}"))?
                .iter()
                .map(|v| v.as_f64().unwrap_or(f64::NAN) as f32)
                .collect();
            if vector.is_empty() || vector.iter().any(|v| v.is_nan()) {
                anyhow::bail!(
                    "embedding {i} from `{}` is empty or contains non-numbers; storing it \
                     would poison every similarity computed against it",
                    self.model
                );
            }
            out.push(vector);
        }
        Ok(out)
    }

    fn supports_embeddings(&self) -> bool {
        true
    }
}

/// Builds [`OllamaBackend`], registered as the `ollama` provider.
pub struct OllamaFactory;

impl crate::backend::BackendFactory for OllamaFactory {
    fn provider(&self) -> &'static str {
        "ollama"
    }

    fn describe(&self) -> &'static str {
        "a local Ollama instance; target is a model tag such as `llama3.2:1b`"
    }

    fn build(&self, target: &str) -> anyhow::Result<std::sync::Arc<dyn InferBackend>> {
        if target.is_empty() {
            anyhow::bail!("an Ollama model name is required, e.g. `llama3.2:1b`");
        }
        // Reachability is deliberately not checked here — see `BackendFactory`.
        Ok(std::sync::Arc::new(OllamaBackend::new(
            OllamaBackend::host_from_env(),
            target,
        )))
    }
}