diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 1a72b8f1d66..656fd36e6d4 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,29 +1,187 @@ +//! Typed OCR domain model. +//! +//! These types mirror the Python OCR surface (`litellm/ocr/main.py` and the +//! `OCRResponse` pydantic models) one-to-one, but as real Rust types rather than +//! free-form JSON. Serde derives produce the exact wire JSON, so no provider +//! ever hand-builds or hand-parses JSON. +//! +//! `serde_json::Value` appears in exactly three leaves, and only where Python's +//! own types are `Any` / `Dict[str, Any]`: [`AnnotationFormat`] (a user-supplied +//! JSON Schema we forward verbatim), [`OcrResponse::document_annotation`] +//! (`Optional[Any]`), and [`OcrPageImage::bbox`] (`Optional[Dict[str, Any]]`). +//! Everywhere Python is concrete, this module is concrete. +//! +//! Unknown upstream fields are dropped: serde ignores unrecognized keys by +//! default, which is the deliberate "fully strict" behavior (vs Python's +//! `model_config = {"extra": "allow"}` passthrough). + +use std::collections::BTreeMap; +use std::time::Duration; + use serde::{Deserialize, Serialize}; use serde_json::Value; -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct OcrRequestData { - pub data: Value, - pub files: Option, +/// Which OCR provider to dispatch to. +/// +/// The Python shell resolves `(custom_llm_provider, model)` into a precise +/// variant before crossing the bridge — e.g. an `azure_ai` model that is a +/// Document Intelligence model becomes [`OcrProvider::AzureDocumentIntelligence`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum OcrProvider { + Mistral, + AzureAi, + AzureDocumentIntelligence, + VertexAi, + VertexDeepseek, + Reducto, } +/// The document to OCR. +/// +/// `type: "file"` inputs are converted to a base64 data-URI in the Python shell, +/// so the core only ever sees a URL or image URL. Serializes to the provider +/// wire shape, e.g. `{"type": "document_url", "document_url": "…"}`. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct OcrResponseData { - pub pages: Vec, +#[serde(tag = "type", rename_all = "snake_case")] +pub enum OcrDocument { + DocumentUrl { document_url: String }, + ImageUrl { image_url: String }, +} + +/// Extra HTTP headers (Python `extra_headers`). Header values are strings on the +/// wire; `BTreeMap` keeps a stable, testable ordering. +pub type OcrHeaders = BTreeMap; + +/// Table output format (Mistral `table_format`). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TableFormat { + Markdown, + Html, +} + +/// Confidence-score granularity (Mistral `confidence_scores_granularity`). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ConfidenceGranularity { + Word, + Page, +} + +/// A provider "response format" (a user-supplied JSON Schema, forwarded verbatim +/// to the provider). A JSON Schema is arbitrary JSON by definition, so +/// `serde_json::Value` is its precise type — mirrors Python's `Dict[str, Any]`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct AnnotationFormat(pub Value); + +/// OCR options — the superset of provider params (Python `**kwargs`). Every field +/// is optional; serde skips `None`, so each provider's request body carries only +/// the keys that were actually set. +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrParams { + #[serde(skip_serializing_if = "Option::is_none")] + pub pages: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub include_image_base64: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub image_limit: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub image_min_size: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub bbox_annotation_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub document_annotation_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub document_annotation_prompt: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extract_header: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub extract_footer: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub table_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub confidence_scores_granularity: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, +} + +/// A fully-typed OCR request — maps argument-for-argument to `litellm.ocr()` / +/// `litellm.aocr()`, grouped into a struct (idiomatic Rust; also keeps the entry +/// point under clippy's `too_many_arguments` ceiling). +#[derive(Clone, Debug)] +pub struct OcrRequest { + pub provider: OcrProvider, pub model: String, + pub document: OcrDocument, + pub api_key: Option, + pub api_base: Option, + pub extra_headers: OcrHeaders, + pub timeout: Option, + pub params: OcrParams, +} + +/// Page dimensions (Python `OCRPageDimensions`). +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrPageDimensions { + #[serde(skip_serializing_if = "Option::is_none")] + pub dpi: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub height: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub width: Option, +} + +/// An image extracted from a page (Python `OCRPageImage`). `bbox` is +/// `Optional[Dict[str, Any]]` in Python, so it stays a `Value`. +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrPageImage { + #[serde(skip_serializing_if = "Option::is_none")] + pub image_base64: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub bbox: Option, +} + +/// A single OCR page (Python `OCRPage`). +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrPage { + pub index: u32, + pub markdown: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub images: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub dimensions: Option, +} + +/// Usage information (Python `OCRUsageInfo`). +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrUsageInfo { + #[serde(skip_serializing_if = "Option::is_none")] + pub pages_processed: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub credits: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub doc_size_bytes: Option, +} + +fn default_ocr_object() -> String { + "ocr".to_string() +} + +/// The standardized OCR response (Python `OCRResponse`). `document_annotation` is +/// `Optional[Any]` in Python (it depends on the user's annotation schema), so it +/// stays a `Value`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrResponse { + #[serde(default)] + pub pages: Vec, + #[serde(default)] + pub model: String, + #[serde(default, skip_serializing_if = "Option::is_none")] pub document_annotation: Option, - pub usage_info: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usage_info: Option, + #[serde(default = "default_ocr_object")] pub object: String, } - -impl OcrResponseData { - pub fn into_json(self) -> Value { - serde_json::json!({ - "pages": self.pages, - "model": self.model, - "document_annotation": self.document_annotation, - "usage_info": self.usage_info, - "object": self.object, - }) - } -}