diff --git a/crates/agentic-server-core/src/storage/file_search.rs b/crates/agentic-server-core/src/storage/file_search.rs index bface853..7609d25b 100644 --- a/crates/agentic-server-core/src/storage/file_search.rs +++ b/crates/agentic-server-core/src/storage/file_search.rs @@ -47,6 +47,8 @@ pub(crate) struct StoredChunk { pub filename: String, pub chunk_index: usize, pub text: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub embedding_text: Option, pub embedding: Option>, pub attributes: crate::types::file_search::FileAttributes, } @@ -494,6 +496,7 @@ mod tests { dimensions: i64::try_from(dimensions).unwrap(), chunks: (0..count) .map(|chunk_index| StoredChunk { + embedding_text: None, file_id: file.id.clone(), filename: file.filename.clone(), chunk_index, diff --git a/crates/agentic-server-core/src/tool/file_search/embeddings.rs b/crates/agentic-server-core/src/tool/file_search/embeddings.rs index b6e8707f..18f34850 100644 --- a/crates/agentic-server-core/src/tool/file_search/embeddings.rs +++ b/crates/agentic-server-core/src/tool/file_search/embeddings.rs @@ -2,8 +2,10 @@ use std::{sync::Arc, time::Duration}; +#[cfg(test)] +use crate::types::retrieval_models::EmbeddingData; +use crate::types::retrieval_models::{EmbeddingRequest, EmbeddingResponse}; use futures::StreamExt; -use serde::{Deserialize, Serialize}; use crate::types::file_search::{FileSearchConfig, FileSearchError}; @@ -17,24 +19,7 @@ pub(super) struct Embeddings { endpoint: reqwest::Url, model: String, api_key: Option, -} - -#[derive(Serialize)] -struct EmbeddingRequest<'a> { - model: &'a str, - input: &'a [String], - encoding_format: &'static str, -} - -#[derive(Deserialize)] -struct EmbeddingResponse { - model: String, - data: Vec, -} -#[derive(Deserialize)] -struct EmbeddingData { - index: usize, - embedding: Vec, + dimensions: Option, } impl Embeddings { @@ -42,6 +27,24 @@ impl Embeddings { client: Arc, config: &FileSearchConfig, ) -> Result, FileSearchError> { + if let Some(model) = &config.vector_stores.default_embedding_model { + let (provider, name) = config.vector_stores.resolve(None, Some(model))?; + if config.embedding_base_url.is_some() + || config.embedding_model.is_some() + || config.embedding_api_key.is_some() + { + return Err(FileSearchError::InvalidRequest( + "configure either grouped or legacy embeddings, not both".into(), + )); + } + return Ok(Some(Self { + client, + endpoint: provider.endpoint("embeddings")?, + model: name.into(), + api_key: provider.api_key.clone(), + dimensions: model.embedding_dimensions, + })); + } let (Some(base_url), Some(model)) = (&config.embedding_base_url, &config.embedding_model) else { if config.embedding_base_url.is_some() || config.embedding_model.is_some() @@ -77,6 +80,7 @@ impl Embeddings { endpoint, model: model.clone(), api_key: config.embedding_api_key.clone(), + dimensions: None, })) } @@ -90,7 +94,12 @@ impl Embeddings { expected_dimensions: Option, ) -> Result>, FileSearchError> { let mut embeddings = Vec::with_capacity(texts.len()); - let mut dimensions = expected_dimensions; + let mut dimensions = expected_dimensions.or(self.dimensions); + if self.dimensions.is_some_and(|configured| dimensions != Some(configured)) { + return Err(FileSearchError::InvalidRequest( + "stored and configured embedding dimensions differ".into(), + )); + } for input in texts.chunks(BATCH_SIZE) { let mut request = self .client @@ -101,6 +110,7 @@ impl Embeddings { model: &self.model, input, encoding_format: "float", + dimensions: self.dimensions, })?); if let Some(key) = &self.api_key { request = request.bearer_auth(key); diff --git a/crates/agentic-server-core/src/tool/file_search/handler.rs b/crates/agentic-server-core/src/tool/file_search/handler.rs index c1a12f18..e11025b5 100644 --- a/crates/agentic-server-core/src/tool/file_search/handler.rs +++ b/crates/agentic-server-core/src/tool/file_search/handler.rs @@ -80,7 +80,7 @@ impl FileSearchHandler { let arguments = parse_arguments(arguments)?; let request = search_request(params, arguments.queries); let result = service - .search(params.vector_store_ids.as_deref().unwrap_or_default(), &request) + .search_for_tool(params.vector_store_ids.as_deref().unwrap_or_default(), &request) .await .map_err(ToolError::FileSearch)?; let output = FileSearchToolOutput { @@ -224,7 +224,8 @@ fn search_request(params: &FileSearchToolParam, queries: Vec) -> SearchR max_num_results: params.max_num_results, filters: params.filters.clone(), ranking_options: params.ranking_options.clone(), - ..SearchRequest::default() + search_mode: params.search_mode, + rewrite_query: params.rewrite_query, } } diff --git a/crates/agentic-server-core/src/tool/file_search/ingest.rs b/crates/agentic-server-core/src/tool/file_search/ingest.rs index edd4f3a7..c0ce48b1 100644 --- a/crates/agentic-server-core/src/tool/file_search/ingest.rs +++ b/crates/agentic-server-core/src/tool/file_search/ingest.rs @@ -12,6 +12,13 @@ const TOKENIZATION_BLOCK_BYTES: usize = 256; pub(super) fn chunking_config(strategy: &ChunkingStrategy) -> Result { let config = match strategy { + ChunkingStrategy::Contextual { contextual } => { + contextual.validate()?; + return Ok(StaticChunking { + max_chunk_size_tokens: contextual.max_chunk_size_tokens, + chunk_overlap_tokens: contextual.chunk_overlap_tokens, + }); + } ChunkingStrategy::Auto => StaticChunking::default(), ChunkingStrategy::Static { config } => config.clone(), }; @@ -77,13 +84,18 @@ pub(super) fn validate_content_type(filename: &str, content_type: &str) -> Resul invalid("Unsupported file type; upload UTF-8 text or a PDF containing extractable text") } +pub(super) struct ExtractedDocument { + pub text: String, + pub chunks: Vec, +} + pub(super) fn extract_and_chunk( bytes: Vec, filename: &str, content_type: &str, chunking: &StaticChunking, cancelled: &AtomicBool, -) -> Result, FileSearchError> { +) -> Result { validate_content_type(filename, content_type)?; let content_type = content_type.split(';').next().unwrap_or_default().trim(); let text = if is_pdf(filename, content_type) { @@ -101,7 +113,8 @@ pub(super) fn extract_and_chunk( if text.contains('\0') { return invalid("The file contains binary content instead of text"); } - chunks(&text, chunking, cancelled) + let chunks = chunks(&text, chunking, cancelled)?; + Ok(ExtractedDocument { text, chunks }) } #[cfg(not(feature = "file-search-pdf"))] @@ -219,10 +232,56 @@ fn chunks(text: &str, config: &StaticChunking, cancelled: &AtomicBool) -> Result Ok(chunks) } +/// Keep complete source chunks that fit the remaining model-context budget. +pub(super) fn limit_context( + results: Vec, + mut budget: usize, + cancelled: &AtomicBool, +) -> Result, FileSearchError> { + if cancelled.load(Ordering::Relaxed) { + return Err(FileSearchError::Unavailable( + "File search context preparation was cancelled".into(), + )); + } + let tokenizer = tiktoken_rs::cl100k_base_singleton(); + let mut selected = Vec::with_capacity(results.len()); + 'passages: for result in results { + let mut tokens = 0usize; + for content in &result.content { + let mut text = content.text.as_str(); + while !text.is_empty() { + if cancelled.load(Ordering::Relaxed) { + return Err(FileSearchError::Unavailable( + "File search context preparation was cancelled".into(), + )); + } + let mut end = text.len().min(TOKENIZATION_BLOCK_BYTES); + while !text.is_char_boundary(end) { + end -= 1; + } + tokens = tokens.saturating_add(tokenizer.encode_ordinary(&text[..end]).len()); + if tokens > budget { + continue 'passages; + } + text = &text[end..]; + } + } + budget -= tokens; + selected.push(result); + } + Ok(selected) +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn cancelled_context_preparation_exits_before_tokenizing() { + let cancelled = AtomicBool::new(true); + assert!(limit_context(Vec::new(), 4000, &cancelled).is_err()); + } + #[test] fn unicode_tokens_never_split_scalars_or_drop_text() { let text = "A coral 🪸 reef conserves biodiversity. 日本語の文章。 café ".repeat(90); @@ -237,7 +296,8 @@ mod tests { &config, &AtomicBool::new(false), ) - .unwrap(); + .unwrap() + .chunks; assert!(chunks.len() > 1); assert_eq!(chunks.concat(), text); assert!(chunks.iter().all(|chunk| !chunk.contains('\u{fffd}'))); diff --git a/crates/agentic-server-core/src/tool/file_search/mod.rs b/crates/agentic-server-core/src/tool/file_search/mod.rs index ae9aa354..ca2bb7a3 100644 --- a/crates/agentic-server-core/src/tool/file_search/mod.rs +++ b/crates/agentic-server-core/src/tool/file_search/mod.rs @@ -3,6 +3,7 @@ mod embeddings; pub(crate) mod handler; mod ingest; +mod models; mod ranking; mod service; diff --git a/crates/agentic-server-core/src/tool/file_search/models.rs b/crates/agentic-server-core/src/tool/file_search/models.rs new file mode 100644 index 00000000..0b8cd354 --- /dev/null +++ b/crates/agentic-server-core/src/tool/file_search/models.rs @@ -0,0 +1,336 @@ +//! Cancellation-safe bounded text generation and vLLM reranking. +use crate::types::{ + file_search::{ + ContextualChunking, FileSearchError, ModelProvider, ScoreInterpretation, SearchResult, VectorStoresConfig, + invalid, + }, + retrieval_models::{ + RetrievalChatRequest, RetrievalChatResponse, RetrievalMessage, RetrievalRole, RetrievalText, RetrievalTextPart, + TextRerankRequest, TextRerankResponse, + }, +}; +use futures::{StreamExt, TryStreamExt}; +use serde::{Serialize, de::DeserializeOwned}; +use std::{sync::Arc, time::Duration}; + +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; +const MAX_REQUEST_BYTES: usize = 32 * 1024 * 1024; + +#[derive(Clone)] +pub(super) struct Models { + client: Arc, + config: Arc, + contextual_workers: Arc, +} +impl Models { + pub(super) fn new(client: Arc, config: Arc) -> Self { + let concurrency = config.contextual_retrieval_params.default_max_concurrency; + Self { + client, + config, + contextual_workers: Arc::new(tokio::sync::Semaphore::new(concurrency)), + } + } + pub(super) async fn rewrite(&self, queries: &[String]) -> Result { + let params = self.config.rewrite_query_params.as_ref().ok_or_else(|| { + FileSearchError::InvalidRequest("rewrite_query requires configured rewrite_query_params".into()) + })?; + let (provider, model) = self.config.resolve(None, params.model.as_ref())?; + let prompt = expand_prompt(¶ms.prompt, "{query}", &queries.join(" "))?; + self.chat( + provider, + &RetrievalChatRequest { + model, + messages: vec![RetrievalMessage { + role: RetrievalRole::User, + content: &prompt, + }], + stream: false, + temperature: params.temperature, + max_tokens: params.max_tokens, + }, + 45, + 4096, + ) + .await + } + pub(super) async fn contextualize( + &self, + document: &str, + chunks: &[String], + params: &ContextualChunking, + ) -> Result, FileSearchError> { + params.validate()?; + let defaults = &self.config.contextual_retrieval_params; + let (provider, model) = self + .config + .resolve(params.model_id.as_deref(), defaults.model.as_ref())?; + if document.chars().count() / 4 > defaults.max_document_tokens { + return invalid("contextual document exceeds max_document_tokens"); + } + let (prefix, suffix) = params + .context_prompt + .split_once("{{CHUNK_CONTENT}}") + .ok_or_else(|| FileSearchError::InvalidRequest("contextual chunk placeholder is missing".into()))?; + let prefix = expand_prompt(prefix, "{{WHOLE_DOCUMENT}}", document)?; + let concurrency = params + .max_concurrency + .unwrap_or(defaults.default_max_concurrency) + .min(defaults.default_max_concurrency); + let prefix = prefix.as_str(); + // Buffered futures are dropped on cancellation or any failed chunk; no detached workers or partial publication. + futures::stream::iter((0..chunks.len()).map(|index| async move { + let chunk = &chunks[index]; + let _permit = self + .contextual_workers + .acquire() + .await + .map_err(|_| FileSearchError::Unavailable("contextual retrieval is shutting down".into()))?; + let user_message = format!("{chunk}{suffix}"); + let context = self + .chat( + provider, + &RetrievalChatRequest { + model, + messages: vec![ + RetrievalMessage { + role: RetrievalRole::System, + content: prefix.trim_end(), + }, + RetrievalMessage { + role: RetrievalRole::User, + content: &user_message, + }, + ], + stream: false, + temperature: 0.0, + max_tokens: 256, + }, + params.timeout_seconds.unwrap_or(defaults.default_timeout_seconds), + 8192, + ) + .await?; + Ok(format!("{context}\n\n{chunk}")) + })) + .buffered(concurrency) + .try_collect() + .await + } + async fn chat( + &self, + provider: &ModelProvider, + request: &RetrievalChatRequest<'_>, + timeout: u64, + max_text: usize, + ) -> Result { + let response: RetrievalChatResponse = self.post(provider, "chat/completions", request, timeout).await?; + if response.choices.len() != 1 { + return Err(FileSearchError::ProviderProtocol); + } + let text = match response + .choices + .into_iter() + .next() + .and_then(|choice| choice.message.content) + { + Some(RetrievalText::Text(text)) => text, + Some(RetrievalText::Parts(parts)) => parts + .into_iter() + .map(|RetrievalTextPart::Text { text }| text) + .collect::>() + .join(" "), + None => return Err(FileSearchError::ProviderProtocol), + }; + let text = text.trim(); + if text.is_empty() || text.len() > max_text { + return Err(FileSearchError::ProviderProtocol); + } + Ok(text.into()) + } + pub(super) async fn rerank( + &self, + query: &str, + mut candidates: Vec, + selection: Option<&str>, + ) -> Result, FileSearchError> { + let (provider, model) = self + .config + .resolve(selection, self.config.default_reranker_model.as_ref())?; + if candidates.is_empty() { + return Ok(candidates); + } + let documents = candidates + .iter() + .map(|result| result.content[0].text.as_str()) + .collect(); + // Score every bounded candidate, then validate the complete permutation before publishing any result. + let request = TextRerankRequest { + model, + query, + documents, + top_n: candidates.len(), + }; + let response: TextRerankResponse = self.post(provider, "rerank", &request, 45).await?; + if response.results.len() != candidates.len() { + return Err(FileSearchError::ProviderProtocol); + } + let mut seen = vec![false; candidates.len()]; + for result in response.results { + if result.index >= candidates.len() || seen[result.index] || !result.relevance_score.is_finite() { + return Err(FileSearchError::ProviderProtocol); + } + seen[result.index] = true; + let score = match provider.score_interpretation { + ScoreInterpretation::Probability if (0.0..=1.0).contains(&result.relevance_score) => { + result.relevance_score + } + ScoreInterpretation::Probability => return Err(FileSearchError::ProviderProtocol), + ScoreInterpretation::Logit if result.relevance_score >= 0.0 => { + 1.0 / (1.0 + (-result.relevance_score).exp()) + } + ScoreInterpretation::Logit => { + let exp = result.relevance_score.exp(); + exp / (1.0 + exp) + } + }; + candidates[result.index].score = score; + } + candidates.sort_by(|left, right| right.score.total_cmp(&left.score)); + Ok(candidates) + } + async fn post( + &self, + provider: &ModelProvider, + operation: &str, + request: &T, + timeout: u64, + ) -> Result { + let mut body = LimitedRequestBody::default(); + if let Err(error) = serde_json::to_writer(&mut body, request) { + if body.exceeded { + return invalid("model request exceeds 32 MiB"); + } + return Err(error.into()); + } + let mut request = self + .client + .post(provider.endpoint(operation)?) + .timeout(Duration::from_secs(timeout)) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(body.bytes); + if let Some(key) = &provider.api_key { + request = request.bearer_auth(key); + } + let response = request + .send() + .await + .map_err(|error| FileSearchError::Provider(error.without_url()))? + .error_for_status() + .map_err(|error| FileSearchError::Provider(error.without_url()))?; + if response.content_length().is_some_and(|n| n > MAX_RESPONSE_BYTES as u64) { + return Err(FileSearchError::ProviderProtocol); + } + let mut bytes = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(part) = stream.next().await { + let part = part.map_err(|error| FileSearchError::Provider(error.without_url()))?; + if bytes.len().saturating_add(part.len()) > MAX_RESPONSE_BYTES { + return Err(FileSearchError::ProviderProtocol); + } + bytes.extend_from_slice(&part); + } + serde_json::from_slice(&bytes).map_err(FileSearchError::ProviderDecode) + } +} + +/// Preflight interpolation before allocating expanded text, including query-list rewrite templates. +fn expand_prompt(template: &str, placeholder: &str, value: &str) -> Result { + let occurrences = template.matches(placeholder).count(); + let expanded = value + .len() + .checked_mul(occurrences) + .and_then(|replacements| (template.len() - placeholder.len() * occurrences).checked_add(replacements)); + if expanded.is_none_or(|size| size > MAX_REQUEST_BYTES) { + return invalid("expanded model prompt exceeds 32 MiB"); + } + Ok(template.replace(placeholder, value)) +} + +/// Reject writes before growing the request buffer beyond the transport limit. +#[derive(Default)] +struct LimitedRequestBody { + bytes: Vec, + exceeded: bool, +} +impl std::io::Write for LimitedRequestBody { + fn write(&mut self, input: &[u8]) -> std::io::Result { + if input.len() > MAX_REQUEST_BYTES - self.bytes.len() { + self.exceeded = true; + return Err(std::io::Error::other("model request exceeds byte limit")); + } + let required = self.bytes.len() + input.len(); + if required > self.bytes.capacity() { + let capacity = self + .bytes + .capacity() + .max(4096) + .saturating_mul(2) + .max(required) + .min(MAX_REQUEST_BYTES); + self.bytes.reserve_exact(capacity - self.bytes.len()); + } + self.bytes.extend_from_slice(input); + Ok(input.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde::ser::SerializeSeq; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[test] + fn query_template_expansion_is_checked_before_allocating() { + assert!(expand_prompt(&"{query}".repeat(2000), "{query}", &"x".repeat(64 * 1024)).is_err()); + assert_eq!( + expand_prompt("Query: {query}", "{query}", "coral").unwrap(), + "Query: coral" + ); + } + + struct CountedRequest<'a>(&'a AtomicUsize); + impl Serialize for CountedRequest<'_> { + fn serialize(&self, serializer: S) -> Result { + let mut sequence = serializer.serialize_seq(Some(33_000))?; + let text = "x".repeat(1024); + for _ in 0..33_000 { + self.0.fetch_add(1, Ordering::Relaxed); + sequence.serialize_element(&text)?; + } + sequence.end() + } + } + + #[tokio::test] + async fn oversized_request_stops_serialization_at_the_byte_limit() { + let models = Models::new( + Arc::new(reqwest::Client::new()), + Arc::new(VectorStoresConfig::default()), + ); + let provider: ModelProvider = + serde_json::from_str(r#"{"base_url":"http://127.0.0.1:9/v1","models":["unused"]}"#).unwrap(); + let serialized = AtomicUsize::new(0); + let result: Result = models + .post(&provider, "chat/completions", &CountedRequest(&serialized), 1) + .await; + assert_eq!(result.err().unwrap().status_code(), 400); + assert!( + serialized.load(Ordering::Relaxed) < 33_000, + "serialization must stop at the limit, before traversing the entire payload" + ); + } +} diff --git a/crates/agentic-server-core/src/tool/file_search/ranking.rs b/crates/agentic-server-core/src/tool/file_search/ranking.rs index 857fab27..a58a365d 100644 --- a/crates/agentic-server-core/src/tool/file_search/ranking.rs +++ b/crates/agentic-server-core/src/tool/file_search/ranking.rs @@ -4,7 +4,7 @@ use std::collections::{HashMap, HashSet}; use crate::{ storage::file_search::StoredChunk, - types::file_search::{SearchContent, SearchMode, SearchRequest, SearchResult}, + types::file_search::{ChunkRetrievalParams, Ranker, SearchContent, SearchMode, SearchRequest, SearchResult}, }; pub(super) fn rank( @@ -13,6 +13,8 @@ pub(super) fn rank( query_embeddings: &[Vec], mode: SearchMode, request: &SearchRequest, + ranker: Ranker, + defaults: &ChunkRetrievalParams, ) -> Vec { let chunks: Vec<_> = chunks .into_iter() @@ -51,21 +53,10 @@ pub(super) fn rank( let scores = match mode { SearchMode::Keyword => keyword, SearchMode::Semantic => semantic, - SearchMode::Hybrid => { - let weights = request - .ranking_options - .as_ref() - .and_then(|options| options.hybrid_search.as_ref()); - fuse( - &semantic, - &keyword, - weights.and_then(|weights| weights.embedding_weight).unwrap_or(1.0), - weights.and_then(|weights| weights.text_weight).unwrap_or(1.0), - ) - } + SearchMode::Hybrid => hybrid_scores(&semantic, &keyword, request, ranker, defaults), }; for (index, score) in scores.into_iter().enumerate() { - if score <= 0.0 || score < threshold { + if !ranker.uses_model() && (score <= 0.0 || score < threshold) { continue; } let key = (chunks[index].file_id.as_str(), chunks[index].text.as_str()); @@ -101,6 +92,43 @@ pub(super) fn rank( .collect() } +fn hybrid_scores( + semantic: &[f64], + keyword: &[f64], + request: &SearchRequest, + ranker: Ranker, + defaults: &ChunkRetrievalParams, +) -> Vec { + let options = request.ranking_options.as_ref(); + let weights = options.and_then(|options| options.hybrid_search.as_ref()); + let alpha = options + .and_then(|options| options.alpha) + .unwrap_or(defaults.weighted_search_alpha); + let embedding = options + .and_then(|options| options.weights.as_ref()) + .map(|weights| weights.vector) + .or_else(|| weights.and_then(|weights| weights.embedding_weight)) + .unwrap_or(if weights.is_some() { 1.0 } else { alpha }); + let text = options + .and_then(|options| options.weights.as_ref()) + .map(|weights| weights.keyword) + .or_else(|| weights.and_then(|weights| weights.text_weight)) + .unwrap_or(if weights.is_some() { 1.0 } else { 1.0 - alpha }); + if ranker == Ranker::Weighted { + weighted(semantic, keyword, embedding, text) + } else { + fuse( + semantic, + keyword, + embedding, + text, + options + .and_then(|options| options.impact_factor) + .unwrap_or(defaults.rrf_impact_factor), + ) + } +} + fn cosine(left: &[f64], right: &[f64]) -> f64 { if left.len() != right.len() { return 0.0; @@ -181,7 +209,7 @@ fn bm25(chunks: &[StoredChunk], query: &str) -> Vec { } #[allow(clippy::cast_precision_loss, reason = "at most 10000 chunks are ranked")] -fn fuse(semantic: &[f64], keyword: &[f64], embedding_weight: f64, text_weight: f64) -> Vec { +fn fuse(semantic: &[f64], keyword: &[f64], embedding_weight: f64, text_weight: f64, impact_factor: f64) -> Vec { let mut scores = vec![0.0; semantic.len()]; for (input, weight) in [(semantic, embedding_weight), (keyword, text_weight)] { let mut ranked: Vec<_> = input @@ -194,12 +222,44 @@ fn fuse(semantic: &[f64], keyword: &[f64], embedding_weight: f64, text_weight: f right_score.total_cmp(left_score).then_with(|| left.cmp(right)) }); for (rank, (index, _)) in ranked.into_iter().enumerate() { - scores[index] += (weight / (embedding_weight + text_weight)) * 61.0 / (61.0 + rank as f64); + scores[index] += (weight / (embedding_weight + text_weight)) * (impact_factor + 1.0) + / (impact_factor + 1.0 + rank as f64); } } scores } +#[allow(clippy::float_cmp, reason = "equal extrema mean the score list is exactly constant")] +fn weighted(semantic: &[f64], keyword: &[f64], embedding_weight: f64, text_weight: f64) -> Vec { + fn normalize(input: &[f64]) -> Vec { + let min = input + .iter() + .copied() + .filter(|score| *score > 0.0) + .fold(f64::INFINITY, f64::min); + let max = input.iter().copied().fold(0.0, f64::max); + input + .iter() + .map(|score| { + if *score <= 0.0 { + 0.0 + } else if max == min { + 1.0 + } else { + (*score - min) / (max - min) + } + }) + .collect() + } + normalize(semantic) + .into_iter() + .zip(normalize(keyword)) + .map(|(semantic, keyword)| { + (semantic * embedding_weight + keyword * text_weight) / (embedding_weight + text_weight) + }) + .collect() +} + #[cfg(test)] mod tests { use super::cosine; diff --git a/crates/agentic-server-core/src/tool/file_search/service.rs b/crates/agentic-server-core/src/tool/file_search/service.rs index ffcf4470..cf1c0006 100644 --- a/crates/agentic-server-core/src/tool/file_search/service.rs +++ b/crates/agentic-server-core/src/tool/file_search/service.rs @@ -12,7 +12,7 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio_util::sync::CancellationToken; -use super::{embeddings::Embeddings, ingest, ranking}; +use super::{embeddings::Embeddings, ingest, models::Models, ranking}; use crate::{ storage::{ DbPool, @@ -23,8 +23,9 @@ use crate::{ }, types::file_search::{ AttachFileRequest, ChunkingStrategy, CreateVectorStoreRequest, DeleteObject, FileAttributes, FileCounts, - FileObject, FileSearchConfig, FileSearchError, ListParams, ListResponse, SearchMode, SearchQuery, - SearchRequest, SearchResponse, VectorStoreFileObject, VectorStoreObject, invalid, validate_attributes, + FileObject, FileSearchConfig, FileSearchError, ListParams, ListResponse, Ranker, RankingOptions, SearchMode, + SearchQuery, SearchRequest, SearchResponse, VectorStoreFileObject, VectorStoreObject, invalid, + validate_attributes, }, }; @@ -46,6 +47,8 @@ pub struct FileSearchService { files: LocalFiles, embeddings: Option, workers: Arc, + models: Models, + config: Arc, } impl std::fmt::Debug for FileSearchService { @@ -71,6 +74,7 @@ impl FileSearchService { client: Arc, config: FileSearchConfig, ) -> Result { + config.vector_stores.validate()?; let directory = match &config.files_storage_dir { Some(directory) => directory.clone(), None => crate::config::agentic_api_home() @@ -79,13 +83,18 @@ impl FileSearchService { }; if !matches!(config.backend, crate::types::file_search::FileSearchBackend::Exact) && config.embedding_base_url.is_none() + && config.vector_stores.default_embedding_model.is_none() { return invalid("pgvector requires configured embeddings"); } + let embeddings = Embeddings::from_config(client.clone(), &config)?; + let vector_config = Arc::new(config.vector_stores); Ok(Self { storage: FileSearchStorage::with_backend(pool, &config.backend)?, files: LocalFiles::new(directory)?, - embeddings: Embeddings::from_config(client, &config)?, + embeddings, + models: Models::new(client, vector_config.clone()), + config: vector_config, workers: Arc::new(Semaphore::new(4)), }) } @@ -94,6 +103,21 @@ impl FileSearchService { self.storage.initialize().await } + async fn prepare_context( + &self, + data: Vec, + permit: Arc, + ) -> Result, FileSearchError> { + let budget = self.config.chunk_retrieval_params.max_tokens_in_context; + let cancelled = Arc::new(AtomicBool::new(false)); + let _cancel_on_drop = CancelIngestionOnDrop(cancelled.clone()); + tokio::task::spawn_blocking(move || { + let _permit = permit; + ingest::limit_context(data, budget, &cancelled) + }) + .await? + } + fn permit(&self) -> Result, FileSearchError> { self.workers .clone() @@ -301,7 +325,9 @@ impl FileSearchService { dimensions = usize::try_from(prepared.dimensions).ok().filter(|value| *value != 0); for chunk in &prepared.chunks { total_bytes = total_bytes.saturating_add( - chunk.text.len() + chunk.embedding.as_ref().map_or(0, |embedding| embedding.len() * 8), + chunk.text.len() + + chunk.embedding_text.as_ref().map_or(0, String::len) + + chunk.embedding.as_ref().map_or(0, |embedding| embedding.len() * 8), ); } if total_bytes > 64 * 1024 * 1024 { @@ -387,18 +413,46 @@ impl FileSearchService { .read_content(&request.file_id, file.bytes, uploaded.content_base64, permit.clone()) .await?; let strategy = request.chunking_strategy.unwrap_or_default(); - let chunking = ingest::chunking_config(&strategy)?; + let chunking = if matches!(strategy, ChunkingStrategy::Auto) { + crate::types::file_search::StaticChunking { + max_chunk_size_tokens: self.config.file_ingestion_params.default_chunk_size_tokens, + chunk_overlap_tokens: self.config.file_ingestion_params.default_chunk_overlap_tokens, + } + } else { + ingest::chunking_config(&strategy)? + }; + if let ChunkingStrategy::Contextual { contextual } = &strategy { + if self.embeddings.is_none() { + return invalid("contextual ingestion requires configured embeddings"); + } + self.config.resolve( + contextual.model_id.as_deref(), + self.config.contextual_retrieval_params.model.as_ref(), + )?; + } let filename = file.filename.clone(); let worker_permit = permit.clone(); let cancelled = Arc::new(AtomicBool::new(false)); let _cancel_on_drop = CancelIngestionOnDrop(cancelled.clone()); - let texts = tokio::task::spawn_blocking(move || { + let document = tokio::task::spawn_blocking(move || { let _permit = worker_permit; ingest::extract_and_chunk(bytes, &filename, &uploaded.content_type, &chunking, &cancelled) }) .await??; + let contextual = if let ChunkingStrategy::Contextual { contextual } = &strategy { + Some( + self.models + .contextualize(&document.text, &document.chunks, contextual) + .await?, + ) + } else { + None + }; + let texts = document.chunks; let vectors = if let Some(embeddings) = &self.embeddings { - embeddings.embed(&texts, dimensions).await? + embeddings + .embed(contextual.as_deref().unwrap_or(&texts), dimensions) + .await? } else { Vec::new() }; @@ -410,18 +464,24 @@ impl FileSearchService { let embedding_dimensions = vectors.first().map_or(0, Vec::len); let mut usage_bytes = 0usize; let mut embeddings = vectors.into_iter(); + let mut contextual = contextual.unwrap_or_default().into_iter(); let chunks = texts .into_iter() .enumerate() .map(|(chunk_index, text)| { let embedding = embeddings.next(); - usage_bytes = usage_bytes - .saturating_add(text.len() + embedding.as_ref().map_or(0, |embedding| embedding.len() * 8)); + let embedding_text = contextual.next(); + usage_bytes = usage_bytes.saturating_add( + text.len() + + embedding_text.as_ref().map_or(0, String::len) + + embedding.as_ref().map_or(0, |embedding| embedding.len() * 8), + ); StoredChunk { file_id: file.id.clone(), filename: file.filename.clone(), chunk_index, text, + embedding_text, embedding, attributes: request.attributes.clone(), } @@ -488,38 +548,35 @@ impl FileSearchService { Ok(deleted(file_id, "vector_store.file.deleted")) } - /// Retrieves and globally ranks deduplicated chunks across selected vector stores. - /// - /// # Errors - /// Returns validation, configuration, provider, resource-limit, or storage errors. - pub async fn search( - &self, - store_ids: &[String], - request: &SearchRequest, - ) -> Result { - request.validate()?; - if store_ids.is_empty() || store_ids.len() > 16 || store_ids.iter().any(|id| id.is_empty() || id.len() > 128) { - return invalid("search requires 1 to 16 vector store IDs of at most 128 bytes each"); - } - let permit = self.permit()?; - let mode = request.search_mode.unwrap_or(if self.embeddings.is_some() { - SearchMode::Hybrid - } else { - SearchMode::Keyword - }); + fn search_mode(&self, request: &SearchRequest) -> Result { + let mode = request + .search_mode + .or(self.config.chunk_retrieval_params.default_search_mode) + .unwrap_or(if self.embeddings.is_some() { + SearchMode::Hybrid + } else { + SearchMode::Keyword + }); if mode != SearchMode::Keyword && self.embeddings.is_none() { return Err(FileSearchError::Unavailable( "Semantic and hybrid search require configured embeddings; use keyword search".into(), )); } if mode != SearchMode::Hybrid - && request - .ranking_options - .as_ref() - .is_some_and(|options| options.hybrid_search.is_some()) + && request.ranking_options.as_ref().is_some_and(|options| { + options.hybrid_search.is_some() || options.weights.is_some() || options.alpha.is_some() + }) { return invalid("hybrid_search weights require hybrid search mode"); } + Ok(mode) + } + + async fn search_dimensions( + &self, + store_ids: &[String], + mode: SearchMode, + ) -> Result<(Option, bool), FileSearchError> { let mut dimensions = None; for id in store_ids { let store = self.storage.store(id).await?; @@ -544,10 +601,73 @@ impl FileSearchService { } dimensions = Some(expected); } - let queries = match &request.query { + Ok((dimensions, has_embeddings)) + } + + fn ranking_strategy(&self, options: &RankingOptions) -> Result { + let mut ranker = options + .ranker + .as_deref() + .map(str::parse::) + .transpose()? + .unwrap_or(self.config.chunk_retrieval_params.default_reranker_strategy); + if ranker == Ranker::Auto { + ranker = self.config.chunk_retrieval_params.default_reranker_strategy; + if ranker == Ranker::Auto { + ranker = Ranker::Rrf; + } + } + if ranker.uses_model() { + self.config + .resolve(options.model.as_deref(), self.config.default_reranker_model.as_ref())?; + } else if options.model.is_some() { + return invalid("model selector requires a model ranker"); + } + Ok(ranker) + } + + /// Retrieves and globally ranks deduplicated chunks across selected vector stores. + /// + /// # Errors + /// Returns validation, configuration, provider, resource-limit, or storage errors. + pub async fn search( + &self, + store_ids: &[String], + request: &SearchRequest, + ) -> Result { + self.search_impl(store_ids, request, false).await + } + + pub(super) async fn search_for_tool( + &self, + store_ids: &[String], + request: &SearchRequest, + ) -> Result { + self.search_impl(store_ids, request, true).await + } + + async fn search_impl( + &self, + store_ids: &[String], + request: &SearchRequest, + prepare_context: bool, + ) -> Result { + request.validate()?; + if store_ids.is_empty() || store_ids.len() > 16 || store_ids.iter().any(|id| id.is_empty() || id.len() > 128) { + return invalid("search requires 1 to 16 vector store IDs of at most 128 bytes each"); + } + let permit = self.permit()?; + let mode = self.search_mode(request)?; + let (dimensions, has_embeddings) = self.search_dimensions(store_ids, mode).await?; + let mut queries = match &request.query { SearchQuery::Text(query) => vec![query.clone()], SearchQuery::Texts(queries) => queries.clone(), }; + let options = request.ranking_options.clone().unwrap_or_default(); + let ranker = self.ranking_strategy(&options)?; + if request.rewrite_query { + queries = vec![self.models.rewrite(&queries).await?]; + } let vectors = if mode == SearchMode::Keyword || !has_embeddings { Vec::new() } else { @@ -580,12 +700,45 @@ impl FileSearchService { } } let worker_queries = queries.clone(); - let request = request.clone(); - let data = tokio::task::spawn_blocking(move || { - let _permit = permit; - ranking::rank(chunks, &worker_queries, &vectors, mode, &request) + let mut candidate_request = request.clone(); + let limit = request.max_num_results.unwrap_or(10); + candidate_request.max_num_results = Some(if ranker.uses_model() { + limit * self.config.chunk_retrieval_params.chunk_multiplier + } else { + limit + }); + if ranker.uses_model() { + candidate_request + .ranking_options + .get_or_insert_with(Default::default) + .score_threshold = None; + } + let config = self.config.clone(); + let worker_permit = permit.clone(); + let mut data = tokio::task::spawn_blocking(move || { + let _permit = worker_permit; + ranking::rank( + chunks, + &worker_queries, + &vectors, + mode, + &candidate_request, + ranker, + &config.chunk_retrieval_params, + ) }) .await?; + if ranker.uses_model() { + data = self + .models + .rerank(&queries.join(" "), data, options.model.as_deref()) + .await?; + data.retain(|result| result.score >= options.score_threshold.unwrap_or(0.0)); + } + data.truncate(limit); + if prepare_context { + data = self.prepare_context(data, permit).await?; + } Ok(SearchResponse { object: "vector_store.search_results.page".into(), search_query: queries, @@ -654,6 +807,73 @@ fn page(mut data: Vec, params: &ListParams, id: impl Fn(&T) -> &str) -> Li mod local_file_tests { use super::*; + #[test] + fn cancelled_context_jobs_retain_all_four_slots_until_blocking_work_exits() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap(); + runtime.block_on(async { + let directory = tempfile::tempdir().unwrap(); + let pool = crate::storage::create_pool_with_schema(Some("sqlite::memory:")) + .await + .unwrap(); + let service = FileSearchService::new( + pool, + Arc::new(reqwest::Client::new()), + FileSearchConfig { + files_storage_dir: Some(directory.path().to_owned()), + ..FileSearchConfig::default() + }, + ) + .unwrap(); + let (release, resume) = std::sync::mpsc::channel(); + let (started, waiting) = tokio::sync::oneshot::channel(); + let blocker = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + resume.recv().unwrap(); + }); + waiting.await.unwrap(); + let mut jobs = Vec::new(); + for _ in 0..4 { + let mut job = Box::pin(service.prepare_context(Vec::new(), service.permit().unwrap())); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(20), &mut job) + .await + .is_err() + ); + jobs.push(job); + } + let admitted_slots = service.workers.available_permits(); + drop(jobs); + let cancelled_slots = service.workers.available_permits(); + let fifth_is_busy = service.permit().is_err(); + // Always release the blocker before assertions, including on the red run. + release.send(()).unwrap(); + blocker.await.unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(3), async { + while service.workers.available_permits() != 4 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_eq!( + admitted_slots, 0, + "queued context jobs must occupy the four operation slots" + ); + assert_eq!( + cancelled_slots, 0, + "cancelling the caller cannot release a queued blocking job's slot" + ); + assert!( + fifth_is_busy, + "new work must be rejected until cancelled blocking work exits" + ); + }); + } + #[test] fn cancelled_local_read_retains_capacity_until_filesystem_io_finishes() { let runtime = tokio::runtime::Builder::new_current_thread() diff --git a/crates/agentic-server-core/src/tool/normalize.rs b/crates/agentic-server-core/src/tool/normalize.rs index 7df35e55..b3e937a3 100644 --- a/crates/agentic-server-core/src/tool/normalize.rs +++ b/crates/agentic-server-core/src/tool/normalize.rs @@ -129,7 +129,7 @@ mod file_search_tests { serde_json::json!({"type": "file_search", "vector_store_ids": [" "]}), serde_json::json!({"type": "file_search", "vector_store_ids": ["vs_1"], "max_num_results": 0}), serde_json::json!({"type": "file_search", "vector_store_ids": ["vs_1"], "max_num_results": 51}), - serde_json::json!({"type": "file_search", "vector_store_ids": ["vs_1"], "ranking_options": {"ranker":"neural"}}), + serde_json::json!({"type": "file_search", "vector_store_ids": ["vs_1"], "ranking_options": {"ranker":"unknown"}}), serde_json::json!({"type": "file_search", "vector_store_ids": ["vs_1"], "ranking_options": {"score_threshold":2.0}}), ] { let tool: ResponsesTool = serde_json::from_value(declaration).unwrap(); diff --git a/crates/agentic-server-core/src/types/file_search.rs b/crates/agentic-server-core/src/types/file_search.rs index 6a11d45e..03603f5b 100644 --- a/crates/agentic-server-core/src/types/file_search.rs +++ b/crates/agentic-server-core/src/types/file_search.rs @@ -4,10 +4,14 @@ use std::{collections::BTreeMap, fmt, path::PathBuf}; use serde::{Deserialize, Serialize}; +pub use super::vector_stores::*; + /// Deployment-controlled embedding connection. An absent connection selects keyword retrieval. #[derive(Clone, Default, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct FileSearchConfig { + #[serde(default)] + pub vector_stores: VectorStoresConfig, #[serde(default)] pub backend: FileSearchBackend, pub files_storage_dir: Option, @@ -20,6 +24,7 @@ impl fmt::Debug for FileSearchConfig { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("FileSearchConfig") .field("backend", &self.backend) + .field("vector_stores", &self.vector_stores) .field("files_storage_dir", &self.files_storage_dir) .field("embedding_configured", &self.embedding_base_url.is_some()) .field("embedding_model", &self.embedding_model) @@ -120,11 +125,11 @@ pub enum FileSearchError { Storage(#[from] sqlx::Error), #[error("stored file search data could not be decoded")] Serialization(#[from] serde_json::Error), - #[error("embedding provider request failed")] + #[error("file search model provider request failed")] Provider(#[source] reqwest::Error), - #[error("embedding provider returned an invalid response")] + #[error("file search model provider returned an invalid response")] ProviderProtocol, - #[error("embedding provider returned malformed JSON")] + #[error("file search model provider returned malformed JSON")] ProviderDecode(#[source] serde_json::Error), #[error("file search worker failed")] Worker(#[from] tokio::task::JoinError), @@ -156,7 +161,7 @@ impl FileSearchError { | Self::Conflict(message) | Self::Unavailable(message) => message.clone(), Self::Provider(_) | Self::ProviderProtocol | Self::ProviderDecode(_) => { - "Embedding service request failed".into() + "File search model service request failed".into() } Self::Storage(_) | Self::Serialization(_) | Self::Worker(_) => "File search operation failed".into(), Self::FileStorage { .. } => { @@ -244,8 +249,25 @@ pub struct CompoundFilter { #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct RankingOptions { pub ranker: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub alpha: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub impact_factor: Option, pub score_threshold: Option, pub hybrid_search: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub weights: Option, +} + +/// Explicit vector/keyword fusion proportions. Neural scores replace fused scores. +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub struct FusionWeights { + pub vector: f64, + pub keyword: f64, } #[derive(Clone, Debug, Default, Serialize, Deserialize)] @@ -339,12 +361,66 @@ pub struct FileObject { pub enum ChunkingStrategy { #[default] Auto, + Contextual { + contextual: ContextualChunking, + }, Static { #[serde(rename = "static")] config: StaticChunking, }, } +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub struct ContextualChunking { + pub model_id: Option, + pub max_chunk_size_tokens: usize, + pub chunk_overlap_tokens: usize, + pub timeout_seconds: Option, + pub max_concurrency: Option, + pub context_prompt: String, +} +impl Default for ContextualChunking { + fn default() -> Self { + Self { + model_id: None, max_chunk_size_tokens: 700, chunk_overlap_tokens: 400, + timeout_seconds: None, max_concurrency: None, + context_prompt: "\n{{WHOLE_DOCUMENT}}\n\nHere is the chunk we want to situate within the whole document\n\n{{CHUNK_CONTENT}}\n\nPlease give a short succinct description to situate this chunk of text within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct description and nothing else.".into(), + } + } +} +impl ContextualChunking { + /// # Errors + /// Rejects unsupported bounds and malformed context templates. + pub fn validate(&self) -> Result<(), FileSearchError> { + if !(100..=4096).contains(&self.max_chunk_size_tokens) + || self.chunk_overlap_tokens >= self.max_chunk_size_tokens + || self.timeout_seconds.is_some_and(|n| !(1..=600).contains(&n)) + || self.max_concurrency.is_some_and(|n| !(1..=32).contains(&n)) + || self + .model_id + .as_ref() + .is_some_and(|model| model.trim().is_empty() || model.len() > 385) + || self.context_prompt.len() > 16384 + { + return invalid("invalid contextual chunk size, overlap, model, timeout, or concurrency"); + } + if self.context_prompt.matches("{{WHOLE_DOCUMENT}}").count() != 1 + || self.context_prompt.matches("{{CHUNK_CONTENT}}").count() != 1 + { + return invalid("context_prompt requires each document and chunk placeholder exactly once"); + } + match ( + self.context_prompt.find("{{WHOLE_DOCUMENT}}"), + self.context_prompt.find("{{CHUNK_CONTENT}}"), + ) { + (Some(document), Some(chunk)) if document < chunk => Ok(()), + _ => invalid("context_prompt requires {{WHOLE_DOCUMENT}} before {{CHUNK_CONTENT}}"), + } + } +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] @@ -486,9 +562,6 @@ impl SearchRequest { if !(1..=50).contains(&self.max_num_results.unwrap_or(10)) { return invalid("max_num_results must be between 1 and 50"); } - if self.rewrite_query { - return invalid("rewrite_query is not supported by in-tree file search"); - } if let Some(options) = &self.ranking_options { options.validate()?; } @@ -503,12 +576,27 @@ impl RankingOptions { /// # Errors /// Returns an invalid-request error for unsupported ranking or invalid weights. pub fn validate(&self) -> Result<(), FileSearchError> { + if let Some(ranker) = &self.ranker { + ranker.parse::()?; + } if self - .ranker - .as_deref() - .is_some_and(|ranker| !matches!(ranker, "auto" | "none")) + .model + .as_ref() + .is_some_and(|model| model.trim().is_empty() || model.len() > 385) { - return invalid("ranker must be auto or none; neural and classifier reranking are not supported"); + return invalid("model must contain 1 to 385 bytes"); + } + if self + .alpha + .is_some_and(|value| !value.is_finite() || !(0.0..=1.0).contains(&value)) + { + return invalid("alpha must be between 0 and 1"); + } + if self + .impact_factor + .is_some_and(|value| !value.is_finite() || !(0.0..=10000.0).contains(&value)) + { + return invalid("impact_factor must be between 0 and 10000"); } if self .score_threshold @@ -516,6 +604,19 @@ impl RankingOptions { { return invalid("score_threshold must be between 0 and 1"); } + if let Some(weights) = &self.weights { + if self.hybrid_search.is_some() + || !weights.vector.is_finite() + || !weights.keyword.is_finite() + || weights.vector < 0.0 + || weights.keyword < 0.0 + || (weights.vector + weights.keyword - 1.0).abs() > 1e-6 + { + return invalid( + "weights must be nonnegative vector/keyword proportions summing to 1; do not combine with hybrid_search", + ); + } + } if let Some(weights) = &self.hybrid_search { let embedding = weights.embedding_weight.unwrap_or(1.0); let text = weights.text_weight.unwrap_or(1.0); diff --git a/crates/agentic-server-core/src/types/mod.rs b/crates/agentic-server-core/src/types/mod.rs index 009964e4..16a86708 100644 --- a/crates/agentic-server-core/src/types/mod.rs +++ b/crates/agentic-server-core/src/types/mod.rs @@ -26,3 +26,6 @@ pub use tools::{ ShellEnvironment, ShellToolParam, ToolSearchExecution, ToolSearchStatus, ToolSearchToolParam, WebSearchContextSize, WebSearchFilters, WebSearchToolParam, WebSearchUserLocation, }; + +pub mod retrieval_models; +pub mod vector_stores; diff --git a/crates/agentic-server-core/src/types/retrieval_models.rs b/crates/agentic-server-core/src/types/retrieval_models.rs new file mode 100644 index 00000000..d102a3c4 --- /dev/null +++ b/crates/agentic-server-core/src/types/retrieval_models.rs @@ -0,0 +1,78 @@ +//! Typed text-only model transport for in-tree file search. +use serde::{Deserialize, Serialize}; +#[derive(Serialize)] +pub struct EmbeddingRequest<'a> { + pub model: &'a str, + pub input: &'a [String], + pub encoding_format: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + pub dimensions: Option, +} +#[derive(Deserialize)] +pub struct EmbeddingResponse { + pub model: String, + pub data: Vec, +} +#[derive(Deserialize)] +pub struct EmbeddingData { + pub index: usize, + pub embedding: Vec, +} +#[derive(Serialize)] +pub struct RetrievalChatRequest<'a> { + pub model: &'a str, + pub messages: Vec>, + pub stream: bool, + pub temperature: f64, + pub max_tokens: usize, +} +#[derive(Serialize)] +pub struct RetrievalMessage<'a> { + pub role: RetrievalRole, + pub content: &'a str, +} +#[derive(Serialize)] +#[serde(rename_all = "snake_case")] +pub enum RetrievalRole { + System, + User, +} +#[derive(Deserialize)] +pub struct RetrievalChatResponse { + pub choices: Vec, +} +#[derive(Deserialize)] +pub struct RetrievalChoice { + pub message: RetrievalChatMessage, +} +#[derive(Deserialize)] +pub struct RetrievalChatMessage { + pub content: Option, +} +#[derive(Deserialize)] +#[serde(untagged)] +pub enum RetrievalText { + Text(String), + Parts(Vec), +} +#[derive(Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum RetrievalTextPart { + Text { text: String }, +} +#[derive(Serialize)] +pub struct TextRerankRequest<'a> { + pub model: &'a str, + pub query: &'a str, + pub documents: Vec<&'a str>, + pub top_n: usize, +} +#[derive(Deserialize)] +pub struct TextRerankResponse { + pub results: Vec, +} +#[derive(Deserialize)] +pub struct TextRerankResult { + pub index: usize, + pub relevance_score: f64, +} diff --git a/crates/agentic-server-core/src/types/tools/params.rs b/crates/agentic-server-core/src/types/tools/params.rs index 5c710ded..8ab36c97 100644 --- a/crates/agentic-server-core/src/types/tools/params.rs +++ b/crates/agentic-server-core/src/types/tools/params.rs @@ -293,6 +293,10 @@ pub struct WebSearchToolParam { #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] #[serde(deny_unknown_fields)] pub struct FileSearchToolParam { + #[serde(skip_serializing_if = "Option::is_none")] + pub search_mode: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub rewrite_query: bool, pub vector_store_ids: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub max_num_results: Option, @@ -862,7 +866,7 @@ mod tests { assert_eq!(serde_json::to_value(tool).unwrap(), wire); assert!( serde_json::from_value::(serde_json::json!({ - "type":"file_search","vector_store_ids":["vs_1"],"rewrite_query":true + "type":"file_search","vector_store_ids":["vs_1"],"model_url":"https://unconfigured.invalid" })) .is_err() ); diff --git a/crates/agentic-server-core/src/types/vector_stores.rs b/crates/agentic-server-core/src/types/vector_stores.rs new file mode 100644 index 00000000..becdec7c --- /dev/null +++ b/crates/agentic-server-core/src/types/vector_stores.rs @@ -0,0 +1,344 @@ +//! Deployment-selected models and bounded retrieval settings. No request controls endpoints or secrets. +use super::file_search::{FileSearchError, SearchMode, invalid}; +use serde::{Deserialize, Serialize}; +use std::{collections::BTreeMap, fmt}; + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct VectorStoresConfig { + pub providers: BTreeMap, + pub default_provider_id: Option, + pub default_embedding_model: Option, + pub default_reranker_model: Option, + pub file_ingestion_params: FileIngestionParams, + pub chunk_retrieval_params: ChunkRetrievalParams, + pub contextual_retrieval_params: ContextualRetrievalParams, + pub rewrite_query_params: Option, + pub file_batch_params: FileBatchParams, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct QualifiedModel { + pub provider_id: String, + pub model_id: String, + pub embedding_dimensions: Option, +} + +#[derive(Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ModelProvider { + pub base_url: String, + pub models: Vec, + pub api_key_env: Option, + #[serde(skip)] + pub api_key: Option, + #[serde(default)] + pub protocol: ModelProtocol, + #[serde(default)] + pub score_interpretation: ScoreInterpretation, +} +impl fmt::Debug for ModelProvider { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ModelProvider") + .field("models", &self.models) + .field("protocol", &self.protocol) + .field("score_interpretation", &self.score_interpretation) + .field("api_key", &self.api_key.as_ref().map(|_| "[REDACTED]")) + .finish_non_exhaustive() + } +} +#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ModelProtocol { + #[default] + Vllm, +} +#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ScoreInterpretation { + #[default] + Probability, + Logit, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub enum Ranker { + Auto, + None, + #[default] + Rrf, + Weighted, + Normalized, + Neural, + Classifier, + #[serde(rename = "default-2024-11-15")] + Default20241115, + #[serde(rename = "default-2024-08-21", alias = "default_2024_08_21")] + Default20240821, +} +impl Ranker { + #[must_use] + pub const fn uses_model(self) -> bool { + matches!( + self, + Self::Neural | Self::Classifier | Self::Default20241115 | Self::Default20240821 + ) + } +} +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ChunkRetrievalParams { + pub chunk_multiplier: usize, + pub max_tokens_in_context: usize, + pub default_reranker_strategy: Ranker, + pub rrf_impact_factor: f64, + pub weighted_search_alpha: f64, + pub default_search_mode: Option, +} +impl Default for ChunkRetrievalParams { + fn default() -> Self { + Self { + chunk_multiplier: 5, + max_tokens_in_context: 4000, + default_reranker_strategy: Ranker::Rrf, + rrf_impact_factor: 60.0, + weighted_search_alpha: 0.5, + default_search_mode: None, + } + } +} +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct FileIngestionParams { + pub default_chunk_size_tokens: usize, + pub default_chunk_overlap_tokens: usize, +} +impl Default for FileIngestionParams { + fn default() -> Self { + Self { + default_chunk_size_tokens: 800, + default_chunk_overlap_tokens: 400, + } + } +} +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ContextualRetrievalParams { + pub model: Option, + pub default_timeout_seconds: u64, + pub default_max_concurrency: usize, + pub max_document_tokens: usize, +} +impl Default for ContextualRetrievalParams { + fn default() -> Self { + Self { + model: None, + default_timeout_seconds: 120, + default_max_concurrency: 3, + max_document_tokens: 100_000, + } + } +} +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct RewriteQueryParams { + pub model: Option, + pub max_tokens: usize, + pub temperature: f64, + pub prompt: String, +} +impl Default for RewriteQueryParams { + fn default() -> Self { + Self{model:None,max_tokens:100,temperature:0.3,prompt:"Expand this query with relevant synonyms and related terms. Return only the improved query, no explanations:\n\n{query}\n\nImproved query:".into()} + } +} +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct FileBatchParams { + pub max_concurrent_files_per_batch: usize, + pub file_batch_chunk_size: usize, + pub cleanup_interval_seconds: u64, +} +impl Default for FileBatchParams { + fn default() -> Self { + Self { + max_concurrent_files_per_batch: 3, + file_batch_chunk_size: 10, + cleanup_interval_seconds: 86400, + } + } +} + +impl VectorStoresConfig { + /// Checks all deployment bounds and registered model references before serving traffic. + /// # Errors + /// Returns an invalid-request configuration error. + pub fn validate(&self) -> Result<(), FileSearchError> { + if self.providers.len() > 32 { + return invalid("at most 32 model providers may be configured"); + } + for (id, provider) in &self.providers { + if id.is_empty() + || id.len() > 128 + || id.contains('/') + || provider.models.is_empty() + || provider.models.len() > 128 + || provider + .models + .iter() + .any(|model| model.trim().is_empty() || model.len() > 256) + { + return invalid("provider IDs and allowlisted models must be nonempty and bounded"); + } + provider.endpoint("chat/completions")?; + if provider.api_key_env.as_ref().is_some_and(|key| { + key.is_empty() || !key.bytes().all(|byte| byte.is_ascii_alphanumeric() || byte == b'_') + }) { + return invalid("provider api_key_env must be an environment variable name"); + } + } + if self + .default_provider_id + .as_ref() + .is_some_and(|id| !self.providers.contains_key(id)) + { + return invalid("default_provider_id is not configured"); + } + for model in [ + self.default_embedding_model.as_ref(), + self.default_reranker_model.as_ref(), + self.contextual_retrieval_params.model.as_ref(), + self.rewrite_query_params + .as_ref() + .and_then(|rewrite| rewrite.model.as_ref()), + ] + .into_iter() + .flatten() + { + self.resolve(None, Some(model))?; + if model.embedding_dimensions.is_some_and(|n| !(1..=4096).contains(&n)) { + return invalid("embedding_dimensions must be 1 to 4096"); + } + } + let retrieval = &self.chunk_retrieval_params; + if !(1..=20).contains(&retrieval.chunk_multiplier) + || !(1..=32768).contains(&retrieval.max_tokens_in_context) + || !retrieval.rrf_impact_factor.is_finite() + || !(0.0..=10000.0).contains(&retrieval.rrf_impact_factor) + || !retrieval.weighted_search_alpha.is_finite() + || !(0.0..=1.0).contains(&retrieval.weighted_search_alpha) + { + return invalid("invalid chunk retrieval bounds, alpha, or RRF impact factor"); + } + let ingestion = &self.file_ingestion_params; + if !(100..=4096).contains(&ingestion.default_chunk_size_tokens) + || ingestion.default_chunk_overlap_tokens > ingestion.default_chunk_size_tokens / 2 + { + return invalid("ingestion chunk size must be 100 to 4096 and overlap at most half"); + } + let context = &self.contextual_retrieval_params; + if !(10..=600).contains(&context.default_timeout_seconds) + || !(1..=32).contains(&context.default_max_concurrency) + || !(1000..=1_000_000).contains(&context.max_document_tokens) + { + return invalid("invalid contextual timeout, concurrency, or document bound"); + } + if let Some(rewrite) = &self.rewrite_query_params { + if !(1..=4096).contains(&rewrite.max_tokens) + || !rewrite.temperature.is_finite() + || !(0.0..=2.0).contains(&rewrite.temperature) + || !rewrite.prompt.contains("{query}") + || rewrite.prompt.len() > 16384 + { + return invalid("invalid query rewrite token limit, temperature, or prompt (requires {query})"); + } + } + let batch = &self.file_batch_params; + if !(1..=32).contains(&batch.max_concurrent_files_per_batch) + || !(1..=1000).contains(&batch.file_batch_chunk_size) + || !(1..=604_800).contains(&batch.cleanup_interval_seconds) + { + return invalid("invalid file batch concurrency, chunk size, or cleanup interval"); + } + Ok(()) + } + /// Resolves a registered provider and its allowlisted provider-local model. + /// # Errors + /// Fails closed on unknown providers/models or absent defaults. + pub fn resolve<'a>( + &'a self, + selection: Option<&'a str>, + default: Option<&'a QualifiedModel>, + ) -> Result<(&'a ModelProvider, &'a str), FileSearchError> { + let (provider_id, model_id) = if let Some(selection) = selection { + match selection.split_once('/') { + Some(parts) => parts, + None => ( + self.default_provider_id.as_deref().ok_or_else(|| { + FileSearchError::InvalidRequest("unqualified model requires default_provider_id".into()) + })?, + selection, + ), + } + } else if let Some(default) = default { + (default.provider_id.as_str(), default.model_id.as_str()) + } else { + return invalid("requested model operation requires a configured default model or model selector"); + }; + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| FileSearchError::InvalidRequest("model provider is not configured".into()))?; + if !provider.models.iter().any(|model| model == model_id) { + return invalid("model is not in the provider allowlist"); + } + Ok((provider, model_id)) + } +} +impl ModelProvider { + /// Constructs a validated provider endpoint without routing request data as URLs. + /// # Errors + /// Returns invalid-request for malformed or credential-bearing base URLs. + pub fn endpoint(&self, operation: &str) -> Result { + let mut url = reqwest::Url::parse(&self.base_url) + .map_err(|_| FileSearchError::InvalidRequest("invalid model provider base_url".into()))?; + if !matches!(url.scheme(), "http" | "https") + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + { + return invalid("model provider base_url must be HTTP(S) without credentials, query, or fragment"); + } + let base = url.path().trim_end_matches('/'); + let base = if operation == "rerank" { + base.strip_suffix("/v1").unwrap_or(base) + } else { + base + }; + url.set_path(&format!("{base}/{operation}")); + Ok(url) + } +} + +impl std::str::FromStr for Ranker { + type Err = FileSearchError; + fn from_str(value: &str) -> Result { + match value { + "auto" => Ok(Self::Auto), + "none" => Ok(Self::None), + "rrf" => Ok(Self::Rrf), + "weighted" => Ok(Self::Weighted), + "normalized" => Ok(Self::Normalized), + "neural" => Ok(Self::Neural), + "classifier" => Ok(Self::Classifier), + "default-2024-11-15" => Ok(Self::Default20241115), + "default-2024-08-21" | "default_2024_08_21" => Ok(Self::Default20240821), + _ => invalid("unsupported file search ranker"), + } + } +} diff --git a/crates/agentic-server-core/tests/file_search_models.rs b/crates/agentic-server-core/tests/file_search_models.rs new file mode 100644 index 00000000..6661389f --- /dev/null +++ b/crates/agentic-server-core/tests/file_search_models.rs @@ -0,0 +1,871 @@ +//! Local HTTP contracts for model-assisted retrieval; never calls external models. +use agentic_core::{storage::create_pool_with_schema, tool::file_search::FileSearchService, types::file_search::*}; +use axum::{Json, extract::State, http::StatusCode}; +use serde_json::{Value, json}; +use std::sync::{Arc, Mutex}; + +#[derive(Clone, Default)] +struct Fixture { + requests: Arc>>, + rerank_response: Arc>>, + chat_fail: Arc>, + chat_fail_after: Arc>>, + chat_response: Arc>>, + chat_wait: Arc>, + chat_started: Arc, + authorization: Arc>>, +} +async fn chat( + State(state): State, + headers: axum::http::HeaderMap, + Json(input): Json, +) -> (StatusCode, Json) { + state.authorization.lock().unwrap().push( + headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_owned(), + ); + state.requests.lock().unwrap().push(("chat".into(), input.clone())); + state.chat_started.notify_one(); + let wait = *state.chat_wait.lock().unwrap(); + if wait { + std::future::pending::<()>().await; + } + let count = state + .requests + .lock() + .unwrap() + .iter() + .filter(|(kind, _)| kind == "chat") + .count(); + if state.chat_fail_after.lock().unwrap().is_some_and(|limit| count > limit) { + return (StatusCode::BAD_GATEWAY, Json(json!({"error":"private-secret"}))); + } + if let Some(response) = state.chat_response.lock().unwrap().clone() { + return (StatusCode::OK, Json(response)); + } + if *state.chat_fail.lock().unwrap() { + return (StatusCode::BAD_GATEWAY, Json(json!({"error":"private-secret"}))); + } + let context = input["messages"].as_array().unwrap().len() == 2; + ( + StatusCode::OK, + Json(json!({"choices":[{"message":{"content": if context {"aquatic context"} else {"coral"}}}]})), + ) +} +async fn embed(State(state): State, headers: axum::http::HeaderMap, Json(input): Json) -> Json { + state.authorization.lock().unwrap().push( + headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_owned(), + ); + state.requests.lock().unwrap().push(("embed".into(), input.clone())); + let data: Vec<_> = input["input"].as_array().unwrap().iter().enumerate().map(|(index,text)| { + let text = text.as_str().unwrap(); + json!({"index":index,"embedding":if text.contains("aquatic") || text=="ocean coral" {vec![1.0,0.0]} else {vec![0.0,1.0]}}) + }).collect(); + Json(json!({"model":input["model"],"data":data})) +} +async fn rerank(State(state): State, headers: axum::http::HeaderMap, Json(input): Json) -> Json { + state.authorization.lock().unwrap().push( + headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_owned(), + ); + state.requests.lock().unwrap().push(("rerank".into(), input.clone())); + if let Some(output) = state.rerank_response.lock().unwrap().clone() { + return Json(output); + } + let data: Vec<_> = input["documents"].as_array().unwrap().iter().enumerate().map(|(index,text)| { + json!({"index":index,"relevance_score": if text.as_str().unwrap().contains("preferred") {0.95} else {0.1}}) + }).collect(); + Json(json!({"results":data})) +} +struct Setup { + service: FileSearchService, + state: Fixture, + config: FileSearchConfig, + pool: Arc, + _files: tempfile::TempDir, + task: tokio::task::JoinHandle<()>, +} +impl Drop for Setup { + fn drop(&mut self) { + self.task.abort(); + } +} +async fn setup(embedding: bool) -> Setup { + let state = Fixture::default(); + let app = axum::Router::new() + .route("/v1/chat/completions", axum::routing::post(chat)) + .route("/v1/embeddings", axum::routing::post(embed)) + .route("/rerank", axum::routing::post(rerank)) + .with_state(state.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}/v1", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let files = tempfile::tempdir().unwrap(); + let mut config: FileSearchConfig = serde_json::from_value(json!({"vector_stores":{ + "default_provider_id":"local", + "providers":{"local":{"base_url":base,"models":["embed","chat","org/rerank"]}}, + "default_reranker_model":{"provider_id":"local","model_id":"org/rerank"}, + "contextual_retrieval_params":{"model":{"provider_id":"local","model_id":"chat"}}, + "rewrite_query_params":{"model":{"provider_id":"local","model_id":"chat"},"temperature":0.0} + }})) + .unwrap(); + config.files_storage_dir = Some(files.path().to_owned()); + if embedding { + // Exercise grouped embedding routing, including provider-local model names. + let mut wire = serde_json::to_value(&config).unwrap(); + wire["vector_stores"]["default_embedding_model"] = + json!({"provider_id":"local","model_id":"embed","embedding_dimensions":2}); + config = serde_json::from_value(wire).unwrap(); + } + let pool = create_pool_with_schema(Some("sqlite::memory:")).await.unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new( + reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .unwrap(), + ), + config.clone(), + ) + .unwrap(); + Setup { + service, + state, + config, + pool, + _files: files, + task, + } +} +async fn attach(service: &FileSearchService, store: &str, text: &str, strategy: Option) -> String { + let file = service + .upload_file("source.txt", "text/plain", "assistants", text.as_bytes().to_vec()) + .await + .unwrap(); + service + .attach_file( + store, + AttachFileRequest { + file_id: file.id.clone(), + chunking_strategy: strategy, + ..Default::default() + }, + ) + .await + .unwrap(); + file.id +} +fn query(text: &str) -> SearchRequest { + SearchRequest { + query: SearchQuery::Text(text.into()), + ..Default::default() + } +} +fn contextual() -> ChunkingStrategy { + serde_json::from_value(json!({"type":"contextual","contextual":{}})).unwrap() +} + +#[tokio::test] +async fn contextual_embedding_changes_retrieval_but_preserves_original_source() { + let setup = setup(true).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let id = attach(&setup.service, &store.id, "unadorned source", Some(contextual())).await; + let result = setup + .service + .search( + &[store.id], + &SearchRequest { + search_mode: Some(SearchMode::Semantic), + ..query("aquatic") + }, + ) + .await + .unwrap(); + assert_eq!(result.data[0].file_id, id); + assert_eq!(result.data[0].content[0].text, "unadorned source"); + { + let requests = setup.state.requests.lock().unwrap(); + let embedded = requests.iter().find(|(kind, _)| kind == "embed").unwrap(); + assert_eq!(embedded.1["input"][0], "aquatic context\n\nunadorned source"); + } + let persisted: String = sqlx::query_scalar("SELECT data FROM file_search_chunks LIMIT 1") + .fetch_one(setup.pool.as_ref()) + .await + .unwrap(); + let persisted: Value = serde_json::from_str(&persisted).unwrap(); + assert_eq!(persisted["text"], "unadorned source"); + assert_eq!(persisted["embedding_text"], "aquatic context\n\nunadorned source"); +} + +#[tokio::test] +async fn rewrite_changes_retrieval_and_reported_query_preserving_zero_temperature() { + let setup = setup(false).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let id = attach(&setup.service, &store.id, "coral reefs", None).await; + assert!( + setup + .service + .search(std::slice::from_ref(&store.id), &query("ocean")) + .await + .unwrap() + .data + .is_empty() + ); + let result = setup + .service + .search( + &[store.id], + &SearchRequest { + rewrite_query: true, + query: SearchQuery::Texts(vec!["ocean".into(), "habitat".into()]), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(result.search_query, vec!["coral"]); + assert_eq!(result.data[0].file_id, id); + let requests = setup.state.requests.lock().unwrap(); + assert_eq!(requests[0].1["temperature"], 0.0); + assert!( + requests[0].1["messages"][0]["content"] + .as_str() + .unwrap() + .contains("ocean habitat") + ); +} + +#[tokio::test] +async fn reranker_sees_candidates_before_truncation_and_none_bypasses_model() { + let setup = setup(false).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + attach(&setup.service, &store.id, "coral coral coral", None).await; + let preferred = attach( + &setup.service, + &store.id, + "coral preferred with many other words to reduce lexical score", + None, + ) + .await; + for ranker in [ + "neural", + "classifier", + "default-2024-11-15", + "default-2024-08-21", + "default_2024_08_21", + ] { + let request: SearchRequest = + serde_json::from_value(json!({"query":"coral","max_num_results":1,"ranking_options":{"ranker":ranker}})) + .unwrap(); + let result = setup + .service + .search(std::slice::from_ref(&store.id), &request) + .await + .unwrap(); + assert_eq!(result.data[0].file_id, preferred); + assert!((result.data[0].score - 0.95).abs() < f64::EPSILON); + } + let request: SearchRequest = + serde_json::from_value(json!({"query":"coral","max_num_results":1,"ranking_options":{"ranker":"none"}})) + .unwrap(); + setup.service.search(&[store.id], &request).await.unwrap(); + let requests = setup.state.requests.lock().unwrap(); + assert_eq!(requests.len(), 5); + assert_eq!(requests[0].1["model"], "org/rerank"); + assert_eq!(requests[0].1["documents"].as_array().unwrap().len(), 2); +} + +#[tokio::test] +async fn failed_context_is_atomic_and_rerank_protocol_errors_are_not_partial_results() { + let setup = setup(true).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + *setup.state.chat_fail.lock().unwrap() = true; + let file = setup + .service + .upload_file("failed.txt", "text/plain", "assistants", b"failed source".to_vec()) + .await + .unwrap(); + let error = setup + .service + .attach_file( + &store.id, + AttachFileRequest { + file_id: file.id, + chunking_strategy: Some(contextual()), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert_eq!(error.status_code(), 502); + assert!(!error.public_message().contains("private-secret")); + assert_eq!( + setup + .service + .get_vector_store(&store.id) + .await + .unwrap() + .file_counts + .total, + 0 + ); + assert!( + !setup + .state + .requests + .lock() + .unwrap() + .iter() + .any(|(kind, _)| kind == "embed") + ); + *setup.state.chat_fail.lock().unwrap() = false; + attach(&setup.service, &store.id, "coral one", None).await; + attach(&setup.service, &store.id, "coral two", None).await; + for results in [ + json!([]), + json!([{"index":0,"relevance_score":0.9},{"index":0,"relevance_score":0.2}]), + json!([{"index":0,"relevance_score":0.9},{"index":2,"relevance_score":0.2}]), + json!([{"index":0,"relevance_score":2.0},{"index":1,"relevance_score":0.2}]), + json!([{"index":0,"relevance_score":-0.1},{"index":1,"relevance_score":0.2}]), + ] { + *setup.state.rerank_response.lock().unwrap() = Some(json!({"results":results})); + let request: SearchRequest = serde_json::from_value( + json!({"query":"coral","search_mode":"keyword","ranking_options":{"ranker":"neural"}}), + ) + .unwrap(); + assert_eq!( + setup + .service + .search(std::slice::from_ref(&store.id), &request) + .await + .unwrap_err() + .status_code(), + 502 + ); + } +} + +#[tokio::test] +async fn configured_ranker_mode_logit_scores_and_request_overrides_apply_consistently() { + let setup = setup(false).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let id = attach(&setup.service, &store.id, "coral", None).await; + let mut config = setup.config.clone(); + config.vector_stores.chunk_retrieval_params.default_search_mode = Some(SearchMode::Keyword); + config.vector_stores.chunk_retrieval_params.default_reranker_strategy = Ranker::Neural; + config + .vector_stores + .providers + .get_mut("local") + .unwrap() + .score_interpretation = ScoreInterpretation::Logit; + let service = FileSearchService::new(setup.pool.clone(), Arc::new(reqwest::Client::new()), config).unwrap(); + for (raw, expected) in [(0.0, 0.5), (1000.0, 1.0), (-1000.0, 0.0)] { + *setup.state.rerank_response.lock().unwrap() = Some(json!({"results":[{"index":0,"relevance_score":raw}]})); + let result = service + .search(std::slice::from_ref(&store.id), &query("coral")) + .await + .unwrap(); + assert_eq!(result.data[0].file_id, id); + assert!((result.data[0].score - expected).abs() < f64::EPSILON); + } + let none: SearchRequest = + serde_json::from_value(json!({"query":"coral","ranking_options":{"ranker":"none"}})).unwrap(); + service.search(std::slice::from_ref(&store.id), &none).await.unwrap(); + assert_eq!(setup.state.requests.lock().unwrap().len(), 3); + let filtered: SearchRequest = serde_json::from_value( + json!({"query":"coral","ranking_options":{"ranker":"classifier","score_threshold":0.1}}), + ) + .unwrap(); + assert!(service.search(&[store.id], &filtered).await.unwrap().data.is_empty()); +} + +#[tokio::test] +async fn unknown_model_selection_fails_before_provider_calls() { + let setup = setup(false).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + attach(&setup.service, &store.id, "coral", None).await; + for model in ["foreign/org/rerank", "local/missing", "https://127.0.0.1/evil"] { + let request: SearchRequest = + serde_json::from_value(json!({"query":"coral","ranking_options":{"ranker":"neural","model":model}})) + .unwrap(); + assert_eq!( + setup + .service + .search(std::slice::from_ref(&store.id), &request) + .await + .unwrap_err() + .status_code(), + 400 + ); + } + assert!(setup.state.requests.lock().unwrap().is_empty()); +} + +#[test] +fn grouped_configuration_and_contextual_bounds_fail_closed() { + for params in [ + json!({"chunk_retrieval_params":{"chunk_multiplier":0}}), + json!({"chunk_retrieval_params":{"weighted_search_alpha":2.0}}), + json!({"file_batch_params":{"max_concurrent_files_per_batch":0}}), + json!({"file_batch_params":{"file_batch_chunk_size":0}}), + json!({"contextual_retrieval_params":{"default_max_concurrency":0}}), + json!({"rewrite_query_params":{"temperature":-1.0}}), + json!({"rewrite_query_params":{"prompt":"missing placeholder"}}), + json!({"default_provider_id":"missing"}), + ] { + let config: VectorStoresConfig = serde_json::from_value(params).unwrap(); + assert!(config.validate().is_err()); + } + for params in [ + json!({"chunk_overlap_tokens":700}), + json!({"max_concurrency":0}), + json!({"timeout_seconds":0}), + json!({"context_prompt":"{{CHUNK_CONTENT}} {{WHOLE_DOCUMENT}}"}), + ] { + let config: ContextualChunking = serde_json::from_value(params).unwrap(); + assert!(config.validate().is_err()); + } + ContextualChunking::default().validate().unwrap(); + assert!( + serde_json::from_value::( + json!({"providers":{"local":{"base_url":"http://localhost/v1","models":["model"],"api_key":"secret"}}}) + ) + .is_err() + ); +} + +#[tokio::test] +async fn deployment_and_request_fusion_weights_choose_different_evidence() { + let setup = setup(true).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let semantic = attach(&setup.service, &store.id, "aquatic habitat", None).await; + let lexical = attach(&setup.service, &store.id, "coral reefs", None).await; + let mut config = setup.config.clone(); + config.vector_stores.chunk_retrieval_params.default_reranker_strategy = Ranker::Weighted; + config.vector_stores.chunk_retrieval_params.weighted_search_alpha = 1.0; + let service = FileSearchService::new(setup.pool.clone(), Arc::new(reqwest::Client::new()), config).unwrap(); + let result = service + .search(std::slice::from_ref(&store.id), &query("ocean coral")) + .await + .unwrap(); + assert_eq!(result.data[0].file_id, semantic); + let request: SearchRequest = serde_json::from_value( + json!({"query":"ocean coral","ranking_options":{"ranker":"weighted","weights":{"vector":0.0,"keyword":1.0}}}), + ) + .unwrap(); + let result = service.search(&[store.id], &request).await.unwrap(); + assert_eq!(result.data[0].file_id, lexical); +} + +#[tokio::test] +async fn responses_tool_honors_rewrite_ranking_defaults_and_context_token_budget() { + use agentic_core::tool::{ + GatewayExecutor, + file_search::{FileSearchExecutionParams, FileSearchHandler}, + }; + let setup = setup(false).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + attach( + &setup.service, + &store.id, + "coral preferred source with enough words to exceed the configured context budget", + None, + ) + .await; + let mut config = setup.config.clone(); + config.vector_stores.chunk_retrieval_params.default_reranker_strategy = Ranker::Neural; + config.vector_stores.chunk_retrieval_params.max_tokens_in_context = 2; + let service = FileSearchService::new(setup.pool.clone(), Arc::new(reqwest::Client::new()), config).unwrap(); + let handler = FileSearchHandler::new(service); + let params = FileSearchExecutionParams { + declaration: serde_json::from_value( + json!({"vector_store_ids":[store.id],"rewrite_query":true,"search_mode":"keyword"}), + ) + .unwrap(), + include_results: true, + }; + let output = handler + .execute("call_1", "file_search", r#"{"queries":["ocean"]}"#, ¶ms) + .await + .unwrap(); + let output: Value = serde_json::from_str(&output.output).unwrap(); + assert_eq!(output["queries"], json!(["coral"])); + assert!(output["retrieved_passages"].as_array().unwrap().is_empty()); + assert_eq!(setup.state.requests.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn partial_context_failure_timeout_and_cancellation_publish_nothing() { + let setup = setup(true).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let file = setup + .service + .upload_file( + "multi.txt", + "text/plain", + "assistants", + "word ".repeat(250).into_bytes(), + ) + .await + .unwrap(); + let strategy:ChunkingStrategy=serde_json::from_value(json!({"type":"contextual","contextual":{"max_chunk_size_tokens":100,"chunk_overlap_tokens":0,"max_concurrency":1,"timeout_seconds":1}})).unwrap(); + let request = AttachFileRequest { + file_id: file.id, + chunking_strategy: Some(strategy), + ..Default::default() + }; + *setup.state.chat_fail_after.lock().unwrap() = Some(1); + assert_eq!( + setup + .service + .attach_file(&store.id, request.clone()) + .await + .unwrap_err() + .status_code(), + 502 + ); + assert_eq!( + setup + .service + .get_vector_store(&store.id) + .await + .unwrap() + .file_counts + .total, + 0 + ); + assert!( + !setup + .state + .requests + .lock() + .unwrap() + .iter() + .any(|(kind, _)| kind == "embed") + ); + *setup.state.chat_fail_after.lock().unwrap() = None; + *setup.state.chat_wait.lock().unwrap() = true; + assert_eq!( + setup + .service + .attach_file(&store.id, request.clone()) + .await + .unwrap_err() + .status_code(), + 502 + ); + setup.state.chat_started.notified().await; + // Cancel while a context request is pending; the service must release permits and publish no attachment. + let service = setup.service.clone(); + let store_id = store.id.clone(); + let work_request = request.clone(); + let work = tokio::spawn(async move { service.attach_file(&store_id, work_request).await }); + tokio::time::timeout(std::time::Duration::from_secs(3), setup.state.chat_started.notified()) + .await + .unwrap(); + work.abort(); + assert!(work.await.unwrap_err().is_cancelled()); + assert_eq!( + setup + .service + .get_vector_store(&store.id) + .await + .unwrap() + .file_counts + .total, + 0 + ); + *setup.state.chat_wait.lock().unwrap() = false; + setup.service.attach_file(&store.id, request).await.unwrap(); + assert_eq!( + setup + .service + .get_vector_store(&store.id) + .await + .unwrap() + .file_counts + .completed, + 1 + ); +} + +#[tokio::test] +async fn malformed_or_oversized_rewrite_outputs_fail_without_fallback() { + let setup = setup(false).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + attach(&setup.service, &store.id, "coral", None).await; + for response in [ + json!({"choices":[]}), + json!({"choices":[{"message":{"content":null}}]}), + json!({"choices":[{"message":{"content":" "}}]}), + json!({"choices":[{"message":{"content":"a".repeat(4097)}}]}), + json!({"choices":[{"message":{"content":"a".repeat(1024*1024)}}]}), + ] { + *setup.state.chat_response.lock().unwrap() = Some(response); + assert_eq!( + setup + .service + .search( + std::slice::from_ref(&store.id), + &SearchRequest { + rewrite_query: true, + ..query("coral") + } + ) + .await + .unwrap_err() + .status_code(), + 502 + ); + } + *setup.state.chat_response.lock().unwrap() = None; + *setup.state.chat_fail.lock().unwrap() = true; + assert_eq!( + setup + .service + .search( + &[store.id], + &SearchRequest { + rewrite_query: true, + ..query("coral") + } + ) + .await + .unwrap_err() + .status_code(), + 502 + ); +} + +#[tokio::test] +async fn neural_ranker_promotes_zero_similarity_before_applying_final_threshold() { + let setup = setup(true).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let preferred = attach(&setup.service, &store.id, "coral preferred", None).await; + attach(&setup.service, &store.id, "aquatic", None).await; + let request:SearchRequest=serde_json::from_value(json!({"query":"aquatic","search_mode":"semantic","max_num_results":1,"ranking_options":{"ranker":"neural","score_threshold":0.9}})).unwrap(); + let result = setup.service.search(&[store.id], &request).await.unwrap(); + assert_eq!(result.data.len(), 1); + assert_eq!(result.data[0].file_id, preferred); + assert!((result.data[0].score - 0.95).abs() < f64::EPSILON); +} + +#[tokio::test] +#[ignore = "requires isolated TEST_POSTGRES_URL with pgvector"] +async fn postgres_pgvector_contextual_and_neural_ranking_preserve_source_and_candidate_bounds() { + let setup = setup(true).await; + let url = std::env::var("TEST_POSTGRES_URL").unwrap(); + let pool = create_pool_with_schema(Some(&url)).await.unwrap(); + let mut config = setup.config.clone(); + config.backend = FileSearchBackend::Pgvector { + dimensions: 2, + index: PgvectorIndex::Hnsw { + m: 16, + ef_construction: 64, + ef_search: 100, + }, + candidate_limit: 50, + }; + let service = FileSearchService::new(pool, Arc::new(reqwest::Client::new()), config).unwrap(); + let store = service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let contextual_id = attach(&service, &store.id, "original source", Some(contextual())).await; + let preferred = attach(&service, &store.id, "coral preferred", None).await; + let result = service + .search( + std::slice::from_ref(&store.id), + &SearchRequest { + search_mode: Some(SearchMode::Semantic), + ..query("aquatic") + }, + ) + .await + .unwrap(); + assert_eq!(result.data[0].file_id, contextual_id); + assert_eq!(result.data[0].content[0].text, "original source"); + let request:SearchRequest=serde_json::from_value(json!({"query":"aquatic","search_mode":"semantic","max_num_results":1,"ranking_options":{"ranker":"neural","score_threshold":0.9}})).unwrap(); + let result = service.search(std::slice::from_ref(&store.id), &request).await.unwrap(); + assert_eq!(result.data[0].file_id, preferred); + let mut extra_files = Vec::new(); + for index in 0..60 { + extra_files.push(attach(&service, &store.id, &format!("aquatic evidence {index}"), None).await); + } + let mut bounded_request = request; + bounded_request.max_num_results = Some(50); + service + .search(std::slice::from_ref(&store.id), &bounded_request) + .await + .unwrap(); + { + let requests = setup.state.requests.lock().unwrap(); + let last = requests.iter().rev().find(|(kind, _)| kind == "rerank").unwrap(); + assert_eq!(last.1["documents"].as_array().unwrap().len(), 50); + } + for file in extra_files { + service.delete_file(&file).await.unwrap(); + } + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&contextual_id).await.unwrap(); + service.delete_file(&preferred).await.unwrap(); +} + +#[tokio::test] +async fn provider_routing_keeps_embedding_and_generation_credentials_independent() { + let embedding = setup(true).await; + let generation = setup(false).await; + let mut config = embedding.config.clone(); + config.vector_stores.providers.get_mut("local").unwrap().api_key = Some("embedding-secret".into()); + let mut other = generation.config.vector_stores.providers["local"].clone(); + other.api_key = Some("generation-secret".into()); + config.vector_stores.providers.insert("other".into(), other); + config + .vector_stores + .default_reranker_model + .as_mut() + .unwrap() + .provider_id = "other".into(); + config + .vector_stores + .rewrite_query_params + .as_mut() + .unwrap() + .model + .as_mut() + .unwrap() + .provider_id = "other".into(); + let service = FileSearchService::new(embedding.pool.clone(), Arc::new(reqwest::Client::new()), config).unwrap(); + let store = service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + attach(&service, &store.id, "coral preferred", None).await; + let request: SearchRequest = + serde_json::from_value(json!({"query":"ocean","rewrite_query":true,"ranking_options":{"ranker":"neural"}})) + .unwrap(); + assert_eq!(service.search(&[store.id], &request).await.unwrap().data.len(), 1); + assert_eq!( + *embedding.state.authorization.lock().unwrap(), + vec!["Bearer embedding-secret", "Bearer embedding-secret"] + ); + assert_eq!( + *generation.state.authorization.lock().unwrap(), + vec!["Bearer generation-secret", "Bearer generation-secret"] + ); +} + +#[tokio::test] +async fn repeated_context_placeholders_are_rejected_before_expansion_or_model_calls() { + let setup = setup(true).await; + let contextual = ContextualChunking { + context_prompt: format!("{}{{{{CHUNK_CONTENT}}}}", "{{WHOLE_DOCUMENT}}".repeat(900)), + ..ContextualChunking::default() + }; + // Validate before allocating the large document: the old implementation would amplify it 900-fold. + assert!(contextual.validate().is_err()); + let file = setup + .service + .upload_file( + "amplification.txt", + "text/plain", + "assistants", + "word ".repeat(79_000).into_bytes(), + ) + .await + .unwrap(); + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let error = setup + .service + .attach_file( + &store.id, + AttachFileRequest { + file_id: file.id, + chunking_strategy: Some(ChunkingStrategy::Contextual { contextual }), + ..AttachFileRequest::default() + }, + ) + .await + .unwrap_err(); + assert_eq!(error.status_code(), 400); + assert!(setup.state.requests.lock().unwrap().is_empty()); + assert_eq!( + setup + .service + .get_vector_store(&store.id) + .await + .unwrap() + .file_counts + .total, + 0 + ); +} + +#[test] +fn hyphenated_openai_ranker_is_accepted_in_deployment_configuration() { + let ranker: Ranker = serde_json::from_str("\"default-2024-08-21\"").unwrap(); + assert_eq!(ranker, Ranker::Default20240821); + assert_eq!(serde_json::to_string(&ranker).unwrap(), "\"default-2024-08-21\""); +} diff --git a/crates/agentic-server/src/config_file.rs b/crates/agentic-server/src/config_file.rs index 480813c3..e9c57049 100644 --- a/crates/agentic-server/src/config_file.rs +++ b/crates/agentic-server/src/config_file.rs @@ -24,6 +24,8 @@ impl FilesFileConfig { #[derive(Debug, Default, Deserialize, Serialize)] #[serde(default, deny_unknown_fields)] pub(crate) struct FileSearchFileConfig { + #[serde(skip_serializing_if = "Option::is_none")] + pub vector_stores: Option, #[serde(default)] pub backend: agentic_core::types::file_search::FileSearchBackend, #[serde(skip_serializing_if = "Option::is_none")] @@ -36,7 +38,8 @@ pub(crate) struct FileSearchFileConfig { impl FileSearchFileConfig { fn is_empty(&self) -> bool { - self.backend == agentic_core::types::file_search::FileSearchBackend::Exact + self.vector_stores.is_none() + && self.backend == agentic_core::types::file_search::FileSearchBackend::Exact && self.embedding_base_url.is_none() && self.embedding_model.is_none() && self.api_key_env.is_none() @@ -219,6 +222,11 @@ impl FileConfig { } fn validate(&self, path: &Path) -> Result<(), Error> { + if let Some(config) = &self.file_search.vector_stores { + config + .validate() + .map_err(|error| Error::Config(format!("invalid file_search.vector_stores: {error}")))?; + } if self .files .storage_dir @@ -375,6 +383,43 @@ ef_search = 100 ); } + #[test] + fn grouped_models_load_with_secret_names_and_reject_raw_secrets() { + let home = tempdir().unwrap(); + let configuration = r#" +[file_search.vector_stores] +default_provider_id="local" +[file_search.vector_stores.providers.local] +base_url="http://localhost:8001/v1" +models=["embed", "chat", "org/rerank"] +api_key_env="LOCAL_RETRIEVAL_KEY" +[file_search.vector_stores.default_reranker_model] +provider_id="local" +model_id="org/rerank" +[file_search.vector_stores.rewrite_query_params] +temperature=0.0 +"#; + fs::write(home.path().join("config.toml"), configuration).unwrap(); + let config = FileConfig::load(home.path()).unwrap().unwrap(); + let mut grouped = config.file_search.vector_stores.unwrap(); + grouped.providers.get_mut("local").unwrap().api_key = Some("never-persist-secret".into()); + assert!(!format!("{grouped:?}").contains("never-persist-secret")); + assert!(!toml::to_string(&grouped).unwrap().contains("never-persist-secret")); + assert!(grouped.rewrite_query_params.unwrap().temperature.abs() < f64::EPSILON); + fs::write( + home.path().join("config.toml"), + configuration.replace("api_key_env=", "api_key="), + ) + .unwrap(); + assert!(FileConfig::load(home.path()).is_err()); + fs::write( + home.path().join("config.toml"), + configuration.replace("org/rerank\"\n", "unknown\"\n"), + ) + .unwrap(); + assert!(FileConfig::load(home.path()).is_err()); + } + #[test] fn rejects_blank_embedding_api_key_environment_name() { let home = tempdir().expect("temp home"); diff --git a/crates/agentic-server/src/main.rs b/crates/agentic-server/src/main.rs index 3838d91d..c276046d 100644 --- a/crates/agentic-server/src/main.rs +++ b/crates/agentic-server/src/main.rs @@ -287,6 +287,13 @@ fn build_config(llm_api_base: String, common: &CommonArgs, file: &FileConfig) -> sqlite, tools: ToolRuntimeConfig { file_search: agentic_core::types::file_search::FileSearchConfig { + vector_stores: { + let mut config = file.file_search.vector_stores.clone().unwrap_or_default(); + for provider in config.providers.values_mut() { + provider.api_key = provider.api_key_env.as_deref().and_then(environment_value); + } + config + }, backend: file.file_search.backend.clone(), files_storage_dir: std::env::var_os("AGENTIC_FILES_STORAGE_DIR") .map(std::path::PathBuf::from) diff --git a/docs/api/file-search.md b/docs/api/file-search.md index c75fece2..0b413106 100644 --- a/docs/api/file-search.md +++ b/docs/api/file-search.md @@ -49,6 +49,158 @@ A vector store records its embedding endpoint, model, and vector dimensions. Changing that configuration requires a new store and reingestion for semantic search. Existing stores remain available for keyword search. +## Model-assisted retrieval configuration + +Use `file_search.vector_stores` to group independently configured model providers. +Requests select allowlisted models; they cannot supply endpoints or credentials. +Qualified selectors split once into provider ID and provider-local model name, so +`local/org/rerank` sends model `org/rerank` to provider `local`. Unqualified selectors +use `default_provider_id`. Unknown providers and models fail before model calls. + +```toml +[file_search.vector_stores] +default_provider_id = "local" + +[file_search.vector_stores.providers.local] +base_url = "http://localhost:8001/v1" +models = ["embedding-model", "context-model", "org/rerank"] +api_key_env = "LOCAL_RETRIEVAL_KEY" +protocol = "vllm" +score_interpretation = "probability" + +[file_search.vector_stores.default_embedding_model] +provider_id = "local" +model_id = "embedding-model" +embedding_dimensions = 768 + +[file_search.vector_stores.default_reranker_model] +provider_id = "local" +model_id = "org/rerank" + +[file_search.vector_stores.contextual_retrieval_params] +model = { provider_id = "local", model_id = "context-model" } +default_timeout_seconds = 120 +default_max_concurrency = 3 +max_document_tokens = 100000 + +[file_search.vector_stores.rewrite_query_params] +model = { provider_id = "local", model_id = "context-model" } +max_tokens = 100 +temperature = 0.0 + +[file_search.vector_stores.chunk_retrieval_params] +chunk_multiplier = 5 +max_tokens_in_context = 4000 +default_reranker_strategy = "rrf" +rrf_impact_factor = 60.0 +weighted_search_alpha = 0.5 +# Optional: default_search_mode = "vector" +``` + +Grouped embeddings and the legacy embedding connection are alternative forms; +configuring both is an error. Existing legacy environment overrides still apply. +Each provider's secret is resolved from its `api_key_env`, redacted in diagnostics, +and omitted from serialized configuration. Credentials are never inherited from +the Responses model. Different providers can serve each operation. The supported +`vllm` text protocol uses embeddings and Chat Completions below the base URL and +`/rerank` after removing a trailing `/v1`. + +Absent explicit defaults preserve keyword-only local setup and hybrid search when +embeddings exist. `file_ingestion_params.default_chunk_size_tokens` and +`default_chunk_overlap_tokens` configure auto ingestion; defaults preserve native +800/400 behavior. Static overlap remains limited to half the chunk size. +`file_batch_params` validates `max_concurrent_files_per_batch` (default 3, range +1–32), `file_batch_chunk_size` (10, 1–1000), and `cleanup_interval_seconds` (86400, +1–604800); asynchronous workers are a later layer. + +### Contextual ingestion + +```json +{ + "file_id": "file-...", + "chunking_strategy": { + "type": "contextual", + "contextual": { + "model_id": "local/context-model", + "max_chunk_size_tokens": 700, + "chunk_overlap_tokens": 400 + } + } +} +``` + +The nested `contextual` object is required; all its fields have defaults. Chunk +size is 100–4096; contextual overlap must be strictly less than chunk size. +`model_id` overrides the configured contextual model. `timeout_seconds` (1–600) +overrides the deployment timeout (10–600, default 120). `max_concurrency` (1–32) +can reduce per-ingestion concurrency, still capped by the deployment's shared +contextual semaphore (1–32, default 3). Document size uses a character-count/4 token +estimate, bounded by `max_document_tokens` (1000–1000000) and existing extraction +byte limits. Optional `context_prompt` must contain `{{WHOLE_DOCUMENT}}` exactly +once before `{{CHUNK_CONTENT}}`, also exactly once, and fit 16 KiB. + +The model receives the document as a shared system prefix and the chunk in a user +message, with temperature zero and at most 256 output tokens. Nonempty context is +prepended to the embedding input and stored separately as `embedding_text`. +Returned/cited source text stays original. Every context call must succeed before +embedding or publishing the attachment; partial failure, timeout, cancellation, +or malformed/empty output publishes nothing. Contextual ingestion requires +embeddings. Calls have no retries or silent fallback. Prompt expansion is +size-checked before allocation. JSON request bodies are bounded to 32 MiB while +being serialized, chat/rerank responses to 1 MiB, contextual output to 8 KiB, and +rewritten queries to 4096 bytes. Rewrite and rerank calls time out after 45 seconds. + +### Rewrite and ranking semantics + +Direct search and the Responses tool declaration accept `rewrite_query` and +`search_mode`. Rewriting joins input queries with spaces and makes one model call +before retrieval. `search_query` (tool `queries`) contains the single rewritten +query; without rewriting, established multi-query retrieval and deduplication +apply. Rewrite `temperature` includes explicit zero (range 0–2), `max_tokens` is +1–4096, and optional `prompt` must include `{query}` and fit 16 KiB. Missing +configuration, failed calls, and empty output fail the search. + +`ranking_options.ranker` accepts `auto`, `none`, `rrf`, `normalized`, `weighted`, +`neural`, and `classifier`. `auto` uses the configured strategy; `none` bypasses +model reranking while retaining the selected retrieval mode's base ranking. +`normalized` aliases normalized RRF. `weighted` min-max normalizes each present +score list and combines vector proportion `alpha` with keyword proportion +`1-alpha`. Equal nonempty scores normalize to one; absent scores normalize to zero. +RRF uses one-based ranks with configurable `impact_factor` (0–10000) and scales +scores to [0,1]. Explicit `weights` contains nonnegative `vector` and `keyword` +values summing to one. Existing `hybrid_search.embedding_weight`/`text_weight` +remain supported; do not combine these two weight forms. Explicit parameters +override deployment defaults. Fusion weights require hybrid mode. Neural score +blending is not supported. + +OpenAI selectors `default-2024-11-15` and `default-2024-08-21` select the configured +model ranker; `default_2024_08_21` remains accepted as a compatibility alias. These +selectors do not reproduce OpenAI's hosted models or scores. `ranking_options.model` +overrides `default_reranker_model` for both neural and classifier ranking. Missing +model configuration is an error. Both use vLLM/Cohere-style text reranking +(`documents`, `top_n`, `results`) before final truncation. The deployment +`chunk_multiplier` (1–20) expands the initial result count within existing +pgvector aggregate candidate limits. No additional corpus data is loaded to +refill results. Filters and store isolation apply before provider calls. + +Reranker indexes must form a complete unique permutation of bounded candidates. +Scores must be finite. `score_interpretation = "probability"` requires [0,1]; +`"logit"` applies a numerically stable sigmoid. Arbitrary scores are never clipped. +Model scores are implementation-specific: select the interpretation for the model +being deployed. `score_threshold` (0–1) applies to final model scores for both +neural and classifier ranking, before the final result limit. Zero initial +similarity does not exclude an otherwise selected model-reranking candidate. +Without model reranking, the threshold applies to base retrieval scores. +Provider errors return no partial search response. + +The Responses tool keeps whole source chunks within `max_tokens_in_context` +(1–32768, default 4000), using bounded `cl100k_base` token counting. Chunks that do +not fit are omitted. Direct search is independent of that context budget. +Search and context preparation share the service admission slot; queued or running +tokenization retains that slot until it exits, including after caller cancellation. +Cancellation is checked between bounded tokenization blocks. Citation item types +and file IDs remain unchanged. + ## Select PostgreSQL indexed retrieval The default `exact` backend works with SQLite and plain PostgreSQL. To use @@ -249,11 +401,9 @@ unreferenced files. Automatic orphan-file cleanup is not included. A missing or damaged file referenced by metadata returns a storage error instead of a partial download. Uploads stored inline by an earlier draft remain readable. -The routes use the gateway's configured authentication policy. This first -implementation uses bounded exact retrieval over SQL storage. It is intended for -small corpora and returns explicit capacity errors when limits are exceeded. +The routes use the gateway's configured authentication policy. Retrieval uses bounded exact SQL or the configured pgvector backend and returns +explicit capacity errors when limits are exceeded. Each store is limited to 10,000 chunks and 64 MiB of serialized chunk data, including embeddings; ingestion enforces these limits before publication. Searching multiple stores shares the same aggregate retrieval budget. -It does not implement OGX's provider catalog, contextual chunking, query rewriting, -neural reranking, or asynchronous file batches. +It does not expose OGX's provider catalog or asynchronous file batches.