use std::borrow::Cow;
use std::sync::LazyLock;
use async_trait::async_trait;
use base64::Engine;
use liter_llm::types::ContentPart;
use liter_llm::{ChatCompletionRequest, ImageUrl, LlmClient, Message, UserContent, UserMessage};
use regex::Regex;
use crate::core::config::LlmConfig;
use crate::plugins::{OcrBackend, OcrBackendType, Plugin};
use crate::types::{FormatMetadata, Formula, Metadata, OcrMetadata, Table};
use crate::ocr_metadata_keys::OCR_PROCESSED_IMAGE_HEIGHT_METADATA_KEY as PROCESSED_HEIGHT_KEY;
use crate::ocr_metadata_keys::OCR_PROCESSED_IMAGE_WIDTH_METADATA_KEY as PROCESSED_WIDTH_KEY;
const DEFAULT_VLM_TIMEOUT_SECS: u64 = 300;
fn effective_vlm_config(config: &LlmConfig) -> Cow<'_, LlmConfig> {
if config.timeout_secs.is_none() {
let mut owned = config.clone();
owned.timeout_secs = Some(DEFAULT_VLM_TIMEOUT_SECS);
Cow::Owned(owned)
} else {
Cow::Borrowed(config)
}
}
#[cfg_attr(alef, alef(skip))]
pub struct VlmOcrBackend;
impl Plugin for VlmOcrBackend {
fn name(&self) -> &str {
"vlm"
}
fn version(&self) -> String {
env!("CARGO_PKG_VERSION").to_string()
}
fn initialize(&self) -> crate::Result<()> {
Ok(())
}
fn shutdown(&self) -> crate::Result<()> {
Ok(())
}
}
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
impl OcrBackend for VlmOcrBackend {
async fn process_image(
&self,
image_bytes: &[u8],
config: &crate::OcrConfig,
) -> crate::Result<crate::ExtractedDocument> {
let vlm_config = config
.vlm_config
.as_ref()
.ok_or_else(|| crate::XbergError::validation("VLM OCR requires vlm_config to be set"))?;
let mime = infer::get(image_bytes).map(|t| t.mime_type()).unwrap_or("image/png");
let languages = config.effective_languages();
let lang_str = languages[0].as_str();
let (text, usage) = vlm_ocr(image_bytes, mime, lang_str, vlm_config, config.vlm_prompt.as_deref()).await?;
let formulas = extract_formulas(&text);
let tables = extract_gfm_tables(&text);
let metadata = build_metadata(image_bytes, tables.len() as u32);
Ok(crate::ExtractedDocument {
content: text,
mime_type: Cow::Borrowed("text/plain"),
llm_usage: usage.map(|u| vec![u]),
formulas,
metadata,
tables,
detected_languages: Some(languages),
..Default::default()
})
}
fn supports_language(&self, _lang: &str) -> bool {
true
}
#[cfg_attr(alef, alef(skip))]
fn backend_type(&self) -> OcrBackendType {
OcrBackendType::Custom
}
fn confidence_semantics(&self) -> crate::plugins::ConfidenceSemantics {
crate::plugins::ConfidenceSemantics::None
}
#[cfg_attr(alef, alef(skip))]
fn probe(&self, config: &crate::OcrConfig) -> crate::doctor::DoctorCheck {
use crate::doctor::DoctorCheck;
let Some(vlm_config) = config.vlm_config.as_ref() else {
return DoctorCheck::fail(
"ocr.vlm",
"ocr.backend = \"vlm\" requires vlm_config (model endpoint configuration)",
);
};
if let Err(e) = super::client::create_client(vlm_config) {
return DoctorCheck::fail("ocr.vlm", format!("{e}"));
}
DoctorCheck::skip(
"ocr.vlm",
"configuration and credentials are valid; endpoint reachability was not checked",
)
}
}
#[allow(clippy::field_reassign_with_default)]
pub(crate) async fn vlm_ocr(
image_bytes: &[u8],
image_mime_type: &str,
language: &str,
config: &LlmConfig,
vlm_prompt: Option<&str>,
) -> crate::Result<(String, Option<crate::types::LlmUsage>)> {
let effective_config = effective_vlm_config(config);
let client = super::client::create_client(&effective_config)?;
let b64 = base64::engine::general_purpose::STANDARD.encode(image_bytes);
let data_url = format!("data:{image_mime_type};base64,{b64}");
let template = vlm_prompt.unwrap_or(super::prompts::VLM_OCR_TEMPLATE);
let ctx = minijinja::context! { language => language };
let prompt = super::prompts::render_template(template, &ctx)?;
let message = Message::User(UserMessage {
content: UserContent::Parts(vec![
ContentPart::Text { text: prompt },
ContentPart::ImageUrl {
image_url: ImageUrl {
url: data_url,
detail: None,
},
},
]),
name: None,
});
let mut request = ChatCompletionRequest::default();
request.model = normalize_vlm_model(&config.model, config.base_url.as_deref());
request.messages = vec![message];
super::client::apply_request_time_params(&mut request, config)?;
let response = client.chat(request).await.map_err(|e| {
crate::XbergError::ocr(format!(
"VLM OCR request failed: model={}, language={}, image_size={}KB: {e}",
config.model,
language,
image_bytes.len() / 1024
))
})?;
let usage = super::usage::extract_usage_from_chat(&response, "vlm_ocr");
let text = response
.choices
.first()
.and_then(|choice| choice.message.content.as_ref().and_then(|m| m.as_text()))
.ok_or_else(|| crate::XbergError::ocr(format!("VLM OCR returned no content (model={})", config.model)))?;
Ok((text, usage))
}
fn normalize_vlm_model(model: &str, base_url: Option<&str>) -> String {
if let Some(url) = base_url
&& url.contains("openai.com")
&& let Some(stripped) = model.strip_prefix("openai/")
{
return stripped.to_string();
}
model.to_string()
}
static FORMULA_PATTERN: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?s)\$\$(?P<dollar>.+?)\$\$|\\\[(?P<bracket>.+?)\\\]|\\\((?P<paren>.+?)\\\)|\$(?P<inline>[^\s$](?:[^$\n]*[^\s$])?)\$",
)
.expect("VLM formula regex pattern is valid and should compile")
});
fn build_metadata(image_bytes: &[u8], table_count: u32) -> Metadata {
let mut metadata = Metadata {
format: Some(FormatMetadata::Ocr(OcrMetadata {
table_count,
..Default::default()
})),
ocr_used: true,
..Default::default()
};
if let Some((width, height)) = probe_image_dimensions(image_bytes) {
metadata
.additional
.insert(Cow::Borrowed(PROCESSED_WIDTH_KEY), serde_json::json!(width));
metadata
.additional
.insert(Cow::Borrowed(PROCESSED_HEIGHT_KEY), serde_json::json!(height));
}
metadata
}
fn probe_image_dimensions(image_bytes: &[u8]) -> Option<(u32, u32)> {
crate::extraction::image_decode::probe_standard_image_with_default_security_limits(image_bytes)
.ok()
.map(|(width, height, _)| (width, height))
}
fn extract_gfm_tables(text: &str) -> Vec<Table> {
use pulldown_cmark::{Event, Options, Parser, Tag, TagEnd};
let mut tables = Vec::new();
let mut in_table = false;
let mut rows: Vec<Vec<String>> = Vec::new();
let mut current_row: Vec<String> = Vec::new();
let mut current_cell = String::new();
let mut in_cell = false;
for event in Parser::new_ext(text, Options::ENABLE_TABLES) {
match event {
Event::Start(Tag::Table(_)) => {
in_table = true;
rows.clear();
}
Event::End(TagEnd::Table) if in_table => {
in_table = false;
if !rows.is_empty() {
let cells = std::mem::take(&mut rows);
let markdown = crate::rendering::common::render_table_markdown(&cells);
tables.push(Table {
cells,
markdown,
page_number: 1,
..Default::default()
});
}
}
Event::Start(Tag::TableHead | Tag::TableRow) if in_table => {
current_row.clear();
}
Event::End(TagEnd::TableHead | TagEnd::TableRow) if in_table && !current_row.is_empty() => {
rows.push(std::mem::take(&mut current_row));
}
Event::Start(Tag::TableCell) if in_table => {
in_cell = true;
current_cell.clear();
}
Event::End(TagEnd::TableCell) if in_table => {
in_cell = false;
current_row.push(current_cell.trim().to_string());
current_cell.clear();
}
Event::Text(cell_text) | Event::Code(cell_text) if in_table && in_cell => {
current_cell.push_str(&cell_text);
}
_ => {}
}
}
tables
}
fn extract_formulas(text: &str) -> Vec<Formula> {
FORMULA_PATTERN
.captures_iter(text)
.filter_map(|caps| {
caps.name("dollar")
.or_else(|| caps.name("bracket"))
.or_else(|| caps.name("paren"))
.or_else(|| caps.name("inline"))
})
.map(|m| m.as_str().trim())
.filter(|latex| !latex.is_empty())
.map(|latex| Formula {
latex: latex.to_string(),
bbox: None,
page: None,
})
.collect()
}
#[cfg(test)]
mod tests {
fn render_ocr_prompt(language: &str) -> String {
let ctx = minijinja::context! { language => language };
super::super::prompts::render_template(super::super::prompts::VLM_OCR_TEMPLATE, &ctx).unwrap()
}
#[test]
fn test_vlm_ocr_prompt_non_english_includes_language() {
let prompt = render_ocr_prompt("deu");
assert!(prompt.contains("language: deu"));
}
#[test]
fn test_vlm_ocr_prompt_english_no_language_hint() {
let prompt = render_ocr_prompt("eng");
assert!(!prompt.contains("language:"));
}
#[test]
fn test_vlm_ocr_prompt_en_no_language_hint() {
let prompt = render_ocr_prompt("en");
assert!(!prompt.contains("language:"));
}
#[test]
fn test_normalize_vlm_model_strips_prefix_for_openai_endpoint() {
assert_eq!(
super::normalize_vlm_model("openai/gpt-4o-mini", Some("https://eu.api.openai.com/v1/")),
"gpt-4o-mini"
);
}
#[test]
fn test_normalize_vlm_model_keeps_prefix_without_base_url() {
assert_eq!(
super::normalize_vlm_model("openai/gpt-4o-mini", None),
"openai/gpt-4o-mini"
);
}
#[test]
fn test_normalize_vlm_model_keeps_prefix_for_gateway_base_url() {
assert_eq!(
super::normalize_vlm_model("openai/gpt-4o", Some("https://openrouter.ai/api/v1")),
"openai/gpt-4o"
);
}
#[test]
fn test_normalize_vlm_model_leaves_unprefixed_model_untouched() {
assert_eq!(
super::normalize_vlm_model("gpt-4o-mini", Some("https://eu.api.openai.com/v1/")),
"gpt-4o-mini"
);
}
#[test]
fn test_effective_vlm_config_applies_default_timeout_when_unset() {
let config = crate::core::config::LlmConfig {
model: "openai/gpt-4o".to_string(),
..Default::default()
};
assert!(config.timeout_secs.is_none());
let effective = super::effective_vlm_config(&config);
assert_eq!(effective.timeout_secs, Some(super::DEFAULT_VLM_TIMEOUT_SECS));
assert_ne!(effective.timeout_secs, Some(60));
}
#[test]
fn test_effective_vlm_config_preserves_explicit_timeout() {
let config = crate::core::config::LlmConfig {
model: "openai/gpt-4o".to_string(),
timeout_secs: Some(1200),
..Default::default()
};
let effective = super::effective_vlm_config(&config);
assert_eq!(effective.timeout_secs, Some(1200));
}
#[test]
fn test_vlm_prompt_custom_template_is_used_issue_760() {
let custom_prompt = "Extract all text from this document image. \
Preserve formatting and use latex for mathematical formulas.";
let ctx = minijinja::context! { language => "eng" };
let prompt = super::super::prompts::render_template(custom_prompt, &ctx).unwrap();
assert!(prompt.contains("latex"), "custom prompt must be used; got: {prompt}");
assert!(
prompt.contains("Preserve formatting"),
"custom prompt must be used; got: {prompt}"
);
assert!(
!prompt.contains("Extract all visible text"),
"default template must NOT be used when custom prompt is set; got: {prompt}"
);
}
#[test]
fn test_extract_formulas_strips_delimiters_from_vlm_response() {
let text = "The quadratic formula is:\n\n$$x = \\frac{-b \\pm \\sqrt{b^2 - 4ac}}{2a}$$\n\nDone.";
let formulas = super::extract_formulas(text);
assert_eq!(formulas.len(), 1, "expected exactly one formula; got: {formulas:?}");
assert_eq!(formulas[0].latex, r"x = \frac{-b \pm \sqrt{b^2 - 4ac}}{2a}");
assert!(!formulas[0].latex.contains("$$"), "delimiters must be stripped");
assert_eq!(formulas[0].bbox, None, "VLM formulas carry no geometry");
assert_eq!(formulas[0].page, None, "VLM formulas carry no page");
}
#[test]
fn test_extract_formulas_handles_multiple_matches() {
let text = "$$a^2 + b^2 = c^2$$ and also $$E = mc^2$$";
let formulas = super::extract_formulas(text);
assert_eq!(formulas.len(), 2, "expected two formulas; got: {formulas:?}");
assert_eq!(formulas[0].latex, "a^2 + b^2 = c^2");
assert_eq!(formulas[1].latex, "E = mc^2");
}
#[test]
fn test_extract_formulas_returns_empty_when_no_math_present() {
let formulas = super::extract_formulas("Just a plain paragraph with no equations.");
assert!(formulas.is_empty(), "expected no formulas; got: {formulas:?}");
}
#[test]
fn test_extract_formulas_extracts_inline_dollar_math() {
let text = "The area is $A = \\pi r^2$ for a circle.";
let formulas = super::extract_formulas(text);
assert_eq!(formulas.len(), 1, "expected exactly one formula; got: {formulas:?}");
assert_eq!(formulas[0].latex, r"A = \pi r^2");
assert_eq!(formulas[0].bbox, None, "VLM formulas carry no geometry");
assert_eq!(formulas[0].page, None, "VLM formulas carry no page");
}
#[test]
fn test_extract_formulas_extracts_inline_paren_math() {
let text = r"Euler's identity: \(e^{i\pi} + 1 = 0\).";
let formulas = super::extract_formulas(text);
assert_eq!(formulas.len(), 1, "expected exactly one formula; got: {formulas:?}");
assert_eq!(formulas[0].latex, r"e^{i\pi} + 1 = 0");
}
#[test]
fn test_extract_formulas_ignores_currency_dollar_signs() {
let text = "The item costs $5 and the other costs $10, for $15 total.";
let formulas = super::extract_formulas(text);
assert!(
formulas.is_empty(),
"currency amounts must not be treated as formulas; got: {formulas:?}"
);
}
#[test]
fn test_extract_formulas_ignores_single_unpaired_dollar() {
let formulas = super::extract_formulas("Prices start at $20 per unit.");
assert!(
formulas.is_empty(),
"unpaired dollar sign must not match; got: {formulas:?}"
);
}
#[test]
fn test_extract_formulas_prefers_display_over_inline_when_both_present() {
let text = "$$x^2 + y^2 = z^2$$ and separately $a+b$.";
let formulas = super::extract_formulas(text);
assert_eq!(
formulas.len(),
2,
"expected one display + one inline; got: {formulas:?}"
);
assert_eq!(formulas[0].latex, "x^2 + y^2 = z^2");
assert_eq!(formulas[1].latex, "a+b");
}
#[test]
fn test_extract_gfm_tables_from_vlm_response() {
let text = "Some text.\n\n| Name | Age |\n|------|-----|\n| Alice | 30 |\n\nMore text.";
let tables = super::extract_gfm_tables(text);
assert_eq!(tables.len(), 1, "expected exactly one table; got: {tables:?}");
assert_eq!(
tables[0].cells,
vec![
vec!["Name".to_string(), "Age".to_string()],
vec!["Alice".to_string(), "30".to_string()],
]
);
assert_eq!(tables[0].page_number, 1);
}
#[test]
fn test_build_metadata_reports_ocr_used_and_table_count() {
let metadata = super::build_metadata(&[], 2);
assert!(metadata.ocr_used, "ocr_used must be true for a VLM OCR result");
let Some(super::FormatMetadata::Ocr(ocr_metadata)) = metadata.format else {
panic!("expected FormatMetadata::Ocr; got: {:?}", metadata.format);
};
assert_eq!(ocr_metadata.table_count, 2);
}
#[test]
fn test_vlm_prompt_none_falls_back_to_default() {
let ctx = minijinja::context! { language => "eng" };
let prompt = super::super::prompts::render_template(super::super::prompts::VLM_OCR_TEMPLATE, &ctx).unwrap();
assert!(
prompt.contains("Extract all visible text"),
"default template must be used when vlm_prompt is None; got: {prompt}"
);
}
}