use std::sync::Arc;
use kbolt_types::KboltError;
use serde::Deserialize;
use serde_json::{json, Value};
use crate::config::{Config, ProviderOperation};
use crate::models::chat::{
ChatCompletionOutputMode, ChatCompletionRequestOptions, HttpChatClient,
LlamaCppChatRequestOptions,
};
use crate::models::completion::CompletionClient;
use crate::models::gateway::{
resolve_inference_gateway_bindings, EmbedderBinding, ExpanderBinding, GatewayProviderKind,
InferenceGatewayBindings, ProviderDeployment, RerankerBinding,
};
use crate::models::http::{HttpJsonClient, HttpOperation};
use crate::models::variants_expander::{ChatVariantsExpander, VARIANTS_GRAMMAR};
use crate::models::{Embedder, EmbeddingDocumentSizer, EmbeddingInputKind, Expander, Reranker};
use crate::Result;
#[cfg(test)]
use crate::models::chat::build_chat_payload;
#[cfg(test)]
use crate::models::http::parse_retry_after_seconds;
#[derive(Debug, Clone)]
struct HttpApiEmbedder {
client: HttpJsonClient,
model: String,
batch_size: usize,
endpoint_suffix: &'static str,
}
#[derive(Debug, Clone)]
struct LlamaCppServerDocumentSizer {
client: HttpJsonClient,
endpoint_suffix: &'static str,
}
#[derive(Clone)]
struct ChatBackedReranker {
chat: Arc<dyn CompletionClient>,
}
#[derive(Debug, Clone)]
struct LlamaCppEndpointReranker {
client: HttpJsonClient,
model: String,
endpoint_suffix: &'static str,
}
#[derive(Debug, Clone)]
struct OpenAiCompatibleEndpointReranker {
client: HttpJsonClient,
model: String,
endpoint_suffix: &'static str,
}
pub(crate) struct BuiltInferenceClients {
pub embedder: Option<Arc<dyn Embedder>>,
pub embedding_document_sizer: Option<Arc<dyn EmbeddingDocumentSizer>>,
pub reranker: Option<Arc<dyn Reranker>>,
pub expander: Option<Arc<dyn Expander>>,
}
impl std::fmt::Debug for BuiltInferenceClients {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BuiltInferenceClients")
.field("embedder", &self.embedder.is_some())
.field(
"embedding_document_sizer",
&self.embedding_document_sizer.is_some(),
)
.field("reranker", &self.reranker.is_some())
.field("expander", &self.expander.is_some())
.finish()
}
}
pub(crate) fn build_inference_clients(config: &Config) -> Result<BuiltInferenceClients> {
let bindings = resolve_inference_gateway_bindings(config)?;
build_provider_bound_clients(&bindings)
}
fn build_provider_bound_clients(
bindings: &InferenceGatewayBindings,
) -> Result<BuiltInferenceClients> {
Ok(BuiltInferenceClients {
embedder: bindings
.embedder
.as_ref()
.map(build_embedder_from_binding)
.transpose()?,
embedding_document_sizer: bindings
.embedder
.as_ref()
.map(build_embedding_document_sizer_from_binding)
.transpose()?
.flatten(),
reranker: bindings
.reranker
.as_ref()
.map(build_reranker_from_binding)
.transpose()?,
expander: bindings
.expander
.as_ref()
.map(build_expander_from_binding)
.transpose()?,
})
}
fn build_embedder_from_binding(binding: &EmbedderBinding) -> Result<Arc<dyn Embedder>> {
if binding.deployment.operation != ProviderOperation::Embedding {
return Err(KboltError::Inference(format!(
"provider profile '{}' uses incompatible operation '{}' for embedder bindings",
binding.provider_name,
binding.deployment.operation.as_str()
))
.into());
}
Ok(Arc::new(HttpApiEmbedder {
client: HttpJsonClient::new(
&binding.deployment.base_url,
binding.deployment.api_key_env.as_deref(),
binding.deployment.timeout_ms,
binding.deployment.max_retries,
"embedding",
binding.deployment.kind.as_str(),
),
model: binding.deployment.model.clone(),
batch_size: binding.batch_size,
endpoint_suffix: embedding_endpoint_suffix(binding.deployment.kind),
}))
}
fn build_embedding_document_sizer_from_binding(
binding: &EmbedderBinding,
) -> Result<Option<Arc<dyn EmbeddingDocumentSizer>>> {
if binding.deployment.operation != ProviderOperation::Embedding {
return Err(KboltError::Inference(format!(
"provider profile '{}' uses incompatible operation '{}' for embedder bindings",
binding.provider_name,
binding.deployment.operation.as_str()
))
.into());
}
match binding.deployment.kind {
GatewayProviderKind::LlamaCppServer => Ok(Some(Arc::new(LlamaCppServerDocumentSizer {
client: HttpJsonClient::new(
&binding.deployment.base_url,
binding.deployment.api_key_env.as_deref(),
binding.deployment.timeout_ms,
binding.deployment.max_retries,
"embedding",
binding.deployment.kind.as_str(),
),
endpoint_suffix: llama_cpp_tokenize_endpoint_suffix(),
}))),
GatewayProviderKind::OpenAiCompatible => Ok(None),
}
}
fn build_reranker_from_binding(binding: &RerankerBinding) -> Result<Arc<dyn Reranker>> {
match binding.deployment.operation {
ProviderOperation::Reranking => match binding.deployment.kind {
GatewayProviderKind::LlamaCppServer => Ok(Arc::new(LlamaCppEndpointReranker {
client: HttpJsonClient::new(
&binding.deployment.base_url,
binding.deployment.api_key_env.as_deref(),
binding.deployment.timeout_ms,
binding.deployment.max_retries,
"reranking",
binding.deployment.kind.as_str(),
),
model: binding.deployment.model.clone(),
endpoint_suffix: llama_cpp_reranking_endpoint_suffix(),
})),
GatewayProviderKind::OpenAiCompatible => {
Ok(Arc::new(OpenAiCompatibleEndpointReranker {
client: HttpJsonClient::new(
&binding.deployment.base_url,
binding.deployment.api_key_env.as_deref(),
binding.deployment.timeout_ms,
binding.deployment.max_retries,
"reranking",
binding.deployment.kind.as_str(),
),
model: binding.deployment.model.clone(),
endpoint_suffix: openai_compatible_reranking_endpoint_suffix(),
}))
}
},
ProviderOperation::ChatCompletion => Ok(Arc::new(ChatBackedReranker {
chat: Arc::new(build_chat_client(
&binding.deployment,
match binding.deployment.kind {
GatewayProviderKind::LlamaCppServer => {
structured_llama_cpp_chat_options(ChatCompletionOutputMode::JsonObject)
}
GatewayProviderKind::OpenAiCompatible => {
ChatCompletionRequestOptions::json_object()
}
},
)),
})),
ProviderOperation::Embedding => Err(KboltError::Inference(format!(
"provider profile '{}' uses incompatible operation 'embedding' for reranker bindings",
binding.provider_name
))
.into()),
}
}
fn build_expander_from_binding(binding: &ExpanderBinding) -> Result<Arc<dyn Expander>> {
if binding.deployment.operation != ProviderOperation::ChatCompletion {
return Err(KboltError::Inference(format!(
"provider profile '{}' uses incompatible operation '{}' for expander bindings",
binding.provider_name,
binding.deployment.operation.as_str()
))
.into());
}
let options = match binding.deployment.kind {
GatewayProviderKind::OpenAiCompatible => {
validate_openai_expander_sampling(binding)?;
openai_expander_options(binding)
}
GatewayProviderKind::LlamaCppServer => llama_cpp_expander_options(binding),
};
Ok(Arc::new(ChatVariantsExpander::new(Arc::new(
build_chat_client(&binding.deployment, options),
))))
}
fn build_chat_client(
deployment: &ProviderDeployment,
options: ChatCompletionRequestOptions,
) -> HttpChatClient {
HttpChatClient::new(
&deployment.base_url,
deployment.api_key_env.as_deref(),
deployment.timeout_ms,
deployment.max_retries,
&deployment.model,
chat_completion_endpoint_suffix(deployment.kind),
options,
deployment.kind.as_str(),
)
}
fn embedding_endpoint_suffix(kind: GatewayProviderKind) -> &'static str {
match kind {
GatewayProviderKind::LlamaCppServer => "v1/embeddings",
GatewayProviderKind::OpenAiCompatible => "embeddings",
}
}
fn chat_completion_endpoint_suffix(kind: GatewayProviderKind) -> &'static str {
match kind {
GatewayProviderKind::LlamaCppServer => "v1/chat/completions",
GatewayProviderKind::OpenAiCompatible => "chat/completions",
}
}
fn llama_cpp_reranking_endpoint_suffix() -> &'static str {
"v1/reranking"
}
fn openai_compatible_reranking_endpoint_suffix() -> &'static str {
"rerank"
}
fn llama_cpp_tokenize_endpoint_suffix() -> &'static str {
"tokenize"
}
fn openai_expander_options(binding: &ExpanderBinding) -> ChatCompletionRequestOptions {
ChatCompletionRequestOptions {
output_mode: ChatCompletionOutputMode::Text,
max_tokens: Some(binding.max_tokens),
seed: Some(binding.sampling.seed),
temperature: Some(binding.sampling.temperature),
top_k: None,
top_p: Some(binding.sampling.top_p),
min_p: None,
repeat_last_n: None,
repeat_penalty: None,
frequency_penalty: Some(binding.sampling.frequency_penalty),
presence_penalty: Some(binding.sampling.presence_penalty),
llama_cpp: None,
}
}
fn llama_cpp_expander_options(binding: &ExpanderBinding) -> ChatCompletionRequestOptions {
let mut llama_cpp = LlamaCppChatRequestOptions::non_thinking();
llama_cpp.grammar = Some(VARIANTS_GRAMMAR.to_string());
ChatCompletionRequestOptions {
output_mode: ChatCompletionOutputMode::Text,
max_tokens: Some(binding.max_tokens),
seed: Some(binding.sampling.seed),
temperature: Some(binding.sampling.temperature),
top_k: Some(binding.sampling.top_k),
top_p: Some(binding.sampling.top_p),
min_p: Some(binding.sampling.min_p),
repeat_last_n: Some(binding.sampling.repeat_last_n),
repeat_penalty: Some(binding.sampling.repeat_penalty),
frequency_penalty: Some(binding.sampling.frequency_penalty),
presence_penalty: Some(binding.sampling.presence_penalty),
llama_cpp: Some(llama_cpp),
}
}
fn structured_llama_cpp_chat_options(
output_mode: ChatCompletionOutputMode,
) -> ChatCompletionRequestOptions {
let mut options = match output_mode {
ChatCompletionOutputMode::JsonObject => ChatCompletionRequestOptions::json_object(),
ChatCompletionOutputMode::Text => ChatCompletionRequestOptions {
output_mode: ChatCompletionOutputMode::Text,
max_tokens: None,
seed: None,
temperature: Some(0.0),
top_k: None,
top_p: None,
min_p: None,
repeat_last_n: None,
repeat_penalty: None,
frequency_penalty: None,
presence_penalty: None,
llama_cpp: None,
},
};
options.llama_cpp = Some(LlamaCppChatRequestOptions::non_thinking());
options
}
fn validate_openai_expander_sampling(binding: &ExpanderBinding) -> Result<()> {
let defaults = crate::config::ExpanderSamplingConfig::default();
if binding.sampling.top_k != defaults.top_k
|| binding.sampling.min_p != defaults.min_p
|| binding.sampling.repeat_last_n != defaults.repeat_last_n
|| binding.sampling.repeat_penalty != defaults.repeat_penalty
{
return Err(KboltError::Config(format!(
"provider profile '{}' uses openai_compatible chat completion for expander bindings, which does not support top_k, min_p, repeat_last_n, or repeat_penalty overrides",
binding.provider_name
))
.into());
}
Ok(())
}
impl Embedder for HttpApiEmbedder {
fn embed_batch(&self, _kind: EmbeddingInputKind, texts: &[String]) -> Result<Vec<Vec<f32>>> {
embed_with_http_api(
&self.client,
self.endpoint_suffix,
&self.model,
self.batch_size,
texts,
)
}
}
impl EmbeddingDocumentSizer for LlamaCppServerDocumentSizer {
fn count_document_tokens(&self, text: &str) -> Result<usize> {
let payload = json!({
"content": text,
"add_special": true,
});
let response = self.client.post_json::<TokenizeResponseEnvelope>(
self.endpoint_suffix,
&payload,
HttpOperation::Tokenize,
)?;
response.token_count()
}
}
impl Reranker for ChatBackedReranker {
fn rerank(&self, query: &str, docs: &[String]) -> Result<Vec<f32>> {
if docs.is_empty() {
return Ok(Vec::new());
}
let system = "You are a retrieval reranker. Return JSON only as {\"scores\":[number,...]} with one score per document, each score between 0 and 1.";
let mut user = format!("Query:\n{query}\n\nDocuments:\n");
for (index, doc) in docs.iter().enumerate() {
user.push_str(&format!("[{index}] {doc}\n"));
}
user.push_str("\nRespond with exactly one score per document, in order.");
let content = self.chat.complete(system, &user)?;
let parsed: RerankerResponse = parse_json_payload("reranker response", &content)?;
let mut scores = parsed.into_scores();
if scores.len() != docs.len() {
return Err(KboltError::Inference(format!(
"reranker response size mismatch: expected {}, got {}",
docs.len(),
scores.len()
))
.into());
}
for score in &mut scores {
*score = score.clamp(0.0, 1.0);
}
Ok(scores)
}
}
impl Reranker for LlamaCppEndpointReranker {
fn rerank(&self, query: &str, docs: &[String]) -> Result<Vec<f32>> {
if docs.is_empty() {
return Ok(Vec::new());
}
let payload = json!({
"model": self.model,
"query": query,
"documents": docs,
"top_n": docs.len(),
"return_text": false,
});
let scored = self
.client
.post_json::<LlamaCppRerankResponseEnvelope>(
self.endpoint_suffix,
&payload,
HttpOperation::Reranking,
)?
.into_items()?;
finalize_rerank_scores(scored, docs.len())
}
}
impl Reranker for OpenAiCompatibleEndpointReranker {
fn rerank(&self, query: &str, docs: &[String]) -> Result<Vec<f32>> {
if docs.is_empty() {
return Ok(Vec::new());
}
let payload = json!({
"model": self.model,
"query": query,
"documents": docs,
"top_n": docs.len(),
"return_text": false,
});
let scored = self
.client
.post_json::<OpenAiCompatibleRerankResponseEnvelope>(
self.endpoint_suffix,
&payload,
HttpOperation::Reranking,
)?
.into_items()?;
finalize_rerank_scores(scored, docs.len())
}
}
fn embed_with_http_api(
client: &HttpJsonClient,
endpoint_suffix: &str,
model: &str,
batch_size: usize,
texts: &[String],
) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(Vec::new());
}
let mut vectors = Vec::new();
for batch in texts.chunks(batch_size) {
let payload = json!({
"model": model,
"input": batch,
});
let parsed: EmbeddingResponseEnvelope =
client.post_json(endpoint_suffix, &payload, HttpOperation::Embedding)?;
let response_vectors = parsed.into_vectors(batch.len())?;
vectors.extend(response_vectors);
}
Ok(vectors)
}
fn parse_json_payload<T>(label: &str, content: &str) -> Result<T>
where
T: for<'de> Deserialize<'de>,
{
let trimmed = content.trim();
serde_json::from_str(trimmed).map_err(|err| {
KboltError::Inference(format!(
"failed to parse {label} as JSON: {err}; payload={content}"
))
.into()
})
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum EmbeddingResponseEnvelope {
OpenAiLike { data: Vec<EmbeddingItem> },
VoyageLike { embeddings: Vec<Vec<f32>> },
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum TokenizeResponseEnvelope {
Wrapped { tokens: Vec<Value> },
Direct(Vec<Value>),
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum RerankerResponse {
Scores(Vec<f32>),
Wrapped { scores: Vec<f32> },
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum LlamaCppRerankResponseEnvelope {
Items(Vec<LlamaCppRerankItem>),
Wrapped { results: Vec<LlamaCppRerankItem> },
}
impl LlamaCppRerankResponseEnvelope {
fn into_items(self) -> Result<Vec<RerankScoreItem>> {
let items = match self {
Self::Items(items) => items,
Self::Wrapped { results } => results,
}
.into_iter()
.map(Into::into)
.collect::<Vec<_>>();
validate_rerank_scores(&items)?;
Ok(items)
}
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum OpenAiCompatibleRerankResponseEnvelope {
Items(Vec<OpenAiCompatibleRerankItem>),
Wrapped {
results: Vec<OpenAiCompatibleRerankItem>,
},
}
impl OpenAiCompatibleRerankResponseEnvelope {
fn into_items(self) -> Result<Vec<RerankScoreItem>> {
let items = match self {
Self::Items(items) => items,
Self::Wrapped { results } => results,
}
.into_iter()
.map(Into::into)
.collect::<Vec<_>>();
validate_rerank_scores(&items)?;
Ok(items)
}
}
impl RerankerResponse {
fn into_scores(self) -> Vec<f32> {
match self {
Self::Scores(scores) => scores,
Self::Wrapped { scores } => scores,
}
}
}
impl EmbeddingResponseEnvelope {
fn into_vectors(self, expected_len: usize) -> Result<Vec<Vec<f32>>> {
let mut vectors = match self {
Self::OpenAiLike { data } => {
let all_indexed = data.iter().all(|item| item.index.is_some());
let mut ordered = data
.into_iter()
.enumerate()
.map(|(position, item)| {
let order = item.index.unwrap_or(position);
(order, item.embedding.into_vec())
})
.collect::<Vec<_>>();
if all_indexed {
ordered.sort_by_key(|(order, _)| *order);
}
ordered
.into_iter()
.map(|(_, embedding)| embedding)
.collect::<Vec<_>>()
}
Self::VoyageLike { embeddings } => embeddings,
};
if vectors.len() != expected_len {
return Err(KboltError::Inference(format!(
"embedding response size mismatch: expected {expected_len}, got {}",
vectors.len()
))
.into());
}
for (index, vector) in vectors.iter().enumerate() {
if vector.is_empty() {
return Err(KboltError::Inference(format!(
"embedding response contains empty vector at index {index}"
))
.into());
}
}
Ok(std::mem::take(&mut vectors))
}
}
impl TokenizeResponseEnvelope {
fn token_count(self) -> Result<usize> {
let tokens = match self {
Self::Wrapped { tokens } => tokens,
Self::Direct(tokens) => tokens,
};
Ok(tokens.len())
}
}
#[derive(Debug, Deserialize)]
struct EmbeddingItem {
#[serde(default)]
index: Option<usize>,
embedding: EmbeddingVector,
}
#[derive(Debug, Deserialize)]
struct RerankScoreItem {
index: usize,
score: f32,
}
#[derive(Debug, Deserialize)]
struct LlamaCppRerankItem {
index: usize,
relevance_score: f32,
}
impl From<LlamaCppRerankItem> for RerankScoreItem {
fn from(value: LlamaCppRerankItem) -> Self {
Self {
index: value.index,
score: value.relevance_score,
}
}
}
#[derive(Debug, Deserialize)]
struct OpenAiCompatibleRerankItem {
index: usize,
score: f32,
}
impl From<OpenAiCompatibleRerankItem> for RerankScoreItem {
fn from(value: OpenAiCompatibleRerankItem) -> Self {
Self {
index: value.index,
score: value.score,
}
}
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum EmbeddingVector {
Direct(Vec<f32>),
Wrapped { values: Vec<f32> },
}
impl EmbeddingVector {
fn into_vec(self) -> Vec<f32> {
match self {
Self::Direct(values) => values,
Self::Wrapped { values } => values,
}
}
}
fn validate_rerank_scores(items: &[RerankScoreItem]) -> Result<()> {
if items.iter().any(|item| !item.score.is_finite()) {
return Err(
KboltError::Inference("rerank response contains non-finite score".to_string()).into(),
);
}
Ok(())
}
fn finalize_rerank_scores(
mut scored: Vec<RerankScoreItem>,
expected_len: usize,
) -> Result<Vec<f32>> {
if scored.len() != expected_len {
return Err(KboltError::Inference(format!(
"rerank response size mismatch: expected {expected_len}, got {}",
scored.len()
))
.into());
}
scored.sort_by_key(|item| item.index);
let mut scores = Vec::with_capacity(scored.len());
for (expected, item) in scored.into_iter().enumerate() {
if item.index != expected {
return Err(KboltError::Inference(format!(
"rerank response index mismatch: expected {expected}, got {}",
item.index
))
.into());
}
scores.push(item.score);
}
Ok(scores)
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::path::PathBuf;
use std::sync::mpsc;
use std::time::Duration;
use super::*;
use crate::config::{
ChunkingConfig, Config, EmbedderRoleConfig, ExpanderRoleConfig, ExpanderSamplingConfig,
ProviderProfileConfig, RankingConfig, ReapingConfig, RerankerRoleConfig,
RoleBindingsConfig,
};
fn base_config() -> Config {
Config {
config_dir: PathBuf::from("/tmp/kbolt-config"),
cache_dir: PathBuf::from("/tmp/kbolt-cache"),
default_space: None,
providers: HashMap::new(),
roles: RoleBindingsConfig::default(),
reaping: ReapingConfig { days: 7 },
chunking: ChunkingConfig::default(),
ranking: RankingConfig::default(),
}
}
fn serve_once(status_code: u16, body: &str) -> String {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind test server");
let addr = listener.local_addr().expect("server address");
let payload = body.to_string();
std::thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept client");
let _ = read_full_request(&mut stream);
let status_line = match status_code {
200 => "HTTP/1.1 200 OK",
401 => "HTTP/1.1 401 Unauthorized",
404 => "HTTP/1.1 404 Not Found",
429 => "HTTP/1.1 429 Too Many Requests",
500 => "HTTP/1.1 500 Internal Server Error",
other => panic!("unsupported status code in test server: {other}"),
};
let response = format!(
"{status_line}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
payload.len(),
payload
);
stream
.write_all(response.as_bytes())
.expect("write response");
});
format!("http://{addr}")
}
struct TestResponse {
status_code: u16,
body: &'static str,
retry_after: Option<&'static str>,
}
fn serve_sequence(responses: Vec<TestResponse>) -> String {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind test server");
let addr = listener.local_addr().expect("server address");
std::thread::spawn(move || {
for response_spec in responses {
let (mut stream, _) = listener.accept().expect("accept client");
let _ = read_full_request(&mut stream);
let status_line = match response_spec.status_code {
200 => "HTTP/1.1 200 OK",
429 => "HTTP/1.1 429 Too Many Requests",
500 => "HTTP/1.1 500 Internal Server Error",
other => panic!("unsupported status code in test server: {other}"),
};
let retry_after = response_spec
.retry_after
.map(|value| format!("Retry-After: {value}\r\n"))
.unwrap_or_default();
let response = format!(
"{status_line}\r\nContent-Type: application/json\r\n{retry_after}Content-Length: {}\r\nConnection: close\r\n\r\n{}",
response_spec.body.len(),
response_spec.body
);
stream
.write_all(response.as_bytes())
.expect("write response");
}
});
format!("http://{addr}")
}
fn serve_once_capturing_request(
status_code: u16,
body: &str,
) -> (String, mpsc::Receiver<Vec<u8>>) {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind test server");
let addr = listener.local_addr().expect("server address");
let payload = body.to_string();
let (tx, rx) = mpsc::channel();
std::thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept client");
let request = read_raw_request(&mut stream);
tx.send(request).expect("send captured request");
let status_line = match status_code {
200 => "HTTP/1.1 200 OK",
401 => "HTTP/1.1 401 Unauthorized",
404 => "HTTP/1.1 404 Not Found",
429 => "HTTP/1.1 429 Too Many Requests",
500 => "HTTP/1.1 500 Internal Server Error",
other => panic!("unsupported status code in test server: {other}"),
};
let response = format!(
"{status_line}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
payload.len(),
payload
);
stream
.write_all(response.as_bytes())
.expect("write response");
});
(format!("http://{addr}"), rx)
}
fn read_full_request(stream: &mut std::net::TcpStream) -> Vec<u8> {
let raw = read_raw_request(stream);
let header_end = raw
.windows(4)
.position(|window| window == b"\r\n\r\n")
.map(|offset| offset + 4);
let Some(header_end) = header_end else {
return Vec::new();
};
let headers = String::from_utf8_lossy(&raw[..header_end]).to_ascii_lowercase();
let mut content_length = 0usize;
for line in headers.lines() {
if let Some(value) = line.strip_prefix("content-length:") {
content_length = value.trim().parse::<usize>().unwrap_or(0);
break;
}
}
let body_end = header_end.saturating_add(content_length).min(raw.len());
raw.get(header_end..body_end).unwrap_or_default().to_vec()
}
fn read_raw_request(stream: &mut std::net::TcpStream) -> Vec<u8> {
let mut raw = Vec::new();
let mut header_end = None;
while header_end.is_none() {
let mut chunk = [0_u8; 1024];
let read = stream.read(&mut chunk).expect("read request bytes");
if read == 0 {
break;
}
raw.extend_from_slice(&chunk[..read]);
header_end = raw.windows(4).position(|window| window == b"\r\n\r\n");
}
let Some(header_end) = header_end else {
return Vec::new();
};
let header_end = header_end + 4;
let headers = String::from_utf8_lossy(&raw[..header_end]).to_ascii_lowercase();
let mut content_length = 0usize;
for line in headers.lines() {
if let Some(value) = line.strip_prefix("content-length:") {
content_length = value.trim().parse::<usize>().unwrap_or(0);
break;
}
}
let already_read_body = raw.len().saturating_sub(header_end);
let mut remaining = content_length.saturating_sub(already_read_body);
while remaining > 0 {
let mut chunk = [0_u8; 1024];
let read = stream.read(&mut chunk).expect("read request body");
if read == 0 {
break;
}
raw.extend_from_slice(&chunk[..read]);
remaining = remaining.saturating_sub(read);
}
raw
}
#[test]
fn provider_profile_embedder_builds_and_embeds() {
let body = r#"{"data":[{"index":0,"embedding":[0.1,0.2]}]}"#;
let mut config = base_config();
config.providers.insert(
"remote_embed".to_string(),
ProviderProfileConfig::OpenAiCompatible {
operation: ProviderOperation::Embedding,
base_url: serve_once(200, body),
model: "embed-model".to_string(),
api_key_env: None,
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.embedder = Some(EmbedderRoleConfig {
provider: "remote_embed".to_string(),
batch_size: 64,
});
let built = build_inference_clients(&config).expect("build clients");
let embedder = built.embedder.expect("embedder should exist");
let vectors = embedder
.embed_batch(EmbeddingInputKind::Document, &["hello".to_string()])
.expect("embed");
assert_eq!(vectors, vec![vec![0.1, 0.2]]);
}
#[test]
fn llama_cpp_embedder_builds_document_sizer_against_tokenize_endpoint() {
let (base_url, requests) = serve_once_capturing_request(200, r#"{"tokens":[1,2,3]}"#);
let mut config = base_config();
config.providers.insert(
"local_embed".to_string(),
ProviderProfileConfig::LlamaCppServer {
operation: ProviderOperation::Embedding,
base_url,
model: "embeddinggemma".to_string(),
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.embedder = Some(EmbedderRoleConfig {
provider: "local_embed".to_string(),
batch_size: 64,
});
let built = build_inference_clients(&config).expect("build clients");
let sizer = built
.embedding_document_sizer
.expect("document sizer should exist");
let token_count = sizer
.count_document_tokens("hello world")
.expect("count tokens");
assert_eq!(token_count, 3);
let raw_request = requests
.recv_timeout(Duration::from_secs(1))
.expect("captured request");
let request = String::from_utf8(raw_request).expect("utf8 request");
assert!(
request.starts_with("POST /tokenize HTTP/1.1\r\n"),
"unexpected request line: {request}"
);
assert!(
request.contains("\"content\":\"hello world\""),
"missing content payload: {request}"
);
assert!(
request.contains("\"add_special\":true"),
"missing add_special payload: {request}"
);
}
#[test]
fn openai_compatible_embedder_does_not_build_document_sizer() {
let body = r#"{"data":[{"index":0,"embedding":[0.1,0.2]}]}"#;
let mut config = base_config();
config.providers.insert(
"remote_embed".to_string(),
ProviderProfileConfig::OpenAiCompatible {
operation: ProviderOperation::Embedding,
base_url: serve_once(200, body),
model: "embed-model".to_string(),
api_key_env: None,
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.embedder = Some(EmbedderRoleConfig {
provider: "remote_embed".to_string(),
batch_size: 64,
});
let built = build_inference_clients(&config).expect("build clients");
assert!(built.embedding_document_sizer.is_none());
}
#[test]
fn provider_profile_reranker_supports_native_rerank_endpoint() {
let body = r#"{"model":"qwen3-reranker","object":"list","usage":{"prompt_tokens":153,"total_tokens":153},"results":[{"index":1,"relevance_score":0.9},{"index":0,"relevance_score":0.2}]}"#;
let mut config = base_config();
config.providers.insert(
"local_rerank".to_string(),
ProviderProfileConfig::LlamaCppServer {
operation: ProviderOperation::Reranking,
base_url: serve_once(200, body),
model: "qwen3-reranker".to_string(),
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.reranker = Some(RerankerRoleConfig {
provider: "local_rerank".to_string(),
});
let built = build_inference_clients(&config).expect("build clients");
let reranker = built.reranker.expect("reranker should exist");
let scores = reranker
.rerank("query", &["doc one".to_string(), "doc two".to_string()])
.expect("rerank");
assert_eq!(scores, vec![0.2, 0.9]);
}
#[test]
fn openai_native_reranker_uses_score_field() {
let body = r#"{"results":[{"index":1,"score":0.9},{"index":0,"score":0.2}]}"#;
let mut config = base_config();
config.providers.insert(
"remote_rerank".to_string(),
ProviderProfileConfig::OpenAiCompatible {
operation: ProviderOperation::Reranking,
base_url: serve_once(200, body),
model: "rerank-model".to_string(),
api_key_env: None,
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.reranker = Some(RerankerRoleConfig {
provider: "remote_rerank".to_string(),
});
let built = build_inference_clients(&config).expect("build clients");
let reranker = built.reranker.expect("reranker should exist");
let scores = reranker
.rerank("query", &["doc one".to_string(), "doc two".to_string()])
.expect("rerank");
assert_eq!(scores, vec![0.2, 0.9]);
}
#[test]
fn provider_profile_reranker_supports_chat_backed_mode() {
let body = r#"{
"choices": [
{
"message": {
"content": "{\"scores\":[0.2,0.9]}"
}
}
]
}"#;
let mut config = base_config();
config.providers.insert(
"remote_rerank".to_string(),
ProviderProfileConfig::OpenAiCompatible {
operation: ProviderOperation::ChatCompletion,
base_url: serve_once(200, body),
model: "gpt-5-mini".to_string(),
api_key_env: None,
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.reranker = Some(RerankerRoleConfig {
provider: "remote_rerank".to_string(),
});
let built = build_inference_clients(&config).expect("build clients");
let reranker = built.reranker.expect("reranker should exist");
let scores = reranker
.rerank("query", &["doc one".to_string(), "doc two".to_string()])
.expect("rerank");
assert_eq!(scores, vec![0.2, 0.9]);
}
#[test]
fn provider_profile_expander_uses_role_sampling() {
let body = r#"{
"choices": [
{
"message": {
"content": "[\"trait object rust\",\"explain rust traits\"]"
}
}
]
}"#;
let mut config = base_config();
config.providers.insert(
"local_expand".to_string(),
ProviderProfileConfig::LlamaCppServer {
operation: ProviderOperation::ChatCompletion,
base_url: serve_once(200, body),
model: "qwen3-1.7b".to_string(),
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.expander = Some(ExpanderRoleConfig {
provider: "local_expand".to_string(),
max_tokens: 256,
sampling: ExpanderSamplingConfig::default(),
});
let built = build_inference_clients(&config).expect("build clients");
let expander = built.expander.expect("expander should exist");
let variants = expander.expand("rust traits", 4).expect("expand");
assert_eq!(
variants,
vec![
"trait object rust".to_string(),
"explain rust traits".to_string(),
]
);
}
#[test]
fn llama_cpp_expander_posts_v1_chat_request_with_non_thinking_grammar() {
let body = r#"{
"choices": [
{
"message": {
"content": "[\"trait object rust\",\"explain rust traits\"]"
}
}
]
}"#;
let (base_url, requests) = serve_once_capturing_request(200, body);
let mut config = base_config();
config.providers.insert(
"local_expand".to_string(),
ProviderProfileConfig::LlamaCppServer {
operation: ProviderOperation::ChatCompletion,
base_url,
model: "qwen3-1.7b".to_string(),
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.expander = Some(ExpanderRoleConfig {
provider: "local_expand".to_string(),
max_tokens: 256,
sampling: ExpanderSamplingConfig::default(),
});
let built = build_inference_clients(&config).expect("build clients");
let expander = built.expander.expect("expander should exist");
let _ = expander.expand("rust traits", 4).expect("expand");
let raw_request = requests
.recv_timeout(Duration::from_secs(1))
.expect("captured request");
let request = String::from_utf8(raw_request).expect("utf8 request");
assert!(
request.starts_with("POST /v1/chat/completions HTTP/1.1\r\n"),
"unexpected request line: {request}"
);
assert!(
request.contains("\"chat_template_kwargs\":{\"enable_thinking\":false}"),
"missing non-thinking kwargs in request: {request}"
);
assert!(
request.contains("array ::= \\\"[\\\" ws elements? ws \\\"]\\\""),
"missing variants grammar in request: {request}"
);
}
#[test]
fn openai_expander_rejects_llama_only_sampling_overrides() {
let mut config = base_config();
config.providers.insert(
"remote_expand".to_string(),
ProviderProfileConfig::OpenAiCompatible {
operation: ProviderOperation::ChatCompletion,
base_url: "https://api.openai.com/v1".to_string(),
model: "gpt-5-mini".to_string(),
api_key_env: Some("OPENAI_API_KEY".to_string()),
timeout_ms: 5_000,
max_retries: 0,
},
);
config.roles.expander = Some(ExpanderRoleConfig {
provider: "remote_expand".to_string(),
max_tokens: 512,
sampling: ExpanderSamplingConfig {
top_k: 99,
..ExpanderSamplingConfig::default()
},
});
let err = build_inference_clients(&config).expect_err("unsupported sampling should fail");
assert!(err.to_string().contains("does not support top_k"));
}
#[test]
fn build_chat_payload_uses_json_response_format_for_reranker() {
let options = ChatCompletionRequestOptions::json_object();
let payload = build_chat_payload("model", "system", "user", &options);
assert_eq!(payload["response_format"]["type"], "json_object");
}
#[test]
fn build_chat_payload_omits_response_format_for_text_generation() {
let options = ChatCompletionRequestOptions::text();
let payload = build_chat_payload("model", "system", "user", &options);
assert!(payload.get("response_format").is_none());
}
#[test]
fn embedder_retries_on_rate_limit_status() {
let base_url = serve_sequence(vec![
TestResponse {
status_code: 429,
body: r#"{"error":"rate limit"}"#,
retry_after: Some("0"),
},
TestResponse {
status_code: 200,
body: r#"{"data":[{"index":0,"embedding":[0.1,0.2]}]}"#,
retry_after: None,
},
]);
let mut config = base_config();
config.providers.insert(
"remote_embed".to_string(),
ProviderProfileConfig::OpenAiCompatible {
operation: ProviderOperation::Embedding,
base_url,
model: "embed-model".to_string(),
api_key_env: None,
timeout_ms: 5_000,
max_retries: 1,
},
);
config.roles.embedder = Some(EmbedderRoleConfig {
provider: "remote_embed".to_string(),
batch_size: 64,
});
let built = build_inference_clients(&config).expect("build clients");
let embedder = built.embedder.expect("embedder should exist");
let vectors = embedder
.embed_batch(EmbeddingInputKind::Document, &["hello".to_string()])
.expect("embed should retry then succeed");
assert_eq!(vectors, vec![vec![0.1, 0.2]]);
}
#[test]
fn parse_retry_after_seconds_accepts_delta_seconds() {
assert_eq!(parse_retry_after_seconds(Some("7")), Some(7));
}
#[test]
fn parse_retry_after_seconds_accepts_http_date() {
let retry_at = std::time::SystemTime::now() + Duration::from_secs(60);
let header = httpdate::fmt_http_date(retry_at);
let parsed = parse_retry_after_seconds(Some(&header));
assert!(parsed.is_some());
let seconds = parsed.expect("parsed seconds");
assert!(seconds <= 60);
}
}