Extract the fields of an invoice

"""
Label each line of the invoice, then read the value out of the line.

Rung N1. N0 asked "what does this line say". This asks "where does this line
sit, and what does it look like": how far down the page, how indented, how
wordy, how many amounts, and whether the value hangs on the right-hand side.

Those features survive a change of supplier, which is exactly what the labels
of N0 do not. Training is a few dozen labelled lines, the model is a few
kilobytes, and nothing is downloaded.
"""

import re

from sklearn.linear_model import LogisticRegression

AMOUNT = r"\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}"
# Month names in any case, accents typed or not.
MONTHS = "janvier|f[ée]vrier|mars|avril|mai|juin|juillet|ao[ûu]t|septembre|octobre|novembre|d[ée]cembre"
DATE = re.compile(rf"\d{{1,2}}/\d{{1,2}}/\d{{2,4}}|\d{{1,2}}\s+(?:{MONTHS})\s+\d{{4}}", re.I)
REFERENCE = r"\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}"


def page_lines(text: str) -> list[str]:
    """The lines that carry something, indentation kept: it is a feature."""
    return [line for line in text.splitlines() if line.strip()]


def line_features(line: str, index: int, count: int) -> list[float]:
    """Where the line sits and what it looks like. Never what it says."""
    text = line.strip()
    amounts = re.findall(AMOUNT, text)
    letters = "".join(c for c in text if c.isalpha())
    return [
        index / max(count - 1, 1),                          # how far down the page
        1.0 if index == count - 1 else 0.0,                 # the very last line
        min(len(line) - len(line.lstrip()), 40) / 40,       # indentation
        min(len(text.split()), 12) / 12,                    # how wordy
        sum(c.isdigit() for c in text) / len(text),         # digit share
        min(len(amounts), 3) / 3,                           # how many amounts
        1.0 if re.search(DATE, text) else 0.0,
        1.0 if re.search(REFERENCE, text) else 0.0,
        # A value hanging on the right of the line, the way a totals block does.
        1.0 if amounts and text.rindex(amounts[-1]) > len(text) / 2 else 0.0,
        1.0 if letters and letters.isupper() else 0.0,
    ]


def train(documents: list[str], labels: list[list[str]]):
    """`labels` carries one label per non-blank line of each document."""
    rows, targets = [], []
    for document, document_labels in zip(documents, labels):
        lines = page_lines(document)
        rows += [line_features(line, i, len(lines)) for i, line in enumerate(lines)]
        targets += list(document_labels)
    model = LogisticRegression(class_weight="balanced", max_iter=1000)
    model.fit(rows, targets)
    return model


def _match(pattern: str):
    """Reads the first value of that shape out of a line, or nothing."""
    def read(line):
        found = re.search(pattern, line)
        return found.group(0) if found else None

    return read


def _amount(line):
    """The rightmost amount: an item line carries a quantity and a unit price."""
    amounts = re.findall(AMOUNT, line)
    if not amounts:
        return None
    cleaned = re.sub(r"[^\d,.]", "", amounts[-1])
    # A comma means French spelling: the dots left are thousands separators.
    if "," in cleaned:
        cleaned = cleaned.replace(".", "").replace(",", ".")
    return float(cleaned)


READERS = {"invoice_number": _match(REFERENCE), "date": _match(DATE), "total": _amount}


def extract_fields(model, text: str) -> dict:
    """
    For each field, walk the lines from the most likely down.

    The model points at a line; reading a date or an amount out of it is still
    ours to do, and a line the model likes but that holds no value is not an
    answer.
    """
    lines = page_lines(text)
    if not lines:
        return {name: None for name in READERS}
    rows = [line_features(line, i, len(lines)) for i, line in enumerate(lines)]
    scores = model.predict_proba(rows)
    classes = list(model.classes_)
    fields = {}
    for name, read in READERS.items():
        column = classes.index(name)
        ranked = sorted(range(len(lines)), key=lambda i: -scores[i][column])
        fields[name] = next((v for v in (read(lines[i]) for i in ranked) if v is not None), None)
    return fields

/**
 * Label each line of the invoice, then read the value out of the line.
 *
 * Rung N1. N0 asked "what does this line say". This asks "where does this
 * line sit, and what does it look like": how far down the page, how indented,
 * how wordy, how many amounts, and whether the value hangs on the right-hand
 * side.
 *
 * Those features survive a change of supplier, which is exactly what the
 * labels of N0 do not. Training is a few dozen labelled lines, and logistic
 * regression is short enough to be read rather than imported.
 */

const AMOUNT = /\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}/g;
// Month names in any case, accents typed or not.
const MONTHS = 'janvier|f[ée]vrier|mars|avril|mai|juin|juillet|ao[ûu]t|septembre|octobre|novembre|d[ée]cembre';
const DATE = new RegExp(`\\d{1,2}/\\d{1,2}/\\d{2,4}|\\d{1,2}\\s+(?:${MONTHS})\\s+\\d{4}`, 'i');
const REFERENCE = /\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}/;

/** The lines that carry something, indentation kept: it is a feature. */
export function pageLines(text) {
  return text.split('\n').filter((line) => line.trim() !== '');
}

/** Where the line sits and what it looks like. Never what it says. */
export function lineFeatures(line, index, count) {
  const text = line.trim();
  const amounts = text.match(AMOUNT) ?? [];
  const letters = text.replace(/[^\p{L}]/gu, '');
  const words = text.split(/\s+/).filter(Boolean);
  return [
    index / Math.max(count - 1, 1), //                    how far down the page
    index === count - 1 ? 1 : 0, //                       the very last line
    Math.min(line.length - line.trimStart().length, 40) / 40, // indentation
    Math.min(words.length, 12) / 12, //                   how wordy
    (text.match(/\d/g) ?? []).length / text.length, //    digit share
    Math.min(amounts.length, 3) / 3, //                   how many amounts
    DATE.test(text) ? 1 : 0,
    REFERENCE.test(text) ? 1 : 0,
    // A value hanging on the right of the line, the way a totals block does.
    amounts.length && text.lastIndexOf(amounts.at(-1)) > text.length / 2 ? 1 : 0,
    letters && letters === letters.toUpperCase() ? 1 : 0,
  ];
}

/**
 * One binary classifier per label, each trained on every line.
 *
 * Each line is weighted so that the lines carrying the label weigh, together,
 * as much as the lines that do not: on a page, they are the few.
 */
function trainOne(rows, targets, label, epochs, rate) {
  const wanted = targets.map((t) => (t === label ? 1 : 0));
  const positives = wanted.reduce((a, b) => a + b, 0);
  const weights = new Float64Array(rows[0].length);
  let bias = 0;
  for (let epoch = 0; epoch < epochs; epoch += 1) {
    for (let i = 0; i < rows.length; i += 1) {
      const balance = rows.length / (2 * (wanted[i] ? positives : rows.length - positives));
      let z = bias;
      for (let j = 0; j < weights.length; j += 1) z += weights[j] * rows[i][j];
      const error = (1 / (1 + Math.exp(-z)) - wanted[i]) * balance;
      for (let j = 0; j < weights.length; j += 1) weights[j] -= rate * error * rows[i][j];
      bias -= rate * error;
    }
  }
  return { weights, bias };
}

/** `labels` carries one label per non-blank line of each document. */
export function train(documents, labels, { epochs = 400, rate = 0.3 } = {}) {
  const rows = [];
  const targets = [];
  documents.forEach((document, d) => {
    const lines = pageLines(document);
    lines.forEach((line, i) => rows.push(lineFeatures(line, i, lines.length)));
    targets.push(...labels[d]);
  });
  const classes = [...new Set(targets)].sort();
  return { classes, models: classes.map((c) => trainOne(rows, targets, c, epochs, rate)) };
}

function score(model, row) {
  let z = model.bias;
  for (let j = 0; j < row.length; j += 1) z += model.weights[j] * row[j];
  return 1 / (1 + Math.exp(-z));
}

/** Reads the first value of that shape out of a line, or nothing. */
const matching = (pattern) => (line) => line.match(pattern)?.[0] ?? null;

/** The rightmost amount: an item line carries a quantity and a unit price. */
function amount(line) {
  const amounts = line.match(AMOUNT);
  if (!amounts) return null;
  let cleaned = amounts.at(-1).replace(/[^\d,.]/g, '');
  // A comma means French spelling: the dots left are thousands separators.
  if (cleaned.includes(',')) cleaned = cleaned.replaceAll('.', '').replace(',', '.');
  return Number(cleaned);
}

const READERS = { invoice_number: matching(REFERENCE), date: matching(DATE), total: amount };

/**
 * For each field, walk the lines from the most likely down.
 *
 * The model points at a line; reading a date or an amount out of it is still
 * ours to do, and a line the model likes but that holds no value is not an
 * answer.
 */
export function extractFields(model, text) {
  const lines = pageLines(text);
  const rows = lines.map((line, i) => lineFeatures(line, i, lines.length));
  const fields = {};
  for (const [name, read] of Object.entries(READERS)) {
    const column = model.classes.indexOf(name);
    const ranked = rows
      .map((row, i) => [score(model.models[column], row), i])
      .sort((a, b) => b[0] - a[0]);
    fields[name] = ranked.map(([, i]) => read(lines[i])).find((value) => value !== null) ?? null;
  }
  return fields;
}

"""
Ask a general-purpose multimodal model to read the invoice.

Rung N3. The model is sent a picture of the page along with its text, and the
picture still holds what the extracted text has lost: the column an amount
sits in.

Note what the code has to do that N0 did not: cap the size of what it sends,
retry on failure, parse an answer that is only probably JSON, and check the
shape of what came back. That plumbing is the real cost of this rung, and it is
the part your tests have to cover, because the model itself is not testable.

One check here is not about shape. The three amounts are asked for, not just
the one that is due, so that the sum EN 16931 makes a rule of — total with VAT
= total without VAT + VAT — can be recomputed on what the model wrote. It
catches nothing the model gets consistently wrong, and it catches a digit read
off the wrong line, which is the failure this entry is about.
"""

from __future__ import annotations

import base64
import json
import math

PROMPT = (
    "Read the invoice below and return its header fields.\n"
    "Answer with JSON only: an object with the keys `invoice_number`, `date`,\n"
    "`total_excluding_vat`, `vat` and `total`. `total` is the amount due, taxes\n"
    "included, as a number.\n"
    "Use null for a field the page does not carry.\n\n"
    "Extracted text:\n{text}"
)

# The provider bills every token of the prompt: the text is capped before the
# call, in characters, not tokens. An image is billed by its dimensions, not its
# bytes, so its cap only bounds what is encoded and uploaded.
MAX_CHARACTERS = 8000
MAX_IMAGE_BYTES = 4_000_000

FIELDS = ("invoice_number", "date", "total_excluding_vat", "vat", "total")
AMOUNTS = ("total_excluding_vat", "vat", "total")

# The provider named here is an example, not a recommendation: the reasoning
# holds for any general-purpose model API, and the client is swappable. Pass
# any object with a `complete(prompt=..., image_url=..., temperature=...)` method.
MODEL = "gpt-4.1-mini"  # an example id: check the parameters your model accepts


class ExtractionUnavailable(Exception):
    """The provider could not be reached, or answered something unusable."""


class ProviderClient:
    """The one call this snippet makes, on top of the provider's SDK."""

    def __init__(self, sdk=None, model: str = MODEL):
        if sdk is None:  # pragma: no cover - needs a key and a network
            from openai import OpenAI

            sdk = OpenAI()
        self.sdk, self.model = sdk, model

    def complete(self, *, prompt: str, image_url: str, temperature: float) -> str:
        response = self.sdk.chat.completions.create(
            model=self.model,
            # The page travels inside the message, as a data URL.
            messages=[{"role": "user", "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": image_url}},
            ]}],
            temperature=temperature,
        )
        return response.choices[0].message.content


def extract_fields(text: str, page_image: bytes, client=None, *, attempts: int = 3) -> dict:
    """
    Read the invoice fields from its text and a picture of the page.

    `client` is injected so this function can be tested without a network call.
    In production it defaults to a real provider client.
    """
    if len(text) > MAX_CHARACTERS:
        raise ValueError(f"text longer than {MAX_CHARACTERS} characters")
    if not page_image or len(page_image) > MAX_IMAGE_BYTES:
        raise ValueError(f"page image empty or larger than {MAX_IMAGE_BYTES} bytes")
    client = client or ProviderClient()

    image_url = "data:image/png;base64," + base64.b64encode(page_image).decode()
    return _decode(_ask(client, text, image_url, attempts))


def _ask(client, text: str, image_url: str, attempts: int) -> str:
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            answer = client.complete(
                prompt=PROMPT.format(text=text),
                image_url=image_url,
                # Temperature zero, because an amount that changes between two
                # identical calls cannot be reconciled with anything.
                temperature=0,
            )
            if isinstance(answer, str):
                return answer
            last_error = ValueError("the model returned no content")  # a refusal: unusable, not empty
        except Exception as error:  # noqa: BLE001 - any provider failure is retried
            last_error = error
    raise ExtractionUnavailable(str(last_error))


def _decode(answer: str) -> dict:
    """
    Turn the answer into the five fields, or refuse it.

    Models like to wrap JSON in a code fence. That is noise, not an error, and
    stripping it is cheaper than another call.
    """
    stripped = answer.strip()
    if stripped.startswith("```") and stripped.endswith("```") and stripped.count("```") == 2:
        stripped = stripped[3:-3].removeprefix("json")
    try:
        parsed = json.loads(stripped)
    except ValueError as error:
        raise ExtractionUnavailable(f"the model did not answer with JSON: {error}")
    if not isinstance(parsed, dict):
        raise ExtractionUnavailable("the model answered something that is not an object")

    fields = {field: parsed.get(field) for field in FIELDS}
    for field in ("invoice_number", "date"):
        if fields[field] is not None and not isinstance(fields[field], str):
            raise ExtractionUnavailable(f"the model answered a {field} that is not text: {fields[field]!r}")
    for field in AMOUNTS:
        value = fields[field]
        # An amount nobody can compute with is worse than none at all: it would
        # travel down the pipeline looking like a number.
        if value is not None and (isinstance(value, bool) or not isinstance(value, (int, float))
                                  or not math.isfinite(value)):
            raise ExtractionUnavailable(f"the model answered a {field} that is not a number: {value!r}")
    return {**fields, "totals_agree": _totals_agree(fields)}


def _totals_agree(fields: dict) -> bool | None:
    """
    BR-CO-15 recomputed on what the model wrote, or None if it wrote too little.

    False is not a reason to throw the answer away: it is the one reason the
    caller has to look at this invoice rather than the other four hundred.
    """
    amounts = [fields[field] for field in AMOUNTS]
    if any(value is None for value in amounts):
        return None
    without, vat, total = amounts
    return abs(without + vat - total) < 0.005

/**
 * Ask a general-purpose multimodal model to read the invoice.
 *
 * Rung N3. The model is sent a picture of the page along with its text, and
 * the picture still holds what the extracted text has lost: the column an
 * amount sits in.
 *
 * Note what the code has to do that N0 did not: cap the size of what it sends,
 * retry on failure, parse an answer that is only probably JSON, and check the
 * shape of what came back. That plumbing is the real cost of this rung, and it
 * is the part your tests have to cover, because the model itself is not
 * testable.
 *
 * One check here is not about shape. The three amounts are asked for, not just
 * the one that is due, so that the sum EN 16931 makes a rule of — total with
 * VAT = total without VAT + VAT — can be recomputed on what the model wrote.
 * It catches nothing the model gets consistently wrong, and it catches a digit
 * read off the wrong line, which is the failure this entry is about.
 */

const PROMPT = [
  'Read the invoice below and return its header fields.',
  'Answer with JSON only: an object with the keys `invoice_number`, `date`,',
  '`total_excluding_vat`, `vat` and `total`. `total` is the amount due, taxes',
  'included, as a number.',
  'Use null for a field the page does not carry.',
  '',
  'Extracted text:',
].join('\n');

// The provider bills every token of the prompt: the text is capped before the
// call, in characters (code points, as in Python), not tokens. An image is
// billed by its dimensions, not its bytes, so its cap only bounds what is
// encoded and uploaded.
export const MAX_CHARACTERS = 8000;
export const MAX_IMAGE_BYTES = 4_000_000;

const FIELDS = ['invoice_number', 'date', 'total_excluding_vat', 'vat', 'total'];
const AMOUNTS = ['total_excluding_vat', 'vat', 'total'];

// The provider named here is an example, not a recommendation: the reasoning
// holds for any general-purpose model API, and the client is swappable. Pass
// any object with a `complete({ prompt, imageUrl, temperature })` method.
export const MODEL = 'gpt-4.1-mini'; // an example id: check the parameters your model accepts

export async function providerClient(sdk, model = MODEL) {
  if (!sdk) {
    const { OpenAI } = await import('openai');
    sdk = new OpenAI();
  }
  return {
    async complete({ prompt, imageUrl, temperature }) {
      const response = await sdk.chat.completions.create({
        model,
        // The page travels inside the message, as a data URL.
        messages: [{ role: 'user', content: [
          { type: 'text', text: prompt },
          { type: 'image_url', image_url: { url: imageUrl } },
        ] }],
        temperature,
      });
      return response.choices[0].message.content;
    },
  };
}

export class ExtractionUnavailable extends Error {}

/**
 * Read the invoice fields from its text and a picture of the page.
 *
 * @param {string} text the invoice, already extracted
 * @param {Uint8Array} pageImage the rendered page; this snippet opens no file
 * @param {object} options
 * @param {{complete: Function}} [options.client] injected so this can be
 *   tested without a network call; defaults to a real provider client
 * @param {number} [options.attempts]
 */
export async function extractFields(text, pageImage, { client, attempts = 3 } = {}) {
  if (text.length > MAX_CHARACTERS && [...text].length > MAX_CHARACTERS) {
    throw new RangeError(`text longer than ${MAX_CHARACTERS} characters`);
  }
  if (pageImage.length === 0 || pageImage.length > MAX_IMAGE_BYTES) {
    throw new RangeError(`page image empty or larger than ${MAX_IMAGE_BYTES} bytes`);
  }
  client ??= await providerClient();

  const imageUrl = `data:image/png;base64,${Buffer.from(pageImage).toString('base64')}`;
  return decode(await ask(client, text, imageUrl, attempts));
}

async function ask(client, text, imageUrl, attempts) {
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      const answer = await client.complete({
        prompt: `${PROMPT}\n${text}`,
        imageUrl,
        // Temperature zero, because an amount that changes between two
        // identical calls cannot be reconciled with anything.
        temperature: 0,
      });
      if (typeof answer === 'string') return answer;
      lastError = new Error('the model returned no content'); // a refusal: unusable, not empty
    } catch (error) {
      lastError = error;
    }
  }
  throw new ExtractionUnavailable(String(lastError));
}

/**
 * Turn the answer into the three fields, or refuse it.
 *
 * Models like to wrap JSON in a code fence. That is noise, not an error, and
 * stripping it is cheaper than another call.
 */
/**
 * BR-CO-15 recomputed on what the model wrote, or null if it wrote too little.
 *
 * False is not a reason to throw the answer away: it is the one reason the
 * caller has to look at this invoice rather than the other four hundred.
 */
function totalsAgree(fields) {
  const [without, vat, total] = AMOUNTS.map((field) => fields[field]);
  if (without === null || vat === null || total === null) return null;
  return Math.abs(without + vat - total) < 0.005;
}

function decode(answer) {
  const trimmed = answer.trim();
  const fenced = trimmed.startsWith('```') && trimmed.endsWith('```') && trimmed.split('```').length === 3;
  const stripped = fenced ? trimmed.slice(3, -3).replace(/^json/, '') : trimmed;
  let parsed;
  try {
    parsed = JSON.parse(stripped);
  } catch (error) {
    throw new ExtractionUnavailable(`the model did not answer with JSON: ${error.message}`);
  }
  if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) {
    throw new ExtractionUnavailable('the model answered something that is not an object');
  }

  const fields = Object.fromEntries(FIELDS.map((field) => [field, parsed[field] ?? null]));
  for (const field of ['invoice_number', 'date']) {
    if (fields[field] !== null && typeof fields[field] !== 'string') {
      throw new ExtractionUnavailable(`the model answered a ${field} that is not text: ${fields[field]}`);
    }
  }
  for (const field of AMOUNTS) {
    // An amount nobody can compute with is worse than none at all: it would
    // travel down the pipeline looking like a number.
    if (fields[field] !== null && (typeof fields[field] !== 'number' || !Number.isFinite(fields[field]))) {
      throw new ExtractionUnavailable(`the model answered a ${field} that is not a number: ${fields[field]}`);
    }
  }
  return { ...fields, totals_agree: totalsAgree(fields) };
}