feat(rust): add OCR compiler plans and dialect profiles

litellm-rust/Cargo.lock

@@ -1428,6 +1428,7 @@ dependencies = [
  "aws-sigv4",
  "aws-smithy-runtime-api",
  "aws-types",
+ "base64",
  "bytes",
  "futures-util",
  "mime",

litellm-rust/crates/core/Cargo.toml

@@ -6,6 +6,7 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
+base64.workspace = true
 futures-util.workspace = true
 bytes.workspace = true
 mime.workspace = true

litellm-rust/crates/core/src/ocr/compiler.rs

@@ -0,0 +1,218 @@
+use std::collections::BTreeMap;
+
+use serde_json::Value;
+use thiserror::Error;
+use url::Url;
+
+use super::canonical::{CanonicalOcrRequest, DocumentKind, OcrDocument};
+use super::plan::{CompletionPlan, DocumentPlan};
+use super::policy::OcrParameterPolicy;
+use super::response::NormalizedOcr;
+use super::types::OcrDialectId;
+pub use super::wire::{MultipartBodyPlan, MultipartPart, OcrJsonValue, OcrWireBody, OcrWireError};
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+pub enum CompileError {
+    #[error("invalid OCR parameter {field}: {reason}")]
+    InvalidParameter {
+        field: &'static str,
+        reason: &'static str,
+    },
+    #[error("OCR dialect is not compiled yet: {0:?}")]
+    UnsupportedDialect(OcrDialectId),
+}
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+pub enum NormalizeError {
+    #[error("invalid terminal OCR response: {0}")]
+    InvalidPayload(&'static str),
+}
+
+#[derive(Clone, Default, PartialEq, Eq)]
+pub struct OcrCredentials {
+    api_key: Option<String>,
+    oauth_token: Option<String>,
+}
+
+impl OcrCredentials {
+    pub fn new(api_key: Option<String>, oauth_token: Option<String>) -> Self {
+        Self {
+            api_key,
+            oauth_token,
+        }
+    }
+
+    pub fn api_key(&self) -> Option<&str> {
+        self.api_key.as_deref()
+    }
+
+    pub fn oauth_token(&self) -> Option<&str> {
+        self.oauth_token.as_deref()
+    }
+}
+
+#[derive(Clone, PartialEq)]
+pub struct ResolvedOcrTarget {
+    dialect: OcrDialectId,
+    api_base: Url,
+    credentials: OcrCredentials,
+}
+
+impl ResolvedOcrTarget {
+    pub fn new(dialect: OcrDialectId, api_base: Url, credentials: OcrCredentials) -> Self {
+        Self {
+            dialect,
+            api_base,
+            credentials,
+        }
+    }
+
+    pub fn dialect(&self) -> OcrDialectId {
+        self.dialect
+    }
+
+    pub fn api_base(&self) -> &Url {
+        &self.api_base
+    }
+
+    pub fn credentials(&self) -> &OcrCredentials {
+        &self.credentials
+    }
+}
+
+#[derive(Clone, PartialEq)]
+pub enum ProviderDocument {
+    RemoteUrl {
+        kind: DocumentKind,
+        url: Url,
+    },
+    Inline {
+        kind: DocumentKind,
+        media_type: mime::Mime,
+        bytes: bytes::Bytes,
+    },
+    Reference {
+        id: String,
+    },
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum HttpMethod {
+    Get,
+    Post,
+}
+
+#[derive(Clone, PartialEq)]
+pub struct CompiledHttpRequest {
+    pub method: HttpMethod,
+    pub url: Url,
+    pub headers: BTreeMap<String, String>,
+    pub body: OcrWireBody,
+}
+
+#[derive(Clone, PartialEq)]
+pub struct ProviderPayload(Value);
+
+impl ProviderPayload {
+    pub fn new(value: Value) -> Self {
+        Self(value)
+    }
+
+    pub fn as_value(&self) -> &Value {
+        &self.0
+    }
+
+    pub fn into_value(self) -> Value {
+        self.0
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum OcrDocumentPolicy {
+    Ready,
+    FetchRemoteUrlAndInline,
+    UploadUnlessProviderReference,
+}
+
+pub trait OcrDialectCompiler: Send + Sync {
+    fn parameter_policy(&self) -> &'static OcrParameterPolicy;
+
+    fn prepare_document(
+        &self,
+        document: &OcrDocument,
+        target: &ResolvedOcrTarget,
+    ) -> Result<DocumentPlan, CompileError>;
+
+    fn compile_submit(
+        &self,
+        request: &CanonicalOcrRequest,
+        document: ProviderDocument,
+        target: &ResolvedOcrTarget,
+    ) -> Result<CompiledHttpRequest, CompileError>;
+
+    fn completion_plan(&self) -> CompletionPlan;
+
+    fn normalize(
+        &self,
+        terminal_response: ProviderPayload,
+    ) -> Result<NormalizedOcr, NormalizeError>;
+}
+
+#[cfg(test)]
+mod tests {
+    use bytes::Bytes;
+    use mime::Mime;
+    use serde_json::json;
+
+    use super::*;
+
+    #[test]
+    fn compilation_preserves_the_inline_media_allocation() {
+        let source = Bytes::from_static(b"pdf payload");
+        let source_pointer = source.as_ptr();
+        let canonical_document = OcrDocument::Inline {
+            kind: DocumentKind::Pdf,
+            media_type: "application/pdf".parse::<Mime>().expect("valid MIME type"),
+            bytes: source.clone(),
+        };
+        let OcrDocument::Inline {
+            kind,
+            media_type,
+            bytes,
+        } = &canonical_document
+        else {
+            panic!("inline canonical document expected");
+        };
+        let provider_document = ProviderDocument::Inline {
+            kind: *kind,
+            media_type: media_type.clone(),
+            bytes: bytes.clone(),
+        };
+        let ProviderDocument::Inline {
+            media_type, bytes, ..
+        } = provider_document
+        else {
+            panic!("inline document expected");
+        };
+        let request = CompiledHttpRequest {
+            method: HttpMethod::Post,
+            url: Url::parse("https://example.com/ocr").expect("valid URL"),
+            headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
+            body: OcrWireBody::JsonWithMedia(OcrJsonValue::Object(BTreeMap::from([
+                (
+                    "document".to_string(),
+                    OcrJsonValue::InlineDataUri { media_type, bytes },
+                ),
+                ("model".to_string(), OcrJsonValue::Value(json!("ocr-model"))),
+            ]))),
+        };
+
+        let OcrWireBody::JsonWithMedia(OcrJsonValue::Object(fields)) = &request.body else {
+            panic!("media JSON body expected");
+        };
+        let OcrJsonValue::InlineDataUri { bytes, .. } = &fields["document"] else {
+            panic!("inline media expected");
+        };
+        assert_eq!(bytes.as_ptr(), source_pointer);
+    }
+}

litellm-rust/crates/core/src/ocr/mod.rs

@@ -1,5 +1,9 @@
 pub mod canonical;
+pub mod compiler;
+pub mod plan;
 pub mod policy;
+pub mod profile;
 pub mod response;
 pub mod transformation;
 pub mod types;
+pub mod wire;

litellm-rust/crates/core/src/ocr/plan.rs

@@ -0,0 +1,40 @@
+use std::time::Duration;
+
+use url::Url;
+
+use super::canonical::{DocumentKind, OcrDocument};
+use super::compiler::ProviderDocument;
+use super::types::OcrDialectId;
+
+#[derive(Clone, PartialEq)]
+pub enum DocumentPlan {
+    Ready(ProviderDocument),
+    FetchAndInline(FetchPlan),
+    Upload(UploadPlan),
+}
+
+#[derive(Clone, PartialEq)]
+pub struct FetchPlan {
+    pub kind: DocumentKind,
+    pub url: Url,
+    pub max_bytes: usize,
+}
+
+#[derive(Clone, PartialEq)]
+pub struct UploadPlan {
+    pub dialect: OcrDialectId,
+    pub document: OcrDocument,
+}
+
+#[derive(Clone, PartialEq)]
+pub enum CompletionPlan {
+    Immediate,
+    Poll(PollPlan),
+}
+
+#[derive(Clone, PartialEq, Eq)]
+pub struct PollPlan {
+    pub operation_location_header: &'static str,
+    pub interval: Duration,
+    pub timeout: Duration,
+}

litellm-rust/crates/core/src/ocr/profile.rs

@@ -0,0 +1,189 @@
+use super::compiler::OcrDocumentPolicy;
+use super::policy::{OcrParameterPolicy, ParameterDisposition};
+use super::types::OcrDialectId;
+
+pub const MISTRAL_OCR_PARAMETER_POLICY: OcrParameterPolicy = OcrParameterPolicy {
+    pages: ParameterDisposition::Forward,
+    include_image_base64: ParameterDisposition::Forward,
+    image_limit: ParameterDisposition::Forward,
+    image_min_size: ParameterDisposition::Forward,
+    bbox_annotation_format: ParameterDisposition::Forward,
+    document_annotation_format: ParameterDisposition::Forward,
+    document_annotation_prompt: ParameterDisposition::Forward,
+    extract_header: ParameterDisposition::Forward,
+    extract_footer: ParameterDisposition::Forward,
+    table_format: ParameterDisposition::Forward,
+    confidence_scores_granularity: ParameterDisposition::Forward,
+    include_blocks: ParameterDisposition::Forward,
+    request_id: ParameterDisposition::Forward,
+};
+
+pub const AZURE_DOCUMENT_INTELLIGENCE_PARAMETER_POLICY: OcrParameterPolicy = OcrParameterPolicy {
+    pages: ParameterDisposition::Transform,
+    include_image_base64: ParameterDisposition::Reject,
+    image_limit: ParameterDisposition::Reject,
+    image_min_size: ParameterDisposition::Reject,
+    bbox_annotation_format: ParameterDisposition::Reject,
+    document_annotation_format: ParameterDisposition::Reject,
+    document_annotation_prompt: ParameterDisposition::Reject,
+    extract_header: ParameterDisposition::Reject,
+    extract_footer: ParameterDisposition::Reject,
+    table_format: ParameterDisposition::Reject,
+    confidence_scores_granularity: ParameterDisposition::Reject,
+    include_blocks: ParameterDisposition::Reject,
+    request_id: ParameterDisposition::Reject,
+};
+
+pub const REJECT_CANONICAL_OCR_PARAMETER_POLICY: OcrParameterPolicy = OcrParameterPolicy {
+    pages: ParameterDisposition::Reject,
+    include_image_base64: ParameterDisposition::Reject,
+    image_limit: ParameterDisposition::Reject,
+    image_min_size: ParameterDisposition::Reject,
+    bbox_annotation_format: ParameterDisposition::Reject,
+    document_annotation_format: ParameterDisposition::Reject,
+    document_annotation_prompt: ParameterDisposition::Reject,
+    extract_header: ParameterDisposition::Reject,
+    extract_footer: ParameterDisposition::Reject,
+    table_format: ParameterDisposition::Reject,
+    confidence_scores_granularity: ParameterDisposition::Reject,
+    include_blocks: ParameterDisposition::Reject,
+    request_id: ParameterDisposition::Reject,
+};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct OcrPollingProfile {
+    pub operation_location_header: &'static str,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct OcrDialectProfile {
+    pub dialect: OcrDialectId,
+    pub parameter_policy: &'static OcrParameterPolicy,
+    pub document_policy: OcrDocumentPolicy,
+    pub polling: Option<OcrPollingProfile>,
+}
+
+pub const OCR_DIALECT_PROFILES: [OcrDialectProfile; 7] = [
+    OcrDialectProfile {
+        dialect: OcrDialectId::Mistral,
+        parameter_policy: &MISTRAL_OCR_PARAMETER_POLICY,
+        document_policy: OcrDocumentPolicy::Ready,
+        polling: None,
+    },
+    OcrDialectProfile {
+        dialect: OcrDialectId::AzureFoundryMistral,
+        parameter_policy: &MISTRAL_OCR_PARAMETER_POLICY,
+        document_policy: OcrDocumentPolicy::FetchRemoteUrlAndInline,
+        polling: None,
+    },
+    OcrDialectProfile {
+        dialect: OcrDialectId::AzureDocumentIntelligence,
+        parameter_policy: &AZURE_DOCUMENT_INTELLIGENCE_PARAMETER_POLICY,
+        document_policy: OcrDocumentPolicy::Ready,
+        polling: Some(OcrPollingProfile {
+            operation_location_header: "operation-location",
+        }),
+    },
+    OcrDialectProfile {
+        dialect: OcrDialectId::VertexMistral,
+        parameter_policy: &MISTRAL_OCR_PARAMETER_POLICY,
+        document_policy: OcrDocumentPolicy::FetchRemoteUrlAndInline,
+        polling: None,
+    },
+    OcrDialectProfile {
+        dialect: OcrDialectId::VertexDeepSeek,
+        parameter_policy: &REJECT_CANONICAL_OCR_PARAMETER_POLICY,
+        document_policy: OcrDocumentPolicy::Ready,
+        polling: None,
+    },
+    OcrDialectProfile {
+        dialect: OcrDialectId::ReductoV3,
+        parameter_policy: &REJECT_CANONICAL_OCR_PARAMETER_POLICY,
+        document_policy: OcrDocumentPolicy::UploadUnlessProviderReference,
+        polling: None,
+    },
+    OcrDialectProfile {
+        dialect: OcrDialectId::ReductoLegacy,
+        parameter_policy: &REJECT_CANONICAL_OCR_PARAMETER_POLICY,
+        document_policy: OcrDocumentPolicy::UploadUnlessProviderReference,
+        polling: None,
+    },
+];
+
+pub const fn ocr_dialect_profile(dialect: OcrDialectId) -> &'static OcrDialectProfile {
+    match dialect {
+        OcrDialectId::Mistral => &OCR_DIALECT_PROFILES[0],
+        OcrDialectId::AzureFoundryMistral => &OCR_DIALECT_PROFILES[1],
+        OcrDialectId::AzureDocumentIntelligence => &OCR_DIALECT_PROFILES[2],
+        OcrDialectId::VertexMistral => &OCR_DIALECT_PROFILES[3],
+        OcrDialectId::VertexDeepSeek => &OCR_DIALECT_PROFILES[4],
+        OcrDialectId::ReductoV3 => &OCR_DIALECT_PROFILES[5],
+        OcrDialectId::ReductoLegacy => &OCR_DIALECT_PROFILES[6],
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ocr::policy::{OcrCanonicalField, ParameterDisposition};
+
+    #[test]
+    fn mistral_compatible_dialects_share_parameter_rules_but_not_document_rules() {
+        let mistral = ocr_dialect_profile(OcrDialectId::Mistral);
+        let foundry = ocr_dialect_profile(OcrDialectId::AzureFoundryMistral);
+        let vertex = ocr_dialect_profile(OcrDialectId::VertexMistral);
+
+        for field in OcrCanonicalField::ALL {
+            let expected = mistral.parameter_policy.disposition(field);
+            assert_eq!(foundry.parameter_policy.disposition(field), expected);
+            assert_eq!(vertex.parameter_policy.disposition(field), expected);
+        }
+        assert_eq!(mistral.document_policy, OcrDocumentPolicy::Ready);
+        assert_eq!(
+            foundry.document_policy,
+            OcrDocumentPolicy::FetchRemoteUrlAndInline
+        );
+        assert_eq!(
+            vertex.document_policy,
+            OcrDocumentPolicy::FetchRemoteUrlAndInline
+        );
+    }
+
+    #[test]
+    fn non_mistral_profiles_preserve_provider_specific_boundaries() {
+        let azure = ocr_dialect_profile(OcrDialectId::AzureDocumentIntelligence);
+        assert_eq!(
+            azure.parameter_policy.disposition(OcrCanonicalField::Pages),
+            ParameterDisposition::Transform
+        );
+        assert_eq!(
+            azure.polling,
+            Some(OcrPollingProfile {
+                operation_location_header: "operation-location"
+            })
+        );
+
+        for dialect in [OcrDialectId::ReductoV3, OcrDialectId::ReductoLegacy] {
+            let reducto = ocr_dialect_profile(dialect);
+            assert_eq!(
+                reducto.document_policy,
+                OcrDocumentPolicy::UploadUnlessProviderReference
+            );
+            assert!(OcrCanonicalField::ALL.iter().all(|field| {
+                reducto.parameter_policy.disposition(*field) == ParameterDisposition::Reject
+            }));
+        }
+    }
+
+    #[test]
+    fn every_dialect_has_exactly_one_profile() {
+        for (index, profile) in OCR_DIALECT_PROFILES.iter().enumerate() {
+            assert_eq!(ocr_dialect_profile(profile.dialect), profile);
+            assert!(
+                OCR_DIALECT_PROFILES[index + 1..]
+                    .iter()
+                    .all(|other| other.dialect != profile.dialect)
+            );
+        }
+    }
+}

litellm-rust/crates/core/src/ocr/wire.rs

@@ -0,0 +1,358 @@
+use std::collections::BTreeMap;
+use std::io::{self, Write};
+use std::str::Utf8Error;
+
+use base64::Engine;
+use base64::engine::general_purpose::STANDARD;
+use bytes::Bytes;
+use mime::Mime;
+use serde_json::Value;
+use thiserror::Error;
+
+const BASE64_INPUT_CHUNK_SIZE: usize = 48 * 1024;
+
+#[derive(Debug, Error)]
+pub enum OcrWireError {
+    #[error("failed to encode OCR JSON body: {0}")]
+    Json(#[from] serde_json::Error),
+    #[error("failed to write OCR body: {0}")]
+    Io(#[from] io::Error),
+    #[error("encoded OCR data URI is not UTF-8: {0}")]
+    InvalidDataUri(#[from] Utf8Error),
+}
+
+#[derive(Clone, PartialEq)]
+pub enum OcrJsonValue {
+    Value(Value),
+    Array(Vec<Self>),
+    Object(BTreeMap<String, Self>),
+    InlineDataUri { media_type: Mime, bytes: Bytes },
+    EncodedDataUri(Bytes),
+}
+
+impl OcrJsonValue {
+    fn write_to(&self, writer: &mut impl Write) -> Result<(), OcrWireError> {
+        match self {
+            Self::Value(value) => serde_json::to_writer(writer, value).map_err(Into::into),
+            Self::Array(values) => {
+                writer.write_all(b"[")?;
+                for (index, value) in values.iter().enumerate() {
+                    if index != 0 {
+                        writer.write_all(b",")?;
+                    }
+                    value.write_to(writer)?;
+                }
+                writer.write_all(b"]")?;
+                Ok(())
+            }
+            Self::Object(fields) => {
+                writer.write_all(b"{")?;
+                for (index, (key, value)) in fields.iter().enumerate() {
+                    if index != 0 {
+                        writer.write_all(b",")?;
+                    }
+                    serde_json::to_writer(&mut *writer, key)?;
+                    writer.write_all(b":")?;
+                    value.write_to(writer)?;
+                }
+                writer.write_all(b"}")?;
+                Ok(())
+            }
+            Self::InlineDataUri { media_type, bytes } => {
+                writer.write_all(b"\"data:")?;
+                writer.write_all(media_type.as_ref().as_bytes())?;
+                writer.write_all(b";base64,")?;
+                for chunk in bytes.chunks(BASE64_INPUT_CHUNK_SIZE) {
+                    let encoded = STANDARD.encode(chunk);
+                    writer.write_all(encoded.as_bytes())?;
+                }
+                writer.write_all(b"\"")?;
+                Ok(())
+            }
+            Self::EncodedDataUri(data_uri) => {
+                let data_uri = std::str::from_utf8(data_uri)?;
+                serde_json::to_writer(writer, data_uri).map_err(Into::into)
+            }
+        }
+    }
+}
+
+#[derive(Clone, PartialEq)]
+pub enum MultipartPart {
+    Text {
+        name: String,
+        value: String,
+    },
+    Json {
+        name: String,
+        value: Value,
+    },
+    File {
+        name: String,
+        file_name: String,
+        media_type: Mime,
+        bytes: Bytes,
+    },
+}
+
+#[derive(Clone, PartialEq)]
+pub struct MultipartBodyPlan {
+    boundary: String,
+    parts: Vec<MultipartPart>,
+}
+
+impl MultipartBodyPlan {
+    pub fn new(boundary: impl Into<String>, parts: Vec<MultipartPart>) -> Self {
+        Self {
+            boundary: boundary.into(),
+            parts,
+        }
+    }
+
+    pub fn boundary(&self) -> &str {
+        &self.boundary
+    }
+
+    pub fn parts(&self) -> &[MultipartPart] {
+        &self.parts
+    }
+
+    fn write_to(&self, writer: &mut impl Write) -> Result<(), OcrWireError> {
+        for part in &self.parts {
+            write!(writer, "--{}\r\n", self.boundary)?;
+            match part {
+                MultipartPart::Text { name, value } => {
+                    write!(
+                        writer,
+                        "Content-Disposition: form-data; name=\"{name}\"\r\n\r\n{value}\r\n"
+                    )?;
+                }
+                MultipartPart::Json { name, value } => {
+                    write!(
+                        writer,
+                        "Content-Disposition: form-data; name=\"{name}\"\r\nContent-Type: application/json\r\n\r\n"
+                    )?;
+                    serde_json::to_writer(&mut *writer, value)?;
+                    writer.write_all(b"\r\n")?;
+                }
+                MultipartPart::File {
+                    name,
+                    file_name,
+                    media_type,
+                    bytes,
+                } => {
+                    write!(
+                        writer,
+                        "Content-Disposition: form-data; name=\"{name}\"; filename=\"{file_name}\"\r\nContent-Type: {media_type}\r\n\r\n"
+                    )?;
+                    writer.write_all(bytes)?;
+                    writer.write_all(b"\r\n")?;
+                }
+            }
+        }
+        write!(writer, "--{}--\r\n", self.boundary)?;
+        Ok(())
+    }
+}
+
+#[derive(Clone, PartialEq)]
+pub enum OcrWireBody {
+    Json(Value),
+    JsonWithMedia(OcrJsonValue),
+    Multipart(MultipartBodyPlan),
+}
+
+impl OcrWireBody {
+    pub fn content_type(&self) -> String {
+        match self {
+            Self::Json(_) | Self::JsonWithMedia(_) => "application/json".to_string(),
+            Self::Multipart(plan) => {
+                format!("multipart/form-data; boundary={}", plan.boundary())
+            }
+        }
+    }
+
+    pub fn write_to(&self, mut writer: impl Write) -> Result<(), OcrWireError> {
+        match self {
+            Self::Json(value) => serde_json::to_writer(writer, value).map_err(Into::into),
+            Self::JsonWithMedia(value) => value.write_to(&mut writer),
+            Self::Multipart(plan) => plan.write_to(&mut writer),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io;
+    use std::sync::Arc;
+
+    use serde_json::json;
+
+    use super::*;
+
+    #[test]
+    fn ordinary_json_body_writes_without_a_media_plan() {
+        let body = OcrWireBody::Json(json!({"model": "ocr-model", "pages": [0, 2]}));
+        let mut encoded = Vec::new();
+
+        body.write_to(&mut encoded).expect("body writes");
+
+        assert_eq!(body.content_type(), "application/json");
+        assert_eq!(
+            serde_json::from_slice::<Value>(&encoded).expect("valid JSON"),
+            json!({"model": "ocr-model", "pages": [0, 2]})
+        );
+    }
+
+    #[test]
+    fn json_with_media_streams_raw_bytes_as_a_data_uri() {
+        let owner: Arc<[u8]> = vec![b'x'; BASE64_INPUT_CHUNK_SIZE + 1].into();
+        let bytes = Bytes::from_owner(Arc::clone(&owner));
+        let source_pointer = bytes.as_ptr();
+        let body = OcrWireBody::JsonWithMedia(OcrJsonValue::Object(BTreeMap::from([
+            (
+                "document".to_string(),
+                OcrJsonValue::InlineDataUri {
+                    media_type: "application/pdf".parse().expect("valid MIME type"),
+                    bytes,
+                },
+            ),
+            ("model".to_string(), OcrJsonValue::Value(json!("ocr-model"))),
+        ])));
+
+        let OcrWireBody::JsonWithMedia(OcrJsonValue::Object(fields)) = &body else {
+            panic!("media JSON body must retain its typed representation");
+        };
+        let OcrJsonValue::InlineDataUri { bytes, .. } = &fields["document"] else {
+            panic!("document must remain shared binary media");
+        };
+        assert_eq!(bytes.as_ptr(), source_pointer);
+        assert_eq!(Arc::strong_count(&owner), 2);
+
+        let mut encoded = Vec::new();
+        body.write_to(&mut encoded).expect("body writes");
+        let expected_data_uri = format!(
+            "data:application/pdf;base64,{}",
+            STANDARD.encode(owner.as_ref())
+        );
+        let expected = json!({"document": expected_data_uri, "model": "ocr-model"});
+        assert_eq!(
+            serde_json::from_slice::<Value>(&encoded).expect("valid JSON"),
+            expected
+        );
+    }
+
+    #[test]
+    fn raw_media_encoding_uses_bounded_writes() {
+        let bytes = Bytes::from(vec![b'x'; BASE64_INPUT_CHUNK_SIZE * 3 + 1]);
+        let body = OcrWireBody::JsonWithMedia(OcrJsonValue::InlineDataUri {
+            media_type: "application/pdf".parse().expect("valid MIME type"),
+            bytes,
+        });
+        let mut sink = BoundedSink {
+            maximum_write: BASE64_INPUT_CHUNK_SIZE * 4 / 3,
+            written: 0,
+        };
+
+        body.write_to(&mut sink).expect("writes remain bounded");
+
+        assert!(sink.written > BASE64_INPUT_CHUNK_SIZE * 4);
+    }
+
+    #[test]
+    fn encoded_data_uri_is_retained_without_decoding_or_copying() {
+        let data_uri = Bytes::from_static(b"data:image/png;base64,aGVsbG8=");
+        let source_pointer = data_uri.as_ptr();
+        let body = OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(data_uri));
+
+        let OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(retained)) = &body else {
+            panic!("encoded data URI must remain bytes");
+        };
+        assert_eq!(retained.as_ptr(), source_pointer);
+
+        let mut encoded = Vec::new();
+        body.write_to(&mut encoded).expect("body writes");
+        assert_eq!(encoded, br#""data:image/png;base64,aGVsbG8=""#);
+    }
+
+    #[test]
+    fn encoded_data_uri_is_json_escaped_without_changing_its_allocation() {
+        let data_uri = Bytes::from_static(b"data:text/plain,quoted%20\"value\"");
+        let source_pointer = data_uri.as_ptr();
+        let body = OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(data_uri));
+
+        let mut encoded = Vec::new();
+        body.write_to(&mut encoded).expect("body writes");
+
+        let OcrWireBody::JsonWithMedia(OcrJsonValue::EncodedDataUri(retained)) = &body else {
+            panic!("encoded data URI expected");
+        };
+        assert_eq!(retained.as_ptr(), source_pointer);
+        assert_eq!(
+            serde_json::from_slice::<String>(&encoded).expect("valid JSON string"),
+            "data:text/plain,quoted%20\"value\""
+        );
+    }
+
+    #[test]
+    fn multipart_file_is_replayable_and_retains_shared_bytes() {
+        let file = Bytes::from_static(b"large-pdf-payload");
+        let source_pointer = file.as_ptr();
+        let body = OcrWireBody::Multipart(MultipartBodyPlan::new(
+            "ocr-boundary",
+            vec![MultipartPart::File {
+                name: "file".to_string(),
+                file_name: "document.pdf".to_string(),
+                media_type: "application/pdf".parse().expect("valid MIME type"),
+                bytes: file,
+            }],
+        ));
+
+        let OcrWireBody::Multipart(plan) = &body else {
+            panic!("multipart body expected");
+        };
+        let MultipartPart::File { bytes, .. } = &plan.parts()[0] else {
+            panic!("file part expected");
+        };
+        assert_eq!(bytes.as_ptr(), source_pointer);
+
+        let mut first = Vec::new();
+        let mut retry = Vec::new();
+        body.write_to(&mut first).expect("first write succeeds");
+        body.write_to(&mut retry).expect("retry write succeeds");
+        assert_eq!(first, retry);
+        assert!(
+            first
+                .windows(file_name_marker().len())
+                .any(|window| window == file_name_marker())
+        );
+        assert!(
+            first
+                .windows(bytes.len())
+                .any(|window| window == bytes.as_ref())
+        );
+    }
+
+    fn file_name_marker() -> &'static [u8] {
+        b"filename=\"document.pdf\""
+    }
+
+    struct BoundedSink {
+        maximum_write: usize,
+        written: usize,
+    }
+
+    impl Write for BoundedSink {
+        fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
+            if buffer.len() > self.maximum_write {
+                return Err(io::Error::other("write exceeded bound"));
+            }
+            self.written += buffer.len();
+            Ok(buffer.len())
+        }
+
+        fn flush(&mut self) -> io::Result<()> {
+            Ok(())
+        }
+    }
+}