From d7f2d16878486b8ed487517705b2b8fc2acd5f2b Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 4 Oct 2026 08:06:49 +0300 Subject: [PATCH 1/4] feat(llm): expose native input transport support and preserve Responses media Co-authored-by: Medulla --- .../src/catalog/mod_tests.rs | 15 ++ .../tinyinference-llm/src/catalog/parsing.rs | 3 + crates/tinyinference-llm/src/catalog/types.rs | 7 +- .../tinyinference-llm/src/model/decorators.rs | 36 +++++ .../src/model/discover/cache.rs | 2 + .../src/model/discover/cache_tests.rs | 1 + .../src/model/discover/mod.rs | 4 +- .../src/model/discover/parse.rs | 40 ++++- .../src/model/discover/parse_tests.rs | 51 ++++++ .../src/model/discover/types.rs | 15 +- .../src/model/migration_tests.rs | 33 ++++ crates/tinyinference-llm/src/model/types.rs | 40 +++++ .../src/providers/anthropic/mod.rs | 21 +++ .../src/providers/anthropic/mod_tests.rs | 58 +++++++ .../src/providers/anthropic/request.rs | 59 +++++++ .../src/providers/openai/convert.rs | 63 ++++++-- .../src/providers/openai/mod_tests.rs | 85 ++++++++++ .../src/providers/openai/responses.rs | 145 ++++++++++++++---- .../src/providers/openai/responses_tests.rs | 76 ++++++++- .../src/providers/openai/transport.rs | 48 +++++- .../src/providers/openai/wire_tests.rs | 32 ++++ 21 files changed, 767 insertions(+), 67 deletions(-) diff --git a/crates/tinyinference-llm/src/catalog/mod_tests.rs b/crates/tinyinference-llm/src/catalog/mod_tests.rs index f1b8bbe2..91197f0d 100644 --- a/crates/tinyinference-llm/src/catalog/mod_tests.rs +++ b/crates/tinyinference-llm/src/catalog/mod_tests.rs @@ -87,6 +87,7 @@ fn null_is_empty_only_for_success_envelopes() { #[test] fn codex_hints_are_merged_without_duplicates() { let mut models = vec![ModelInfo { + input_modalities: None, id: "gpt-5.4".into(), owned_by: None, context_window: None, @@ -117,3 +118,17 @@ fn blank_preferred_identifier_falls_back_to_the_next_usable_field() { assert_eq!(models[0].id, "usable-slug"); assert_eq!(models[1].id, "usable-name"); } + +#[test] +fn catalog_preserves_advertised_input_modalities() { + let models = parse_models_response(&serde_json::json!({"data":[ + {"id":"vision","architecture":{"input_modalities":["text","image"]}}, + {"id":"unknown"} + ]})) + .unwrap(); + assert_eq!( + models[0].input_modalities, + Some(vec!["text".into(), "image".into()]) + ); + assert_eq!(models[1].input_modalities, None); +} diff --git a/crates/tinyinference-llm/src/catalog/parsing.rs b/crates/tinyinference-llm/src/catalog/parsing.rs index 691cadf5..61d802a8 100644 --- a/crates/tinyinference-llm/src/catalog/parsing.rs +++ b/crates/tinyinference-llm/src/catalog/parsing.rs @@ -128,6 +128,7 @@ pub fn merge_openai_codex_model_hints(models: &mut Vec) { id: (*id).to_string(), owned_by: Some("openai-codex".to_string()), context_window: None, + input_modalities: None, display_name: None, input_per_1m: None, output_per_1m: None, @@ -142,6 +143,7 @@ fn model_info_from_catalog_item(item: &serde_json::Value) -> Option { id: id.to_string(), owned_by: None, context_window: None, + input_modalities: None, display_name: None, input_per_1m: None, output_per_1m: None, @@ -184,6 +186,7 @@ fn model_info_from_catalog_item(item: &serde_json::Value) -> Option { id, owned_by, context_window, + input_modalities: crate::model::discover::input_modalities_from_entry(item), display_name, input_per_1m: price("inputPer1M"), output_per_1m: price("outputPer1M"), diff --git a/crates/tinyinference-llm/src/catalog/types.rs b/crates/tinyinference-llm/src/catalog/types.rs index f9baa46b..59ac1925 100644 --- a/crates/tinyinference-llm/src/catalog/types.rs +++ b/crates/tinyinference-llm/src/catalog/types.rs @@ -1,8 +1,8 @@ //! `ModelInfo`: the typed representation of one `/models` catalog entry. -use serde::Serialize; +use serde::{Deserialize, Serialize}; -#[derive(Debug, Clone, PartialEq, Serialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] /// Normalized metadata for one model advertised by a provider catalog. pub struct ModelInfo { /// Provider model identifier used in inference requests. @@ -10,6 +10,9 @@ pub struct ModelInfo { /// Provider or organization that owns the model, when advertised. #[serde(skip_serializing_if = "Option::is_none")] pub owned_by: Option, + /// Advertised input modalities; `None` means unknown, never assumed vision. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_modalities: Option>, /// Maximum context length in tokens, when advertised. #[serde(skip_serializing_if = "Option::is_none")] pub context_window: Option, diff --git a/crates/tinyinference-llm/src/model/decorators.rs b/crates/tinyinference-llm/src/model/decorators.rs index bd646028..fbda09d1 100644 --- a/crates/tinyinference-llm/src/model/decorators.rs +++ b/crates/tinyinference-llm/src/model/decorators.rs @@ -93,6 +93,15 @@ impl ChatModel for ProfileOverrideModel { Some(&self.profile) } + fn supports_input( + &self, + modality: super::InputModality, + mime_type: &str, + source: super::InputSource, + ) -> bool { + self.inner.supports_input(modality, mime_type, source) + } + fn cache_identity(&self) -> Option { self.inner.cache_identity() } @@ -144,6 +153,15 @@ impl ChatModel for MaxTokensModel { self.inner.profile() } + fn supports_input( + &self, + modality: super::InputModality, + mime_type: &str, + source: super::InputSource, + ) -> bool { + self.inner.supports_input(modality, mime_type, source) + } + fn cache_identity(&self) -> Option { self.inner.cache_identity() } @@ -201,6 +219,15 @@ impl ChatModel for RouteRecordingModel { self.inner.profile() } + fn supports_input( + &self, + modality: super::InputModality, + mime_type: &str, + source: super::InputSource, + ) -> bool { + self.inner.supports_input(modality, mime_type, source) + } + fn cache_identity(&self) -> Option { self.inner.cache_identity() } @@ -362,6 +389,15 @@ impl ChatModel for ObservingModel { self.inner.profile() } + fn supports_input( + &self, + modality: super::InputModality, + mime_type: &str, + source: super::InputSource, + ) -> bool { + self.inner.supports_input(modality, mime_type, source) + } + fn cache_identity(&self) -> Option { self.inner.cache_identity() } diff --git a/crates/tinyinference-llm/src/model/discover/cache.rs b/crates/tinyinference-llm/src/model/discover/cache.rs index e043021d..5f60f31c 100644 --- a/crates/tinyinference-llm/src/model/discover/cache.rs +++ b/crates/tinyinference-llm/src/model/discover/cache.rs @@ -74,6 +74,7 @@ impl CachedLimits { Some(ModelLimits { context_window: Some(learned), max_output_tokens: limits.max_output_tokens, + input_modalities: limits.input_modalities, source: LimitSource::LearnedFromOverflow, }) } @@ -81,6 +82,7 @@ impl CachedLimits { (None, Some(learned)) => Some(ModelLimits { context_window: Some(learned), max_output_tokens: None, + input_modalities: None, source: LimitSource::LearnedFromOverflow, }), (None, None) => None, diff --git a/crates/tinyinference-llm/src/model/discover/cache_tests.rs b/crates/tinyinference-llm/src/model/discover/cache_tests.rs index f845c1eb..51183773 100644 --- a/crates/tinyinference-llm/src/model/discover/cache_tests.rs +++ b/crates/tinyinference-llm/src/model/discover/cache_tests.rs @@ -7,6 +7,7 @@ fn limits(window: u64) -> ModelLimits { ModelLimits { context_window: Some(window), max_output_tokens: Some(8_192), + input_modalities: None, source: LimitSource::ProviderListing, } } diff --git a/crates/tinyinference-llm/src/model/discover/mod.rs b/crates/tinyinference-llm/src/model/discover/mod.rs index bc2bfc6c..68ccca56 100644 --- a/crates/tinyinference-llm/src/model/discover/mod.rs +++ b/crates/tinyinference-llm/src/model/discover/mod.rs @@ -32,8 +32,8 @@ pub use fetch::{ ModelListingFetcher, ReqwestListingFetcher, discover_model_limits, discover_model_limits_with, }; pub use parse::{ - limits_from_entry, model_ids_match, parse_listing_limits, parse_model_limits, - parse_openrouter_endpoint_limits, pinned_openrouter_providers, + input_modalities_from_entry, limits_from_entry, model_ids_match, parse_listing_limits, + parse_model_limits, parse_openrouter_endpoint_limits, pinned_openrouter_providers, }; pub use types::{DiscoveryRequest, LimitSource, ModelLimits}; diff --git a/crates/tinyinference-llm/src/model/discover/parse.rs b/crates/tinyinference-llm/src/model/discover/parse.rs index 66515bb6..f49426ad 100644 --- a/crates/tinyinference-llm/src/model/discover/parse.rs +++ b/crates/tinyinference-llm/src/model/discover/parse.rs @@ -60,6 +60,21 @@ fn entry_id(item: &Value) -> Option<&str> { .find(|id| !id.is_empty()) } +/// Reads advertised input modalities from OpenRouter, models.dev, or managed catalogs. +/// Missing or malformed metadata remains unknown rather than asserting support. +#[must_use] +pub fn input_modalities_from_entry(item: &Value) -> Option> { + let values = item + .get("input_modalities") + .or_else(|| item.pointer("/architecture/input_modalities")) + .or_else(|| item.pointer("/modalities/input"))? + .as_array()?; + values + .iter() + .map(|value| value.as_str().map(|s| s.trim().to_ascii_lowercase())) + .collect() +} + /// Limits one listing entry advertises, or `None` when it advertises neither a /// window nor an output cap. /// @@ -69,15 +84,19 @@ fn entry_id(item: &Value) -> Option<&str> { #[must_use] pub fn limits_from_entry(item: &Value) -> Option { let top_provider = item.get("top_provider"); - let context_window = first_key(item, CONTEXT_KEYS).or_else(|| { - top_provider.and_then(|top| first_key(top, &["context_length", "context_window"])) - }); + let context_window = first_key(item, CONTEXT_KEYS) + .or_else(|| item.pointer("/limit/context").and_then(positive_u64)) + .or_else(|| { + top_provider.and_then(|top| first_key(top, &["context_length", "context_window"])) + }); let max_output_tokens = top_provider .and_then(|top| first_key(top, OUTPUT_KEYS)) - .or_else(|| first_key(item, OUTPUT_KEYS)); + .or_else(|| first_key(item, OUTPUT_KEYS)) + .or_else(|| item.pointer("/limit/output").and_then(positive_u64)); let limits = ModelLimits { context_window, max_output_tokens, + input_modalities: input_modalities_from_entry(item), source: LimitSource::ProviderListing, }; (!limits.is_empty()).then_some(limits) @@ -228,6 +247,19 @@ pub fn parse_openrouter_endpoint_limits(body: &Value, providers: &[String]) -> O let limits = ModelLimits { context_window, max_output_tokens, + input_modalities: matched + .iter() + .try_fold(None::>, |common, endpoint| { + let reported = input_modalities_from_entry(endpoint)?; + Some(Some(match common { + None => reported, + Some(mut values) => { + values.retain(|m| reported.contains(m)); + values + } + })) + }) + .flatten(), source: LimitSource::ProviderEndpoint { provider }, }; (!limits.is_empty()).then_some(limits) diff --git a/crates/tinyinference-llm/src/model/discover/parse_tests.rs b/crates/tinyinference-llm/src/model/discover/parse_tests.rs index 69d2063a..320776fa 100644 --- a/crates/tinyinference-llm/src/model/discover/parse_tests.rs +++ b/crates/tinyinference-llm/src/model/discover/parse_tests.rs @@ -205,3 +205,54 @@ fn pinned_providers_from_routing_options() { assert!(pinned_openrouter_providers(&json!({ "provider": { "sort": "price" } })).is_empty()); assert!(pinned_openrouter_providers(&json!({})).is_empty()); } + +#[test] +fn discovery_retains_known_modalities_without_inventing_unknown_support() { + assert_eq!(limits_from_entry(&serde_json::json!({"context_length":1024,"architecture":{"input_modalities":["text","image"]}})).unwrap().input_modalities, Some(vec!["text".into(), "image".into()])); + assert_eq!( + limits_from_entry( + &serde_json::json!({"limit":{"context":1024},"modalities":{"input":["text","audio"]}}) + ) + .unwrap() + .input_modalities, + Some(vec!["text".into(), "audio".into()]) + ); + assert_eq!( + limits_from_entry(&serde_json::json!({"context_length":1024})) + .unwrap() + .input_modalities, + None + ); +} + +#[test] +fn pinned_endpoint_modalities_are_intersected_and_unknown_remains_unknown() { + let body = serde_json::json!({"data":{"endpoints":[ + {"provider_name":"a","context_length":1024,"input_modalities":["text","image"]}, + {"provider_name":"b","context_length":1024,"input_modalities":["text"]} + ]}}); + let providers = vec!["a".into(), "b".into()]; + assert_eq!( + parse_openrouter_endpoint_limits(&body, &providers) + .unwrap() + .input_modalities, + Some(vec!["text".into()]) + ); + let mut unknown = body; + unknown["data"]["endpoints"][1] + .as_object_mut() + .unwrap() + .remove("input_modalities"); + assert_eq!( + parse_openrouter_endpoint_limits(&unknown, &providers) + .unwrap() + .input_modalities, + None + ); +} + +#[test] +fn legacy_discovery_metadata_deserializes_without_modality_assumptions() { + let limits: ModelLimits = serde_json::from_value(serde_json::json!({"context_window":1024,"max_output_tokens":128,"source":{"kind":"provider_listing"}})).unwrap(); + assert_eq!(limits.input_modalities, None); +} diff --git a/crates/tinyinference-llm/src/model/discover/types.rs b/crates/tinyinference-llm/src/model/discover/types.rs index 08be38ac..c0301b6f 100644 --- a/crates/tinyinference-llm/src/model/discover/types.rs +++ b/crates/tinyinference-llm/src/model/discover/types.rs @@ -2,10 +2,10 @@ use std::time::Duration; -use serde::Serialize; +use serde::{Deserialize, Serialize}; /// Where a [`ModelLimits`] value came from. -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] pub enum LimitSource { /// The model's entry in the provider's `/models` listing (or its @@ -34,21 +34,26 @@ impl LimitSource { } /// The token limits a provider reports for one model. -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct ModelLimits { /// Maximum context window (prompt plus completion) in tokens. pub context_window: Option, /// Maximum completion tokens the provider allows per response. pub max_output_tokens: Option, + /// Advertised input modalities; `None` means unknown, never assumed vision. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_modalities: Option>, /// Where the values came from. pub source: LimitSource, } impl ModelLimits { - /// Whether the record carries no usable limit at all. + /// Whether the record carries neither usable limits nor modality metadata. #[must_use] pub fn is_empty(&self) -> bool { - self.context_window.is_none() && self.max_output_tokens.is_none() + self.context_window.is_none() + && self.max_output_tokens.is_none() + && self.input_modalities.is_none() } } diff --git a/crates/tinyinference-llm/src/model/migration_tests.rs b/crates/tinyinference-llm/src/model/migration_tests.rs index 5f33349f..e8b5f933 100644 --- a/crates/tinyinference-llm/src/model/migration_tests.rs +++ b/crates/tinyinference-llm/src/model/migration_tests.rs @@ -24,6 +24,17 @@ impl RecordingModel { #[async_trait] impl ChatModel<()> for RecordingModel { + fn supports_input( + &self, + modality: InputModality, + mime_type: &str, + source: InputSource, + ) -> bool { + modality == InputModality::Image + && mime_type == "image/png" + && source == InputSource::Base64 + } + async fn invoke(&self, _state: &(), request: ModelRequest) -> crate::Result { self.requests.lock().unwrap().push(request); Ok(self.response.clone()) @@ -496,3 +507,25 @@ async fn observer_streams_report_one_terminal_outcome_with_stream_metadata() { if observed.as_ref() == Some(&correlation) && message == "mock returned: provider failed" )); } + +#[test] +fn transport_support_defaults_closed_and_all_decorators_forward() { + assert!(!FailingModel.supports_input(InputModality::Image, "image/png", InputSource::Base64)); + let inner: Arc> = + Arc::new(RecordingModel::new(ModelResponse::assistant("ok"))); + let inner: Arc> = + Arc::new(ProfileOverrideModel::new(inner, ModelProfile::default())); + let inner: Arc> = Arc::new(MaxTokensModel::new(inner, 10)); + let inner: Arc> = Arc::new(RouteRecordingModel::new( + inner, + ResolvedModelRoute::new("test", "model", "route"), + )); + let model = ObservingModel::new(inner, Arc::new(RecordingObserver::default())); + assert!(model.supports_input(InputModality::Image, "image/png", InputSource::Base64)); + assert!(!model.supports_input(InputModality::Image, "image/png", InputSource::Path)); + assert!(!model.supports_input( + InputModality::Document, + "application/pdf", + InputSource::Base64 + )); +} diff --git a/crates/tinyinference-llm/src/model/types.rs b/crates/tinyinference-llm/src/model/types.rs index 6a98de39..29bbd5ff 100644 --- a/crates/tinyinference-llm/src/model/types.rs +++ b/crates/tinyinference-llm/src/model/types.rs @@ -1254,6 +1254,32 @@ impl Stream for ModelStream { } } +/// Native input modality, independent of a model's advertised capabilities. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum InputModality { + /// Still images. + Image, + /// Audio recordings. + Audio, + /// Video recordings. + Video, + /// Documents such as PDFs. + Document, +} + +/// Reference representation accepted by an inference transport. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum InputSource { + /// Remote HTTP(S) URL. + Url, + /// Inline base64 bytes (including image data URIs). + Base64, + /// Unresolved local path; providers never read these. + Path, +} + /// A provider-neutral chat model. /// /// Generic over the application `State`. @@ -1269,6 +1295,20 @@ pub trait ChatModel: Send + Sync { None } + /// Whether this transport can serialize this modality, MIME type, and source. + /// + /// This does not assert that the selected model understands the input: + /// hosts must also check effective model capabilities. Unknown transports + /// decline by default, keeping existing implementations source-compatible. + fn supports_input( + &self, + _modality: InputModality, + _mime_type: &str, + _source: InputSource, + ) -> bool { + false + } + /// Returns a stable, credential-safe identity for response-cache scoping. /// /// The default declines identity. Implementations must never include raw diff --git a/crates/tinyinference-llm/src/providers/anthropic/mod.rs b/crates/tinyinference-llm/src/providers/anthropic/mod.rs index 5fdb7bbe..6b6e2fc8 100644 --- a/crates/tinyinference-llm/src/providers/anthropic/mod.rs +++ b/crates/tinyinference-llm/src/providers/anthropic/mod.rs @@ -295,6 +295,7 @@ impl AnthropicModel { } async fn post(&self, request: &ModelRequest, streaming: bool) -> Result { + request::validate_media(request)?; crate::network_guard::ensure_network_models_allowed()?; let endpoint = reqwest::Url::parse(&self.endpoint()) .map_err(|error| Error::Validation(format!("invalid Anthropic base URL: {error}")))?; @@ -438,6 +439,26 @@ impl ChatModel for AnthropicModel { Some(&self.profile) } + fn supports_input( + &self, + modality: crate::model::InputModality, + mime_type: &str, + source: crate::model::InputSource, + ) -> bool { + use crate::model::{InputModality, InputSource}; + if source == InputSource::Path { + return false; + } + match modality { + InputModality::Image => matches!( + mime_type.trim().to_ascii_lowercase().as_str(), + "image/png" | "image/jpeg" | "image/webp" | "image/gif" + ), + InputModality::Document => mime_type.eq_ignore_ascii_case("application/pdf"), + InputModality::Audio | InputModality::Video => false, + } + } + fn cache_identity(&self) -> Option { Some(format!("anthropic:{}:{}", self.base_url, self.model)) } diff --git a/crates/tinyinference-llm/src/providers/anthropic/mod_tests.rs b/crates/tinyinference-llm/src/providers/anthropic/mod_tests.rs index a0f8a878..145a8f47 100644 --- a/crates/tinyinference-llm/src/providers/anthropic/mod_tests.rs +++ b/crates/tinyinference-llm/src/providers/anthropic/mod_tests.rs @@ -1203,3 +1203,61 @@ async fn overflow_error_records_the_stated_window() { .expect("overflow window recorded"); assert_eq!(learned.context_window, Some(200_000)); } + +#[test] +fn anthropic_transport_media_support_and_validation_match_native_wire() { + use crate::message::MediaRef; + use crate::model::{ChatModel, InputModality as M, InputSource as S}; + let model = AnthropicModel::new("key"); + assert!(>::supports_input( + &model, + M::Image, + "image/jpeg", + S::Url + )); + assert!(>::supports_input( + &model, + M::Document, + "application/pdf", + S::Base64 + )); + for (modality, mime, source) in [ + (M::Video, "video/mp4", S::Base64), + (M::Audio, "audio/wav", S::Base64), + (M::Document, "application/msword", S::Base64), + (M::Document, "application/pdf", S::Path), + ] { + assert!(!>::supports_input( + &model, modality, mime, source + )); + } + for block in [ + ContentBlock::Audio(MediaRef::base64("QQ==", "audio/wav")), + ContentBlock::Document(MediaRef::path("/private.pdf")), + ContentBlock::Document(MediaRef::base64("QQ==", "application/msword")), + ] { + let request = ModelRequest::new(vec![Message::User(crate::message::UserMessage { + content: vec![block], + })]); + assert!(super::request::validate_media(&request).is_err()); + } +} + +#[tokio::test] +async fn unsupported_anthropic_media_fails_before_invoke_or_stream_network_io() { + let model = AnthropicModel::with_base_url("key", "http://127.0.0.1:1"); + let request = ModelRequest::new(vec![Message::User(crate::message::UserMessage { + content: vec![ContentBlock::Video(crate::message::MediaRef::base64( + "QQ==", + "video/mp4", + ))], + })]); + assert!(matches!( + model.invoke(&(), request.clone()).await, + Err(crate::Error::Validation(_)) + )); + assert!(matches!( + model.stream(&(), request).await, + Err(crate::Error::Validation(_)) + )); +} diff --git a/crates/tinyinference-llm/src/providers/anthropic/request.rs b/crates/tinyinference-llm/src/providers/anthropic/request.rs index 075fcd3d..49dcb7c0 100644 --- a/crates/tinyinference-llm/src/providers/anthropic/request.rs +++ b/crates/tinyinference-llm/src/providers/anthropic/request.rs @@ -351,3 +351,62 @@ fn unsupported_media_placeholder(kind: &str, media: &crate::message::MediaRef) - "text": format!("[{kind} attachment omitted: {descriptor}]"), }) } + +/// Rejects media the native endpoint cannot carry before performing any I/O. +pub(super) fn validate_media(request: &ModelRequest) -> crate::Result<()> { + use crate::message::MediaRef; + for message in &request.messages { + let blocks = match message { + Message::User(m) => &m.content, + Message::Tool(m) => &m.content, + Message::Assistant(m) => &m.content, + Message::System(m) => &m.content, + Message::Custom(_) => continue, + }; + for block in blocks { + let valid = match block { + ContentBlock::Image(image) => { + let remote = reqwest::Url::parse(&image.url).is_ok_and(|url| { + matches!(url.scheme(), "http" | "https") && url.host_str().is_some() + }); + let inline_mime = image + .url + .strip_prefix("data:") + .and_then(|rest| rest.split_once(";base64,").map(|(mime, _)| mime)); + let mime_valid = + inline_mime + .or(image.mime_type.as_deref()) + .is_none_or(|mime| { + matches!( + mime.to_ascii_lowercase().as_str(), + "image/png" | "image/jpeg" | "image/webp" | "image/gif" + ) + }); + (remote || inline_mime.is_some()) + && mime_valid + && !matches!(message, Message::System(_) | Message::Assistant(_)) + } + ContentBlock::Document(media) => { + let source_valid = match media { + MediaRef::Base64 { .. } => true, + MediaRef::Url { url, .. } => reqwest::Url::parse(url).is_ok_and(|url| { + matches!(url.scheme(), "http" | "https") && url.host_str().is_some() + }), + MediaRef::Path { .. } => false, + }; + source_valid + && media + .media_type() + .is_none_or(|mime| mime.eq_ignore_ascii_case("application/pdf")) + && !matches!(message, Message::System(_) | Message::Assistant(_)) + } + ContentBlock::Audio(_) | ContentBlock::Video(_) => false, + _ => true, + }; + if !valid { + return Err(crate::Error::Validation("Anthropic endpoint cannot represent this media type or source; resolve local paths in the host".into())); + } + } + } + Ok(()) +} diff --git a/crates/tinyinference-llm/src/providers/openai/convert.rs b/crates/tinyinference-llm/src/providers/openai/convert.rs index 69493ab5..11a87f37 100644 --- a/crates/tinyinference-llm/src/providers/openai/convert.rs +++ b/crates/tinyinference-llm/src/providers/openai/convert.rs @@ -151,11 +151,14 @@ pub(super) fn translate_user_content(blocks: &[ContentBlock]) -> Result parts.push(ContentPartWire::ImageUrl { - image_url: ImageUrlWire { - url: image.url.clone(), - }, - }), + ContentBlock::Image(image) => { + validate_image_ref(image)?; + parts.push(ContentPartWire::ImageUrl { + image_url: ImageUrlWire { + url: image.url.clone(), + }, + }); + } ContentBlock::Audio(media) => parts.push(input_audio_part(media)?), // See the string-rendering arm above: reasoning blocks have no // OpenAI representation and are dropped, not failed. @@ -211,8 +214,8 @@ fn mark_last_text_part(message: &mut ChatMessageWire) { /// request. Failing closed keeps the block from being silently dropped. pub(super) fn unrepresentable_block_error() -> Error { Error::Validation( - "OpenAI request cannot represent a provider-extension content block; \ - remove it or target the originating provider" + "OpenAI Chat Completions cannot represent document, video, or provider-extension content blocks; \ + remove the block or target a transport that supports it" .to_string(), ) } @@ -230,7 +233,7 @@ fn input_audio_part(media: &crate::message::MediaRef) -> Result MediaRef::Base64 { data, media_type } => Ok(ContentPartWire::InputAudio { input_audio: InputAudioWire { data: data.clone(), - format: audio_format_from_media_type(media_type), + format: audio_format_from_media_type(media_type)?, }, }), MediaRef::Url { .. } | MediaRef::Path { .. } => Err(Error::Validation( @@ -241,15 +244,16 @@ fn input_audio_part(media: &crate::message::MediaRef) -> Result } } -/// Derives the OpenAI `input_audio.format` token (`"wav"`, `"mp3"`, …) from a -/// MIME type such as `audio/wav`, defaulting to `"wav"` when unrecognized. -fn audio_format_from_media_type(media_type: &str) -> String { - media_type - .rsplit('/') - .next() - .filter(|format| !format.is_empty()) - .unwrap_or("wav") - .to_string() +/// Maps supported WAV and MP3 MIME types onto OpenAI audio format tokens. +/// Unrecognized MIME types fail instead of emitting a malformed request. +fn audio_format_from_media_type(media_type: &str) -> Result { + match media_type.trim().to_ascii_lowercase().as_str() { + "audio/wav" | "audio/x-wav" => Ok("wav".into()), + "audio/mpeg" | "audio/mp3" => Ok("mp3".into()), + _ => Err(Error::Validation( + "OpenAI input_audio supports only WAV or MP3 MIME types".into(), + )), + } } /// Translates a [`ToolChoice`] into the OpenAI `tool_choice` JSON value. @@ -592,3 +596,28 @@ pub(super) fn delta_reasoning_text(delta: &mut ChunkDeltaWire) -> String { } text } + +/// Whether a URL names a remote HTTP(S) resource, never a local file. +pub(super) fn is_remote_url(url: &str) -> bool { + reqwest::Url::parse(url) + .is_ok_and(|url| matches!(url.scheme(), "https" | "http") && url.host_str().is_some()) +} + +/// Checks image source and MIME without fetching any bytes. +pub(super) fn validate_image_ref(image: &crate::message::ImageRef) -> Result<()> { + let data_mime = image + .url + .strip_prefix("data:") + .and_then(|rest| rest.split_once(";base64,").map(|(mime, _)| mime)); + let mime = data_mime.or(image.mime_type.as_deref()); + if mime.is_some_and(|mime| { + !matches!( + mime.to_ascii_lowercase().as_str(), + "image/png" | "image/jpeg" | "image/webp" | "image/gif" + ) + }) || (data_mime.is_none() && !is_remote_url(&image.url)) + { + return Err(Error::Validation("Image input requires PNG/JPEG/WebP/GIF inline base64 or a remote HTTP(S) URL; resolve local paths in the host".into())); + } + Ok(()) +} diff --git a/crates/tinyinference-llm/src/providers/openai/mod_tests.rs b/crates/tinyinference-llm/src/providers/openai/mod_tests.rs index c8963c30..ea01616f 100644 --- a/crates/tinyinference-llm/src/providers/openai/mod_tests.rs +++ b/crates/tinyinference-llm/src/providers/openai/mod_tests.rs @@ -3242,3 +3242,88 @@ async fn streamed_terminal_response_carries_origin() { assert_eq!(origin.api, "chat_completions"); assert_eq!(origin.model, "gpt-4.1-mini"); } + +#[test] +fn input_audio_normalizes_mpeg_and_rejects_unknown_mime() { + use crate::message::{ContentBlock, MediaRef}; + let content = super::convert::translate_user_content(&[ContentBlock::Audio(MediaRef::base64( + "QQ==", + "audio/mpeg", + ))]) + .unwrap(); + assert_eq!( + serde_json::to_value(content).unwrap()[0]["input_audio"]["format"], + "mp3" + ); + assert!( + super::convert::translate_user_content(&[ContentBlock::Audio(MediaRef::base64( + "QQ==", + "audio/flac" + ))]) + .is_err() + ); +} + +#[test] +fn native_transport_support_tracks_selected_wire_and_reference_source() { + use crate::model::{ChatModel, InputModality as M, InputSource as S}; + let chat = OpenAiModel::new("key") + .with_base_url("http://localhost/v1") + .with_model("test"); + assert!(>::supports_input( + &chat, + M::Image, + "image/png", + S::Base64 + )); + assert!(>::supports_input( + &chat, + M::Audio, + "audio/mpeg", + S::Base64 + )); + assert!(!>::supports_input( + &chat, + M::Audio, + "audio/wav", + S::Url + )); + assert!(!>::supports_input( + &chat, + M::Audio, + "audio/flac", + S::Base64 + )); + assert!(!>::supports_input( + &chat, + M::Document, + "application/pdf", + S::Base64 + )); + let responses = chat.with_responses_api_primary(); + assert!(!>::supports_input( + &responses, + M::Audio, + "audio/wav", + S::Base64 + )); + assert!(!>::supports_input( + &responses, + M::Document, + "application/pdf", + S::Base64 + )); + let responses = responses.with_responses_document_input(true); + assert!(>::supports_input( + &responses, + M::Document, + "application/pdf", + S::Base64 + )); + assert!(!>::supports_input( + &responses, + M::Document, + "application/pdf", + S::Path + )); +} diff --git a/crates/tinyinference-llm/src/providers/openai/responses.rs b/crates/tinyinference-llm/src/providers/openai/responses.rs index 0c600e0e..d6f688f8 100644 --- a/crates/tinyinference-llm/src/providers/openai/responses.rs +++ b/crates/tinyinference-llm/src/providers/openai/responses.rs @@ -182,7 +182,10 @@ pub(super) struct ResponsesInput { pub(super) struct ResponsesContentPart { #[serde(rename = "type")] pub(super) kind: String, + #[serde(skip_serializing_if = "String::is_empty")] pub(super) text: String, + #[serde(flatten)] + pub(super) extra: serde_json::Map, } #[derive(Debug, Deserialize)] @@ -309,53 +312,127 @@ fn normalize_role(message: &Message) -> &'static str { /// the content-part `kind` tracks the *normalized* role (`output_text` for /// assistant/tool, `input_text` otherwise) — the API rejects `input_text` on an /// assistant item. -pub(super) fn build_responses_input(messages: &[Message]) -> (Option, Vec) { +pub(super) fn build_responses_input( + messages: &[Message], + allow_documents: bool, +) -> crate::Result<(Option, Vec)> { + use crate::message::MediaRef; let mut instructions_parts = Vec::new(); let mut input = Vec::new(); - for message in messages { - let text = match message { + let blocks = match message { Message::System(m) => { - let t = message_text(&m.content); - if !t.trim().is_empty() { - instructions_parts.push(t); + if m.content.iter().any(|b| { + matches!( + b, + ContentBlock::Image(_) + | ContentBlock::Audio(_) + | ContentBlock::Video(_) + | ContentBlock::Document(_) + ) + }) { + return Err(crate::Error::Validation( + "Responses instructions cannot contain media".into(), + )); } - continue; - } - Message::User(m) => message_text(&m.content), - Message::Assistant(m) => message_text(&m.content), - // Keep the call id visible so the model can tell *which* call this - // answers. Folding it into an anonymous assistant turn lost that. - Message::Tool(m) => { - let body = message_text(&m.content); - if body.trim().is_empty() { - String::new() - } else { - format!("[tool_result id={} ]\n{body}", m.tool_call_id) + let text = message_text(&m.content); + if !text.trim().is_empty() { + instructions_parts.push(text); } + continue; } - // Host-side out-of-band record; never sent to the provider. + Message::User(m) => &m.content, + Message::Assistant(m) => &m.content, + Message::Tool(m) => &m.content, Message::Custom(_) => continue, }; - if text.trim().is_empty() { - continue; - } let role = normalize_role(message); - input.push(ResponsesInput { - role: role.to_string(), - content: vec![ResponsesContentPart { - kind: if role == "assistant" { - "output_text".to_string() - } else { - "input_text".to_string() - }, + let mut content = Vec::new(); + for block in blocks { + let mut extra = serde_json::Map::new(); + let (kind, text) = match block { + ContentBlock::Text(t) => ( + if role == "assistant" { + "output_text" + } else { + "input_text" + }, + t.clone(), + ), + ContentBlock::Json(v) => ( + if role == "assistant" { + "output_text" + } else { + "input_text" + }, + v.to_string(), + ), + ContentBlock::Image(image) if matches!(message, Message::User(_)) => { + super::convert::validate_image_ref(image)?; + extra.insert("image_url".into(), Value::String(image.url.clone())); + ("input_image", String::new()) + } + ContentBlock::Document(media) + if matches!(message, Message::User(_)) && allow_documents => + { + if media + .media_type() + .is_none_or(|mime| !mime.eq_ignore_ascii_case("application/pdf")) + { + return Err(crate::Error::Validation( + "Responses document input requires application/pdf MIME type".into(), + )); + } + match media { + MediaRef::Base64 { data, .. } => { + extra.insert("filename".into(), Value::String("document.pdf".into())); + extra.insert("file_data".into(), Value::String(format!("data:application/pdf;base64,{data}"))); + } + MediaRef::Url { url, .. } if super::convert::is_remote_url(url) => { + extra.insert("file_url".into(), Value::String(url.clone())); + } + _ => return Err(crate::Error::Validation("Responses document requires inline base64 or a remote HTTPS URL; resolve local paths in the host".into())), + } + ("input_file", String::new()) + } + ContentBlock::Thinking { .. } | ContentBlock::RedactedThinking { .. } => continue, + _ => { + return Err(crate::Error::Validation( + "Responses endpoint cannot represent this content block or source".into(), + )); + } + }; + if text.trim().is_empty() && extra.is_empty() { + continue; + } + content.push(ResponsesContentPart { + kind: kind.into(), text, - }], - }); + extra, + }); + } + if let Message::Tool(m) = message { + // Preserve the historical causal tool-result marker and merged text. + let body = content.iter().map(|p| p.text.as_str()).collect::(); + content = if body.trim().is_empty() { + Vec::new() + } else { + vec![ResponsesContentPart { + kind: "input_text".into(), + text: format!("[tool_result id={} ]\n{body}", m.tool_call_id), + extra: serde_json::Map::new(), + }] + }; + } + if !content.is_empty() { + input.push(ResponsesInput { + role: role.into(), + content, + }); + } } - let instructions = (!instructions_parts.is_empty()).then(|| instructions_parts.join("\n\n")); - (instructions, input) + Ok((instructions, input)) } /// Extracts the assistant text from a Responses body: the convenience diff --git a/crates/tinyinference-llm/src/providers/openai/responses_tests.rs b/crates/tinyinference-llm/src/providers/openai/responses_tests.rs index 857bc538..4b18a385 100644 --- a/crates/tinyinference-llm/src/providers/openai/responses_tests.rs +++ b/crates/tinyinference-llm/src/providers/openai/responses_tests.rs @@ -11,7 +11,7 @@ fn build_input_folds_system_into_instructions_and_keys_roles() { Message::assistant("hello"), Message::user(" "), // empty → skipped ]; - let (instructions, input) = build_responses_input(&messages); + let (instructions, input) = build_responses_input(&messages, false).unwrap(); assert_eq!(instructions.as_deref(), Some("be terse\n\nand correct")); assert_eq!(input.len(), 2); assert_eq!(input[0].role, "user"); @@ -34,7 +34,7 @@ fn build_input_skips_custom_messages() { }), Message::assistant("hello"), ]; - let (_, input) = build_responses_input(&messages); + let (_, input) = build_responses_input(&messages, false).unwrap(); assert_eq!(input.len(), 2); assert_eq!(input[0].content[0].text, "hi"); assert_eq!(input[1].content[0].text, "hello"); @@ -206,3 +206,75 @@ async fn responses_invoke_avoids_codex_stream_required_400() { assert_eq!(response.text(), "ok"); server.join().unwrap(); } + +#[test] +fn native_inputs_preserve_order_and_emit_responses_wire_shapes() { + use crate::message::{ContentBlock, ImageRef, MediaRef, UserMessage}; + let messages = vec![Message::User(UserMessage { + content: vec![ + ContentBlock::Text("before".into()), + ContentBlock::Image(ImageRef { + url: "data:image/png;base64,QQ==".into(), + mime_type: Some("image/png".into()), + }), + ContentBlock::Text("after".into()), + ContentBlock::Document(MediaRef::base64("QQ==", "application/pdf")), + ], + })]; + let (_, input) = build_responses_input(&messages, true).unwrap(); + assert_eq!( + serde_json::to_value(&input).unwrap(), + json!([{"role":"user","content":[ + {"type":"input_text","text":"before"}, + {"type":"input_image","image_url":"data:image/png;base64,QQ=="}, + {"type":"input_text","text":"after"}, + {"type":"input_file","filename":"document.pdf","file_data":"data:application/pdf;base64,QQ=="} + ]}]) + ); +} + +#[test] +fn responses_native_inputs_reject_unsupported_sources_and_modalities() { + use crate::message::{ContentBlock, ImageRef, MediaRef, UserMessage}; + for block in [ + ContentBlock::Audio(MediaRef::base64("QQ==", "audio/wav")), + ContentBlock::Video(MediaRef::base64("QQ==", "video/mp4")), + ContentBlock::Document(MediaRef::path("/secret.pdf")), + ContentBlock::Document(MediaRef::base64("QQ==", "application/msword")), + ContentBlock::Image(ImageRef { + url: "file:///secret.png".into(), + mime_type: Some("image/png".into()), + }), + ] { + let messages = [Message::User(UserMessage { + content: vec![block], + })]; + assert!(build_responses_input(&messages, true).is_err()); + } + let pdf = [Message::User(UserMessage { + content: vec![ContentBlock::Document(MediaRef::base64( + "QQ==", + "application/pdf", + ))], + })]; + assert!( + build_responses_input(&pdf, false).is_err(), + "Codex endpoint must decline files" + ); +} + +#[test] +fn responses_pdf_url_uses_file_url() { + use crate::message::{ContentBlock, MediaRef, UserMessage}; + let messages = [Message::User(UserMessage { + content: vec![ContentBlock::Document(MediaRef::Url { + url: "https://example.com/a.pdf".into(), + media_type: Some("application/pdf".into()), + })], + })]; + let (_, input) = build_responses_input(&messages, true).unwrap(); + assert_eq!( + serde_json::to_value(input).unwrap()[0]["content"][0], + json!({"type":"input_file","file_url":"https://example.com/a.pdf"}) + ); +} diff --git a/crates/tinyinference-llm/src/providers/openai/transport.rs b/crates/tinyinference-llm/src/providers/openai/transport.rs index 4e23c896..bec34d51 100644 --- a/crates/tinyinference-llm/src/providers/openai/transport.rs +++ b/crates/tinyinference-llm/src/providers/openai/transport.rs @@ -90,6 +90,7 @@ pub struct OpenAiModel { /// When `true`, calls go to the OpenAI **Responses API** (`/v1/responses`) /// instead of Chat Completions. See [`Self::with_responses_api_primary`]. responses_api_primary: bool, + responses_document_input: bool, /// When `true` (Responses path only), omit `max_output_tokens` from the wire /// body — the OpenAI Codex OAuth backend rejects it. See /// [`Self::with_responses_omit_max_output_tokens`]. @@ -384,6 +385,7 @@ impl OpenAiModel { explicit_cache_control: false, anthropic_cache_control: true, responses_api_primary: false, + responses_document_input: false, responses_omit_max_output_tokens: false, extra_query_params: Vec::new(), user_agent: None, @@ -1458,13 +1460,24 @@ impl OpenAiModel { } } + /// Enables PDF inputs for a Responses endpoint known to accept `input_file`. + /// + /// Disabled by default because compatible gateways (including Codex OAuth) + /// do not all accept document inputs. This never changes model capability. + #[must_use] + pub fn with_responses_document_input(mut self, enabled: bool) -> Self { + self.responses_document_input = enabled; + self + } + /// Builds the `/v1/responses` request body from a provider-neutral request. pub(super) fn translate_responses_request( &self, request: &ModelRequest, ) -> Result { let model = request.model.clone().unwrap_or_else(|| self.model.clone()); - let (instructions, input) = responses::build_responses_input(&request.messages); + let (instructions, input) = + responses::build_responses_input(&request.messages, self.responses_document_input)?; let max_output_tokens = if self.responses_omit_max_output_tokens { None } else { @@ -2117,6 +2130,39 @@ impl ChatModel for OpenAiModel { Some(&self.profile) } + fn supports_input( + &self, + modality: crate::model::InputModality, + mime_type: &str, + source: crate::model::InputSource, + ) -> bool { + use crate::model::{InputModality, InputSource}; + if source == InputSource::Path { + return false; + } + let mime = mime_type.trim().to_ascii_lowercase(); + match modality { + InputModality::Image => matches!( + mime.as_str(), + "image/png" | "image/jpeg" | "image/webp" | "image/gif" + ), + InputModality::Audio => { + !self.responses_api_primary + && source == InputSource::Base64 + && matches!( + mime.as_str(), + "audio/wav" | "audio/x-wav" | "audio/mpeg" | "audio/mp3" + ) + } + InputModality::Document => { + self.responses_api_primary + && self.responses_document_input + && mime == "application/pdf" + } + InputModality::Video => false, + } + } + /// Invokes the OpenAI Chat Completions endpoint and maps the response into a /// [`ModelResponse`]. /// diff --git a/crates/tinyinference-llm/src/providers/openai/wire_tests.rs b/crates/tinyinference-llm/src/providers/openai/wire_tests.rs index 81101124..53f05e77 100644 --- a/crates/tinyinference-llm/src/providers/openai/wire_tests.rs +++ b/crates/tinyinference-llm/src/providers/openai/wire_tests.rs @@ -794,3 +794,35 @@ fn debug_output_reports_a_bearer_source_without_reading_it() { assert!(debug.contains("bearer_source: true"), "{debug}"); assert!(!debug.contains("static-key"), "{debug}"); } + +#[tokio::test] +async fn responses_http_sends_native_image_and_opted_in_pdf_in_order() { + use crate::message::{ContentBlock, ImageRef, MediaRef, UserMessage}; + let (base, server) = serve_sequence(vec![(200, json!({"output_text":"accepted"}).to_string())]); + let model = OpenAiModel::new("k") + .with_base_url(&base) + .with_responses_api_primary() + .with_responses_document_input(true); + let request = ModelRequest::new(vec![Message::User(UserMessage { + content: vec![ + ContentBlock::Text("inspect".into()), + ContentBlock::Image(ImageRef { + url: "data:image/png;base64,QQ==".into(), + mime_type: Some("image/png".into()), + }), + ContentBlock::Document(MediaRef::base64("QQ==", "application/pdf")), + ], + })]); + assert_eq!(model.invoke(&(), request).await.unwrap().text(), "accepted"); + let requests = server.join().unwrap(); + assert!(requests[0].starts_with("POST /v1/responses ")); + let body = request_body(&requests[0]); + assert_eq!( + body["input"][0]["content"], + json!([ + {"type":"input_text","text":"inspect"}, + {"type":"input_image","image_url":"data:image/png;base64,QQ=="}, + {"type":"input_file","filename":"document.pdf","file_data":"data:application/pdf;base64,QQ=="} + ]) + ); +} From c76ed9eac0325b8b00b9d91cf3f8f90d053d960f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 4 Oct 2026 08:16:50 +0300 Subject: [PATCH 2/4] fix(llm): validate inline media and reject assistant media loss Co-authored-by: Medulla --- Cargo.lock | 1 + crates/tinyinference-llm/Cargo.toml | 1 + .../src/providers/anthropic/request.rs | 34 ++- .../tinyinference-llm/src/providers/media.rs | 18 ++ .../src/providers/media_tests.rs | 259 ++++++++++++++++++ crates/tinyinference-llm/src/providers/mod.rs | 5 + .../src/providers/openai/convert.rs | 38 ++- .../src/providers/openai/responses.rs | 1 + 8 files changed, 335 insertions(+), 22 deletions(-) create mode 100644 crates/tinyinference-llm/src/providers/media.rs create mode 100644 crates/tinyinference-llm/src/providers/media_tests.rs diff --git a/Cargo.lock b/Cargo.lock index d9a7fec9..cd994c01 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1320,6 +1320,7 @@ name = "tinyinference-llm" version = "0.3.0" dependencies = [ "async-trait", + "base64 0.22.1", "bytes", "futures", "reqwest", diff --git a/crates/tinyinference-llm/Cargo.toml b/crates/tinyinference-llm/Cargo.toml index 75da70a2..21cfa895 100644 --- a/crates/tinyinference-llm/Cargo.toml +++ b/crates/tinyinference-llm/Cargo.toml @@ -12,6 +12,7 @@ keywords = ["llm", "inference", "streaming", "openai", "anthropic"] categories = ["api-bindings", "asynchronous"] [dependencies] +base64 = "0.22" async-trait = { workspace = true } bytes = { workspace = true } futures = { workspace = true } diff --git a/crates/tinyinference-llm/src/providers/anthropic/request.rs b/crates/tinyinference-llm/src/providers/anthropic/request.rs index 49dcb7c0..35973c3f 100644 --- a/crates/tinyinference-llm/src/providers/anthropic/request.rs +++ b/crates/tinyinference-llm/src/providers/anthropic/request.rs @@ -224,7 +224,7 @@ fn text_only_blocks(content: &[ContentBlock]) -> Vec { /// User-side content: text, images, and documents (Anthropic's native /// `document` block). Audio and video have no Messages API representation -/// and are rendered as placeholder text rather than silently dropped. +/// and are rejected by public invoke/stream validation before conversion. /// Thinking blocks never appear in user content; provider extensions have no /// faithful representation and are dropped. fn content_blocks(content: &[ContentBlock]) -> Vec { @@ -316,8 +316,8 @@ fn image_block(image: &ImageRef) -> Value { /// Renders a document reference as Anthropic's `document` content block. /// `MediaRef::Path` has no wire representation (the harness never reads -/// local files) and falls back to a placeholder text block instead of being -/// silently dropped. +/// local files) and is rejected by public invoke/stream validation. The +/// internal renderer retains a diagnostic fallback for unchecked conversion. fn document_block(media: &crate::message::MediaRef) -> Value { use crate::message::MediaRef; match media { @@ -372,23 +372,29 @@ pub(super) fn validate_media(request: &ModelRequest) -> crate::Result<()> { let inline_mime = image .url .strip_prefix("data:") - .and_then(|rest| rest.split_once(";base64,").map(|(mime, _)| mime)); - let mime_valid = - inline_mime - .or(image.mime_type.as_deref()) - .is_none_or(|mime| { - matches!( - mime.to_ascii_lowercase().as_str(), - "image/png" | "image/jpeg" | "image/webp" | "image/gif" - ) - }); + .and_then(|rest| rest.split_once(";base64,")); + if let Some((_, data)) = inline_mime { + crate::providers::media::validate_base64(data)?; + } + let mime_valid = inline_mime + .map(|(mime, _)| mime) + .or(image.mime_type.as_deref()) + .is_none_or(|mime| { + matches!( + mime.to_ascii_lowercase().as_str(), + "image/png" | "image/jpeg" | "image/webp" | "image/gif" + ) + }); (remote || inline_mime.is_some()) && mime_valid && !matches!(message, Message::System(_) | Message::Assistant(_)) } ContentBlock::Document(media) => { let source_valid = match media { - MediaRef::Base64 { .. } => true, + MediaRef::Base64 { data, .. } => { + crate::providers::media::validate_base64(data)?; + true + } MediaRef::Url { url, .. } => reqwest::Url::parse(url).is_ok_and(|url| { matches!(url.scheme(), "http" | "https") && url.host_str().is_some() }), diff --git a/crates/tinyinference-llm/src/providers/media.rs b/crates/tinyinference-llm/src/providers/media.rs new file mode 100644 index 00000000..47877524 --- /dev/null +++ b/crates/tinyinference-llm/src/providers/media.rs @@ -0,0 +1,18 @@ +//! Shared validation for inline media carried by hosted transports. + +use base64::Engine; + +/// Rejects empty or malformed standard Base64 before a provider sends it. +/// Media-format recognition remains the receiving provider's responsibility. +pub(super) fn validate_base64(data: &str) -> crate::Result<()> { + if data.is_empty() + || base64::engine::general_purpose::STANDARD + .decode(data) + .is_err() + { + return Err(crate::Error::Validation( + "Inline media requires nonempty, valid standard Base64 data".into(), + )); + } + Ok(()) +} diff --git a/crates/tinyinference-llm/src/providers/media_tests.rs b/crates/tinyinference-llm/src/providers/media_tests.rs new file mode 100644 index 00000000..46a9aeed --- /dev/null +++ b/crates/tinyinference-llm/src/providers/media_tests.rs @@ -0,0 +1,259 @@ +use super::*; +use crate::message::{ImageRef, MediaRef, UserMessage}; +use crate::providers::{anthropic::AnthropicModel, openai::OpenAiModel}; + +async fn assert_rejected(model: &dyn ChatModel<()>, request: ModelRequest, expected: &str) { + match model.invoke(&(), request.clone()).await { + Err(Error::Validation(message)) => assert!(message.contains(expected), "{message}"), + other => panic!("expected validation before I/O, got {other:?}"), + } + match model.stream(&(), request).await { + Err(Error::Validation(message)) => assert!(message.contains(expected), "{message}"), + _ => panic!("expected validation before streaming I/O"), + } +} + +#[tokio::test] +async fn invalid_inline_media_is_rejected_before_public_transport_io() { + let chat = OpenAiModel::compatible("key", "http://127.0.0.1:1/v1", "gpt-4.1"); + let responses = OpenAiModel::compatible("key", "http://127.0.0.1:1/v1", "gpt-4.1") + .with_responses_api_primary() + .with_responses_document_input(true); + let anthropic = + AnthropicModel::with_base_url("key", "http://127.0.0.1:1").with_insecure_http(true); + for data in ["", "%%%", "Q", "QR=="] { + let image = ContentBlock::Image(ImageRef { + url: format!("data:image/png;base64,{data}"), + mime_type: None, + }); + let audio = ContentBlock::Audio(MediaRef::base64(data, "audio/wav")); + let pdf = ContentBlock::Document(MediaRef::base64(data, "application/pdf")); + for (model, block) in [ + (&chat as &dyn ChatModel<()>, image.clone()), + (&chat, audio), + (&responses, image.clone()), + (&responses, pdf.clone()), + (&anthropic, image), + (&anthropic, pdf), + ] { + assert_rejected( + model, + ModelRequest::new(vec![Message::User(UserMessage { + content: vec![block], + })]), + "Base64", + ) + .await; + } + } +} + +#[tokio::test] +async fn assistant_media_is_rejected_before_public_chat_transport_io() { + let model = OpenAiModel::compatible("key", "http://127.0.0.1:1/v1", "gpt-4.1"); + for block in [ + ContentBlock::Image(ImageRef { + url: "data:image/png;base64,QQ==".into(), + mime_type: None, + }), + ContentBlock::Audio(MediaRef::base64("QQ==", "audio/wav")), + ContentBlock::Video(MediaRef::base64("QQ==", "video/mp4")), + ContentBlock::Document(MediaRef::base64("QQ==", "application/pdf")), + ] { + let mut message = Message::assistant("visible text"); + if let Message::Assistant(assistant) = &mut message { + assistant.content.push(block); + } + assert_rejected( + &model, + ModelRequest::new(vec![message]), + "assistant messages", + ) + .await; + } +} + +/// Captures complete outgoing JSON on loopback and returns synthetic responses. +fn wire_fixture( + count: usize, + reply: &'static str, +) -> (String, std::thread::JoinHandle>) { + use std::io::{Read, Write}; + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let endpoint = format!("http://{}/v1", listener.local_addr().unwrap()); + let server = std::thread::spawn(move || { + let mut captured = Vec::new(); + for _ in 0..count { + let (mut socket, _) = listener.accept().unwrap(); + let mut bytes = Vec::new(); + let mut buffer = [0; 4096]; + let (start, length) = loop { + let read = socket.read(&mut buffer).unwrap(); + assert!(read > 0); + bytes.extend_from_slice(&buffer[..read]); + if let Some(split) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { + let headers = std::str::from_utf8(&bytes[..split]).unwrap(); + let length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap(); + break (split + 4, length); + } + }; + while bytes.len() < start + length { + let read = socket.read(&mut buffer).unwrap(); + assert!(read > 0); + bytes.extend_from_slice(&buffer[..read]); + } + let body: Value = serde_json::from_slice(&bytes[start..start + length]).unwrap(); + let (content_type, payload) = if body["stream"] == true && body.get("input").is_some() { + ( + "text/event-stream", + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"output_text\":\"ok\"}}\n\n", + ) + } else if body["stream"] == true { + ( + "text/event-stream", + "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n", + ) + } else { + ("application/json", reply) + }; + captured.push(body); + write!(socket, "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{payload}", payload.len()).unwrap(); + } + captured + }); + (endpoint, server) +} + +const CHAT_REPLY: &str = + r#"{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}"#; + +#[tokio::test] +async fn valid_inline_media_reaches_public_provider_wire() { + let image = ContentBlock::Image(ImageRef { + url: "data:image/png;base64,QQ==".into(), + mime_type: None, + }); + let audio = ContentBlock::Audio(MediaRef::base64("QQ==", "audio/wav")); + let pdf = ContentBlock::Document(MediaRef::base64("QQ==", "application/pdf")); + let request = |blocks| ModelRequest::new(vec![Message::User(UserMessage { content: blocks })]); + + let (endpoint, server) = wire_fixture(1, CHAT_REPLY); + let model = OpenAiModel::compatible("key", endpoint, "gpt-4.1"); + assert_eq!( + model + .invoke(&(), request(vec![image.clone(), audio])) + .await + .unwrap() + .text(), + "ok" + ); + let bodies = server.join().unwrap(); + assert_eq!( + bodies[0]["messages"][0]["content"][0]["image_url"]["url"], + "data:image/png;base64,QQ==" + ); + assert_eq!( + bodies[0]["messages"][0]["content"][1]["input_audio"]["data"], + "QQ==" + ); + + let (endpoint, server) = wire_fixture(1, r#"{"output_text":"ok"}"#); + let model = OpenAiModel::compatible("key", endpoint, "gpt-4.1") + .with_responses_api_primary() + .with_responses_document_input(true); + assert_eq!( + model + .invoke(&(), request(vec![image.clone(), pdf.clone()])) + .await + .unwrap() + .text(), + "ok" + ); + let bodies = server.join().unwrap(); + assert_eq!( + bodies[0]["input"][0]["content"][0]["image_url"], + "data:image/png;base64,QQ==" + ); + assert_eq!( + bodies[0]["input"][0]["content"][1]["file_data"], + "data:application/pdf;base64,QQ==" + ); + + let (endpoint, server) = wire_fixture( + 1, + r#"{"id":"msg","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1},"stop_reason":"end_turn"}"#, + ); + let model = AnthropicModel::with_base_url("key", endpoint).with_insecure_http(true); + assert_eq!( + model + .invoke(&(), request(vec![image, pdf])) + .await + .unwrap() + .text(), + "ok" + ); + let bodies = server.join().unwrap(); + assert_eq!( + bodies[0]["messages"][0]["content"][0]["source"]["data"], + "QQ==" + ); + assert_eq!( + bodies[0]["messages"][0]["content"][1]["source"]["data"], + "QQ==" + ); +} + +#[tokio::test] +async fn assistant_text_thinking_and_tools_keep_public_chat_wire_semantics() { + use futures::StreamExt; + let (endpoint, server) = wire_fixture(2, CHAT_REPLY); + let model = OpenAiModel::compatible("key", endpoint, "gpt-4.1"); + let mut text = Message::assistant("visible text"); + let mut tools = Message::assistant(""); + if let Message::Assistant(assistant) = &mut text { + assistant + .content + .push(ContentBlock::thinking("private reasoning")); + assistant.content.push(ContentBlock::RedactedThinking { + data: "opaque".into(), + }); + } + if let Message::Assistant(assistant) = &mut tools { + assistant.tool_calls.push(ToolCall::new( + "call-1", + "lookup", + serde_json::json!({"key":"value"}), + )); + } + let request = ModelRequest::new(vec![text, tools]); + assert_eq!( + model.invoke(&(), request.clone()).await.unwrap().text(), + "ok" + ); + let items: Vec<_> = model.stream(&(), request).await.unwrap().collect().await; + assert!( + items + .iter() + .any(|item| matches!(item, ModelStreamItem::Completed(_))) + ); + for body in server.join().unwrap() { + assert_eq!(body["messages"][0]["content"], "visible text"); + assert!(body["messages"][1]["content"].is_null()); + assert_eq!(body["messages"][1]["tool_calls"][0]["id"], "call-1"); + assert_eq!( + body["messages"][1]["tool_calls"][0]["function"]["name"], + "lookup" + ); + assert_eq!( + body["messages"][1]["tool_calls"][0]["function"]["arguments"], + r#"{"key":"value"}"# + ); + } +} diff --git a/crates/tinyinference-llm/src/providers/mod.rs b/crates/tinyinference-llm/src/providers/mod.rs index 08e53bbd..56ce46dd 100644 --- a/crates/tinyinference-llm/src/providers/mod.rs +++ b/crates/tinyinference-llm/src/providers/mod.rs @@ -34,6 +34,7 @@ //! // #[cfg(feature = "ollama")] pub mod ollama; //! ``` +mod media; mod types; // --- real provider integrations --- @@ -64,3 +65,7 @@ mod mock; #[cfg(test)] #[path = "mod_tests.rs"] mod test; + +#[cfg(test)] +#[path = "media_tests.rs"] +mod media_tests; diff --git a/crates/tinyinference-llm/src/providers/openai/convert.rs b/crates/tinyinference-llm/src/providers/openai/convert.rs index 11a87f37..6d89bb59 100644 --- a/crates/tinyinference-llm/src/providers/openai/convert.rs +++ b/crates/tinyinference-llm/src/providers/openai/convert.rs @@ -30,6 +30,20 @@ pub(super) fn translate_message(message: &Message) -> Result { tool_call_id: None, }, Message::Assistant(assistant) => { + if assistant.content.iter().any(|block| { + matches!( + block, + ContentBlock::Image(_) + | ContentBlock::Audio(_) + | ContentBlock::Video(_) + | ContentBlock::Document(_) + ) + }) { + return Err(Error::Validation( + "OpenAI Chat Completions cannot represent media content in assistant messages" + .into(), + )); + } let text = message.text(); // OpenAI accepts a null content for tool-call-only assistant turns. let content = if text.is_empty() && !assistant.tool_calls.is_empty() { @@ -230,12 +244,15 @@ pub(super) fn unrepresentable_block_error() -> Error { fn input_audio_part(media: &crate::message::MediaRef) -> Result { use crate::message::MediaRef; match media { - MediaRef::Base64 { data, media_type } => Ok(ContentPartWire::InputAudio { - input_audio: InputAudioWire { - data: data.clone(), - format: audio_format_from_media_type(media_type)?, - }, - }), + MediaRef::Base64 { data, media_type } => { + crate::providers::media::validate_base64(data)?; + Ok(ContentPartWire::InputAudio { + input_audio: InputAudioWire { + data: data.clone(), + format: audio_format_from_media_type(media_type)?, + }, + }) + } MediaRef::Url { .. } | MediaRef::Path { .. } => Err(Error::Validation( "OpenAI input_audio requires inline base64 data; resolve the \ audio reference to bytes before sending it" @@ -608,8 +625,13 @@ pub(super) fn validate_image_ref(image: &crate::message::ImageRef) -> Result<()> let data_mime = image .url .strip_prefix("data:") - .and_then(|rest| rest.split_once(";base64,").map(|(mime, _)| mime)); - let mime = data_mime.or(image.mime_type.as_deref()); + .and_then(|rest| rest.split_once(";base64,")); + if let Some((_, data)) = data_mime { + crate::providers::media::validate_base64(data)?; + } + let mime = data_mime + .map(|(mime, _)| mime) + .or(image.mime_type.as_deref()); if mime.is_some_and(|mime| { !matches!( mime.to_ascii_lowercase().as_str(), diff --git a/crates/tinyinference-llm/src/providers/openai/responses.rs b/crates/tinyinference-llm/src/providers/openai/responses.rs index d6f688f8..92eed216 100644 --- a/crates/tinyinference-llm/src/providers/openai/responses.rs +++ b/crates/tinyinference-llm/src/providers/openai/responses.rs @@ -385,6 +385,7 @@ pub(super) fn build_responses_input( } match media { MediaRef::Base64 { data, .. } => { + crate::providers::media::validate_base64(data)?; extra.insert("filename".into(), Value::String("document.pdf".into())); extra.insert("file_data".into(), Value::String(format!("data:application/pdf;base64,{data}"))); } From f57b99adba3dc9fff43daad3dec63b870e2212c2 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 4 Oct 2026 08:59:40 +0300 Subject: [PATCH 3/4] fix: aggregate conservative media capabilities and remove socket fixtures Co-authored-by: Medulla --- crates/tinyinference-llm/src/lib.rs | 4 +- .../src/model/discover/cache.rs | 35 +++- .../src/model/discover/cache_tests.rs | 101 ++++++++++ .../src/model/migration_tests.rs | 7 + .../src/providers/media_tests.rs | 185 ------------------ .../src/providers/openai/wire_tests.rs | 51 +++++ 6 files changed, 194 insertions(+), 189 deletions(-) diff --git a/crates/tinyinference-llm/src/lib.rs b/crates/tinyinference-llm/src/lib.rs index 6d745a1b..c515068c 100644 --- a/crates/tinyinference-llm/src/lib.rs +++ b/crates/tinyinference-llm/src/lib.rs @@ -26,8 +26,8 @@ pub use failure::{ }; pub use message::{AssistantMessage, ContentBlock, Message, MessageDelta}; pub use model::{ - ChatModel, ModelRequest, ModelResponse, ModelStream, ModelStreamItem, - context_window_for_model_id, model_id_supports_vision, + ChatModel, InputModality, InputSource, ModelRequest, ModelResponse, ModelStream, + ModelStreamItem, context_window_for_model_id, model_id_supports_vision, }; pub use network_guard::{allow_network_models, deny_network_models, network_models_denied}; pub use providers::{MockModel, ProviderKind, ProviderSpec}; diff --git a/crates/tinyinference-llm/src/model/discover/cache.rs b/crates/tinyinference-llm/src/model/discover/cache.rs index 5f60f31c..ada727c5 100644 --- a/crates/tinyinference-llm/src/model/discover/cache.rs +++ b/crates/tinyinference-llm/src/model/discover/cache.rs @@ -142,7 +142,9 @@ impl ModelLimitsCache { /// The fresh cached state for `(endpoint, model)` across every discovery /// variant: the smallest fresh discovered window, lowered by any learned - /// overflow window. Use [`Self::get_variant`] for one request's slot. + /// overflow window. Input modalities intersect fresh variant facts; any + /// unknown variant makes the aggregate unknown. Use [`Self::get_variant`] + /// for one request's slot. #[must_use] pub fn get(&self, endpoint: &str, model: &str) -> CachedLimits { self.get_at(endpoint, model, Instant::now()) @@ -161,9 +163,24 @@ impl ModelLimitsCache { // else a fresh negative, else nothing. let mut best: Option = None; let mut negative = false; + let mut modalities: Option> = None; + let mut saw_modalities = false; + let mut unknown_modalities = false; for discovered in entry.discovered.values() { match self.fresh(discovered, now) { Some(Some(limits)) => { + if let Some(values) = &limits.input_modalities { + if saw_modalities { + if let Some(common) = &mut modalities { + common.retain(|value| values.contains(value)); + } + } else { + modalities = Some(values.clone()); + saw_modalities = true; + } + } else { + unknown_modalities = true; + } let smaller = best.as_ref().is_none_or(|current| { limits.context_window.unwrap_or(u64::MAX) < current.context_window.unwrap_or(u64::MAX) @@ -172,10 +189,24 @@ impl ModelLimitsCache { best = Some(limits); } } - Some(None) => negative = true, + Some(None) => { + negative = true; + unknown_modalities = true; + } None => {} } } + if let Some(best) = &mut best { + best.input_modalities = if unknown_modalities { + None + } else { + modalities.map(|mut values| { + values.sort(); + values.dedup(); + values + }) + }; + } CachedLimits { discovered: best.map(Some).or(negative.then_some(None)), learned_context_window: self.fresh_learned(entry, now), diff --git a/crates/tinyinference-llm/src/model/discover/cache_tests.rs b/crates/tinyinference-llm/src/model/discover/cache_tests.rs index 51183773..d6c22082 100644 --- a/crates/tinyinference-llm/src/model/discover/cache_tests.rs +++ b/crates/tinyinference-llm/src/model/discover/cache_tests.rs @@ -147,3 +147,104 @@ fn clear_drops_everything() { cache.clear(); assert_eq!(cache.get(ENDPOINT, MODEL), CachedLimits::default()); } + +fn variant(window: u64, modalities: Option<&[&str]>) -> ModelLimits { + ModelLimits { + input_modalities: modalities + .map(|values| values.iter().map(|value| (*value).to_owned()).collect()), + ..limits(window) + } +} + +#[test] +fn aggregate_modalities_intersect_independently_of_window_and_ties() { + for windows in [(32_000, 64_000), (64_000, 32_000), (32_000, 32_000)] { + let cache = cache(); + let now = Instant::now(); + cache.insert_discovered_variant_at( + ENDPOINT, + MODEL, + "image", + Some(variant(windows.0, Some(&["text", "image"]))), + now, + ); + cache.insert_discovered_variant_at( + ENDPOINT, + MODEL, + "text", + Some(variant(windows.1, Some(&["text"]))), + now, + ); + let aggregate = cache.get_at(ENDPOINT, MODEL, now).effective().unwrap(); + assert_eq!(aggregate.context_window, Some(windows.0.min(windows.1))); + assert_eq!(aggregate.input_modalities, Some(vec!["text".into()])); + assert_eq!( + cache + .get_variant_at(ENDPOINT, MODEL, "image", now) + .effective() + .unwrap() + .input_modalities, + Some(vec!["text".into(), "image".into()]) + ); + } +} + +#[test] +fn aggregate_modalities_preserve_unknown_negative_empty_and_expiration() { + let cache = cache(); + let now = Instant::now(); + cache.insert_discovered_variant_at( + ENDPOINT, + MODEL, + "image", + Some(variant(32_000, Some(&["image", "text"]))), + now, + ); + cache.insert_discovered_variant_at( + ENDPOINT, + MODEL, + "unknown", + Some(variant(64_000, None)), + now, + ); + assert_eq!( + cache + .get_at(ENDPOINT, MODEL, now) + .effective() + .unwrap() + .input_modalities, + None + ); + cache.insert_discovered_variant_at( + ENDPOINT, + MODEL, + "unknown", + Some(variant(64_000, Some(&[]))), + now, + ); + assert_eq!( + cache + .get_at(ENDPOINT, MODEL, now) + .effective() + .unwrap() + .input_modalities, + Some(vec![]) + ); + cache.insert_discovered_variant_at(ENDPOINT, MODEL, "unknown", None, now); + assert_eq!( + cache + .get_at(ENDPOINT, MODEL, now) + .effective() + .unwrap() + .input_modalities, + None + ); + assert_eq!( + cache + .get_at(ENDPOINT, MODEL, now + Duration::from_secs(10)) + .effective() + .unwrap() + .input_modalities, + Some(vec!["image".into(), "text".into()]) + ); +} diff --git a/crates/tinyinference-llm/src/model/migration_tests.rs b/crates/tinyinference-llm/src/model/migration_tests.rs index e8b5f933..c20ceb6c 100644 --- a/crates/tinyinference-llm/src/model/migration_tests.rs +++ b/crates/tinyinference-llm/src/model/migration_tests.rs @@ -529,3 +529,10 @@ fn transport_support_defaults_closed_and_all_decorators_forward() { InputSource::Base64 )); } + +#[test] +fn input_capability_types_are_available_from_crate_root() { + use crate::{InputModality, InputSource}; + assert_eq!(InputModality::Image, crate::model::InputModality::Image); + assert_eq!(InputSource::Base64, crate::model::InputSource::Base64); +} diff --git a/crates/tinyinference-llm/src/providers/media_tests.rs b/crates/tinyinference-llm/src/providers/media_tests.rs index 46a9aeed..e7f2367a 100644 --- a/crates/tinyinference-llm/src/providers/media_tests.rs +++ b/crates/tinyinference-llm/src/providers/media_tests.rs @@ -72,188 +72,3 @@ async fn assistant_media_is_rejected_before_public_chat_transport_io() { .await; } } - -/// Captures complete outgoing JSON on loopback and returns synthetic responses. -fn wire_fixture( - count: usize, - reply: &'static str, -) -> (String, std::thread::JoinHandle>) { - use std::io::{Read, Write}; - let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - let endpoint = format!("http://{}/v1", listener.local_addr().unwrap()); - let server = std::thread::spawn(move || { - let mut captured = Vec::new(); - for _ in 0..count { - let (mut socket, _) = listener.accept().unwrap(); - let mut bytes = Vec::new(); - let mut buffer = [0; 4096]; - let (start, length) = loop { - let read = socket.read(&mut buffer).unwrap(); - assert!(read > 0); - bytes.extend_from_slice(&buffer[..read]); - if let Some(split) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { - let headers = std::str::from_utf8(&bytes[..split]).unwrap(); - let length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().unwrap()) - }) - .unwrap(); - break (split + 4, length); - } - }; - while bytes.len() < start + length { - let read = socket.read(&mut buffer).unwrap(); - assert!(read > 0); - bytes.extend_from_slice(&buffer[..read]); - } - let body: Value = serde_json::from_slice(&bytes[start..start + length]).unwrap(); - let (content_type, payload) = if body["stream"] == true && body.get("input").is_some() { - ( - "text/event-stream", - "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"output_text\":\"ok\"}}\n\n", - ) - } else if body["stream"] == true { - ( - "text/event-stream", - "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n", - ) - } else { - ("application/json", reply) - }; - captured.push(body); - write!(socket, "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{payload}", payload.len()).unwrap(); - } - captured - }); - (endpoint, server) -} - -const CHAT_REPLY: &str = - r#"{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}"#; - -#[tokio::test] -async fn valid_inline_media_reaches_public_provider_wire() { - let image = ContentBlock::Image(ImageRef { - url: "data:image/png;base64,QQ==".into(), - mime_type: None, - }); - let audio = ContentBlock::Audio(MediaRef::base64("QQ==", "audio/wav")); - let pdf = ContentBlock::Document(MediaRef::base64("QQ==", "application/pdf")); - let request = |blocks| ModelRequest::new(vec![Message::User(UserMessage { content: blocks })]); - - let (endpoint, server) = wire_fixture(1, CHAT_REPLY); - let model = OpenAiModel::compatible("key", endpoint, "gpt-4.1"); - assert_eq!( - model - .invoke(&(), request(vec![image.clone(), audio])) - .await - .unwrap() - .text(), - "ok" - ); - let bodies = server.join().unwrap(); - assert_eq!( - bodies[0]["messages"][0]["content"][0]["image_url"]["url"], - "data:image/png;base64,QQ==" - ); - assert_eq!( - bodies[0]["messages"][0]["content"][1]["input_audio"]["data"], - "QQ==" - ); - - let (endpoint, server) = wire_fixture(1, r#"{"output_text":"ok"}"#); - let model = OpenAiModel::compatible("key", endpoint, "gpt-4.1") - .with_responses_api_primary() - .with_responses_document_input(true); - assert_eq!( - model - .invoke(&(), request(vec![image.clone(), pdf.clone()])) - .await - .unwrap() - .text(), - "ok" - ); - let bodies = server.join().unwrap(); - assert_eq!( - bodies[0]["input"][0]["content"][0]["image_url"], - "data:image/png;base64,QQ==" - ); - assert_eq!( - bodies[0]["input"][0]["content"][1]["file_data"], - "data:application/pdf;base64,QQ==" - ); - - let (endpoint, server) = wire_fixture( - 1, - r#"{"id":"msg","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1},"stop_reason":"end_turn"}"#, - ); - let model = AnthropicModel::with_base_url("key", endpoint).with_insecure_http(true); - assert_eq!( - model - .invoke(&(), request(vec![image, pdf])) - .await - .unwrap() - .text(), - "ok" - ); - let bodies = server.join().unwrap(); - assert_eq!( - bodies[0]["messages"][0]["content"][0]["source"]["data"], - "QQ==" - ); - assert_eq!( - bodies[0]["messages"][0]["content"][1]["source"]["data"], - "QQ==" - ); -} - -#[tokio::test] -async fn assistant_text_thinking_and_tools_keep_public_chat_wire_semantics() { - use futures::StreamExt; - let (endpoint, server) = wire_fixture(2, CHAT_REPLY); - let model = OpenAiModel::compatible("key", endpoint, "gpt-4.1"); - let mut text = Message::assistant("visible text"); - let mut tools = Message::assistant(""); - if let Message::Assistant(assistant) = &mut text { - assistant - .content - .push(ContentBlock::thinking("private reasoning")); - assistant.content.push(ContentBlock::RedactedThinking { - data: "opaque".into(), - }); - } - if let Message::Assistant(assistant) = &mut tools { - assistant.tool_calls.push(ToolCall::new( - "call-1", - "lookup", - serde_json::json!({"key":"value"}), - )); - } - let request = ModelRequest::new(vec![text, tools]); - assert_eq!( - model.invoke(&(), request.clone()).await.unwrap().text(), - "ok" - ); - let items: Vec<_> = model.stream(&(), request).await.unwrap().collect().await; - assert!( - items - .iter() - .any(|item| matches!(item, ModelStreamItem::Completed(_))) - ); - for body in server.join().unwrap() { - assert_eq!(body["messages"][0]["content"], "visible text"); - assert!(body["messages"][1]["content"].is_null()); - assert_eq!(body["messages"][1]["tool_calls"][0]["id"], "call-1"); - assert_eq!( - body["messages"][1]["tool_calls"][0]["function"]["name"], - "lookup" - ); - assert_eq!( - body["messages"][1]["tool_calls"][0]["function"]["arguments"], - r#"{"key":"value"}"# - ); - } -} diff --git a/crates/tinyinference-llm/src/providers/openai/wire_tests.rs b/crates/tinyinference-llm/src/providers/openai/wire_tests.rs index 53f05e77..9e5c258f 100644 --- a/crates/tinyinference-llm/src/providers/openai/wire_tests.rs +++ b/crates/tinyinference-llm/src/providers/openai/wire_tests.rs @@ -826,3 +826,54 @@ async fn responses_http_sends_native_image_and_opted_in_pdf_in_order() { ]) ); } + +#[test] +fn valid_native_image_audio_and_assistant_history_serialize_without_io() { + use crate::message::{ContentBlock, ImageRef, MediaRef, UserMessage}; + use crate::tool::ToolCall; + let model = OpenAiModel::new("k").with_model("gpt-4.1"); + let user = Message::User(UserMessage { + content: vec![ + ContentBlock::Image(ImageRef { + url: "data:image/png;base64,QQ==".into(), + mime_type: None, + }), + ContentBlock::Audio(MediaRef::base64("QQ==", "audio/wav")), + ], + }); + let mut text = Message::assistant("visible text"); + let mut tools = Message::assistant(""); + if let Message::Assistant(assistant) = &mut text { + assistant + .content + .push(ContentBlock::thinking("private reasoning")); + assistant.content.push(ContentBlock::RedactedThinking { + data: "opaque".into(), + }); + } + if let Message::Assistant(assistant) = &mut tools { + assistant + .tool_calls + .push(ToolCall::new("call-1", "lookup", json!({"key":"value"}))); + } + let body = body_for(&model, &ModelRequest::new(vec![user, text, tools])); + assert_eq!( + body["messages"][0]["content"][0]["image_url"]["url"], + "data:image/png;base64,QQ==" + ); + assert_eq!( + body["messages"][0]["content"][1]["input_audio"]["data"], + "QQ==" + ); + assert_eq!(body["messages"][1]["content"], "visible text"); + assert!(body["messages"][2]["content"].is_null()); + assert_eq!(body["messages"][2]["tool_calls"][0]["id"], "call-1"); + assert_eq!( + body["messages"][2]["tool_calls"][0]["function"]["name"], + "lookup" + ); + assert_eq!( + body["messages"][2]["tool_calls"][0]["function"]["arguments"], + r#"{"key":"value"}"# + ); +} From 84d4cafa51eed9a34533774d755a10082c6ecb33 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 4 Oct 2026 09:09:07 +0300 Subject: [PATCH 4/4] fix: supplement modality-only discovery with token limits Co-authored-by: Medulla --- .../src/model/discover/fetch.rs | 26 ++++++- .../src/model/discover/fetch_tests.rs | 73 +++++++++++++++++++ .../src/model/discover/types.rs | 2 +- 3 files changed, 96 insertions(+), 5 deletions(-) diff --git a/crates/tinyinference-llm/src/model/discover/fetch.rs b/crates/tinyinference-llm/src/model/discover/fetch.rs index c52128c9..a094f07a 100644 --- a/crates/tinyinference-llm/src/model/discover/fetch.rs +++ b/crates/tinyinference-llm/src/model/discover/fetch.rs @@ -104,8 +104,9 @@ impl ModelListingFetcher for ReqwestListingFetcher { /// - A fresh cache entry (found or not found) answers without any request. /// - Otherwise the listing is read (`{endpoint}/models` or /// `request.listing_url`); every entry it lists is cached, so one fetch -/// serves every model on the endpoint. When the model is missing and -/// `probe_single_model` is set, `{endpoint}/models/{id}` is tried. When +/// serves every model on the endpoint. When the model or either token limit +/// is missing and `probe_single_model` is set, `{endpoint}/models/{id}` is +/// tried; missing fields are filled without discarding listing modalities. When /// `pinned_providers` is non-empty, OpenRouter's /// `{endpoint}/models/{id}/endpoints` narrows the window to those providers. /// - The whole probe is bounded by `request.timeout`; a timeout or any error @@ -209,10 +210,27 @@ async fn fetch_limits( } }; - if found.is_none() && request.probe_single_model { + if found + .as_ref() + .is_none_or(|limits| limits.context_window.is_none() || limits.max_output_tokens.is_none()) + && request.probe_single_model + { let url = format!("{}/models/{}", request.endpoint, request.model); match fetcher.get_json(&url, &request.headers).await { - Ok(body) => found = parse_model_limits(&body, &request.model), + Ok(body) => { + if let Some(record) = parse_model_limits(&body, &request.model) { + if let Some(listed) = &mut found { + listed.context_window = listed.context_window.or(record.context_window); + listed.max_output_tokens = + listed.max_output_tokens.or(record.max_output_tokens); + if listed.input_modalities.is_none() { + listed.input_modalities = record.input_modalities; + } + } else { + found = Some(record); + } + } + } Err(error) => tracing::debug!( endpoint = %request.endpoint, model = %request.model, diff --git a/crates/tinyinference-llm/src/model/discover/fetch_tests.rs b/crates/tinyinference-llm/src/model/discover/fetch_tests.rs index 7db4f3d1..cd37c3fe 100644 --- a/crates/tinyinference-llm/src/model/discover/fetch_tests.rs +++ b/crates/tinyinference-llm/src/model/discover/fetch_tests.rs @@ -363,3 +363,76 @@ async fn cached_accessor_sees_a_default_discovery() { Some(1_048_576) ); } + +#[tokio::test] +async fn modality_only_listing_merges_single_model_token_limits() { + let fetcher = FakeFetcher::default() + .with( + &format!("{ENDPOINT}/models"), + json!({"data":[{"id":MODEL,"input_modalities":["text","image"]}]}), + ) + .with( + &format!("{ENDPOINT}/models/{MODEL}"), + json!({"id":MODEL,"context_window":64000,"max_output_tokens":8192}), + ); + let cache = cache(); + let request = DiscoveryRequest::new(ENDPOINT, MODEL); + let result = discover_model_limits_with(&fetcher, &cache, &request) + .await + .unwrap(); + assert_eq!(result.context_window, Some(64000)); + assert_eq!(result.max_output_tokens, Some(8192)); + assert_eq!( + result.input_modalities, + Some(vec!["text".into(), "image".into()]) + ); + assert_eq!(fetcher.calls().len(), 2); + assert_eq!( + discover_model_limits_with(&fetcher, &cache, &request).await, + Some(result) + ); + assert_eq!(fetcher.calls().len(), 2); +} + +#[tokio::test] +async fn partial_listing_fills_missing_limits_and_retains_facts_if_probe_unavailable() { + for probe in [ + None, + Some(json!({"id":MODEL,"context_window":96000,"max_output_tokens":4096})), + ] { + let mut fetcher = FakeFetcher::default().with( + &format!("{ENDPOINT}/models"), + json!({"data":[{"id":MODEL,"context_window":32000,"input_modalities":[]}]}), + ); + if let Some(record) = probe { + fetcher = fetcher.with(&format!("{ENDPOINT}/models/{MODEL}"), record); + } + let result = + discover_model_limits_with(&fetcher, &cache(), &DiscoveryRequest::new(ENDPOINT, MODEL)) + .await + .unwrap(); + assert_eq!(result.context_window, Some(32000)); + assert_eq!(result.input_modalities, Some(vec![])); + assert_eq!( + result.max_output_tokens, + fetcher + .routes + .get(&format!("{ENDPOINT}/models/{MODEL}")) + .map(|_| 4096) + ); + assert_eq!(fetcher.calls().len(), 2); + } + let fetcher = FakeFetcher::default().with( + &format!("{ENDPOINT}/models"), + json!({"data":[{"id":MODEL,"input_modalities":["text"]}]}), + ); + let result = discover_model_limits_with( + &fetcher, + &cache(), + &DiscoveryRequest::new(ENDPOINT, MODEL).with_single_model_probe(false), + ) + .await + .unwrap(); + assert_eq!(result.input_modalities, Some(vec!["text".into()])); + assert_eq!(fetcher.calls().len(), 1); +} diff --git a/crates/tinyinference-llm/src/model/discover/types.rs b/crates/tinyinference-llm/src/model/discover/types.rs index c0301b6f..72363c17 100644 --- a/crates/tinyinference-llm/src/model/discover/types.rs +++ b/crates/tinyinference-llm/src/model/discover/types.rs @@ -74,7 +74,7 @@ pub struct DiscoveryRequest { /// provider publishes its catalogue elsewhere. pub listing_url: Option, /// Whether to fall back to `GET {endpoint}/models/{id}` when the listing - /// does not contain the model. + /// does not contain the model or either token limit is missing. pub probe_single_model: bool, /// Headers to send (for example `Authorization`). Never logged. pub headers: Vec<(String, String)>,