kibble 0.1.0

//! Embedding backends: an OpenAI-compatible endpoint plus a deterministic stub
//! for tests. Batching is the caller's job (see `vectors::get_or_embed`);
//! `embed_batch` embeds exactly the slice it is given, in order.

use crate::config::EmbedConfig;
use std::io;

#[allow(async_fn_in_trait)]
pub trait Embedder {
    /// Embed each text; returns one vector per input, in order.
    async fn embed_batch(&self, texts: &[String]) -> io::Result<Vec<Vec<f32>>>;
}

/// Deterministic stub for tests: identical text -> identical vector.
#[cfg(test)]
pub struct StubEmbedder {
    pub dim: usize,
}
#[cfg(test)]
impl StubEmbedder {
    pub fn new() -> Self {
        StubEmbedder { dim: 16 }
    }
    fn vec_for(&self, text: &str) -> Vec<f32> {
        let mut v = vec![0f32; self.dim];
        let mut h: u64 = 0xcbf29ce484222325;
        for b in text.bytes() {
            h ^= b as u64;
            h = h.wrapping_mul(0x100000001b3);
        }
        for (i, slot) in v.iter_mut().enumerate() {
            let hh = h.wrapping_add(i as u64).wrapping_mul(0x9E3779B97F4A7C15);
            *slot = ((hh >> 16) & 0xffff) as f32 / 65535.0;
        }
        v
    }
}
#[cfg(test)]
impl Default for StubEmbedder {
    fn default() -> Self {
        Self::new()
    }
}
#[cfg(test)]
impl Embedder for StubEmbedder {
    async fn embed_batch(&self, texts: &[String]) -> io::Result<Vec<Vec<f32>>> {
        Ok(texts.iter().map(|t| self.vec_for(t)).collect())
    }
}

/// Real backend: POSTs to `{base_url}/embeddings` (OpenAI-compatible).
pub struct EndpointEmbedder {
    client: reqwest::Client,
    base_url: String,
    model: String,
    api_key: String,
}
impl EndpointEmbedder {
    pub fn new(cfg: &EmbedConfig, proxy: Option<&str>) -> io::Result<Self> {
        let client = crate::net::build_client(proxy).map_err(|e| io::Error::other(e.to_string()))?;
        let api_key = std::env::var("EMBED_API_KEY")
            .or_else(|_| std::env::var("OPENAI_API_KEY"))
            .unwrap_or_default();
        Ok(EndpointEmbedder {
            client,
            base_url: cfg.base_url.clone(),
            model: cfg.model.clone(),
            api_key,
        })
    }
}
/// Max chars sent to the embedding endpoint. nomic-embed-v1.5's gguf hard-caps at its
/// 2048-token trained context (llama-server ignores `-c`/`-ub` for embeddings), and an
/// over-long input 500s the whole batch. 2000 chars is < 2048 tokens even at the theoretical
/// 1-char/token max density, so it never trips the cap — plenty of signal for clustering/dedup.
const MAX_EMBED_CHARS: usize = 2000;

/// Truncate an embedding input to `MAX_EMBED_CHARS` on a char boundary.
fn cap_embed_input(text: &str) -> String {
    if text.chars().count() > MAX_EMBED_CHARS {
        text.chars().take(MAX_EMBED_CHARS).collect()
    } else {
        text.to_string()
    }
}

impl Embedder for EndpointEmbedder {
    async fn embed_batch(&self, texts: &[String]) -> io::Result<Vec<Vec<f32>>> {
        // Embedding models have a bounded context (nomic-embed-v1.5 = 2048 tokens); an over-long
        // input makes the server 500 and fails the WHOLE batch, silently disabling clustering/
        // rebalance/classify. Cap each input to a conservative char budget (safe under 2048 tokens
        // even for dense text) — the leading text carries the topic signal for clustering/dedup,
        // and truncation beats losing the entire embed stage.
        let capped: Vec<String> = texts.iter().map(|t| cap_embed_input(t)).collect();
        let body = serde_json::json!({ "model": self.model, "input": capped }).to_string();
        let mut req = self
            .client
            .post(format!("{}/embeddings", self.base_url))
            .header("content-type", "application/json");
        if !self.api_key.is_empty() {
            req = req.header("authorization", format!("Bearer {}", self.api_key));
        }
        let text = req
            .body(body)
            .send()
            .await
            .and_then(|r| r.error_for_status())
            .map_err(|e| io::Error::other(format!("embed endpoint failed: {e}")))?
            .text()
            .await
            .map_err(|e| io::Error::other(format!("embed read failed: {e}")))?;
        let v: serde_json::Value = serde_json::from_str(&text)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
        let data = v
            .get("data")
            .and_then(|d| d.as_array())
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "no data[] in embeddings response"))?;
        if data.len() != texts.len() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("embeddings count mismatch: got {} for {} inputs", data.len(), texts.len()),
            ));
        }
        let mut out = Vec::with_capacity(data.len());
        for item in data {
            let arr = item
                .get("embedding")
                .and_then(|e| e.as_array())
                .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "no embedding[] in data item"))?;
            out.push(arr.iter().map(|x| x.as_f64().unwrap_or(0.0) as f32).collect());
        }
        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cap_embed_input_truncates_long_only() {
        let short = "hello world";
        assert_eq!(cap_embed_input(short), short);
        let long: String = "x".repeat(MAX_EMBED_CHARS + 500);
        assert_eq!(cap_embed_input(&long).chars().count(), MAX_EMBED_CHARS);
    }

    #[tokio::test]
    async fn stub_is_deterministic_and_order_preserving() {
        let e = StubEmbedder::new();
        let a = e.embed_batch(&["foo".into(), "bar".into(), "foo".into()]).await.unwrap();
        assert_eq!(a.len(), 3);
        assert_eq!(a[0].len(), 16);
        assert_eq!(a[0], a[2]); // identical text -> identical vector
        assert_ne!(a[0], a[1]);
    }

    #[tokio::test]
    async fn endpoint_parses_embeddings_response() {
        // Minimal one-shot HTTP mock returning an OpenAI embeddings body.
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let server = tokio::spawn(async move {
            let (mut sock, _) = listener.accept().await.unwrap();
            use tokio::io::{AsyncReadExt, AsyncWriteExt};
            let mut buf = [0u8; 2048];
            let _ = sock.read(&mut buf).await.unwrap();
            let json = r#"{"data":[{"embedding":[0.1,0.2,0.3]},{"embedding":[0.4,0.5,0.6]}]}"#;
            let resp = format!(
                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                json.len(), json
            );
            sock.write_all(resp.as_bytes()).await.unwrap();
        });
        let cfg = EmbedConfig { base_url: format!("http://{addr}"), model: "m".into(), batch_size: 64, store: "x".into() };
        let emb = EndpointEmbedder::new(&cfg, None).unwrap();
        let out = emb.embed_batch(&["a".into(), "b".into()]).await.unwrap();
        server.await.unwrap();
        assert_eq!(out, vec![vec![0.1f32, 0.2, 0.3], vec![0.4, 0.5, 0.6]]);
    }
}