From d5ea8c71767ef90f4db9ffa3085b9c56afa9e3aa Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Tue, 2 Dec 2025 19:41:40 -0800 Subject: [PATCH] [model-gateway] multimodality initialization (#13350) --- sgl-router/Cargo.toml | 9 +- sgl-router/src/lib.rs | 1 + sgl-router/src/multimodal/error.rs | 47 ++ sgl-router/src/multimodal/media.rs | 206 +++++++++ sgl-router/src/multimodal/mod.rs | 11 + sgl-router/src/multimodal/registry.rs | 461 ++++++++++++++++++++ sgl-router/src/multimodal/tracker.rs | 176 ++++++++ sgl-router/src/multimodal/types.rs | 290 ++++++++++++ sgl-router/tests/multimodal_tracker_test.rs | 151 +++++++ 9 files changed, 1351 insertions(+), 1 deletion(-) create mode 100644 sgl-router/src/multimodal/error.rs create mode 100644 sgl-router/src/multimodal/media.rs create mode 100644 sgl-router/src/multimodal/mod.rs create mode 100644 sgl-router/src/multimodal/registry.rs create mode 100644 sgl-router/src/multimodal/tracker.rs create mode 100644 sgl-router/src/multimodal/types.rs create mode 100644 sgl-router/tests/multimodal_tracker_test.rs diff --git a/sgl-router/Cargo.toml b/sgl-router/Cargo.toml index 45db679e2..4c786c21a 100644 --- a/sgl-router/Cargo.toml +++ b/sgl-router/Cargo.toml @@ -4,7 +4,10 @@ version = "0.2.3" edition = "2021" [features] -default = [] +default = ["grpc-client"] +grpc-client = [] +grpc-server = [] + vendored-openssl = ["openssl/vendored"] [lints.rust] @@ -36,6 +39,7 @@ serde_json = { version = "1.0", default-features = false, features = [ "std", "preserve_order", ] } +serde_bytes = "0.11" bytes = "1.8.0" rand = "0.9.2" reqwest = { version = "0.12.8", features = ["stream", "blocking", "json", "rustls-tls"], default-features = false } @@ -86,6 +90,9 @@ oracle = { version = "0.6.3", features = ["chrono"] } subtle = "2.6" rustpython-parser = "0.4.0" num-traits = "0.2" +image = { version = "0.25.4", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] } +ndarray = "0.16" +base64 = "0.22" openai-harmony = { git = "https://github.com/openai/harmony", tag = "v0.0.4" } openmetrics-parser = "0.4.4" diff --git a/sgl-router/src/lib.rs b/sgl-router/src/lib.rs index c22064586..1b0fe03b9 100644 --- a/sgl-router/src/lib.rs +++ b/sgl-router/src/lib.rs @@ -8,6 +8,7 @@ pub mod grpc_client; pub mod mcp; pub mod metrics; pub mod middleware; +pub mod multimodal; pub mod policies; pub mod protocols; pub mod reasoning_parser; diff --git a/sgl-router/src/multimodal/error.rs b/sgl-router/src/multimodal/error.rs new file mode 100644 index 000000000..0d73d3c3d --- /dev/null +++ b/sgl-router/src/multimodal/error.rs @@ -0,0 +1,47 @@ +use std::time::Duration; + +use thiserror::Error; + +use super::types::Modality; + +pub type MultiModalResult = Result; + +#[derive(Debug, Error)] +pub enum MediaConnectorError { + #[error("unsupported media scheme: {0}")] + UnsupportedScheme(String), + #[error("invalid media URL: {0}")] + InvalidUrl(String), + #[error("media domain '{0}' is not in the allow list")] + DisallowedDomain(String), + #[error("local media path is not allowed: {0}")] + DisallowedLocalPath(String), + #[error("HTTP error while fetching media: {0}")] + Http(#[from] reqwest::Error), + #[error("I/O error while reading media: {0}")] + Io(#[from] std::io::Error), + #[error("base64 decode error: {0}")] + Base64Decode(#[from] base64::DecodeError), + #[error("data URL parse error: {0}")] + DataUrl(String), + #[error("media decode task failed: {0}")] + Blocking(#[from] tokio::task::JoinError), + #[error("image decode error: {0}")] + Image(#[from] image::ImageError), + #[error("media fetch timed out after {0:?}")] + Timeout(Duration), +} + +#[derive(Debug, Error)] +pub enum MultiModalError { + #[error(transparent)] + Media(#[from] MediaConnectorError), + #[error("unsupported content part: {0}")] + UnsupportedContent(&'static str), + #[error("too many {modality:?} items provided. limit={limit}")] + ModalityLimit { modality: Modality, limit: usize }, + #[error("tracker task join error: {0}")] + Join(#[from] tokio::task::JoinError), + #[error("tracker validation error: {0}")] + Validation(String), +} diff --git a/sgl-router/src/multimodal/media.rs b/sgl-router/src/multimodal/media.rs new file mode 100644 index 000000000..d44633128 --- /dev/null +++ b/sgl-router/src/multimodal/media.rs @@ -0,0 +1,206 @@ +use std::{collections::HashSet, path::PathBuf, sync::Arc, time::Duration}; + +use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; +use bytes::Bytes; +use image::DynamicImage; +use reqwest::Client; +use tokio::{fs, task}; +use url::Url; + +use super::{ + error::MediaConnectorError, + types::{ImageDetail, ImageFrame, ImageSource}, +}; + +#[derive(Clone)] +pub struct MediaConnectorConfig { + pub allowed_domains: Option>, + pub allowed_local_media_path: Option, + pub fetch_timeout: Duration, +} + +impl Default for MediaConnectorConfig { + fn default() -> Self { + Self { + allowed_domains: None, + allowed_local_media_path: None, + fetch_timeout: Duration::from_secs(10), + } + } +} + +#[derive(Clone, Copy, Debug)] +pub struct ImageFetchConfig { + pub detail: ImageDetail, +} + +impl Default for ImageFetchConfig { + fn default() -> Self { + Self { + detail: ImageDetail::Auto, + } + } +} + +#[derive(Debug, Clone)] +pub enum MediaSource { + Url(String), + DataUrl(String), + InlineBytes(Vec), + File(PathBuf), +} + +#[derive(Clone)] +pub struct MediaConnector { + client: Client, + allowed_domains: Option>, + allowed_local_media_path: Option, + fetch_timeout: Duration, +} + +impl MediaConnector { + pub fn new(client: Client, config: MediaConnectorConfig) -> Result { + let allowed_domains = config.allowed_domains.map(|domains| { + domains + .into_iter() + .map(|d| d.to_ascii_lowercase()) + .collect::>() + }); + + let allowed_local_media_path = if let Some(path) = config.allowed_local_media_path { + Some(std::fs::canonicalize(path)?) + } else { + None + }; + + Ok(Self { + client, + allowed_domains, + allowed_local_media_path, + fetch_timeout: config.fetch_timeout, + }) + } + + pub async fn fetch_image( + &self, + source: MediaSource, + cfg: ImageFetchConfig, + ) -> Result, MediaConnectorError> { + match source { + MediaSource::Url(url) => self.fetch_http_image(url, cfg).await, + MediaSource::DataUrl(data_url) => self.fetch_data_url(data_url, cfg).await, + MediaSource::InlineBytes(bytes) => { + self.decode_image(bytes.into(), cfg.detail, ImageSource::InlineBytes) + .await + } + MediaSource::File(path) => self.fetch_file(path, cfg).await, + } + } + + async fn fetch_http_image( + &self, + url: String, + cfg: ImageFetchConfig, + ) -> Result, MediaConnectorError> { + let parsed = Url::parse(&url).map_err(|_| MediaConnectorError::InvalidUrl(url.clone()))?; + self.ensure_domain_allowed(&parsed)?; + + let mut req = self.client.get(parsed.as_str()); + if self.fetch_timeout > Duration::ZERO { + req = req.timeout(self.fetch_timeout); + } + + let resp = req.send().await.map_err(|err| { + if err.is_timeout() { + MediaConnectorError::Timeout(self.fetch_timeout) + } else { + MediaConnectorError::Http(err) + } + })?; + + let resp = resp.error_for_status()?; + let bytes = resp.bytes().await?; + self.decode_image( + bytes, + cfg.detail, + ImageSource::Url { + url: parsed.to_string(), + }, + ) + .await + } + + async fn fetch_data_url( + &self, + data_url: String, + cfg: ImageFetchConfig, + ) -> Result, MediaConnectorError> { + let (metadata, data) = data_url + .split_once(',') + .ok_or_else(|| MediaConnectorError::DataUrl("missing comma in data url".into()))?; + + if !metadata.ends_with(";base64") { + return Err(MediaConnectorError::DataUrl( + "only base64 encoded data URLs are supported".into(), + )); + } + + let data = data.trim(); + let decoded = BASE64_STANDARD.decode(data)?; + self.decode_image(decoded.into(), cfg.detail, ImageSource::DataUrl) + .await + } + + async fn fetch_file( + &self, + path: PathBuf, + cfg: ImageFetchConfig, + ) -> Result, MediaConnectorError> { + let allowed_root = self + .allowed_local_media_path + .as_ref() + .ok_or_else(|| MediaConnectorError::DisallowedLocalPath(path.display().to_string()))?; + + let canonical = fs::canonicalize(&path).await?; + if !canonical.starts_with(allowed_root) { + return Err(MediaConnectorError::DisallowedLocalPath( + path.display().to_string(), + )); + } + + let bytes = fs::read(&canonical).await?; + self.decode_image( + bytes.into(), + cfg.detail, + ImageSource::File { path: canonical }, + ) + .await + } + + fn ensure_domain_allowed(&self, url: &Url) -> Result<(), MediaConnectorError> { + if let Some(allowed) = &self.allowed_domains { + let host = url + .host_str() + .map(|h| h.to_ascii_lowercase()) + .ok_or_else(|| MediaConnectorError::InvalidUrl(url.to_string()))?; + if !allowed.contains(&host) { + return Err(MediaConnectorError::DisallowedDomain(host)); + } + } + Ok(()) + } + + async fn decode_image( + &self, + bytes: Bytes, + detail: ImageDetail, + source: ImageSource, + ) -> Result, MediaConnectorError> { + let raw: Arc> = Arc::new(bytes.to_vec()); + let raw_clone = raw.clone(); + let image: DynamicImage = + task::spawn_blocking(move || image::load_from_memory(&raw_clone)).await??; + + Ok(Arc::new(ImageFrame::new(image, raw, detail, source))) + } +} diff --git a/sgl-router/src/multimodal/mod.rs b/sgl-router/src/multimodal/mod.rs new file mode 100644 index 000000000..76ee5d1c6 --- /dev/null +++ b/sgl-router/src/multimodal/mod.rs @@ -0,0 +1,11 @@ +pub mod error; +pub mod media; +pub mod registry; +pub mod tracker; +pub mod types; + +pub use error::{MediaConnectorError, MultiModalError, MultiModalResult}; +pub use media::{ImageFetchConfig, MediaConnector, MediaConnectorConfig, MediaSource}; +pub use registry::{ModelProcessorSpec, ModelRegistry}; +pub use tracker::{AsyncMultiModalTracker, TrackerConfig, TrackerOutput}; +pub use types::*; diff --git a/sgl-router/src/multimodal/registry.rs b/sgl-router/src/multimodal/registry.rs new file mode 100644 index 000000000..eb5c30d33 --- /dev/null +++ b/sgl-router/src/multimodal/registry.rs @@ -0,0 +1,461 @@ +use std::collections::HashMap; + +use once_cell::sync::Lazy; +use serde_json::{json, Value}; +use thiserror::Error; + +use super::types::{ImageSize, Modality, PromptReplacement, TokenId}; +use crate::tokenizer::traits::Tokenizer as TokenizerTrait; + +#[derive(Debug, Error)] +pub enum ModelRegistryError { + #[error("unsupported model: {0}")] + UnsupportedModel(String), + #[error("token '{token}' not found in tokenizer vocabulary")] + TokenNotFound { token: String }, + #[error("missing config field '{field}'")] + MissingConfigField { field: String }, +} + +pub type RegistryResult = Result; + +/// Metadata about the current model used to derive tokenizer/config dependent fields. +pub struct ModelMetadata<'a> { + pub model_id: &'a str, + pub tokenizer: &'a dyn TokenizerTrait, + pub config: &'a Value, +} + +impl<'a> ModelMetadata<'a> { + pub fn token_id(&self, token: &str) -> RegistryResult { + self.tokenizer + .token_to_id(token) + .map(|id| id as TokenId) + .ok_or_else(|| ModelRegistryError::TokenNotFound { + token: token.to_string(), + }) + } + + pub fn config_u32(&self, path: &[&str]) -> Option { + Self::find_value(self.config, path).and_then(|value| value.as_u64().map(|v| v as u32)) + } + + fn find_value<'v>(value: &'v Value, path: &[&str]) -> Option<&'v Value> { + let mut current = value; + for key in path { + current = current.get(*key)?; + } + Some(current) + } +} + +pub trait ModelProcessorSpec: Send + Sync { + fn name(&self) -> &'static str; + fn matches(&self, metadata: &ModelMetadata) -> bool; + fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult; + fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult; + fn modality_limits(&self, metadata: &ModelMetadata) + -> RegistryResult>; + fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult; + fn prompt_replacements( + &self, + metadata: &ModelMetadata, + image_sizes: &[ImageSize], + ) -> RegistryResult>; +} + +pub struct ModelRegistry { + specs: Vec, +} + +impl ModelRegistry { + pub fn new() -> Self { + Self { + specs: vec![ + LazySpec::new("llava", || Box::new(LlavaSpec)), + LazySpec::new("qwen_vl", || Box::new(QwenVLVisionSpec)), + LazySpec::new("phi3_v", || Box::new(Phi3VisionSpec)), + ], + } + } + + pub fn lookup<'a>(&'a self, metadata: &ModelMetadata) -> Option<&'a dyn ModelProcessorSpec> { + for spec in &self.specs { + let spec_ref = spec.get(); + if spec_ref.matches(metadata) { + return Some(spec_ref); + } + } + None + } +} + +impl Default for ModelRegistry { + fn default() -> Self { + Self::new() + } +} + +struct LazySpec { + inner: Lazy>, +} + +impl LazySpec { + fn new(_id: &'static str, factory: fn() -> Box) -> Self { + Self { + inner: Lazy::new(factory), + } + } + + fn get(&self) -> &dyn ModelProcessorSpec { + self.inner.as_ref() + } +} + +struct LlavaSpec; + +impl LlavaSpec { + fn patch_size(metadata: &ModelMetadata) -> u32 { + metadata + .config_u32(&["vision_config", "patch_size"]) + .unwrap_or(14) + } + + fn tokens_per_image(metadata: &ModelMetadata, size: &ImageSize) -> usize { + let patch = Self::patch_size(metadata); + let cols = size.width.div_ceil(patch) as usize; + let rows = size.height.div_ceil(patch) as usize; + cols * rows + } +} + +impl ModelProcessorSpec for LlavaSpec { + fn name(&self) -> &'static str { + "llava" + } + + fn matches(&self, metadata: &ModelMetadata) -> bool { + metadata.model_id.to_ascii_lowercase().contains("llava") + } + + fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok("".to_string()) + } + + fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { + if let Some(value) = metadata.config_u32(&["image_token_index"]) { + return Ok(value as TokenId); + } + metadata.token_id("") + } + + fn modality_limits( + &self, + _metadata: &ModelMetadata, + ) -> RegistryResult> { + Ok(HashMap::from([(Modality::Image, 4)])) + } + + fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(json!({})) + } + + fn prompt_replacements( + &self, + metadata: &ModelMetadata, + image_sizes: &[ImageSize], + ) -> RegistryResult> { + let token_id = self.placeholder_token_id(metadata)?; + let token = self.placeholder_token(metadata)?; + Ok(image_sizes + .iter() + .map(|size| { + let count = Self::tokens_per_image(metadata, size); + PromptReplacement::repeated(Modality::Image, &token, token_id, count) + }) + .collect()) + } +} + +struct QwenVLVisionSpec; + +impl QwenVLVisionSpec { + fn pad_token_id(metadata: &ModelMetadata) -> RegistryResult { + metadata + .config_u32(&["vision_token_id"]) + .map(|v| v as TokenId) + .ok_or_else(|| ModelRegistryError::MissingConfigField { + field: "vision_token_id".to_string(), + }) + } + + fn start_token_id(metadata: &ModelMetadata) -> RegistryResult { + metadata + .config_u32(&["vision_start_token_id"]) + .map(|v| v as TokenId) + .ok_or_else(|| ModelRegistryError::MissingConfigField { + field: "vision_start_token_id".to_string(), + }) + } + + fn patch_grid(metadata: &ModelMetadata, size: &ImageSize) -> (usize, usize) { + let patch = metadata + .config_u32(&["vision_config", "patch_size"]) + .unwrap_or(14); + let cols = size.width.div_ceil(patch) as usize; + let rows = size.height.div_ceil(patch) as usize; + (rows, cols) + } +} + +impl ModelProcessorSpec for QwenVLVisionSpec { + fn name(&self) -> &'static str { + "qwen_vl" + } + + fn matches(&self, metadata: &ModelMetadata) -> bool { + metadata.model_id.to_ascii_lowercase().contains("qwen") + && metadata.model_id.to_ascii_lowercase().contains("vl") + } + + fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok("".to_string()) + } + + fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { + metadata + .config_u32(&["image_token_id"]) + .map(|v| v as TokenId) + .ok_or_else(|| ModelRegistryError::MissingConfigField { + field: "image_token_id".to_string(), + }) + } + + fn modality_limits( + &self, + _metadata: &ModelMetadata, + ) -> RegistryResult> { + Ok(HashMap::from([(Modality::Image, 10)])) + } + + fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(json!({})) + } + + fn prompt_replacements( + &self, + metadata: &ModelMetadata, + image_sizes: &[ImageSize], + ) -> RegistryResult> { + let start_token_id = Self::start_token_id(metadata)?; + let pad_token_id = Self::pad_token_id(metadata)?; + let placeholder_token = self.placeholder_token(metadata)?; + Ok(image_sizes + .iter() + .map(|size| { + let (rows, cols) = Self::patch_grid(metadata, size); + let pad_len = rows * cols; + let mut tokens = Vec::with_capacity(pad_len + 1); + tokens.push(start_token_id); + tokens.extend(std::iter::repeat_n(pad_token_id, pad_len)); + PromptReplacement::sequence(Modality::Image, &placeholder_token, tokens) + }) + .collect()) + } +} + +struct Phi3VisionSpec; + +impl Phi3VisionSpec { + fn tokens_per_image(metadata: &ModelMetadata) -> usize { + metadata + .config_u32(&["img_processor", "num_img_tokens"]) + .unwrap_or(256) as usize + } +} + +impl ModelProcessorSpec for Phi3VisionSpec { + fn name(&self) -> &'static str { + "phi3_v" + } + + fn matches(&self, metadata: &ModelMetadata) -> bool { + let id = metadata.model_id.to_ascii_lowercase(); + id.contains("phi") && id.contains("vision") + } + + fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok("".to_owned()) + } + + fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { + metadata.token_id("") + } + + fn modality_limits( + &self, + _metadata: &ModelMetadata, + ) -> RegistryResult> { + Ok(HashMap::from([(Modality::Image, 4)])) + } + + fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(json!({})) + } + + fn prompt_replacements( + &self, + metadata: &ModelMetadata, + image_sizes: &[ImageSize], + ) -> RegistryResult> { + let token_id = self.placeholder_token_id(metadata)?; + let token = self.placeholder_token(metadata)?; + let count = Self::tokens_per_image(metadata); + Ok(image_sizes + .iter() + .map(|_| PromptReplacement::repeated(Modality::Image, &token, token_id, count)) + .collect()) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use serde_json::json; + + use super::*; + use crate::tokenizer::traits::{ + Decoder, Encoder, Encoding, SpecialTokens, Tokenizer as TokenizerTrait, + }; + + struct TestTokenizer { + vocab: HashMap, + } + + impl TestTokenizer { + fn new(pairs: &[(&str, u32)]) -> Self { + let vocab = pairs + .iter() + .map(|(token, id)| ((*token).to_string(), *id)) + .collect(); + Self { vocab } + } + } + + impl Encoder for TestTokenizer { + fn encode(&self, _input: &str) -> anyhow::Result { + Ok(Encoding::Sp(Vec::new())) + } + + fn encode_batch(&self, inputs: &[&str]) -> anyhow::Result> { + inputs.iter().map(|_| self.encode("")).collect() + } + } + + impl Decoder for TestTokenizer { + fn decode(&self, _token_ids: &[u32], _skip_special_tokens: bool) -> anyhow::Result { + Ok(String::new()) + } + } + + impl TokenizerTrait for TestTokenizer { + fn vocab_size(&self) -> usize { + self.vocab.len() + } + + fn get_special_tokens(&self) -> &SpecialTokens { + static TOKENS: Lazy = Lazy::new(|| SpecialTokens { + bos_token: None, + eos_token: None, + unk_token: None, + sep_token: None, + pad_token: None, + cls_token: None, + mask_token: None, + additional_special_tokens: vec![], + }); + &TOKENS + } + + fn token_to_id(&self, token: &str) -> Option { + self.vocab.get(token).copied() + } + + fn id_to_token(&self, _id: u32) -> Option { + None + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + #[test] + fn llava_prompt_replacement_uses_config_ids() { + let tokenizer = TestTokenizer::new(&[("", 32000)]); + let config = json!({ + "model_type": "llava", + "image_token_index": 32000, + "vision_config": {"patch_size": 14} + }); + let metadata = ModelMetadata { + model_id: "llava-v1.5", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).expect("llava spec"); + let replacements = spec + .prompt_replacements(&metadata, &[ImageSize::new(336, 336)]) + .unwrap(); + assert_eq!(replacements[0].tokens.len(), 576); + } + + #[test] + fn qwen_vision_uses_config_token_ids() { + let tokenizer = TestTokenizer::new(&[("", 999)]); + let config = json!({ + "model_type": "qwen2_vl", + "vision_start_token_id": 151652, + "vision_token_id": 151654, + "image_token_id": 151655, + "vision_config": {"patch_size": 14} + }); + let metadata = ModelMetadata { + model_id: "Qwen2-VL-7B", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).expect("qwen spec"); + let replacements = spec + .prompt_replacements(&metadata, &[ImageSize::new(448, 448)]) + .unwrap(); + // 448/14 = 32 patches => 1024 pad tokens + 1 start token + assert_eq!(replacements[0].tokens.len(), 1025); + assert_eq!(replacements[0].tokens[0], 151652); + assert_eq!(replacements[0].tokens[1], 151654); + } + + #[test] + fn phi3_uses_num_img_tokens() { + let tokenizer = TestTokenizer::new(&[("", 555)]); + let config = json!({ + "model_type": "phi3_v", + "img_processor": {"num_img_tokens": 144} + }); + let metadata = ModelMetadata { + model_id: "Phi-3-vision", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).expect("phi3 spec"); + let replacements = spec + .prompt_replacements(&metadata, &[ImageSize::new(336, 336)]) + .unwrap(); + assert_eq!(replacements[0].tokens.len(), 144); + assert_eq!(replacements[0].tokens[0], 555); + } +} diff --git a/sgl-router/src/multimodal/tracker.rs b/sgl-router/src/multimodal/tracker.rs new file mode 100644 index 000000000..25a5bc668 --- /dev/null +++ b/sgl-router/src/multimodal/tracker.rs @@ -0,0 +1,176 @@ +use std::{collections::HashMap, sync::Arc}; + +use tokio::task::JoinHandle; + +use super::{ + error::{MultiModalError, MultiModalResult}, + media::{ImageFetchConfig, MediaConnector, MediaSource}, + types::{ + ChatContentPart, ConversationSegment, ImageDetail, Modality, MultiModalData, + MultiModalUUIDs, PlaceholderHandle, PlaceholderMap, TrackedMedia, DEFAULT_PLACEHOLDERS, + }, +}; + +type PendingTask = JoinHandle>; + +#[derive(Debug, Clone)] +pub struct TrackerConfig { + pub placeholder_tokens: HashMap, + pub modality_limits: HashMap, +} + +impl Default for TrackerConfig { + fn default() -> Self { + Self { + placeholder_tokens: DEFAULT_PLACEHOLDERS + .iter() + .map(|(k, v)| (*k, (*v).to_string())) + .collect(), + modality_limits: HashMap::new(), + } + } +} + +#[derive(Debug)] +pub struct TrackerOutput { + pub conversation: Vec, + pub data: MultiModalData, + pub uuids: MultiModalUUIDs, + pub placeholders: PlaceholderMap, +} + +pub struct AsyncMultiModalTracker { + media_connector: Arc, + config: TrackerConfig, + pending: HashMap>, + placeholders: PlaceholderMap, + conversation: Vec, + uuids: MultiModalUUIDs, + counts: HashMap, +} + +impl AsyncMultiModalTracker { + pub fn new(media_connector: Arc, config: TrackerConfig) -> Self { + Self { + media_connector, + config, + pending: HashMap::new(), + placeholders: PlaceholderMap::new(), + conversation: Vec::new(), + uuids: HashMap::new(), + counts: HashMap::new(), + } + } + + pub fn push_part(&mut self, part: ChatContentPart) -> MultiModalResult<()> { + match part { + ChatContentPart::Text { text } => { + if !text.is_empty() { + self.conversation.push(ConversationSegment::text(text)); + } + Ok(()) + } + ChatContentPart::ImageUrl { url, detail, uuid } => { + self.enqueue_image(MediaSource::Url(url), detail.unwrap_or_default(), uuid) + } + ChatContentPart::ImageData { + data, + mime_type: _, + uuid, + detail, + } => self.enqueue_image( + MediaSource::InlineBytes(data), + detail.unwrap_or_default(), + uuid, + ), + ChatContentPart::ImageEmbeds { .. } => { + Err(MultiModalError::UnsupportedContent("image_embeds")) + } + } + } + + pub async fn finalize(mut self) -> MultiModalResult { + let mut data = MultiModalData::new(); + for (modality, tasks) in self.pending.drain() { + let mut items = Vec::with_capacity(tasks.len()); + for task in tasks { + let media = task.await??; + items.push(media); + } + data.insert(modality, items); + } + + Ok(TrackerOutput { + conversation: self.conversation, + data, + uuids: self.uuids, + placeholders: self.placeholders, + }) + } + + fn placeholder_token(&self, modality: Modality) -> &str { + if let Some(token) = self.config.placeholder_tokens.get(&modality) { + token.as_str() + } else { + DEFAULT_PLACEHOLDERS + .get(&modality) + .copied() + .unwrap_or("") + } + } + + fn next_index(&mut self, modality: Modality) -> MultiModalResult { + let count = self.counts.entry(modality).or_insert(0); + let next = *count; + let limit = self.config.modality_limits.get(&modality).copied(); + if let Some(limit) = limit { + if next >= limit { + return Err(MultiModalError::ModalityLimit { modality, limit }); + } + } + *count += 1; + Ok(next) + } + + fn record_uuid(&mut self, modality: Modality, uuid: Option) { + self.uuids.entry(modality).or_default().push(uuid); + } + + fn add_placeholder(&mut self, token: String, handle: PlaceholderHandle) { + self.placeholders.entry(token).or_default().push(handle); + } + + fn enqueue_image( + &mut self, + source: MediaSource, + detail: ImageDetail, + uuid: Option, + ) -> MultiModalResult<()> { + let modality = Modality::Image; + let idx = self.next_index(modality)?; + let token = self.placeholder_token(modality).to_string(); + let text_position = self.conversation.len(); + self.conversation + .push(ConversationSegment::placeholder(token.clone())); + self.add_placeholder( + token.clone(), + PlaceholderHandle { + modality, + item_index: idx, + text_position, + }, + ); + self.record_uuid(modality, uuid); + + let connector = Arc::clone(&self.media_connector); + let handle = tokio::spawn(async move { + let frame = connector + .fetch_image(source, ImageFetchConfig { detail }) + .await?; + Ok(TrackedMedia::Image(frame)) + }); + + self.pending.entry(modality).or_default().push(handle); + Ok(()) + } +} diff --git a/sgl-router/src/multimodal/types.rs b/sgl-router/src/multimodal/types.rs new file mode 100644 index 000000000..fe3bb1dc9 --- /dev/null +++ b/sgl-router/src/multimodal/types.rs @@ -0,0 +1,290 @@ +use std::{ + collections::{BTreeMap, HashMap}, + fmt, + path::PathBuf, + sync::{Arc, LazyLock}, +}; + +use image::DynamicImage; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// Supported multimodal modalities. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Modality { + Image, + ImageEmbeds, + Audio, + Video, +} + +impl fmt::Display for Modality { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Modality::Image => write!(f, "image"), + Modality::ImageEmbeds => write!(f, "image_embeds"), + Modality::Audio => write!(f, "audio"), + Modality::Video => write!(f, "video"), + } + } +} + +/// Simple helper used for default placeholder tokens per modality. +pub static DEFAULT_PLACEHOLDERS: LazyLock> = LazyLock::new(|| { + HashMap::from([ + (Modality::Image, ""), + (Modality::ImageEmbeds, ""), + (Modality::Audio, "