263 lines
8.1 KiB
Rust
263 lines
8.1 KiB
Rust
|
use anyhow::Result;
|
||
|
|
use mem_llm::EmbeddingsClient;
|
||
|
|
use serde_json::json;
|
||
|
|
use wiremock::matchers::{header, method, path};
|
||
|
|
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a1_batches_at_32() -> Result<()> {
|
||
|
|
// Test: 100 inputs produce exactly 4 requests (32+32+32+4)
|
||
|
|
let server = MockServer::start().await;
|
||
|
|
|
||
|
|
// Mock POST /v1/embeddings to count calls and return 768-dim vectors
|
||
|
|
let mut call_count = 0;
|
||
|
|
Mock::given(method("POST"))
|
||
|
|
.and(path("/v1/embeddings"))
|
||
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||
|
|
"object": "list",
|
||
|
|
"data": (0..32).map(|i| json!({
|
||
|
|
"embedding": vec![0.1; 768],
|
||
|
|
"index": i
|
||
|
|
})).collect::<Vec<_>>(),
|
||
|
|
"usage": {"prompt_tokens": 1, "completion_tokens": 1}
|
||
|
|
})))
|
||
|
|
.mount(&server)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let client = EmbeddingsClient::from_env();
|
||
|
|
let mut client = client?;
|
||
|
|
client.base_url = server.uri();
|
||
|
|
|
||
|
|
// 100 inputs
|
||
|
|
let texts: Vec<String> = (0..100).map(|i| format!("text {}", i)).collect();
|
||
|
|
let result = client.embed(&texts).await?;
|
||
|
|
|
||
|
|
assert_eq!(result.len(), 100, "Should return 100 vectors for 100 inputs");
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a2_order_preserved() -> Result<()> {
|
||
|
|
// Test: order is preserved across batch boundaries
|
||
|
|
// Mock returns vectors with distinct values based on input index
|
||
|
|
let server = MockServer::start().await;
|
||
|
|
|
||
|
|
Mock::given(method("POST"))
|
||
|
|
.and(path("/v1/embeddings"))
|
||
|
|
.respond_with(|req: &wiremock::Request| {
|
||
|
|
// Parse request body to extract input texts
|
||
|
|
let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap_or_default();
|
||
|
|
let input = body["input"].as_array().unwrap_or(&vec![]);
|
||
|
|
|
||
|
|
let data: Vec<_> = input
|
||
|
|
.iter()
|
||
|
|
.enumerate()
|
||
|
|
.map(|(idx, text)| {
|
||
|
|
let text_str = text.as_str().unwrap_or("");
|
||
|
|
// Extract the number from "text N" to create distinguishable vectors
|
||
|
|
let marker = text_str
|
||
|
|
.split_whitespace()
|
||
|
|
.last()
|
||
|
|
.and_then(|s| s.parse::<f32>().ok())
|
||
|
|
.unwrap_or(0.0);
|
||
|
|
|
||
|
|
json!({
|
||
|
|
"embedding": vec![marker; 768], // Distinctive marker value
|
||
|
|
"index": idx
|
||
|
|
})
|
||
|
|
})
|
||
|
|
.collect();
|
||
|
|
|
||
|
|
ResponseTemplate::new(200).set_body_json(json!({
|
||
|
|
"object": "list",
|
||
|
|
"data": data,
|
||
|
|
"usage": {}
|
||
|
|
}))
|
||
|
|
})
|
||
|
|
.mount(&server)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let mut client = EmbeddingsClient::from_env()?;
|
||
|
|
client.base_url = server.uri();
|
||
|
|
|
||
|
|
// 70 inputs to cross batch boundary (32 + 32 + 6)
|
||
|
|
let texts: Vec<String> = (0..70).map(|i| format!("text {}", i)).collect();
|
||
|
|
let result = client.embed(&texts).await?;
|
||
|
|
|
||
|
|
// Verify order: each vector's first element should match input index
|
||
|
|
for (i, vec) in result.iter().enumerate() {
|
||
|
|
let first_val = vec.as_ref()[0];
|
||
|
|
let expected = i as f32;
|
||
|
|
assert!(
|
||
|
|
(first_val - expected).abs() < 0.01,
|
||
|
|
"Vector {} has marker {}, expected {}",
|
||
|
|
i,
|
||
|
|
first_val,
|
||
|
|
expected
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a3_dimension_asserted() -> Result<()> {
|
||
|
|
// Test: 512-dim response triggers error naming the model
|
||
|
|
let server = MockServer::start().await;
|
||
|
|
|
||
|
|
Mock::given(method("POST"))
|
||
|
|
.and(path("/v1/embeddings"))
|
||
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||
|
|
"object": "list",
|
||
|
|
"data": [{
|
||
|
|
"embedding": vec![0.1; 512], // Wrong dimension!
|
||
|
|
"index": 0
|
||
|
|
}],
|
||
|
|
"usage": {}
|
||
|
|
})))
|
||
|
|
.mount(&server)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let mut client = EmbeddingsClient::from_env()?;
|
||
|
|
client.base_url = server.uri();
|
||
|
|
|
||
|
|
let result = client.embed(&["test".to_string()]).await;
|
||
|
|
|
||
|
|
assert!(result.is_err(), "Should error on dimension mismatch");
|
||
|
|
let err_msg = format!("{:?}", result.unwrap_err());
|
||
|
|
assert!(
|
||
|
|
err_msg.contains("768") && err_msg.contains("512"),
|
||
|
|
"Error should name both dimensions: {}",
|
||
|
|
err_msg
|
||
|
|
);
|
||
|
|
assert!(
|
||
|
|
err_msg.contains("nomic-ai/nomic-embed-text-v2-moe"),
|
||
|
|
"Error should name the model: {}",
|
||
|
|
err_msg
|
||
|
|
);
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a4_apikey_sent() -> Result<()> {
|
||
|
|
// Test: apikey header is present in request
|
||
|
|
let server = MockServer::start().await;
|
||
|
|
|
||
|
|
Mock::given(method("POST"))
|
||
|
|
.and(path("/v1/embeddings"))
|
||
|
|
.and(header("apikey", wiremock::matchers::Matcher::regex(".*"))) // Match any apikey value
|
||
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||
|
|
"object": "list",
|
||
|
|
"data": [{
|
||
|
|
"embedding": vec![0.1; 768],
|
||
|
|
"index": 0
|
||
|
|
}],
|
||
|
|
"usage": {}
|
||
|
|
})))
|
||
|
|
.mount(&server)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let mut client = EmbeddingsClient::from_env()?;
|
||
|
|
client.base_url = server.uri();
|
||
|
|
client.api_key = "test-key-12345".to_string();
|
||
|
|
|
||
|
|
let result = client.embed(&["test".to_string()]).await?;
|
||
|
|
|
||
|
|
assert_eq!(result.len(), 1, "Should return 1 vector");
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
#[ignore] // Live test against real gateway
|
||
|
|
async fn a5_live_dims() -> Result<()> {
|
||
|
|
// Test: Real gateway returns 768-dim vectors
|
||
|
|
// Run with: cargo test a5_live_dims -- --ignored
|
||
|
|
let client = EmbeddingsClient::from_env()?;
|
||
|
|
|
||
|
|
let result = client.embed(&["hello world".to_string()]).await?;
|
||
|
|
|
||
|
|
assert_eq!(result.len(), 1, "Should return 1 vector");
|
||
|
|
assert_eq!(result[0].as_ref().len(), 768, "Should be 768-dim");
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_empty_input() -> Result<()> {
|
||
|
|
let client = EmbeddingsClient::from_env()?;
|
||
|
|
let result = client.embed(&[]).await?;
|
||
|
|
assert_eq!(result.len(), 0, "Empty input should return empty output");
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_batch_boundary_32() -> Result<()> {
|
||
|
|
// Exact boundary: 32 inputs should fit in 1 batch
|
||
|
|
let server = MockServer::start().await;
|
||
|
|
|
||
|
|
Mock::given(method("POST"))
|
||
|
|
.and(path("/v1/embeddings"))
|
||
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||
|
|
"object": "list",
|
||
|
|
"data": (0..32).map(|i| json!({
|
||
|
|
"embedding": vec![0.1; 768],
|
||
|
|
"index": i
|
||
|
|
})).collect::<Vec<_>>(),
|
||
|
|
"usage": {}
|
||
|
|
})))
|
||
|
|
.mount(&server)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let mut client = EmbeddingsClient::from_env()?;
|
||
|
|
client.base_url = server.uri();
|
||
|
|
|
||
|
|
let texts: Vec<String> = (0..32).map(|i| format!("text {}", i)).collect();
|
||
|
|
let result = client.embed(&texts).await?;
|
||
|
|
|
||
|
|
assert_eq!(result.len(), 32, "32 inputs = 1 batch");
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_batch_boundary_33() -> Result<()> {
|
||
|
|
// Over boundary: 33 inputs should need 2 batches (32+1)
|
||
|
|
let server = MockServer::start().await;
|
||
|
|
|
||
|
|
Mock::given(method("POST"))
|
||
|
|
.and(path("/v1/embeddings"))
|
||
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||
|
|
"object": "list",
|
||
|
|
"data": (0..32).map(|i| json!({
|
||
|
|
"embedding": vec![0.1; 768],
|
||
|
|
"index": i
|
||
|
|
})).collect::<Vec<_>>(),
|
||
|
|
"usage": {}
|
||
|
|
})))
|
||
|
|
.mount(&server)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
// Add a separate mock for the 1-element batch
|
||
|
|
Mock::given(method("POST"))
|
||
|
|
.and(path("/v1/embeddings"))
|
||
|
|
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||
|
|
"object": "list",
|
||
|
|
"data": [{
|
||
|
|
"embedding": vec![0.1; 768],
|
||
|
|
"index": 0
|
||
|
|
}],
|
||
|
|
"usage": {}
|
||
|
|
})))
|
||
|
|
.mount(&server)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let mut client = EmbeddingsClient::from_env()?;
|
||
|
|
client.base_url = server.uri();
|
||
|
|
|
||
|
|
let texts: Vec<String> = (0..33).map(|i| format!("text {}", i)).collect();
|
||
|
|
let result = client.embed(&texts).await?;
|
||
|
|
|
||
|
|
assert_eq!(result.len(), 33, "33 inputs = 2 batches");
|
||
|
|
Ok(())
|
||
|
|
}
|