feat: add contextual retrieval, reranking and vector store configuration

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<String>,
     pub embedding: Option<Vec<f64>>,
     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,

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,31 +19,32 @@ pub(super) struct Embeddings {
     endpoint: reqwest::Url,
     model: String,
     api_key: Option<String>,
-}
-
-#[derive(Serialize)]
-struct EmbeddingRequest<'a> {
-    model: &'a str,
-    input: &'a [String],
-    encoding_format: &'static str,
-}
-
-#[derive(Deserialize)]
-struct EmbeddingResponse {
-    model: String,
-    data: Vec<EmbeddingData>,
-}
-#[derive(Deserialize)]
-struct EmbeddingData {
-    index: usize,
-    embedding: Vec<f64>,
+    dimensions: Option<usize>,
 }
 
 impl Embeddings {
     pub(super) fn from_config(
         client: Arc<reqwest::Client>,
         config: &FileSearchConfig,
     ) -> Result<Option<Self>, 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<usize>,
     ) -> Result<Vec<Vec<f64>>, 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);

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<String>) -> 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,
     }
 }
 

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<StaticChunking, FileSearchError> {
     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<String>,
+}
+
 pub(super) fn extract_and_chunk(
     bytes: Vec<u8>,
     filename: &str,
     content_type: &str,
     chunking: &StaticChunking,
     cancelled: &AtomicBool,
-) -> Result<Vec<String>, FileSearchError> {
+) -> Result<ExtractedDocument, FileSearchError> {
     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<crate::types::file_search::SearchResult>,
+    mut budget: usize,
+    cancelled: &AtomicBool,
+) -> Result<Vec<crate::types::file_search::SearchResult>, 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}')));

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;