tawasuyu/brahman

1//! Modo embeddings: usa fastembed-rs (ONNX Runtime) para producir
2//! vectores reales de text-embedding.
3//!
4//! Modelo default: `all-MiniLM-L6-v2` (384-d). Se descarga al primer
5//! arranque a `~/.cache/fastembed` y queda cacheado.
6//!
7//! ## Mapeo del contrato
8//!
9//! - `EmbedText`: pasa el texto al modelo, devuelve el vector 384-d.
10//! - `EmbedFile`: lee hasta los primeros 8 KiB del archivo, los
11//!   interpreta como UTF-8 con replacement-char, y los embeda como
12//!   texto. Para archivos binarios el resultado no es semánticamente
13//!   útil — caller decide qué hacer.
14//! - `Ping`: devuelve `model_id` y `embed_dim` reales.
15
16use std::fs::File;
17use std::io::Read;
18use std::path::PathBuf;
19use std::sync::Arc;
20use std::time::Instant;
21
22use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
23use nouser_nous::{
24    EmbedFilePayload, EmbedRequest, EmbedResponse, EmbedTextPayload, ErrorResponse, PingResponse,
25    RequestKind,
26};
27use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
28use tokio::net::UnixStream;
29use tracing::{info, warn};
30
31use crate::cache::EmbedCache;
32
33const MAX_FILE_BYTES: usize = 8192;
34
35/// Backend concreto: posee el modelo cargado.
36pub struct Backend {
37    model: TextEmbedding,
38}
39
40impl Backend {
41    pub fn init() -> Result<Self, String> {
42        info!("cargando modelo all-MiniLM-L6-v2 (puede descargar ~80MB la primera vez)");
43        let opts = InitOptions::new(EmbeddingModel::AllMiniLML6V2)
44            .with_show_download_progress(true);
45        let model = TextEmbedding::try_new(opts).map_err(|e| format!("fastembed init: {e}"))?;
46        info!("modelo listo");
47        Ok(Self { model })
48    }
49
50    fn embed_one(&self, text: &str) -> Result<Vec<f32>, String> {
51        let out = self
52            .model
53            .embed(vec![text], None)
54            .map_err(|e| format!("embed: {e}"))?;
55        out.into_iter()
56            .next()
57            .ok_or_else(|| "fastembed devolvió 0 vectores".to_string())
58    }
59}
60
61pub async fn handle_conn(
62    stream: UnixStream,
63    backend: Arc<Backend>,
64    cache: Option<EmbedCache>,
65) -> std::io::Result<()> {
66    let mut reader = BufReader::new(stream);
67    let mut line = String::new();
68    let n = reader.read_line(&mut line).await?;
69    if n == 0 {
70        return Ok(());
71    }
72
73    let req: EmbedRequest = match serde_json::from_str(&line) {
74        Ok(r) => r,
75        Err(e) => {
76            return write_error(reader.into_inner(), format!("JSON inválido: {e}")).await;
77        }
78    };
79
80    let started = Instant::now();
81    let result = match req.kind {
82        RequestKind::EmbedFile => handle_file(req.payload, &backend, cache.as_ref(), started),
83        RequestKind::EmbedText => handle_text(req.payload, &backend, started),
84        RequestKind::Ping => handle_ping(),
85    };
86
87    let mut stream = reader.into_inner();
88    match result {
89        Ok(json) => {
90            stream.write_all(json.as_bytes()).await?;
91            stream.write_all(b"\n").await?;
92        }
93        Err(msg) => return write_error(stream, msg).await,
94    }
95    stream.shutdown().await?;
96    Ok(())
97}
98
99fn handle_text(
100    payload: serde_json::Value,
101    backend: &Backend,
102    started: Instant,
103) -> Result<String, String> {
104    let p: EmbedTextPayload =
105        serde_json::from_value(payload).map_err(|e| format!("payload: {e}"))?;
106    info!(text_len = p.text.len(), "embed_text");
107    let v = backend.embed_one(&p.text)?;
108    let resp = EmbedResponse {
109        embedding: v,
110        model: super::model_id().to_string(),
111        elapsed_ms: started.elapsed().as_millis() as u64,
112    };
113    serde_json::to_string(&resp).map_err(|e| format!("encode: {e}"))
114}
115
116fn handle_file(
117    payload: serde_json::Value,
118    backend: &Backend,
119    cache: Option<&EmbedCache>,
120    started: Instant,
121) -> Result<String, String> {
122    let p: EmbedFilePayload =
123        serde_json::from_value(payload).map_err(|e| format!("payload: {e}"))?;
124
125    let path = PathBuf::from(&p.path);
126    let mut file = File::open(&path).map_err(|e| format!("abrir archivo: {e}"))?;
127    let mut buf = vec![0u8; MAX_FILE_BYTES];
128    let n = file.read(&mut buf).map_err(|e| format!("leer archivo: {e}"))?;
129    buf.truncate(n);
130
131    let model_id = super::model_id();
132    // Hash de los bytes que el modelo realmente verá. Si el archivo
133    // crece pasada la ventana MAX_FILE_BYTES sin modificar la cabeza,
134    // el hash NO cambia — el embedding cacheado sigue siendo válido
135    // bajo la semántica del proveedor (el modelo nunca vio los bytes
136    // adicionales). Si la cabeza cambia, el hash cambia y caemos a
137    // re-embed naturalmente.
138    let file_sha = ente_cas::sha256_of(&buf);
139
140    if let Some(cache) = cache {
141        if let Some(cached) = cache.get(&file_sha, model_id) {
142            info!(
143                path = %p.path,
144                sha = %ente_cas::hex(&file_sha),
145                bytes = n,
146                "embed_file: cache HIT"
147            );
148            let resp = EmbedResponse {
149                embedding: cached,
150                model: model_id.to_string(),
151                elapsed_ms: started.elapsed().as_millis() as u64,
152            };
153            return serde_json::to_string(&resp).map_err(|e| format!("encode: {e}"));
154        }
155    }
156
157    info!(
158        path = %p.path,
159        sha = %ente_cas::hex(&file_sha),
160        bytes = n,
161        "embed_file: cache MISS — invocando modelo"
162    );
163
164    // Write-through al CAS de arje: hacemos la cabeza del archivo
165    // direccionable por contenido. No es la fuente de verdad para
166    // el cache (sled lo es) pero deja un registro consultable por
167    // herramientas como `ente-cas gc` y permite que otros consumers
168    // resuelvan los bytes por hash.
169    if let Err(e) = ente_cas::store(&buf) {
170        // No-fatal: si CAS no escribe, cacheamos el embedding igual.
171        warn!(error = %e, "ente_cas::store falló (no-fatal)");
172    }
173
174    let text = String::from_utf8_lossy(&buf).to_string();
175    let v = backend.embed_one(&text)?;
176
177    if let Some(cache) = cache {
178        cache.put(&file_sha, model_id, &v);
179    }
180
181    let resp = EmbedResponse {
182        embedding: v,
183        model: model_id.to_string(),
184        elapsed_ms: started.elapsed().as_millis() as u64,
185    };
186    serde_json::to_string(&resp).map_err(|e| format!("encode: {e}"))
187}
188
189fn handle_ping() -> Result<String, String> {
190    let resp = PingResponse {
191        model: super::model_id().to_string(),
192        embed_dim: super::embed_dim(),
193    };
194    serde_json::to_string(&resp).map_err(|e| format!("encode: {e}"))
195}
196
197async fn write_error(mut stream: UnixStream, msg: String) -> std::io::Result<()> {
198    warn!(error = %msg, "respuesta de error");
199    let resp = ErrorResponse { error: msg };
200    let json = serde_json::to_string(&resp).unwrap_or_else(|_| "{\"error\":\"encode\"}".into());
201    stream.write_all(json.as_bytes()).await?;
202    stream.write_all(b"\n").await?;
203    stream.shutdown().await?;
204    Ok(())
205}