rust_llm/
ocr.rs

1//! Port of `lib/ruby_llm/ocr.rb` (`RubyLLM.ocr`), `Protocol#ocr`, and
2//! `lib/ruby_llm/providers/mistral/ocr.rb`, the only OCR seam among the providers RustLLM supports.
3//!
4//! ```ruby
5//! ocr = RubyLLM.ocr("contract.pdf")
6//! ocr.markdown             # => "# Contract\n\n..."
7//! ```
8
9use std::sync::Arc;
10
11use serde_json::{Value, json};
12
13use crate::attachment::{Attachment, AttachmentType};
14use crate::chat::{failure_tokens, resolve_model};
15use crate::config::Config;
16use crate::cost::{Cost, Tier};
17use crate::error::{Error, Result};
18use crate::message::{Operation, UsageEntry, UsageStatus};
19use crate::providers::Provider;
20use crate::tokens::Tokens;
21use crate::transport::Connection;
22
23/// One page of an OCR result (`RubyLLM::OCR::Page`).
24#[derive(Debug, Clone, PartialEq)]
25pub struct Page {
26    /// The zero-based page index.
27    pub index: i64,
28    pub markdown: Option<String>,
29    /// The images the provider reports for the page, or `None`.
30    pub images: Option<Value>,
31    /// The tables the provider reports for the page, or `None`.
32    pub tables: Option<Value>,
33    /// The provider's unmodified page hash.
34    pub raw: Value,
35}
36
37/// The text a document AI model extracted from a file (`RubyLLM::OCR`).
38#[derive(Debug, Clone)]
39pub struct Ocr {
40    pub pages: Vec<Page>,
41    /// The id of the model that performed the OCR.
42    pub model: String,
43    /// The provider's usage block, such as `pages_processed`, or `None`.
44    pub usage: Option<Value>,
45    /// The provider's raw response.
46    pub raw: Value,
47    /// One entry per provider attempt.
48    pub usage_entries: Vec<UsageEntry>,
49}
50
51impl Ocr {
52    /// `OCR.new(pages:, model:, usage:, raw:)`: page hashes become [`Page`]s, indexed by position
53    /// when the provider omits `index`.
54    pub fn new(pages: &[Value], model: impl Into<String>, usage: Option<Value>, raw: Value) -> Ocr {
55        let pages = pages
56            .iter()
57            .enumerate()
58            .map(|(position, page)| Page {
59                index: page
60                    .get("index")
61                    .and_then(Value::as_i64)
62                    .unwrap_or(position as i64),
63                markdown: page
64                    .get("markdown")
65                    .and_then(Value::as_str)
66                    .map(str::to_string),
67                images: page.get("images").filter(|v| !v.is_null()).cloned(),
68                tables: page.get("tables").filter(|v| !v.is_null()).cloned(),
69                raw: page.clone(),
70            })
71            .collect();
72        Ocr {
73            pages,
74            model: model.into(),
75            usage,
76            raw,
77            usage_entries: Vec::new(),
78        }
79    }
80
81    /// The markdown of every page, joined with blank lines.
82    pub fn markdown(&self) -> String {
83        self.pages
84            .iter()
85            .filter_map(|p| p.markdown.as_deref())
86            .collect::<Vec<_>>()
87            .join("\n\n")
88    }
89
90    /// Provider-reported usage across every attempt.
91    pub fn tokens(&self) -> Tokens {
92        Tokens::aggregate(self.usage_entries.iter().map(|e| &e.tokens))
93    }
94
95    /// The OCR cost across every attempt.
96    pub fn cost(&self) -> Cost {
97        let complete = self.usage_entries.iter().all(UsageEntry::cost_available);
98        Cost::aggregate(self.usage_entries.iter().map(|e| &e.cost), complete)
99    }
100}
101
102/// Options for [`ocr`], the keyword arguments of `OCR.ocr`.
103#[derive(Default)]
104pub struct OcrOptions<'a> {
105    /// `model:`, defaulting to `config.default_ocr_model`.
106    pub model: Option<&'a str>,
107    pub provider: Option<&'a str>,
108    pub assume_model_exists: bool,
109    /// `pages:`: zero-based page indexes to read.
110    pub pages: Option<Vec<i64>>,
111    /// `provider_options:`: merged into the request in the provider's vocabulary.
112    pub provider_options: Value,
113    /// `context:`: use this configuration instead of the global one.
114    pub config: Option<Arc<Config>>,
115    /// `metadata:`: added to the `ocr.rust_llm` event payload, never sent to the provider.
116    pub metadata: Option<Value>,
117    /// `owner:`: who the usage is attributed to, such as a user; wins over
118    /// [`crate::accounting::with_usage_owner`].
119    pub owner: Option<crate::accounting::UsageOwner>,
120}
121
122/// `RubyLLM.ocr(file, model:, provider:, pages:, provider_options:, metadata:)`, inside an
123/// `ocr.rust_llm` event. `file` is a path, URL, or [`Attachment`].
124pub async fn ocr(file: impl Into<Attachment>, options: OcrOptions<'_>) -> Result<Ocr> {
125    let config = options.config.clone().unwrap_or_else(crate::config);
126    let model_id = options
127        .model
128        .unwrap_or(&config.default_ocr_model)
129        .to_string();
130    let (model, provider) =
131        resolve_model(&model_id, options.provider, options.assume_model_exists)?;
132    let owner = options.owner.clone();
133    let mut event = crate::instrumentation::Event::start(&config, "ocr.rust_llm", || {
134        crate::instrumentation::payload([
135            ("provider", provider.slug().into()),
136            ("provider_class", provider.display().into()),
137            ("model", model.id.clone().into()),
138            ("pages", json!(options.pages)),
139            ("provider_options", options.provider_options.clone()),
140            (
141                "metadata",
142                crate::instrumentation::metadata(&options.metadata),
143            ),
144        ])
145    });
146    let result = event
147        .instrument(crate::accounting::owned_by(
148            owner,
149            ocr_inner(file.into(), options, config.clone(), model, provider),
150        ))
151        .await;
152    if let Ok(r) = &result {
153        crate::accounting::report(&config, &r.usage_entries).await;
154        event.set(
155            "result",
156            || json!({ "model": r.model, "pages": r.pages.len() }),
157        );
158        event.set("response_model", || r.model.clone().into());
159    }
160    event.finish(result.as_ref().err());
161    result
162}
163
164async fn ocr_inner(
165    file: Attachment,
166    options: OcrOptions<'_>,
167    config: Arc<Config>,
168    model: crate::model::Model,
169    provider: Provider,
170) -> Result<Ocr> {
171    let OcrOptions {
172        pages,
173        provider_options,
174        ..
175    } = options;
176    provider.ensure_configured(&config)?;
177    if provider != Provider::Mistral {
178        return Err(Error::Api(
179            format!("{} doesn't support OCR", provider.display()),
180            None,
181        ));
182    }
183    let connection = Connection::new(provider, config.clone())?;
184    let mut attachment = file;
185    if !attachment.is_url() {
186        attachment.load(connection.client()).await?;
187    }
188    let payload = render_payload(&attachment, &model.id, pages.as_deref(), &provider_options)?;
189
190    // `track_usage(:ocr)`: one entry per HTTP attempt.
191    let mut retried: Vec<Tokens> = Vec::new();
192    let mut on_attempt = |previous: Option<&Error>| {
193        if let Some(e) = previous {
194            retried.push(failure_tokens(e, None));
195        }
196    };
197    let raw = connection
198        .post("ocr", &payload, &[], &mut on_attempt)
199        .await?;
200    let data = raw.body;
201    let pages: Vec<Value> = data
202        .get("pages")
203        .and_then(Value::as_array)
204        .cloned()
205        .unwrap_or_default();
206    let response_model = data
207        .get("model")
208        .and_then(Value::as_str)
209        .unwrap_or(&model.id)
210        .to_string();
211    let usage = data.get("usage_info").filter(|v| !v.is_null()).cloned();
212    let mut result = Ocr::new(&pages, response_model, usage, data);
213    let entry = |status, tokens: Tokens| UsageEntry {
214        id: UsageEntry::next_id(),
215        owner: crate::accounting::usage_owner(),
216        operation: Operation::Ocr,
217        provider: provider.slug().into(),
218        model: model.id.clone(),
219        status,
220        cost: Cost::new(&tokens, Some(&model), Tier::Standard),
221        tokens,
222    };
223    let mut entries: Vec<UsageEntry> = retried
224        .into_iter()
225        .map(|t| entry(UsageStatus::Failed, t))
226        .collect();
227    // OCR responses report pages, not tokens, so the billed attempt carries empty tokens as in Ruby.
228    entries.push(entry(UsageStatus::Succeeded, Tokens::default()));
229    result.usage_entries = entries;
230    Ok(result)
231}
232
233/// `Mistral::OCR#render_ocr_payload`: remote files go as URLs, local ones as data URIs; images
234/// use the `image_url` document variant.
235fn render_payload(
236    attachment: &Attachment,
237    model: &str,
238    pages: Option<&[i64]>,
239    provider_options: &Value,
240) -> Result<Value> {
241    let reference = attachment.url_or_data_uri()?;
242    let document = if attachment.kind() == AttachmentType::Image {
243        json!({ "type": "image_url", "image_url": reference })
244    } else {
245        json!({ "type": "document_url", "document_url": reference })
246    };
247    let mut payload = json!({ "model": model, "document": document });
248    if let Some(pages) = pages {
249        payload["pages"] = json!(pages);
250    }
251    if let (Some(p), Some(options)) = (payload.as_object_mut(), provider_options.as_object()) {
252        p.extend(options.clone());
253    }
254    Ok(payload)
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260
261    // ocr_spec.rb: "continues converting hash pages with their original metadata".
262    #[test]
263    fn pages_without_an_index_take_their_position_and_keep_raw() {
264        let raw = json!({ "markdown": "# Ruby", "images": [], "tables": [] });
265        let ocr = Ocr::new(
266            std::slice::from_ref(&raw),
267            "mistral-ocr-latest",
268            None,
269            Value::Null,
270        );
271        assert_eq!(ocr.pages[0].index, 0);
272        assert_eq!(ocr.pages[0].markdown.as_deref(), Some("# Ruby"));
273        assert_eq!(ocr.pages[0].images, Some(json!([])));
274        assert_eq!(ocr.pages[0].raw, raw);
275        assert_eq!(ocr.markdown(), "# Ruby");
276    }
277}
rust_llm/ocr.rs

rust_llm/
ocr.rs
