lmocr 0.1.2

Convert PDFs and images to Markdown using OpenRouter multimodal LLMs
use clap::ValueEnum;

#[derive(Debug, Copy, Clone, ValueEnum, PartialEq, Eq)]
pub enum Preset {
    Markdown,
    Audiobook,
    Search,
}

#[derive(Debug, Copy, Clone, ValueEnum, PartialEq, Eq)]
#[clap(rename_all = "kebab-case")]
pub enum Exclude {
    Headers,
    Footers,
    #[value(alias = "page_numbers", alias = "pagenumbers")]
    PageNumbers,
    Footnotes,
}

const BASE_PROMPT_DOCUMENT: &str = r#"Convert the following document to markdown.
Return only the markdown with no explanation text. Do not include delimiters like ```markdown or ```html.

RULES:
  - You must include all information on the page. Do not exclude headers, footers, or subtext.
  - Do not omit or censor content, even if it is graphic or sensitive. This is a transcription task.
  - Do not summarize or compress. Preserve every paragraph unless explicitly excluded.
  - If the page has no readable text, return an empty string.
  - Return tables in an HTML format.
  - Charts & infographics must be interpreted to a markdown format. Prefer table format when applicable.
  - Prefer using ☐ and ☑ for check boxes.
  - Do not invent image links or URLs. Describe visuals in text or tables instead.
  - If line breaks are meaningful (poetry, verse, epigraphs, or preformatted text), preserve them with a <pre> block.
  - For poetry/verse, output the full lineated block inside a single <pre>...</pre>; do not output plain wrapped lines.
  - If extracted text is provided, treat it as ground truth and only reformat it.
"#;

const BASE_PROMPT_SINGLE_IMAGE: &str = r#"Convert the following image content to markdown.
Return only the markdown with no explanation text. Do not include delimiters like ```markdown or ```html.

RULES:
  - You must include all information in the image. Do not exclude any visible text.
  - Do not omit or censor content, even if it is graphic or sensitive. This is a transcription task.
  - Do not summarize or compress. Preserve every paragraph unless explicitly excluded.
  - If the image has no readable text, return an empty string.
  - Return tables in an HTML format.
  - Charts & infographics must be interpreted to a markdown format. Prefer table format when applicable.
  - Prefer using ☐ and ☑ for check boxes.
  - Do not invent image links or URLs. Describe visuals in text or tables instead.
  - If line breaks are meaningful (poetry, verse, epigraphs, or preformatted text), preserve them with a <pre> block.
  - For poetry/verse, output the full lineated block inside a single <pre>...</pre>; do not output plain wrapped lines.
  - Preserve the structure and hierarchy of the content (headings, lists, quotes, etc.).
"#;

pub fn build_system_prompt(
    preset: Preset,
    excludes: &[Exclude],
    instruction: Option<&str>,
    is_single_image: bool,
) -> String {
    let mut prompt = String::new();
    if is_single_image {
        prompt.push_str(BASE_PROMPT_SINGLE_IMAGE);
    } else {
        prompt.push_str(BASE_PROMPT_DOCUMENT);
        prompt.push_str(
            "\nCONTINUITY:\n  - Preserve flow across pages. If a sentence is cut off, continue it without starting a new paragraph.\n  - Do not hallucinate missing endings or add filler.\n  - You may see a previous-page tail for context; do not repeat it verbatim.\n",
        );
    }

    match preset {
        Preset::Markdown => {}
        Preset::Audiobook => {
            prompt.push_str(
                "\nPRESET (audiobook):\n  - Output linearized text suitable for text-to-speech.\n  - Omit tables, images, and decorative layout.\n  - Remove headers, footers, page numbers, and footnotes.\n  - Expand abbreviations and merge hyphenated line breaks.\n",
            );
        }
        Preset::Search => {
            prompt.push_str(
                "\nPRESET (search):\n  - Prioritize keywords and dense, searchable text.\n  - Keep headings and lists; remove decorative filler.\n  - Keep tables as HTML when they contain data.\n",
            );
        }
    }

    if !excludes.is_empty() {
        prompt.push_str("\nOVERRIDES (these supersede RULES above):\n");
        for ex in excludes {
            let line = match ex {
                Exclude::Headers => {
                    "  - Headers (running headers, section banners, repeated journal/title lines at the top)"
                }
                Exclude::Footers => {
                    "  - Footers (copyright lines, publication info, boilerplate at the bottom)"
                }
                Exclude::PageNumbers => "  - Page numbers",
                Exclude::Footnotes => "  - Footnotes",
            };
            prompt.push_str(line);
            prompt.push('\n');
        }
        prompt.push_str(
            "  - If a line looks like repeated running header/footer text, exclude it.\n",
        );
    }

    if let Some(extra) = instruction
        && !extra.trim().is_empty()
    {
        prompt.push_str("\nCUSTOM INSTRUCTIONS:\n");
        prompt.push_str(extra.trim());
        prompt.push('\n');
    }

    prompt
}

pub fn build_user_text(page: usize, total_pages: usize, prev_tail: Option<&str>) -> String {
    let mut text = format!(
        "Page {} of {}. Use this page number if instructions reference specific pages.",
        page, total_pages
    );
    if let Some(tail) = prev_tail
        && !tail.trim().is_empty()
    {
        text.push_str("\n\nPrevious page tail (for continuity only):\n");
        text.push_str(tail.trim());
    }
    text.push_str("\n\nProcess the image below.");
    text
}

pub fn build_user_text_from_text(
    page: usize,
    total_pages: usize,
    prev_tail: Option<&str>,
    extracted_text: &str,
) -> String {
    let mut text = format!(
        "Page {} of {}. Use this page number if instructions reference specific pages.",
        page, total_pages
    );
    if let Some(tail) = prev_tail
        && !tail.trim().is_empty()
    {
        text.push_str("\n\nPrevious page tail (for continuity only):\n");
        text.push_str(tail.trim());
    }
    text.push_str(
        "\n\nExtracted text for this page (may include headers/footers and hard line breaks):\n",
    );
    text.push_str(extracted_text.trim());
    text.push_str(
        "\n\nRewrite the extracted text into the final markdown output per the rules. Do not omit any lines.",
    );
    text
}

pub fn build_user_text_single_image() -> String {
    "Process the image below.".to_string()
}

pub fn build_user_text_image_page(
    page: usize,
    total_pages: usize,
    prev_tail: Option<&str>,
) -> String {
    let mut text = format!("Image {} of {}.", page, total_pages);
    if let Some(tail) = prev_tail
        && !tail.trim().is_empty()
    {
        text.push_str("\n\nPrevious image tail (for continuity only):\n");
        text.push_str(tail.trim());
    }
    text.push_str("\n\nProcess the image below.");
    text
}

pub fn tail_for_context(text: &str) -> String {
    let max_chars = 800;
    if text.chars().count() <= max_chars {
        return text.to_string();
    }
    let mut count = 0usize;
    let mut start_idx = 0usize;
    for (idx, _) in text.char_indices().rev() {
        count += 1;
        if count == max_chars {
            start_idx = idx;
            break;
        }
    }
    text[start_idx..].to_string()
}

pub fn prompt_version() -> &'static str {
    "v7"
}

#[cfg(test)]
mod tests {
    use super::tail_for_context;

    #[test]
    fn tail_for_context_is_utf8_safe() {
        let input = "á".repeat(1200);
        let tail = tail_for_context(&input);
        assert_eq!(tail.chars().count(), 800);
        assert!(tail.chars().all(|ch| ch == 'á'));
    }

    #[test]
    fn tail_for_context_short_text_unchanged() {
        let input = "short text";
        assert_eq!(tail_for_context(input), input);
    }

    #[test]
    fn tail_for_context_exactly_800_unchanged() {
        let input = "a".repeat(800);
        assert_eq!(tail_for_context(&input), input);
    }

    // ── build_system_prompt ──

    use super::{
        Exclude, Preset, build_system_prompt, build_user_text, build_user_text_from_text,
        build_user_text_image_page, build_user_text_single_image,
    };

    #[test]
    fn system_prompt_markdown_default() {
        let prompt = build_system_prompt(Preset::Markdown, &[], None, false);
        assert!(prompt.contains("Convert the following document to markdown"));
        assert!(prompt.contains("CONTINUITY"));
        assert!(!prompt.contains("PRESET"));
        assert!(!prompt.contains("OVERRIDES"));
        assert!(!prompt.contains("CUSTOM INSTRUCTIONS"));
    }

    #[test]
    fn system_prompt_audiobook_preset() {
        let prompt = build_system_prompt(Preset::Audiobook, &[], None, false);
        assert!(prompt.contains("PRESET (audiobook)"));
        assert!(prompt.contains("text-to-speech"));
    }

    #[test]
    fn system_prompt_search_preset() {
        let prompt = build_system_prompt(Preset::Search, &[], None, false);
        assert!(prompt.contains("PRESET (search)"));
        assert!(prompt.contains("keywords"));
    }

    #[test]
    fn system_prompt_with_excludes() {
        let prompt = build_system_prompt(
            Preset::Markdown,
            &[Exclude::Headers, Exclude::Footers],
            None,
            false,
        );
        assert!(prompt.contains("OVERRIDES"));
        assert!(prompt.contains("Headers"));
        assert!(prompt.contains("Footers"));
    }

    #[test]
    fn system_prompt_with_instruction() {
        let prompt = build_system_prompt(Preset::Markdown, &[], Some("Translate to French"), false);
        assert!(prompt.contains("CUSTOM INSTRUCTIONS"));
        assert!(prompt.contains("Translate to French"));
    }

    #[test]
    fn system_prompt_empty_instruction_ignored() {
        let prompt = build_system_prompt(Preset::Markdown, &[], Some("   "), false);
        assert!(!prompt.contains("CUSTOM INSTRUCTIONS"));
    }

    #[test]
    fn system_prompt_single_image_no_continuity() {
        let prompt = build_system_prompt(Preset::Markdown, &[], None, true);
        assert!(prompt.contains("Convert the following image content to markdown"));
        assert!(!prompt.contains("CONTINUITY"));
    }

    // ── build_user_text_single_image ──

    #[test]
    fn user_text_single_image() {
        let text = build_user_text_single_image();
        assert!(text.contains("Process the image below"));
    }

    // ── build_user_text_image_page ──

    #[test]
    fn user_text_image_page_no_context() {
        let text = build_user_text_image_page(2, 5, None);
        assert!(text.contains("Image 2 of 5"));
        assert!(text.contains("Process the image below"));
        assert!(!text.contains("Previous image tail"));
    }

    #[test]
    fn user_text_image_page_with_context() {
        let text = build_user_text_image_page(3, 10, Some("end of previous"));
        assert!(text.contains("Image 3 of 10"));
        assert!(text.contains("Previous image tail"));
        assert!(text.contains("end of previous"));
    }

    // ── build_user_text ──

    #[test]
    fn user_text_no_context() {
        let text = build_user_text(3, 10, None);
        assert!(text.contains("Page 3 of 10"));
        assert!(text.contains("Process the image below"));
        assert!(!text.contains("Previous page tail"));
    }

    #[test]
    fn user_text_with_context() {
        let text = build_user_text(3, 10, Some("end of page two"));
        assert!(text.contains("Previous page tail"));
        assert!(text.contains("end of page two"));
    }

    #[test]
    fn user_text_empty_tail_ignored() {
        let text = build_user_text(1, 5, Some("   "));
        assert!(!text.contains("Previous page tail"));
    }

    // ── build_user_text_from_text ──

    #[test]
    fn user_text_from_text_includes_extracted() {
        let text = build_user_text_from_text(1, 5, None, "Hello world");
        assert!(text.contains("Hello world"));
        assert!(text.contains("Extracted text"));
        assert!(text.contains("Rewrite the extracted text"));
    }

    #[test]
    fn user_text_from_text_with_tail() {
        let text = build_user_text_from_text(2, 5, Some("prev tail"), "Content");
        assert!(text.contains("Previous page tail"));
        assert!(text.contains("prev tail"));
    }
}