[model-gateway] multimodality initialization (#13350)

This commit is contained in:
Simo Lin
2025-12-02 19:41:40 -08:00
committed by GitHub
parent 42271376d1
commit d5ea8c7176
9 changed files with 1351 additions and 1 deletions

View File

@@ -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"

View File

@@ -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;

View File

@@ -0,0 +1,47 @@
use std::time::Duration;
use thiserror::Error;
use super::types::Modality;
pub type MultiModalResult<T> = Result<T, MultiModalError>;
#[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),
}

View File

@@ -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<Vec<String>>,
pub allowed_local_media_path: Option<PathBuf>,
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<u8>),
File(PathBuf),
}
#[derive(Clone)]
pub struct MediaConnector {
client: Client,
allowed_domains: Option<HashSet<String>>,
allowed_local_media_path: Option<PathBuf>,
fetch_timeout: Duration,
}
impl MediaConnector {
pub fn new(client: Client, config: MediaConnectorConfig) -> Result<Self, MediaConnectorError> {
let allowed_domains = config.allowed_domains.map(|domains| {
domains
.into_iter()
.map(|d| d.to_ascii_lowercase())
.collect::<HashSet<_>>()
});
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<Arc<ImageFrame>, 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<Arc<ImageFrame>, 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<Arc<ImageFrame>, 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<Arc<ImageFrame>, 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<Arc<ImageFrame>, MediaConnectorError> {
let raw: Arc<Vec<u8>> = 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)))
}
}

View File

@@ -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::*;

View File

@@ -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<T> = Result<T, ModelRegistryError>;
/// 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<TokenId> {
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<u32> {
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<String>;
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId>;
fn modality_limits(&self, metadata: &ModelMetadata)
-> RegistryResult<HashMap<Modality, usize>>;
fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult<Value>;
fn prompt_replacements(
&self,
metadata: &ModelMetadata,
image_sizes: &[ImageSize],
) -> RegistryResult<Vec<PromptReplacement>>;
}
pub struct ModelRegistry {
specs: Vec<LazySpec>,
}
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<Box<dyn ModelProcessorSpec>>,
}
impl LazySpec {
fn new(_id: &'static str, factory: fn() -> Box<dyn ModelProcessorSpec>) -> 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<String> {
Ok("<image>".to_string())
}
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
if let Some(value) = metadata.config_u32(&["image_token_index"]) {
return Ok(value as TokenId);
}
metadata.token_id("<image>")
}
fn modality_limits(
&self,
_metadata: &ModelMetadata,
) -> RegistryResult<HashMap<Modality, usize>> {
Ok(HashMap::from([(Modality::Image, 4)]))
}
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
Ok(json!({}))
}
fn prompt_replacements(
&self,
metadata: &ModelMetadata,
image_sizes: &[ImageSize],
) -> RegistryResult<Vec<PromptReplacement>> {
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<TokenId> {
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<TokenId> {
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<String> {
Ok("<image>".to_string())
}
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
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<HashMap<Modality, usize>> {
Ok(HashMap::from([(Modality::Image, 10)]))
}
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
Ok(json!({}))
}
fn prompt_replacements(
&self,
metadata: &ModelMetadata,
image_sizes: &[ImageSize],
) -> RegistryResult<Vec<PromptReplacement>> {
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<String> {
Ok("<image>".to_owned())
}
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
metadata.token_id("<image>")
}
fn modality_limits(
&self,
_metadata: &ModelMetadata,
) -> RegistryResult<HashMap<Modality, usize>> {
Ok(HashMap::from([(Modality::Image, 4)]))
}
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
Ok(json!({}))
}
fn prompt_replacements(
&self,
metadata: &ModelMetadata,
image_sizes: &[ImageSize],
) -> RegistryResult<Vec<PromptReplacement>> {
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<String, u32>,
}
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<Encoding> {
Ok(Encoding::Sp(Vec::new()))
}
fn encode_batch(&self, inputs: &[&str]) -> anyhow::Result<Vec<Encoding>> {
inputs.iter().map(|_| self.encode("")).collect()
}
}
impl Decoder for TestTokenizer {
fn decode(&self, _token_ids: &[u32], _skip_special_tokens: bool) -> anyhow::Result<String> {
Ok(String::new())
}
}
impl TokenizerTrait for TestTokenizer {
fn vocab_size(&self) -> usize {
self.vocab.len()
}
fn get_special_tokens(&self) -> &SpecialTokens {
static TOKENS: Lazy<SpecialTokens> = 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<u32> {
self.vocab.get(token).copied()
}
fn id_to_token(&self, _id: u32) -> Option<String> {
None
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[test]
fn llava_prompt_replacement_uses_config_ids() {
let tokenizer = TestTokenizer::new(&[("<image>", 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(&[("<image>", 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(&[("<image>", 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);
}
}

View File

@@ -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<MultiModalResult<TrackedMedia>>;
#[derive(Debug, Clone)]
pub struct TrackerConfig {
pub placeholder_tokens: HashMap<Modality, String>,
pub modality_limits: HashMap<Modality, usize>,
}
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<ConversationSegment>,
pub data: MultiModalData,
pub uuids: MultiModalUUIDs,
pub placeholders: PlaceholderMap,
}
pub struct AsyncMultiModalTracker {
media_connector: Arc<MediaConnector>,
config: TrackerConfig,
pending: HashMap<Modality, Vec<PendingTask>>,
placeholders: PlaceholderMap,
conversation: Vec<ConversationSegment>,
uuids: MultiModalUUIDs,
counts: HashMap<Modality, usize>,
}
impl AsyncMultiModalTracker {
pub fn new(media_connector: Arc<MediaConnector>, 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<TrackerOutput> {
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("<mm>")
}
}
fn next_index(&mut self, modality: Modality) -> MultiModalResult<usize> {
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<String>) {
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<String>,
) -> 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(())
}
}

View File

@@ -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<HashMap<Modality, &'static str>> = LazyLock::new(|| {
HashMap::from([
(Modality::Image, "<image>"),
(Modality::ImageEmbeds, "<image_embeds>"),
(Modality::Audio, "<audio>"),
(Modality::Video, "<video>"),
])
});
/// Detail level passed by OpenAI style APIs.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum ImageDetail {
#[default]
Auto,
Low,
High,
}
/// A normalized content part understood by the tracker.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ChatContentPart {
Text {
text: String,
},
ImageUrl {
url: String,
#[serde(skip_serializing_if = "Option::is_none")]
detail: Option<ImageDetail>,
#[serde(skip_serializing_if = "Option::is_none")]
uuid: Option<String>,
},
ImageData {
data: Vec<u8>,
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
uuid: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
detail: Option<ImageDetail>,
},
ImageEmbeds {
payload: Value,
#[serde(skip_serializing_if = "Option::is_none")]
uuid: Option<String>,
},
}
/// Represents one segment in the conversation after placeholders are inserted.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ConversationSegment {
Text(String),
Placeholder { token: String },
}
impl ConversationSegment {
pub fn text<S: Into<String>>(text: S) -> Self {
ConversationSegment::Text(text.into())
}
pub fn placeholder<S: Into<String>>(token: S) -> Self {
ConversationSegment::Placeholder {
token: token.into(),
}
}
}
/// Metadata describing where placeholder tokens were inserted.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PlaceholderHandle {
pub modality: Modality,
pub item_index: usize,
pub text_position: usize,
}
pub type PlaceholderMap = HashMap<String, Vec<PlaceholderHandle>>;
/// Image source metadata (useful for hashing & tracing).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ImageSource {
Url { url: String },
DataUrl,
InlineBytes,
File { path: PathBuf },
}
/// Concrete image payload captured by the media connector.
#[derive(Debug)]
pub struct ImageFrame {
image: DynamicImage,
raw_bytes: Arc<Vec<u8>>,
pub detail: ImageDetail,
pub source: ImageSource,
}
impl ImageFrame {
pub fn new(
image: DynamicImage,
raw_bytes: Arc<Vec<u8>>,
detail: ImageDetail,
source: ImageSource,
) -> Self {
Self {
image,
raw_bytes,
detail,
source,
}
}
pub fn data(&self) -> &DynamicImage {
&self.image
}
pub fn raw_bytes(&self) -> &[u8] {
self.raw_bytes.as_slice()
}
pub fn source(&self) -> &ImageSource {
&self.source
}
pub fn size(&self) -> ImageSize {
ImageSize::new(self.image.width(), self.image.height())
}
}
/// Container for all supported multimodal media objects.
#[derive(Debug, Clone)]
pub enum TrackedMedia {
Image(Arc<ImageFrame>),
/// Placeholder variants for future modalities.
Audio,
Video,
Embeddings,
}
pub type MultiModalData = HashMap<Modality, Vec<TrackedMedia>>;
pub type MultiModalUUIDs = HashMap<Modality, Vec<Option<String>>>;
pub type TokenId = i32;
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub struct ImageSize {
pub width: u32,
pub height: u32,
}
impl ImageSize {
pub fn new(width: u32, height: u32) -> Self {
Self { width, height }
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PlaceholderRange {
pub offset: usize,
pub length: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultiModalTensor {
pub shape: Vec<usize>,
pub dtype: String,
#[serde(with = "serde_bytes")]
pub data: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum MultiModalValue {
Tensor(MultiModalTensor),
Json(Value),
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct MultiModalInputs {
pub prompt_token_ids: Vec<u32>,
#[serde(default)]
pub mm_kwargs: BTreeMap<String, Vec<MultiModalValue>>,
#[serde(default)]
pub mm_hashes: BTreeMap<String, Vec<String>>,
#[serde(default)]
pub mm_placeholders: BTreeMap<String, Vec<PlaceholderRange>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_salt: Option<String>,
}
impl MultiModalInputs {
pub fn new(prompt_token_ids: Vec<u32>) -> Self {
Self {
prompt_token_ids,
..Default::default()
}
}
}
#[derive(Debug, Clone)]
pub struct PromptReplacement {
pub modality: Modality,
pub placeholder_token: String,
pub tokens: Vec<TokenId>,
}
impl PromptReplacement {
pub fn repeated(
modality: Modality,
placeholder_token: &str,
token_id: TokenId,
count: usize,
) -> Self {
Self {
modality,
placeholder_token: placeholder_token.to_string(),
tokens: vec![token_id; count],
}
}
pub fn sequence(modality: Modality, placeholder_token: &str, sequence: Vec<TokenId>) -> Self {
Self {
modality,
placeholder_token: placeholder_token.to_string(),
tokens: sequence,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn multimodal_inputs_defaults() {
let inputs = MultiModalInputs::new(vec![1, 2, 3]);
assert_eq!(inputs.prompt_token_ids, vec![1, 2, 3]);
assert!(inputs.mm_kwargs.is_empty());
}
#[test]
fn placeholder_range_serializes() {
let range = PlaceholderRange {
offset: 10,
length: 4,
};
let json = serde_json::to_string(&range).unwrap();
assert!(json.contains("offset"));
}
#[test]
fn prompt_replacement_builders() {
let rep = PromptReplacement::repeated(Modality::Image, "<image>", 100, 3);
assert_eq!(rep.tokens, vec![100, 100, 100]);
}
}

View File

@@ -0,0 +1,151 @@
use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
use reqwest::Client;
use sgl_model_gateway::multimodal::{
AsyncMultiModalTracker, ChatContentPart, ConversationSegment, ImageFetchConfig, ImageSource,
MediaConnector, MediaConnectorConfig, MediaSource, Modality, TrackerConfig,
};
use tempfile::tempdir;
const TINY_PNG_BASE64: &str =
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNgYAAAAAMAASsJTYQAAAAASUVORK5CYII=";
fn tiny_png_bytes() -> Vec<u8> {
BASE64_STANDARD
.decode(TINY_PNG_BASE64)
.expect("decode tiny png fixture")
}
fn test_connector(allowed_path: Option<PathBuf>) -> MediaConnector {
let client = Client::builder()
.timeout(Duration::from_secs(5))
.no_proxy()
.build()
.expect("client");
MediaConnector::new(
client,
MediaConnectorConfig {
allowed_domains: None,
allowed_local_media_path: allowed_path,
fetch_timeout: Duration::from_secs(5),
},
)
.expect("media connector")
}
#[tokio::test]
async fn fetch_image_from_inline_bytes() {
let connector = test_connector(None);
let bytes = tiny_png_bytes();
let frame = connector
.fetch_image(
MediaSource::InlineBytes(bytes.clone()),
ImageFetchConfig::default(),
)
.await
.expect("inline image");
assert_eq!(frame.data().width(), 1);
assert_eq!(frame.data().height(), 1);
assert_eq!(frame.raw_bytes(), bytes.as_slice());
}
#[tokio::test]
async fn fetch_image_from_data_url() {
let connector = test_connector(None);
let bytes = tiny_png_bytes();
let data_url = format!(
"data:image/png;base64,{}",
BASE64_STANDARD.encode(bytes.clone())
);
let frame = connector
.fetch_image(MediaSource::DataUrl(data_url), ImageFetchConfig::default())
.await
.expect("data url");
assert_eq!(frame.data().width(), 1);
assert_eq!(frame.raw_bytes(), bytes.as_slice());
}
#[tokio::test]
async fn fetch_image_from_file() {
let tmp = tempdir().expect("tempdir");
let allowed_root = std::fs::canonicalize(tmp.path()).expect("canonical tmp path");
let file_path = allowed_root.join("tiny.png");
std::fs::write(&file_path, tiny_png_bytes()).expect("write png");
let connector = test_connector(Some(allowed_root));
let frame = connector
.fetch_image(
MediaSource::File(file_path.clone()),
ImageFetchConfig::default(),
)
.await
.expect("file png");
assert_eq!(frame.data().width(), 1);
let expected = std::fs::canonicalize(&file_path).expect("canonical path");
match frame.source() {
ImageSource::File { path } => assert_eq!(path, &expected),
other => panic!("expected file source, got {:?}", other),
}
}
#[tokio::test]
async fn tracker_collects_conversation_and_placeholders() {
let connector = Arc::new(test_connector(None));
let mut tracker = AsyncMultiModalTracker::new(
connector,
TrackerConfig {
placeholder_tokens: Default::default(),
modality_limits: HashMap::from([(Modality::Image, 2)]),
},
);
tracker
.push_part(ChatContentPart::Text {
text: "before".into(),
})
.expect("text part");
tracker
.push_part(ChatContentPart::ImageData {
data: tiny_png_bytes(),
mime_type: Some("image/png".into()),
uuid: Some("img-1".into()),
detail: None,
})
.expect("image part");
tracker
.push_part(ChatContentPart::Text {
text: "after".into(),
})
.expect("text part");
let output = tracker.finalize().await.expect("tracker finalize");
assert_eq!(output.conversation.len(), 3);
assert!(matches!(
&output.conversation[0],
ConversationSegment::Text(text) if text == "before"
));
assert!(matches!(
&output.conversation[1],
ConversationSegment::Placeholder { token } if token == "<image>"
));
assert!(matches!(
&output.conversation[2],
ConversationSegment::Text(text) if text == "after"
));
let images = output.data.get(&Modality::Image).expect("image entry");
assert_eq!(images.len(), 1);
let uuids = output.uuids.get(&Modality::Image).expect("uuid entry");
assert_eq!(uuids, &vec![Some("img-1".into())]);
let placeholders = output
.placeholders
.get("<image>")
.expect("placeholder entry");
assert_eq!(placeholders.len(), 1);
assert_eq!(placeholders[0].text_position, 1);
assert_eq!(placeholders[0].item_index, 0);
}