123 lines
3.6 KiB
Rust
123 lines
3.6 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
|
|
use super::client::LlmClient;
|
|
use crate::error::AgentError;
|
|
|
|
// ── Embedding types ────────────────────────────────────────────
|
|
|
|
#[derive(Serialize)]
|
|
struct EmbeddingRequest {
|
|
model: String,
|
|
input: Vec<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct EmbeddingResponse {
|
|
data: Vec<EmbeddingData>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct EmbeddingData {
|
|
embedding: Vec<f64>,
|
|
index: usize,
|
|
}
|
|
|
|
/// Max inputs per embedding request. The bge/OpenAI-like backends cap the input
|
|
/// array (bge-multilingual-gemma2 rejects >25 with "batch size overflow"), so we
|
|
/// chunk larger corpora — a whole control catalog (~1.8k) would otherwise 500.
|
|
const EMBED_BATCH_SIZE: usize = 16;
|
|
|
|
// ── Embedding implementation ───────────────────────────────────
|
|
|
|
impl LlmClient {
|
|
pub fn embed_model(&self) -> &str {
|
|
&self.embed_model
|
|
}
|
|
|
|
/// Generate embeddings for a batch of texts, chunking into backend-sized
|
|
/// requests and preserving input order across chunks.
|
|
pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
|
|
if texts.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut out = Vec::with_capacity(texts.len());
|
|
for chunk in texts.chunks(EMBED_BATCH_SIZE) {
|
|
out.extend(self.embed_batch(chunk.to_vec()).await?);
|
|
}
|
|
Ok(out)
|
|
}
|
|
|
|
/// Embed one backend-sized batch (≤ [`EMBED_BATCH_SIZE`]) in a single request.
|
|
async fn embed_batch(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
|
|
let url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));
|
|
|
|
let request_body = EmbeddingRequest {
|
|
model: self.embed_model.clone(),
|
|
input: texts,
|
|
};
|
|
|
|
let mut req = self
|
|
.http
|
|
.post(&url)
|
|
.header("content-type", "application/json")
|
|
.json(&request_body);
|
|
|
|
if let Some(auth) = self.auth_header() {
|
|
req = req.header("Authorization", auth);
|
|
}
|
|
|
|
let resp = req
|
|
.send()
|
|
.await
|
|
.map_err(|e| AgentError::Other(format!("Embedding request failed: {e}")))?;
|
|
|
|
if !resp.status().is_success() {
|
|
let status = resp.status();
|
|
let body = resp.text().await.unwrap_or_default();
|
|
return Err(AgentError::Other(format!(
|
|
"Embedding API returned {status}: {body}"
|
|
)));
|
|
}
|
|
|
|
let body: EmbeddingResponse = resp
|
|
.json()
|
|
.await
|
|
.map_err(|e| AgentError::Other(format!("Failed to parse embedding response: {e}")))?;
|
|
|
|
let mut data = body.data;
|
|
data.sort_by_key(|d| d.index);
|
|
|
|
Ok(data.into_iter().map(|d| d.embedding).collect())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use secrecy::SecretString;
|
|
|
|
fn client() -> LlmClient {
|
|
LlmClient::new(
|
|
"http://unused".into(),
|
|
SecretString::from(String::new()),
|
|
"m".into(),
|
|
"e".into(),
|
|
)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn empty_input_makes_no_request() {
|
|
// Must short-circuit before any HTTP call (base_url is unroutable).
|
|
let out = client().embed(Vec::new()).await.unwrap();
|
|
assert!(out.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn batch_size_is_within_backend_cap() {
|
|
assert!(
|
|
EMBED_BATCH_SIZE <= 25,
|
|
"must stay under the bge 25-input cap"
|
|
);
|
|
}
|
|
}
|