Skip to content

C
coderai
Commit 29233812 authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

chunk: POST /v1/chunk, and the OCR pipeline for all three VLM engines (v0.3.76) 

RETRIEVAL CHUNKING. A new endpoint that splits a document into the pieces an
index stores, with the document's structure decided per request by a named or
inline template. NOT embedding windowing: `embed_long_input` makes ONE input fit
a model window and pools it into ONE vector; this produces MANY chunks, each
indexable, each with its section, parent and references. It never pools and
never drops text.

Pure CPU -- token counts are characters / chars_per_token, so no tokenizer is
loaded and no model, GPU or pod is involved. Every image serves it, light-core
ones included, and a test asserts no torch/transformers import creeps in.

The packing loop is reproduced verbatim, including the parts that look
improvable: paragraphs joined with "\n\n" (hence the `+ 2`), and an overlap that
is a RAW CHARACTER TAIL which may begin mid-word. Changing any of it moves every
existing chunk boundary, and then a re-chunked index cannot be compared with the
one it replaces.

`conclusioni` is its own label and is never folded into `dispositivo`: in a
ruling the dispositivo is what the judge decided, in a filing the conclusioni
are what a party ASKED FOR, and merging them would let a demand surface as a
holding.

TWO DEVIATIONS FROM THE SPEC, both recorded in the source where someone
comparing against the original will look:

  * The ceiling is target + overlap + 2 (3682), not the 3680 the prose states --
    which agrees with the spec's own "observed max is 3682" two lines later.
    Asserting 3680 would fail a correct port.
  * The verbatim patterns cannot match `**FATTO**`, yet the fidelity list asks
    for that case. The line-anchored alternatives are widened at BOTH ends
    (`^[\s*#>_]*fatto[\s*#>_]*$`), which is how the patterns that already coped
    spell it. Purely additive: every line that matched still matches, so no
    existing match moves -- but a document whose headers were emphasised used to
    fall back to flat chunking and now gets sectioned, so ITS chunks change and
    it needs re-chunking.

FORWARDABLE, with the hazard that comes with it. /v1/chunk rides the EMBEDDINGS
capability: chunks are on their way to be embedded, that image is full-core so
it carries the chunker, and a capability of its own would be one nothing
advertises -- or would rent a GPU to run a regex split. It was first classified
as deliberately local on the grounds that it is cheap; that was the wrong lens,
since the orchestrator is the box that was OOM-killed today, so WHERE the CPU is
matters more than how much there is.

But a template registered HERE does not exist on the remote: forwarded by name
it returns 404, or worse matches a DIFFERENT template with the same name and
returns different boundaries silently. So a named non-builtin template is
resolved and INLINED before the request leaves. The hook never raises -- a body
it cannot parse is forwarded untouched -- because a rewrite that can throw would
take down every forwarded request, not just a chunk one.

THE OCR PIPELINE NOW COVERS ALL THREE VLM ENGINES, not just surya. Leaving
chandra and olmOCR on the orchestrator was an asymmetry rather than a design:
surya needed a 1.6 GB client venv baked in, while those two speak only HTTP in
served mode (PIL + requests, already in the main venv), so for them the pod side
costs nothing beyond the rasteriser. Two bugs fixed on the way, both of which
would have produced plausible WRONG text rather than an error: the dispatch
hardcoded `engine: "surya"` in the form it posted, and the pod env was
surya-shaped throughout -- a chandra pod was told to run surya against chandra's
model id. Per-engine tuning and per-engine licence acceptance now travel.

AND THE LOG THAT CAUSED A MISDIAGNOSIS. One pod request walks a candidate list
until a (card, region) pair has capacity, and each attempt printed
"provisioning pod for ...". 312 attempts for 10 pods therefore read as runaway
provisioning and was reported as such. It now says "trying candidate N/M for
ONE pod", so the shape of the list is visible instead of inferred.

What these tests CANNOT do is prove byte-identity with the implementation this
ports -- that needs fixtures from a live corpus. They prove the documented
behaviour on each path; identity is what decides whether a re-chunk is safe and
it is still outstanding.

Suite at the 25 pre-existing failures.
Co-Authored-By: 's avatarClaude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HXRMe9ZPUQCCCtBai9uqt3
parent 35a8304e
Hide whitespace changes
Inline Side-by-side
Showing with 1773 additions and 44 deletions
codai/__init__.py
View file @ 29233812
......@@ -16,7 +16,7 @@
# Canonical product version for CoderAI — single source of truth. Both the API
# metadata and the admin web UI read from here.
__version__ = "0.3.75"
__version__ = "0.3.76"
__tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
# Configure the CUDA caching allocator BEFORE torch is imported anywhere.
......
codai/api/app.py
View file @ 29233812
......@@ -760,6 +760,10 @@ app.include_router(loras_router, tags=["LoRAs"])
app.include_router(environments_router, tags=["Environments"])
app.include_router(spatial_router, tags=["Spatial / 3D"])
app.include_router(ocr_router, tags=["OCR"])
# Retrieval chunking. A pure CPU text operation (no tokenizer, no model,
# no GPU), so every image serves it -- including the light-core ones.
from codai.api.chunking import router as chunking_router # noqa: E402
app.include_router(chunking_router, tags=["Chunking"])
# The address an attached OCR engine drives: one stable URL of ours, forwarded
# per request to the least-busy server that holds the model. Not in the schema —
# it is machine-to-machine, authenticated by a lease rather than an API token.
......
codai/api/chunking.py 0 → 100644
View file @ 29233812
#!/usr/bin/env python3
"""Retrieval chunking — ``POST /v1/chunk``.
Splits a document into the pieces a retrieval index stores, with the document's
structure decided PER REQUEST by a named or inline *structure template*.
NOT the same thing as embedding windowing. ``embed_long_input`` in
``codai/api/embeddings.py`` makes ONE over-long input fit a model's window and
pools the pieces into ONE vector. This produces MANY chunks, each separately
indexable, each carrying its section, parent and references. A caller may use
both: chunk here, then embed each chunk, where windowing is the safety net for a
chunk that still overflows. This endpoint must never pool and never drop text.
WHY STRUCTURE-AWARE. Length-only chunking cuts an argument wherever the
character budget runs out, so a chunk can straddle a section boundary and a
retrieved passage mixes what the parties alleged with what the court held.
Sectioning first and length-chunking INSIDE each section removes that class of
error and gives every chunk a label and a recoverable parent.
A pure CPU text operation: token counts are approximated as characters divided
by ``chars_per_token``, so no tokenizer is loaded, no model is needed and no GPU
or pod is involved. That is deliberate -- it means every image that runs coderai
serves this, including the light-core ones.
"""
import re
from typing import Any, Dict, List, Optional, Tuple
#: Approximation divisor: token counts are characters / this. Keeps the endpoint
#: tokenizer-free, which is what makes it callable with nothing loaded.
DEFAULT_CHARS_PER_TOKEN = 4
DEFAULT_CHUNK_TOKENS = 800
DEFAULT_OVERLAP_TOKENS = 120
#: A pattern longer than this, or more sections than this, is refused as
#: `bad_template`. Patterns come from callers; Python's `re` has NO timeout, so
#: these caps are the only bound available here. They are a blunt guard, not a
#: proof: a pathological pattern within the caps can still backtrack badly. Said
#: plainly rather than implied, because the alternative is believing this is
#: safe when it is only bounded in SIZE.
MAX_PATTERN_CHARS = 2000
MAX_TEMPLATE_SECTIONS = 64
#: Paragraphs are separated by a blank line (possibly with whitespace in it).
_PARA_SPLIT = re.compile(r"\n\s*\n")
#: Sentence boundary used to break a paragraph that exceeds the target on its
#: own. Deliberately simple: it is what the production chunker used, and a
#: "better" splitter moves every existing chunk boundary.
_SENT_SPLIT = re.compile(r"(?<=[.;:!?])\s+")
class ChunkError(Exception):
"""A request-level failure, carrying the status and machine-readable code
the endpoint reports (see docs/chunking.md#errors)."""
def __init__(self, code: str, message: str, status: int = 400):
super().__init__(message)
self.code = code
self.status = status
self.message = message
# ---------------------------------------------------------------------------
# Templates
# ---------------------------------------------------------------------------
#
# `conclusioni` IS ITS OWN LABEL and must never be folded into `dispositivo`.
# In a ruling the dispositivo is what the judge decided; in a filing the
# conclusioni are what a party ASKED FOR. Merging them lets a demand surface as
# if it were a holding, which is the most damaging mistake a legal index can
# make. A template that wants them merged has to say so by giving both patterns
# the same label.
#
# Two marker families, because filings outnumber rulings in a merito corpus:
# with ruling markers alone only 8-10% of 400 measured documents were
# recognised.
#: What may follow a header on its own line. `line_prefix` tolerates decoration
#: BEFORE a header; nothing tolerated it after, so `**FATTO**` was missed while
#: `> fatto` matched -- the end-anchored alternatives (`^\s*fatto\s*$` and
#: friends) require the line to end right after the word, and a trailing `**`
#: defeats that.
#:
#: Widening only the END anchor is PURELY ADDITIVE: every line that matched
#: `\s*$` still matches `[\s*#>_]*$`, so no existing match moves. What changes
#: is which documents get sectioned AT ALL -- one whose headers were markdown-
#: emphasised used to fall back to flat chunking and now gets labels. Those
#: documents' chunks therefore differ from an index built before this, and they
#: need re-chunking; documents already sectioned are unaffected.
#: And the same class BEFORE the word. `^\s*fatto` tolerates only
#: whitespace, so `**FATTO**` failed on the leading asterisks as well as
#: the trailing ones. The patterns that already handled decoration
#: (`premesso che`, `conclusioni`) spell it `^\s*[*#>_ ]*` for exactly this
#: reason; these three simply never did.
_LEAD = r"[\s*#>_]*"
_TRAIL = r"[\s*#>_]*"
IT_LEGAL: Dict[str, Any] = {
"name": "it-legal",
"head_label": "intestazione",
"flags": "im",
"line_prefix": r"[ \t>*_#]*",
"merge_adjacent": True,
"keep_header_line": True,
"sections": [
{"label": "premessa",
"pattern": r"tanto\s+premesso|ci[oò]\s+premesso|^\s*[*#>_ ]*premesso\s+che\b|^" + _LEAD + r"premess[ae]" + _TRAIL + "$"},
{"label": "fatto",
"pattern": r"svolgimento\s+del\s+processo|ragioni\s+di\s+fatto|in\s+fatto\b|fatto\s+e\s+diritto|^" + _LEAD + r"fatto" + _TRAIL + r"$|esposizione\s+del\s+fatto"},
{"label": "motivazione",
"pattern": r"motivi\s+della\s+decisione|ragioni\s+della\s+decisione|motivazione|considerato\s+in\s+diritto|in\s+diritto\b|motivi\s+in\s+diritto|ragioni\s+in\s+diritto"},
{"label": "dispositivo",
"pattern": r"p\.?\s*q\.?\s*m\.?|p\.?\s*t\.?\s*m\.?|per\s+questi\s+motivi|^" + _LEAD + r"dispositivo" + _TRAIL + r"$|così\s+deciso|così\s+provvede"},
{"label": "conclusioni",
"pattern": r"^\s*[*#>_ ]*conclusion[ei]\b|\bconclusioni\s*:|voglia\s+(?:l['’]ill\.?mo\s+)?(?:il\s+)?(?:tribunale|giudice|giudicante)|^\s*[*#>_ ]*chiede\b|salvis\s+iuribus"},
],
}
#: DEVIATION FROM docs/chunking.md, deliberate and recorded here because it
#: affects byte-identity with the implementation this ports.
#:
#: The doc gives the patterns verbatim and asks for a fidelity case with a
#: markdown-decorated header (`**FATTO**`). Those two requirements contradict
#: each other: the verbatim `^\s*fatto\s*$` cannot match `**FATTO**`. The
#: anchors here are widened to `^[\s*#>_]*fatto[\s*#>_]*$` so the documented
#: case works.
#:
#: The widening is purely additive -- no line that matched before stops
#: matching, so no existing match position moves. What changes is which
#: documents are sectioned AT ALL: one whose headers were emphasised used to
#: fall back to flat chunking and now gets labels. Those documents' chunks
#: differ from an index built before this and must be re-chunked; documents
#: already sectioned are identical.
#:
#: Worth confirming against the production implementation: if it DOES match
#: `**FATTO**`, its patterns differ from the table in the doc, and that
#: difference is what matters for identity rather than this note.
BUILTIN_TEMPLATES: Dict[str, Dict[str, Any]] = {"it-legal": IT_LEGAL}
def _flags_from(spec: str) -> int:
out = 0
for ch in (spec or "im"):
if ch == "i":
out |= re.IGNORECASE
elif ch == "m":
out |= re.MULTILINE
elif ch == "s":
out |= re.DOTALL
elif ch == "x":
out |= re.VERBOSE
return out
def resolve_template(structure, registered: Optional[dict] = None) -> Optional[dict]:
"""A template from a name or an inline object.
Inline beats registered beats built-in. An unknown NAME is an error and
never a silent fallback to the built-in: a caller who asked for their own
template and got Italian legal sectioning would index the wrong boundaries
and have no way to tell.
"""
if structure in (None, "", False):
return None
if isinstance(structure, dict):
return structure
name = str(structure).strip()
for src in (registered or {}, BUILTIN_TEMPLATES):
if isinstance(src, dict) and name in src:
tpl = src[name]
if not isinstance(tpl, dict):
raise ChunkError("bad_template",
f"registered template {name!r} is not an object")
return tpl
raise ChunkError("unknown_template", f"no structure template named {name!r}",
status=404)
def compile_template(tpl: dict) -> Tuple[Any, List[str], dict]:
"""(compiled alternation, label-by-group-index, normalised options).
ONE alternation with named groups, scanned left to right, so the marks come
out in document order and the first matching alternative wins at a given
position. Scanning per pattern and merging afterwards changes WHICH header
claims an overlapping position -- a different segmentation for the same
document, which is exactly the kind of divergence a port must not introduce.
"""
if not isinstance(tpl, dict):
raise ChunkError("bad_template", "structure template must be an object")
secs = tpl.get("sections")
if not isinstance(secs, list) or not secs:
raise ChunkError("bad_template", "template has no sections")
if len(secs) > MAX_TEMPLATE_SECTIONS:
raise ChunkError("bad_template",
f"template has {len(secs)} sections; the limit is "
f"{MAX_TEMPLATE_SECTIONS}")
prefix = tpl.get("line_prefix")
prefix = r"[ \t>*_#]*" if prefix is None else str(prefix)
labels: List[str] = []
parts: List[str] = []
for i, sec in enumerate(secs):
if not isinstance(sec, dict):
raise ChunkError("bad_template", f"section {i} is not an object")
pat = sec.get("pattern")
if not isinstance(pat, str) or not pat:
raise ChunkError("bad_template", f"section {i} has no pattern")
if len(pat) > MAX_PATTERN_CHARS:
raise ChunkError("bad_template",
f"section {i} pattern is {len(pat)} chars; the "
f"limit is {MAX_PATTERN_CHARS}")
labels.append(str(sec.get("label") or ""))
# Group names are positional (g0, g1, ...) because a label may be any
# string and need not be a valid Python identifier.
parts.append(f"(?P<g{i}>^{prefix}(?:{pat}))")
try:
rx = re.compile("|".join(parts), _flags_from(str(tpl.get("flags") or "im")))
except re.error as exc:
raise ChunkError("bad_template", f"pattern does not compile: {exc}")
opts = {
"head_label": tpl.get("head_label", "intestazione"),
"merge_adjacent": bool(tpl.get("merge_adjacent", True)),
"keep_header_line": bool(tpl.get("keep_header_line", True)),
"name": str(tpl.get("name") or "") or None,
}
return rx, labels, opts
# ---------------------------------------------------------------------------
# Hygiene
# ---------------------------------------------------------------------------
_BINARY_SIGNATURES = (b"PK\x03\x04", b"%PDF-", b"\x7fELF", b"\xd0\xcf\x11\xe0")
def looks_binary(raw) -> Optional[str]:
"""Why this input is not text, or None.
A chunker that accepts container bytes embeds garbage and the index cannot
tell afterwards: measured on a production corpus, 0.59% of chunks held
control characters and 0.04% raw OOXML -- documents whose CONTAINER was
chunked instead of its extracted text.
"""
if isinstance(raw, str):
data = raw.encode("utf-8", "ignore")
text = raw
else:
data = bytes(raw or b"")
text = data.decode("utf-8", "ignore")
for sig in _BINARY_SIGNATURES:
if data.startswith(sig):
return f"input begins with a {sig!r} container header"
if "[Content_Types].xml" in text[:4096]:
return "input looks like an OOXML container, not its extracted text"
if text:
bad = sum(1 for c in text if ord(c) < 32 and c not in "\t\n\f\r")
if bad / max(1, len(text)) > 0.01:
return (f"{bad} control characters in {len(text)} -- this looks like "
f"container bytes rather than extracted text")
return None
# ---------------------------------------------------------------------------
# Sectioning
# ---------------------------------------------------------------------------
def split_sections(text: str, tpl: Optional[dict]) -> Tuple[List[dict], bool, Optional[str]]:
"""(sections, structural_applied, template_name).
A section runs from its header's match START to the next header's start when
``keep_header_line`` is set, so the header stays inside its own section.
"""
if tpl is None:
return [{"label": "", "text": text}], False, None
rx, labels, opts = compile_template(tpl)
marks = [(m.start(), labels[int(m.lastgroup[1:])])
for m in rx.finditer(text) if m.lastgroup]
if not marks:
# No header anywhere: flat, and SAID so, rather than inventing one
# section per paragraph or pretending the template applied.
return [{"label": "", "text": text}], False, opts["name"]
out: List[dict] = []
head = text[:marks[0][0]]
if head.strip() and opts["head_label"] is not None:
out.append({"label": str(opts["head_label"] or ""), "text": head})
for i, (pos, label) in enumerate(marks):
end = marks[i + 1][0] if i + 1 < len(marks) else len(text)
body = text[pos:end] if opts["keep_header_line"] else text[pos:end]
if opts["merge_adjacent"] and out and out[-1]["label"] == label:
# Filings routinely open the conclusions twice (CONCLUSIONI, then
# "Voglia il Tribunale..."). Left separate, one block becomes two
# and the parent ids multiply, defeating the point of `parent`.
out[-1]["text"] = out[-1]["text"] + "\n" + body
else:
out.append({"label": label, "text": body})
return out, True, opts["name"]
# ---------------------------------------------------------------------------
# Packing
# ---------------------------------------------------------------------------
def _pieces(section_text: str, target: int) -> List[str]:
"""A section's paragraphs, with over-long ones broken down. Nothing is ever
discarded -- a chunker that drops text is worse than one that chunks badly,
because the loss is invisible in the index."""
out: List[str] = []
for para in _PARA_SPLIT.split(section_text):
para = para.strip()
if not para:
continue
if len(para) <= target:
out.append(para)
continue
# Sentence boundaries, greedily refilled to target.
cur = ""
for sent in _SENT_SPLIT.split(para):
if not sent:
continue
if len(sent) > target:
if cur:
out.append(cur)
cur = ""
# A single sentence still over target: fixed slices.
for i in range(0, len(sent), target):
out.append(sent[i:i + target])
continue
if cur and len(cur) + len(sent) + 1 > target:
out.append(cur)
cur = sent
else:
cur = (cur + " " + sent) if cur else sent
if cur:
out.append(cur)
return out
def pack(pieces: List[str], target: int, overlap: int) -> List[str]:
"""Pack pieces into chunks with a carried character tail.
Reproduced verbatim from the production chunker, including the parts that
look like they could be improved:
* paragraphs are joined with "\\n\\n", and the `+ 2` accounts for it;
* the overlap is a RAW CHARACTER TAIL of the emitted chunk, which is not
sentence- or token-aligned and may begin mid-word;
* overlap == 0 starts the next chunk at the piece with no carry.
"Improving" any of those moves every existing chunk boundary, which makes a
re-chunked index incomparable with the one it replaces -- and then retrieval
quality cannot be measured across the change.
A chunk may therefore reach target + overlap characters.
"""
chunks: List[str] = []
cur = ""
for p in pieces:
if cur and len(cur) + len(p) + 2 > target:
chunks.append(cur)
cur = (cur[-overlap:] + "\n\n" + p) if overlap else p
else:
cur = (cur + "\n\n" + p).strip() if cur else p
if cur:
chunks.append(cur)
return chunks
# ---------------------------------------------------------------------------
# Reference extraction
# ---------------------------------------------------------------------------
_ART_SUFFIX = r"(?:\s*(?:bis|ter|quater|quinquies|sexies))?"
_ART_RX = re.compile(
r"\bart(?:t|icol[oi])?\.?\s*(\d+" + _ART_SUFFIX + r")"
r"(?:\s*(?:del\s+|della\s+)?(codciv|c\.?c\.?|c\.?p\.?c\.?|c\.?p\.?p\.?|c\.?p\.?|cost\.?))?",
re.IGNORECASE)
_ACT_RX = re.compile(
r"\b(d\.?\s*lgs\.?|d\.?\s*l\.?|d\.?\s*p\.?\s*r\.?|legge|l\.?)\s*"
r"(?:n\.?\s*)?(\d+)\s*/\s*(\d{4})", re.IGNORECASE)
_CODE_MAP = {"codciv": "c.c.", "cc": "c.c.", "cpc": "c.p.c.", "cpp": "c.p.p.",
"cp": "c.p.", "cost": "cost."}
_KIND_MAP = {"dlgs": "d.lgs.", "dl": "d.l.", "dpr": "d.p.r.", "legge": "l.",
"l": "l."}
def extract_refs(text: str) -> List[str]:
"""Normalised article and act citations, de-duplicated, FIRST-SEEN ORDER kept.
An exact-match signal for hybrid search, and the thing dense embeddings
blur: `art. 2043` and `art. 2054` sit close in vector space and are entirely
different claims.
"""
seen: Dict[str, None] = {}
for m in _ART_RX.finditer(text):
num = re.sub(r"\s+", "", m.group(1) or "").lower()
code_raw = re.sub(r"[.\s]", "", (m.group(2) or "")).lower()
code = _CODE_MAP.get(code_raw, "")
seen.setdefault(f"art.{num} {code}".strip(), None)
for m in _ACT_RX.finditer(text):
kind_raw = re.sub(r"[.\s]", "", (m.group(1) or "")).lower()
kind = _KIND_MAP.get(kind_raw, kind_raw + ".")
seen.setdefault(f"{kind}{m.group(2)}/{m.group(3)}", None)
return list(seen)
# ---------------------------------------------------------------------------
# The whole operation
# ---------------------------------------------------------------------------
def chunk_document(text: str, *, chunk_tokens: int = DEFAULT_CHUNK_TOKENS,
overlap_tokens: int = DEFAULT_OVERLAP_TOKENS,
chars_per_token: int = DEFAULT_CHARS_PER_TOKEN,
structural: bool = True, structure: Any = "it-legal",
refs: Any = "it-legal", include_offsets: bool = False,
require_text: bool = True,
registered: Optional[dict] = None) -> dict:
"""Chunk `text` and return the response document (see docs/chunking.md)."""
if not isinstance(text, str):
text = (text or b"").decode("utf-8", "ignore")
if require_text:
why = looks_binary(text)
if why:
raise ChunkError("not_text", why)
# 1. Normalise.
text = text.replace("\r", "")
if not text.strip():
raise ChunkError("empty_input", "input is empty after normalisation")
target = max(1, int(chunk_tokens) * max(1, int(chars_per_token)))
overlap = max(0, int(overlap_tokens) * max(1, int(chars_per_token)))
tpl = resolve_template(structure, registered) if structural else None
sections, applied, tpl_name = split_sections(text, tpl)
want_refs = refs not in (None, False, "", "false")
ref_tpl = None
if want_refs and isinstance(refs, dict):
ref_tpl = refs # reserved: a template may replace the patterns
chunks: List[dict] = []
sec_out: List[dict] = []
cursor = 0
for parent, sec in enumerate(sections):
body = sec["text"]
label = sec["label"] or ""
sec_out.append({"parent": parent, "section": label or None,
"n_chars": len(body)})
for piece in pack(_pieces(body, target), target, overlap):
obj: Dict[str, Any] = {
"ord": len(chunks), "text": piece, "n_chars": len(piece),
"section": label or None, "parent": parent,
}
if want_refs:
obj["norm_refs"] = extract_refs(piece)
if include_offsets:
# Offsets into the NORMALISED text, found forward from the last
# one so a repeated passage does not resolve to the first copy.
at = text.find(piece, cursor)
if at < 0:
at = text.find(piece)
if at >= 0:
obj["start"], obj["end"] = at, at + len(piece)
cursor = at
chunks.append(obj)
return {
"chunks": chunks,
"n_chunks": len(chunks),
"sections": sec_out,
"structure": tpl_name if applied else None,
"structural_applied": bool(applied),
"params": {"chunk_tokens": int(chunk_tokens),
"overlap_tokens": int(overlap_tokens),
"chars_per_token": int(chars_per_token)},
}
# ---------------------------------------------------------------------------
# The endpoint
# ---------------------------------------------------------------------------
import os # noqa: E402
from fastapi import APIRouter, HTTPException, UploadFile, File, Form, Request # noqa: E402
from pydantic import BaseModel, ConfigDict, Field # noqa: E402
router = APIRouter()
#: Refuse an input above this many characters with 413 rather than spending
#: minutes of CPU on it. Generous: a 400-page filing is ~1 M characters.
DEFAULT_MAX_CHARS = 8_000_000
def _cfg():
try:
from codai.admin.routes import config_manager
return getattr(getattr(config_manager, "config", None), "chunking", None)
except Exception:
return None
def _setting(name: str, env: str, default, cast=int):
"""Per-request > config > env > default, the convention the OCR and
embedding settings already use."""
c = _cfg()
if c is not None:
v = getattr(c, name, None)
if v not in (None, ""):
try:
return cast(v)
except (TypeError, ValueError):
pass
raw = (os.environ.get(env) or "").strip()
if raw:
try:
return cast(raw)
except (TypeError, ValueError):
pass
return default
def _bool_setting(name: str, env: str, default: bool) -> bool:
c = _cfg()
if c is not None:
v = getattr(c, name, None)
if isinstance(v, bool):
return v
raw = (os.environ.get(env) or "").strip().lower()
if raw in ("1", "true", "on", "yes"):
return True
if raw in ("0", "false", "off", "no"):
return False
return default
def _registered_templates() -> dict:
c = _cfg()
t = getattr(c, "templates", None) if c is not None else None
return t if isinstance(t, dict) else {}
class ChunkRequest(BaseModel):
text: Optional[str] = None
chunk_tokens: Optional[int] = Field(None, description="Target size in approximated tokens (chars / chars_per_token).")
overlap_tokens: Optional[int] = Field(None, description="Carried tail in approximated tokens. 0 = no carry.")
chars_per_token: Optional[int] = Field(None, description="The approximation divisor; no tokenizer is loaded.")
structural: Optional[bool] = Field(None, description="false = flat length chunking with no sectioning.")
structure: Optional[Any] = Field(None, description="Template name, or an inline template object.")
extract_refs: Optional[Any] = Field(None, description="Reference extractor: a name, an inline object, or false to skip.")
require_text: Optional[bool] = Field(None, description="Reject binary input instead of chunking container bytes.")
include_offsets: Optional[bool] = Field(None, description="Add char offsets into the normalised source text.")
model_config = ConfigDict(extra="allow")
def _resolve_params(text: str, body) -> dict:
"""Merge request fields over the configured defaults."""
g = (lambda n, e, d, c=int: _setting(n, e, d, c))
return dict(
text=text,
chunk_tokens=(body.chunk_tokens if body.chunk_tokens is not None
else g("chunk_tokens", "CODERAI_CHUNK_TOKENS", DEFAULT_CHUNK_TOKENS)),
overlap_tokens=(body.overlap_tokens if body.overlap_tokens is not None
else g("overlap", "CODERAI_CHUNK_OVERLAP", DEFAULT_OVERLAP_TOKENS)),
chars_per_token=(body.chars_per_token if body.chars_per_token is not None
else g("chars_per_token", "CODERAI_CHUNK_CHARS_PER_TOKEN",
DEFAULT_CHARS_PER_TOKEN)),
structural=(body.structural if body.structural is not None
else _bool_setting("structural", "CODERAI_CHUNK_STRUCTURAL", True)),
structure=(body.structure if body.structure is not None
else (_setting("structure", "CODERAI_CHUNK_STRUCTURE", "it-legal", str))),
refs=(body.extract_refs if body.extract_refs is not None else "it-legal"),
include_offsets=bool(body.include_offsets),
require_text=(body.require_text if body.require_text is not None
else _bool_setting("require_text", "CODERAI_CHUNK_REQUIRE_TEXT", True)),
registered=_registered_templates(),
)
def _run(params: dict) -> dict:
limit = _setting("max_chars", "CODERAI_CHUNK_MAX_CHARS", DEFAULT_MAX_CHARS)
if len(params["text"]) > limit:
raise HTTPException(status_code=413, detail={
"code": "too_large",
"message": f"input is {len(params['text'])} characters; the limit is "
f"{limit} (chunking.max_chars / CODERAI_CHUNK_MAX_CHARS)"})
try:
return chunk_document(**params)
except ChunkError as e:
# The CODE travels, not just the prose: a caller retrying on 400 needs to
# tell "your file is a zip" from "your template does not compile".
raise HTTPException(status_code=e.status,
detail={"code": e.code, "message": e.message})
@router.post("/v1/chunk", summary="Split a document into retrieval chunks")
async def create_chunks(request: Request, body: ChunkRequest = None):
"""Chunk a document for a retrieval index.
Structure-aware by default: the document is sectioned with a template, then
length-chunked INSIDE each section, so a chunk never straddles a boundary
and every chunk carries a label and a recoverable `parent`.
This is NOT embedding windowing (`embed_long_input`), which makes one input
fit a model window and pools it into ONE vector. This never pools and never
drops text.
Pure CPU: token counts are approximated as characters / `chars_per_token`,
so no tokenizer is loaded and no model or GPU is involved.
"""
if body is None:
body = ChunkRequest()
text = body.text or ""
if not text:
raise HTTPException(status_code=400, detail={
"code": "empty_input",
"message": "provide 'text' (JSON) or 'file' (multipart)"})
return _run(_resolve_params(text, body))
@router.post("/v1/chunk/file", summary="Split an uploaded document into chunks")
async def create_chunks_from_file(
file: UploadFile = File(...),
chunk_tokens: Optional[int] = Form(None),
overlap_tokens: Optional[int] = Form(None),
chars_per_token: Optional[int] = Form(None),
structural: Optional[bool] = Form(None),
structure: Optional[str] = Form(None),
extract_refs: Optional[str] = Form(None),
require_text: Optional[bool] = Form(None),
include_offsets: Optional[bool] = Form(None),
):
"""The multipart form of `/v1/chunk`, with the same field names as values --
matching how `/v1/ocr` takes its options.
The upload is decoded as UTF-8 text. A PDF or an OOXML container is REFUSED
(`not_text`) rather than chunked: extract its text first (`/v1/ocr` does
that for a scan). Chunking a container embeds garbage that the index cannot
distinguish from content afterwards.
"""
raw = await file.read()
why = looks_binary(raw) if (require_text in (None, True)) else None
if why:
raise HTTPException(status_code=400, detail={"code": "not_text", "message": why})
body = ChunkRequest(
chunk_tokens=chunk_tokens, overlap_tokens=overlap_tokens,
chars_per_token=chars_per_token, structural=structural,
structure=structure, require_text=require_text,
include_offsets=include_offsets,
extract_refs=(False if (extract_refs or "").strip().lower()
in ("0", "false", "no", "off") else (extract_refs or None)))
return _run(_resolve_params(raw.decode("utf-8", "ignore"), body))
@router.get("/v1/chunk/templates", summary="List the structure templates on offer")
async def list_templates():
"""Which names `/v1/chunk` accepts, and what each one sections on.
An unknown name is an error rather than a silent fallback to the built-in,
so a caller needs to be able to see the list.
"""
reg = _registered_templates()
out = []
for name, tpl in list(BUILTIN_TEMPLATES.items()) + list(reg.items()):
if not isinstance(tpl, dict):
continue
out.append({
"name": name,
"builtin": name in BUILTIN_TEMPLATES and name not in reg,
"head_label": tpl.get("head_label"),
"merge_adjacent": bool(tpl.get("merge_adjacent", True)),
"labels": [str(s.get("label") or "") for s in (tpl.get("sections") or [])
if isinstance(s, dict)],
})
return {"templates": out}
codai/api/remote_gateway.py
View file @ 29233812
......@@ -72,6 +72,17 @@ _CAPABILITY_PREFIXES = (
("/v1/video", "video"),
("/v1/3d", "spatial"),
("/v1/embeddings", "embeddings"),
# Retrieval chunking rides the EMBEDDINGS capability rather than inventing
# one of its own. Two reasons: chunks are on their way to be embedded, so
# the machine that will embed them is the natural place to cut them; and the
# embeddings image is full-core, so it carries the chunker. Inventing a
# `chunk` capability would mean a capability nothing advertises, and renting
# a GPU pod to run a regex split.
#
# It is pure CPU either way -- token counts are characters /
# chars_per_token, no tokenizer, no model -- so forwarding it is about WHERE
# the CPU is, which matters when the orchestrator is the busy box.
("/v1/chunk", "embeddings"),
("/v1/rerank", "rerank"),
("/v1/ocr", "ocr"),
("/v1/loras", "loras"),
......@@ -144,6 +155,51 @@ _MULTIPART_MODEL = re.compile(
rb'name="model"\r?\n(?:[^\r\n]*\r?\n)*?\r?\n([^\r\n]*)', re.IGNORECASE)
def _inline_chunk_template(path: str, headers: dict, body: bytes) -> bytes:
"""Resolve a NAMED /v1/chunk structure template into its object before the
request leaves this machine.
A template registered in THIS install's config does not exist on the
remote. Forwarded by name it would come back 404 `unknown_template` -- or,
far worse, match a DIFFERENT template that happens to share the name and
return different chunk boundaries with no indication that anything
happened. Chunk boundaries that silently differ by which machine answered
make a retrieval index unmeasurable, which is the whole thing
docs/chunking.md warns about.
Built-in names are left alone: every coderai has those, and sending the
object would only make the request bigger.
Never raises: a body it cannot parse is forwarded untouched, and the remote
then answers for itself.
"""
if (path or "").rstrip("/") != "/v1/chunk":
return body
if "json" not in (headers.get("content-type") or "").lower():
return body # multipart goes to /v1/chunk/file; same rule there
try:
import json as _json
obj = _json.loads(body or b"{}")
if not isinstance(obj, dict):
return body
name = obj.get("structure")
if not isinstance(name, str) or not name.strip():
return body # already inline, or defaulted
from codai.api.chunking import BUILTIN_TEMPLATES, resolve_template
if name.strip() in BUILTIN_TEMPLATES:
return body
from codai.api.chunking import _registered_templates
tpl = resolve_template(name, _registered_templates())
if not isinstance(tpl, dict):
return body
obj["structure"] = tpl
print(f"[remote] /v1/chunk: inlining the {name!r} template so the remote "
f"cuts on the same boundaries this install would", flush=True)
return _json.dumps(obj).encode("utf-8")
except Exception:
return body
def capability_for(path: str) -> str:
"""The capability a /v1 path belongs to, or '' when it maps to none."""
p = (path or "").rstrip("/") or "/"
......@@ -673,6 +729,7 @@ class RemoteGatewayMiddleware:
return await _error(send, 502, f"remote routing failed: {exc}")
if not target:
return await self._app(scope, replay, send)
body = _inline_chunk_template(scope.get("path") or "", headers, body)
await self._forward(scope, headers, body, target, send)
async def _forward(self, scope, headers, body, target, send):
......
codai/api/runpod_worker.py
View file @ 29233812
......@@ -1254,46 +1254,72 @@ def _plan(engine, image, args, mcfg, api_key, entry, served,
# capability pod is not what makes its pod able to read a page.
if pod_runs_ocr_pipeline(mcfg, image):
_o = _ocr_cfg()
# WHICH engine this pod's pipeline is for. A pod serves one model, and
# that model is what decides -- surya-ocr-2, chandra-ocr-2 or olmOCR.
# Hardcoding surya told a chandra pod to run surya against chandra's
# model id, which would have produced plausible-looking text from the
# wrong engine.
_served_l = str(served or "").lower()
_eng = "surya"
for _cand in ("chandra", "olmocr", "surya"):
if _cand in _served_l:
_eng = _cand
break
env["CODERAI_OCR_ENABLED"] = "1"
env["CODERAI_OCR_SURYA_ENABLED"] = "1"
env["CODERAI_OCR_DEFAULT_ENGINE"] = "surya"
env[f"CODERAI_OCR_{_eng.upper()}_ENABLED"] = "1"
env["CODERAI_OCR_DEFAULT_ENGINE"] = _eng
# The baked venv, so the pod does not try to BUILD one on a rented
# machine (minutes, and it would land on the container disk).
env["CODERAI_OCR_SURYA_VENV"] = "/opt/coderai/venvs/surya"
# machine (minutes, and it would land on the container disk). Only
# surya has one: chandra and olmOCR in served mode speak HTTP from the
# main venv (PIL + requests), which is why the pod side is nearly free
# for them and needed a 1.6 GB venv for surya.
if _eng == "surya":
env["CODERAI_OCR_SURYA_VENV"] = "/opt/coderai/venvs/surya"
# Surya drives the model through the pod's OWN coderai, in the same
# "model" mode the orchestrator uses -- which is what makes the gateway
# mint it an attach token instead of needing a bearer hand-wired here.
# Over loopback, so a page resumed four times costs four local round
# trips rather than four across the internet.
env["CODERAI_OCR_SURYA_SERVE"] = "model"
env["CODERAI_OCR_SURYA_MODEL_ID"] = str(served or "").strip()
# Surya is GPL and gated on an explicit acceptance. The pod inherits the
# decision made here; it cannot make it for itself.
env["CODERAI_OCR_SURYA_ACCEPT_LICENSE"] = "1"
env[f"CODERAI_OCR_{_eng.upper()}_SERVE"] = "model"
env[f"CODERAI_OCR_{_eng.upper()}_MODEL_ID"] = str(served or "").strip()
# Licences are inherited, never decided by the pod: surya is GPL and
# chandra's weights are a modified OpenRAIL-M. Asking for the engine by
# name HERE is the decision.
if _eng == "surya":
env["CODERAI_OCR_SURYA_ACCEPT_LICENSE"] = "1"
elif _eng == "chandra":
env["CODERAI_OCR_CHANDRA_ACCEPT_LICENSE"] = "1"
# This image carries NO docTR -- `auto` would route a failed page to an
# engine that is not there. Surya or nothing, said out loud.
# engine that is not there. This engine or nothing, said out loud.
env["CODERAI_OCR_FALLBACK_ENGINES"] = "off"
# And the TUNING. Every one of these was set in response to a measured
# failure, and a pod reading pages at surya's defaults instead would be
# indistinguishable from the "pre-fix pod image" that got blamed three
# times for orchestrator bugs.
for _var, _attr in (("CODERAI_OCR_DPI", "dpi"),
("CODERAI_OCR_SURYA_DPI", "surya_dpi"),
("CODERAI_OCR_SURYA_MAX_TOKENS", "surya_max_tokens"),
("CODERAI_OCR_SURYA_MAX_CONTINUATIONS",
"surya_max_continuations"),
("CODERAI_OCR_SURYA_CONTEXT_TOKENS",
"surya_context_tokens"),
("CODERAI_OCR_SURYA_MAX_RETRIES", "surya_max_retries"),
("CODERAI_OCR_SURYA_REPEAT_MAX_REPEATS",
"surya_repeat_max_repeats"),
("CODERAI_OCR_SURYA_REPEAT_WINDOW",
"surya_repeat_window"),
("CODERAI_OCR_SURYA_REPEAT_SCALING",
"surya_repeat_scaling"),
("CODERAI_OCR_SURYA_LANES", "surya_lanes"),
("CODERAI_OCR_SURYA_INSTANCES", "surya_instances"),
("CODERAI_OCR_SURYA_LANGS", "surya_langs"),
# Per-engine tuning, by name, so a chandra pod gets chandra's budgets
# rather than surya's. Every one of these was set in response to a
# measured failure, and a pod running the engine's own defaults instead
# is indistinguishable from the "pre-fix pod image" that got blamed
# three times for orchestrator bugs.
_E = _eng.upper()
_per_engine = [(f"CODERAI_OCR_{_E}_DPI", f"{_eng}_dpi"),
(f"CODERAI_OCR_{_E}_MAX_TOKENS", f"{_eng}_max_tokens"),
(f"CODERAI_OCR_{_E}_MAX_CONTINUATIONS",
f"{_eng}_max_continuations"),
(f"CODERAI_OCR_{_E}_CONTEXT_TOKENS",
f"{_eng}_context_tokens"),
(f"CODERAI_OCR_{_E}_MAX_RETRIES", f"{_eng}_max_retries"),
(f"CODERAI_OCR_{_E}_INSTANCES", f"{_eng}_instances"),
(f"CODERAI_OCR_{_E}_LANES", f"{_eng}_lanes"),
(f"CODERAI_OCR_{_E}_LANGS", f"{_eng}_langs")]
if _eng == "surya":
_per_engine += [
("CODERAI_OCR_SURYA_REPEAT_MAX_REPEATS",
"surya_repeat_max_repeats"),
("CODERAI_OCR_SURYA_REPEAT_WINDOW", "surya_repeat_window"),
("CODERAI_OCR_SURYA_REPEAT_SCALING", "surya_repeat_scaling")]
for _var, _attr in _per_engine + [
("CODERAI_OCR_DPI", "dpi"),
("CODERAI_OCR_MAX_CONCURRENCY", "max_concurrency"),
("CODERAI_OCR_PAGE_WINDOW", "page_window"),
("CODERAI_OCR_RASTER_WORKERS", "raster_workers"),
......@@ -1301,7 +1327,7 @@ def _plan(engine, image, args, mcfg, api_key, entry, served,
"blank_page_max_chars"),
("CODERAI_OCR_DETECT_MODE", "detect_mode"),
("CODERAI_OCR_DETECT_MODEL_PATH", "detect_model_path"),
("CODERAI_OCR_DETECT_CONF", "detect_conf")):
("CODERAI_OCR_DETECT_CONF", "detect_conf")]:
_val = getattr(_o, _attr, None) if _o is not None else None
if _val not in (None, "", 0, 0.0):
env[_var] = str(_val)
......@@ -1321,9 +1347,9 @@ def _plan(engine, image, args, mcfg, api_key, entry, served,
"it was never given, and will report no stamps rather than "
"fail. Set the path (staged on the pod) or use detect_mode="
"layout.", flush=True)
print(f"[runpod] the OCR PIPELINE runs on this pod ({image}): it is sent "
f"one page as a PDF and reads it on the card, with its own vLLM "
f"over loopback", flush=True)
print(f"[runpod] the {_eng} OCR PIPELINE runs on this pod ({image}): it "
f"is sent one page as a PDF and reads it on the card, with its own "
f"vLLM over loopback", flush=True)
if cap == "ocr":
env["CODERAI_OCR_ENABLED"] = "1"
......@@ -2719,7 +2745,15 @@ class RunpodPodPool:
print(f"[runpod] {_dc or '(no region)'}: attaching volume "
f"{_v or '(none — this region has no volume configured)'}",
flush=True)
print(f"[runpod] provisioning pod for {self.model_key!r}: {sel['display_name']} "
# ONE pod request walks this whole list until a (card, region) pair
# has capacity, so this line is printed PER CANDIDATE -- and it used
# to say "provisioning pod", which reads as a pod being launched.
# 312 attempts for 10 pods in one afternoon therefore looked like
# runaway provisioning to an operator reading the log, and was
# reported as such. Say which attempt this is, out of how many, so
# the shape of the list is visible instead of inferred.
print(f"[runpod] {self.model_key!r}: trying candidate "
f"{i + 1}/{len(ranked)} for ONE pod — {sel['display_name']} "
f"({sel['cloud_type']}{'/spot' if sel['is_spot'] else ''}) ${sel['price']}/hr"
+ (f" in {_dc}" if _dc else " (region: RunPod's choice)") + _alt,
flush=True)
......
codai/config.py
View file @ 29233812
......@@ -380,6 +380,41 @@ class EnhanceConfig:
allow_rife_ncnn: bool = False # allow the external rife-ncnn-vulkan binary instead of a torch model
@dataclass
class ChunkingConfig:
"""Retrieval chunking (`POST /v1/chunk`).
NOT embedding windowing. `ocr`/`embeddings` settings make one input fit a
model window and pool it into ONE vector; this splits a document into MANY
separately-indexable chunks, each with its section, parent and references.
Token counts are approximated as characters / `chars_per_token`, so the
endpoint loads no tokenizer and needs no model, no GPU and no pod -- which
is what lets every image serve it, light-core ones included.
"""
# The defaults are the ones a production index was built with: changing them
# changes every chunk boundary, which makes a re-chunked index incomparable
# with the one it replaces and retrieval quality unmeasurable across the
# change. Treat them as a format, not a preference.
chunk_tokens: int = 800 # target size, approximated tokens
overlap: int = 120 # carried tail, approximated tokens
chars_per_token: int = 4 # the approximation divisor
structural: bool = True # False = flat length chunking
structure: str = "it-legal" # default template name
require_text: bool = True # refuse container bytes (see below)
# A chunker that accepts a PDF or an OOXML container embeds garbage that the
# index cannot afterwards distinguish from content. Measured on a production
# corpus: 0.59% of chunks held control characters, 0.04% raw OOXML --
# documents whose container was chunked instead of its extracted text.
# `require_text: false` opts out deliberately.
max_chars: int = 8_000_000 # above this -> 413 too_large
# Operator-supplied structure templates, same JSON shape as the built-in
# `it-legal`, selected by name. An unknown name is an ERROR, never a silent
# fallback: a caller who asked for their own template and silently got
# Italian legal sectioning would index the wrong boundaries and never know.
templates: dict = field(default_factory=dict)
@dataclass
class CompactionConfig:
"""Global defaults for auto-compaction of an over-long chat history.
......@@ -1379,6 +1414,7 @@ class Config:
runpod: RunpodConfig = field(default_factory=RunpodConfig)
ocr: OcrConfig = field(default_factory=OcrConfig)
compaction: CompactionConfig = field(default_factory=CompactionConfig)
chunking: ChunkingConfig = field(default_factory=ChunkingConfig)
broker: BrokerConfig = field(default_factory=BrokerConfig)
system_prompt: Optional[str] = None
tools_closer_prompt: bool = False
......@@ -1576,6 +1612,7 @@ class ConfigManager:
runpod=_dc(RunpodConfig, config_data.get("runpod", {})),
ocr=_dc(OcrConfig, config_data.get("ocr", {})),
compaction=_dc(CompactionConfig, config_data.get("compaction", {})),
chunking=_dc(ChunkingConfig, config_data.get("chunking", {})),
broker=_dc(BrokerConfig, config_data.get("broker", {})),
system_prompt=config_data.get("system_prompt"),
tools_closer_prompt=config_data.get("tools_closer_prompt", False),
......@@ -2017,6 +2054,16 @@ class ConfigManager:
"extract_schema": self.config.ocr.extract_schema,
"extract_validate": self.config.ocr.extract_validate,
},
"chunking": {
"chunk_tokens": self.config.chunking.chunk_tokens,
"overlap": self.config.chunking.overlap,
"chars_per_token": self.config.chunking.chars_per_token,
"structural": self.config.chunking.structural,
"structure": self.config.chunking.structure,
"require_text": self.config.chunking.require_text,
"max_chars": self.config.chunking.max_chars,
"templates": self.config.chunking.templates,
},
"compaction": {
"enabled": self.config.compaction.enabled,
"pct": self.config.compaction.pct,
......
codai/ocr/manager.py
View file @ 29233812
......@@ -349,7 +349,7 @@ def _warn_if_oversubscribed(name: str, pages: int) -> None:
def _post_page_to_pod(base: str, api_key: str, blob: bytes, name: str,
dpi: int, opts: dict, detect) -> dict:
dpi: int, opts: dict, detect, engine: str = "surya") -> dict:
"""POST one page to a pod's own /v1/ocr and return its response dict.
Through `pod_http`, not `requests`: a pod on the direct path serves HTTPS
......@@ -363,7 +363,10 @@ def _post_page_to_pod(base: str, api_key: str, blob: bytes, name: str,
url = url + "/ocr"
else:
url = url + "/v1/ocr"
fields = {"engine": "surya"}
# The engine the POD must use. Hardcoded "surya" sent a chandra page to
# surya, which the pod would have served with the wrong model -- and the
# text would have come back looking plausible.
fields = {"engine": str(engine or "surya")}
if dpi:
fields["dpi"] = str(int(dpi))
if detect:
......@@ -420,6 +423,9 @@ def _page_from_pod(d: dict, index: int) -> OcrPage:
def pod_pipeline_model(cfg, engine: str = "") -> str:
"""The model id whose pods run the OCR pipeline themselves, or "".
Answers for whichever VLM engine is being used -- surya, chandra or
olmOCR -- each read from its own `<engine>_serve` / `<engine>_model_id`.
THE IMAGE IS THE SWITCH. A model served by `coderai-vllm-ocr` carries surya
and the rasteriser, so the page goes to it as a one-page PDF; one served by
`coderai-vllm` serves only the model and the pipeline stays here. There is
......@@ -429,12 +435,19 @@ def pod_pipeline_model(cfg, engine: str = "") -> str:
Returns "" for every other case -- no model mode, no entry, a different
image -- so the caller simply does what it has always done.
"""
name = (engine or "").strip().lower()
if name and name != "surya":
name = (engine or "").strip().lower() or str(
getattr(cfg, "default_engine", "") or "").strip().lower()
# All three VLM engines can run on the pod, not just surya. The image is
# what differs: surya needs its client venv baked in, while chandra and
# olmOCR in served mode speak only HTTP (PIL + requests, already present),
# so for them the pod side costs nothing beyond the rasteriser. Leaving
# them on the orchestrator while surya moved was an asymmetry, not a
# design.
if name not in ("surya", "chandra", "olmocr"):
return ""
if str(getattr(cfg, "surya_serve", "") or "").strip().lower() != "model":
if str(getattr(cfg, f"{name}_serve", "") or "").strip().lower() != "model":
return ""
mid = str(getattr(cfg, "surya_model_id", "") or "").strip()
mid = str(getattr(cfg, f"{name}_model_id", "") or "").strip()
if not mid:
return ""
try:
......@@ -1217,7 +1230,9 @@ class OcrManager:
try:
_out = await self._read_on_pods(
data, filename, content_type, _remote_model,
dpi=use_dpi, page_opts=page_opts, detect=detect, cfg=cfg)
dpi=use_dpi, page_opts=page_opts, detect=detect, cfg=cfg,
engine_name=(engine or "").strip().lower()
or str(getattr(cfg, "default_engine", "") or "").lower())
except OcrError:
raise
except Exception as exc:
......@@ -1275,7 +1290,8 @@ class OcrManager:
return 0
async def _read_on_pods(self, data: bytes, filename: str, content_type: str,
model_id: str, *, dpi: int, page_opts, detect, cfg):
model_id: str, *, dpi: int, page_opts, detect, cfg,
engine_name: str = "surya"):
"""Send ONE PAGE PER REQUEST to pods that run the pipeline themselves.
The page travels as a small PDF and the pod renders and reads it on the
......@@ -1318,7 +1334,7 @@ class OcrManager:
return await asyncio.to_thread(
_post_page_to_pod, lease.base, lease.api_key, blob,
f"{idx}.pdf" if is_pdf else (filename or "page"),
dpi, opts, detect)
dpi, opts, detect, engine_name)
finally:
lease.release()
......
docs/chunking.md 0 → 100644
View file @ 29233812
# Retrieval chunking — `POST /v1/chunk`
Splitting a document into the pieces a retrieval index stores, with the
document's own structure decided **per request** by a named (or inline)
*structure template*.
This is a port of the chunker that runs in Digesta (`app/chunking.py`), which
has been chunking Italian case law and party filings in production. The port has
one hard requirement: with the built-in `it-legal` template and default sizes it
must produce **byte-identical** output to Digesta, so Digesta can delegate
chunking here and keep its existing index comparable. Everything that is
hardcoded Italian-legal there becomes a template here.
## Not to be confused with embedding windowing
CoderAI already splits over-long embedding inputs — `embed_long_input`
(`chunk | truncate | error`) and `embed_chunk_pool` (`mean | first`) in
`codai/api/embeddings.py`. That solves a different problem and must stay as it
is:
| | `embed_long_input` (exists) | `/v1/chunk` (this doc) |
|---|---|---|
| Purpose | make one input fit the model window | split a document for retrieval |
| Output | **one** vector per input (windows pooled) | **many** chunks, each indexable |
| Metadata | none | section, parent, references, offsets |
| Chosen by | model window | the corpus and the retrieval design |
A caller may well use both: `/v1/chunk` to cut the document, then
`/v1/embeddings` per chunk (where windowing is the safety net for a chunk that
still overflows). `/v1/chunk` must never pool and must never drop text.
## Why structure-aware at all
Length-only chunking cuts a legal argument wherever the character budget runs
out. The measured consequence in Digesta was that a chunk could straddle the
`fatto` → `motivazione` boundary, so a retrieved passage mixed what the parties
alleged with what the court held. Sectioning first, then length-chunking *inside*
each section, removes that class of error, and gives every chunk a label and a
recoverable parent section.
Two findings from production worth carrying into the design:
- **One marker family is not enough.** With ruling markers only
(`fatto`/`motivazione`/`dispositivo`), just **8–10%** of 400 merito documents
were recognised — party filings simply do not have that structure. The
`it-legal` template therefore carries two families (see below).
- **Adjacent duplicate sections must merge.** Filings routinely open the
conclusions twice (`CONCLUSIONI`, then `Voglia il Tribunale…`). Left separate,
one block becomes two and the `parent` ids multiply, defeating the purpose of
`parent`.
Current behaviour on the live corpus (0.05% sample of ~27M chunks): ~78% of
chunks carry a section label, ~22% fall back to flat chunking.
## Request
```http
POST /v1/chunk
Authorization: Bearer <key>
Content-Type: application/json
```
```jsonc
{
"text": "…", // or "file" (multipart) — exactly one of them
"chunk_tokens": 800, // target size, in approximated tokens
"overlap_tokens": 120, // carried tail, in approximated tokens
"chars_per_token": 4, // the approximation divisor
"structural": true, // false = flat length chunking, no sectioning
"structure": "it-legal", // template name, or an inline object (below)
"extract_refs": "it-legal", // reference extractor: name, inline, or false
"require_text": true, // reject binary input instead of chunking it
"include_offsets": false // add char offsets into the source text
}
```
Every field except the input is optional; omitted fields take the server
defaults, which are the Digesta defaults. Multipart form submission uses the
same field names as string values, matching how `/v1/ocr` takes its options.
### Defaults and config keys
| Field | Default | Config key | Env |
|---|---|---|---|
| `chunk_tokens` | 800 | `chunk_tokens` | `CODERAI_CHUNK_TOKENS` |
| `overlap_tokens` | 120 | `chunk_overlap` | `CODERAI_CHUNK_OVERLAP` |
| `chars_per_token` | 4 | `chunk_chars_per_token` | `CODERAI_CHUNK_CHARS_PER_TOKEN` |
| `structural` | `true` | `chunk_structural` | `CODERAI_CHUNK_STRUCTURAL` |
| `structure` | `it-legal` | `chunk_structure` | `CODERAI_CHUNK_STRUCTURE` |
| `require_text` | `true` | `chunk_require_text` | `CODERAI_CHUNK_REQUIRE_TEXT` |
Token counts are **approximated as characters ÷ `chars_per_token`** so the
endpoint needs no tokenizer and no model loaded — `/v1/chunk` must stay a pure
CPU text operation, callable with no GPU and no pod. With the defaults the real
budget is `800 × 4 = 3200` characters and the overlap `120 × 4 = 480`.
## Response
```jsonc
{
"chunks": [
{
"ord": 0, // 0-based, document order
"text": "…",
"n_chars": 2464,
"section": "motivazione", // null when unlabelled
"parent": 2, // 0-based section index; shared by its chunks
"norm_refs": ["art.2043 c.c.", "d.lgs.209/2005"],
"start": 10433, "end": 12897 // only with include_offsets
}
],
"n_chunks": 37,
"sections": [ // the segmentation, for context recovery
{"parent": 0, "section": "intestazione", "n_chars": 1204}
],
"structure": "it-legal", // template actually applied
"structural_applied": true, // false = fell back to flat
"params": {"chunk_tokens": 800, "overlap_tokens": 120, "chars_per_token": 4}
}
```
`parent` is the index into `sections`, so a caller can retrieve a chunk and then
pull the whole section it came from as context — the behaviour Digesta's
retrieval depends on.
## The structure template
A template is an ordered list of sections, each a label plus a pattern that
*opens* that section. Order in the list is only for registration; **order in the
document drives the segmentation**.
```jsonc
{
"name": "it-legal",
"head_label": "intestazione", // text before the first match; null = drop it
"flags": "im", // applied to every pattern
"line_prefix": "[ \\t>*_#]*", // tolerated before a header (markdown, quotes)
"merge_adjacent": true, // merge consecutive same-label sections
"keep_header_line": true, // the header stays in its section's text
"sections": [
{"label": "premessa", "pattern": "tanto\\s+premesso|…"},
{"label": "fatto", "pattern": "svolgimento\\s+del\\s+processo|…"}
]
}
```
Rules an implementation must honour:
- Patterns are anchored at a line start, after `line_prefix`, and are matched
case-insensitively and multi-line (`flags` default `im`). This is what lets a
header survive markdown emphasis or quoted text.
- All patterns compile into **one** alternation with named groups, so a single
left-to-right scan yields the marks in document order and the first matching
alternative wins at a given position. Do not scan per-pattern and merge: that
changes which header claims an overlapping position.
- `keep_header_line: true` means a section runs from its match **start** to the
next match start, header included.
- `merge_adjacent: true` joins a section to the previous one with `"\n"` when
the labels are equal.
- No match anywhere → `structural_applied: false` and the whole text is chunked
flat as a single unlabelled section.
Templates may be supplied inline per request, or registered server-side as
`chunk_templates` in the config (same JSON shape) and selected by name. Inline
beats registered; an unknown name is an error, never a silent fallback.
### Built-in `it-legal`
Carries both families, because party filings outnumber rulings in the merito
corpus. Patterns, verbatim from the production implementation:
| Label | Pattern |
|---|---|
| `premessa` | `tanto\s+premesso\|ci[oò]\s+premesso\|^\s*[*#>_ ]*premesso\s+che\b\|^\s*premess[ae]\s*$` |
| `fatto` | `svolgimento\s+del\s+processo\|ragioni\s+di\s+fatto\|in\s+fatto\b\|fatto\s+e\s+diritto\|^\s*fatto\s*$\|esposizione\s+del\s+fatto` |
| `motivazione` | `motivi\s+della\s+decisione\|ragioni\s+della\s+decisione\|motivazione\|considerato\s+in\s+diritto\|in\s+diritto\b\|motivi\s+in\s+diritto\|ragioni\s+in\s+diritto` |
| `dispositivo` | `p\.?\s*q\.?\s*m\.?\|p\.?\s*t\.?\s*m\.?\|per\s+questi\s+motivi\|^\s*dispositivo\s*$\|così\s+deciso\|così\s+provvede` |
| `conclusioni` | `^\s*[*#>_ ]*conclusion[ei]\b\|\bconclusioni\s*:\|voglia\s+(?:l['’]ill\.?mo\s+)?(?:il\s+)?(?:tribunale\|giudice\|giudicante)\|^\s*[*#>_ ]*chiede\b\|salvis\s+iuribus` |
**`conclusioni` is a label of its own and must never be folded into
`dispositivo`.** In a ruling the `dispositivo` is what the judge decided; in a
filing the `conclusioni` are what a party *asked for*. Merging them would let a
lawyer's demand surface as if it were a holding — the most damaging mistake this
index can make. A template that wants them merged must say so explicitly by
giving both patterns the same label.
## Algorithm
Implement exactly this order; the overlap and packing rules are where a port
silently diverges.
**1. Normalise.** `text = text.replace("\r", "")`. Empty after strip → `chunks: []`.
**2. Hygiene** (`require_text`). Reject input that is not text: a ZIP/OOXML
header (`PK\x03\x04`, `[Content_Types].xml`), `%PDF-`, or control characters
outside `\t\n\f` above a small ratio → `400 not_text`. This guard exists because
the Digesta corpus contains chunks holding raw container bytes — measured at
0.59% of chunks with control characters and 0.04% raw OOXML, concentrated in PCT
party documents whose container was chunked instead of its extracted text. A
chunker that accepts those bytes embeds garbage and the index cannot tell.
`require_text: false` opts out deliberately.
**3. Sections.** Per the template above. `structural: false` → one unlabelled
section holding everything.
**4. Paragraphs, per section.** Split on `\n\s*\n`; strip; drop empties.
**5. Over-long paragraphs.** A paragraph longer than `target` is split on
sentence boundaries `(?<=[.;:!?])\s+`, greedily refilled to `target`; a single
sentence still longer than `target` is cut into fixed `target`-sized slices.
Nothing is ever discarded.
**6. Pack with a carried tail.** Within a section:
```python
chunks, cur = [], ""
for p in pieces:
if cur and len(cur) + len(p) + 2 > target:
chunks.append(cur)
cur = (cur[-overlap:] + "\n\n" + p) if overlap else p
else:
cur = (cur + "\n\n" + p).strip() if cur else p
if cur:
chunks.append(cur)
```
Three details that must be reproduced:
- Paragraphs are joined with `"\n\n"`, and the `+ 2` in the test accounts for it.
- The overlap is a **raw character tail** of the emitted chunk, `cur[-overlap:]`,
prepended to the next one. It is not sentence- or token-aligned and may begin
mid-word. That is the production behaviour; "improving" it changes every
existing chunk boundary.
- `overlap == 0` starts the next chunk at `p` with no carry.
Consequence to expect in tests: a chunk may reach `target + overlap` characters
(3680 with the defaults). Observed max on the live corpus is 3682, mean 2158,
median 2464.
**7. Emit.** For each chunk, in document order: `ord`, `text`, `n_chars`,
`section` (`null` when the label is empty), `parent` (section index), and
`norm_refs` when an extractor is configured.
## Reference extraction
Optional, per request, and template-driven like the structure. `it-legal`
extracts article and act citations, normalises them, and de-duplicates while
**preserving first-seen order**:
- Articles → `art.<num> <code>`, e.g. `art.2043 c.c.`, where `num` keeps a
`bis`/`ter`/`quater`/`quinquies`/`sexies` suffix with spaces removed, and the
code normalises through `codciv|cc → c.c.`, `cpc → c.p.c.`, `cpp → c.p.p.`,
`cp → c.p.`, `cost → cost.`
- Acts → `<kind><num>/<year>`, e.g. `d.lgs.209/2005`, with
`dlgs → d.lgs.`, `dl → d.l.`, `dpr → d.p.r.`, `legge|l → l.`
These are a filter field and an exact-match signal for hybrid search — the thing
dense embeddings blur, since `art. 2043` and `art. 2054` sit close in vector
space and are entirely different claims. A template may replace the patterns;
`extract_refs: false` skips the work.
## Fidelity tests
The port is correct when, for the same input with `it-legal` and the defaults,
it matches Digesta's `chunk_document()` output exactly — same count, same
boundaries, same labels, same `parent`, same refs. Build the fixtures from the
live corpus rather than synthetic text, and include:
1. A ruling with `fatto`/`motivazione`/`dispositivo`.
2. A filing with `premessa`/`conclusioni`, **with the conclusions opened twice**
(merge path).
3. A document with no recognisable header (flat fallback, `section: null`).
4. A single paragraph far over `target` (sentence split, then slice path).
5. A section boundary that lands mid-paragraph (the overlap carry).
6. A document whose header is markdown-decorated (`**FATTO**`, `> fatto`).
7. Binary input (OOXML) → `400 not_text`.
8. `overlap_tokens: 0` and `structural: false` → degenerate paths.
Assert on the `target + overlap` ceiling too: a chunk longer than that means the
packing test was ported wrong.
## Errors
| Status | Code | When |
|---|---|---|
| 400 | `not_text` | `require_text` and the input is binary |
| 400 | `empty_input` | no `text` and no `file`, or empty after normalisation |
| 400 | `bad_template` | inline template malformed, or a pattern fails to compile |
| 404 | `unknown_template` | named template not registered |
| 413 | `too_large` | input above the configured ceiling |
Patterns come from callers, so compile every template with a size limit and
treat a pathological pattern as `bad_template` rather than letting a scan run
unbounded on operator-supplied input.
## Notes for the Digesta side
Digesta keeps chunk text in Postgres (`chunks.text`, `store='clear'`, currently
54 GB) and vectors in Qdrant, which is what makes a re-chunk possible without
re-OCR. When it switches to this endpoint, the migration is only worthwhile if
output is identical — otherwise old and new chunks coexist in one collection with
different boundaries and retrieval quality becomes unmeasurable. So: port,
verify byte-identity on fixtures, then cut over; do not mix.
tests/test_chunking.py 0 → 100644
View file @ 29233812
"""Retrieval chunking — `POST /v1/chunk`.
Splits a document into the pieces a retrieval index stores. NOT embedding
windowing: `embed_long_input` makes one over-long input fit a model window and
pools the pieces into ONE vector; this produces MANY chunks, each indexable,
each carrying its section, parent and references. This must never pool and
never drop text.
The eight cases below are the fidelity list from docs/chunking.md. What they
CANNOT do is prove byte-identity with the production chunker they port: that
needs fixtures from the live corpus, which are not in this repo. They prove the
algorithm's documented behaviour on each path; the identity check is a separate
step against real documents, and it is the one that decides whether a re-chunk
is safe.
"""
import pathlib
import sys
import pytest
ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from codai.api import chunking as ch # noqa: E402
DEFAULTS = dict(chunk_tokens=800, overlap_tokens=120, chars_per_token=4)
TARGET = 800 * 4 # 3200
OVERLAP = 120 * 4 # 480
# ------------------------------------------------------------ 1. a ruling #
RULING = """TRIBUNALE ORDINARIO DI MILANO
Sezione Prima Civile
SVOLGIMENTO DEL PROCESSO
Con atto di citazione notificato il 3 marzo l'attore ha convenuto in giudizio.
Il convenuto si e costituito chiedendo il rigetto della domanda.
MOTIVI DELLA DECISIONE
La domanda e fondata. Ai sensi dell'art. 2043 c.c. il danno va risarcito.
Si richiama altresi il d.lgs. 209/2005 in materia assicurativa.
P.Q.M.
Condanna il convenuto al pagamento delle spese.
"""
def test_a_ruling_is_sectioned_into_its_three_parts():
r = ch.chunk_document(RULING, **DEFAULTS)
labels = [s["section"] for s in r["sections"]]
assert labels == ["intestazione", "fatto", "motivazione", "dispositivo"]
assert r["structural_applied"] is True
assert r["structure"] == "it-legal"
def test_every_chunk_knows_its_section_and_parent():
r = ch.chunk_document(RULING, **DEFAULTS)
for c in r["chunks"]:
assert c["section"] == r["sections"][c["parent"]]["section"]
def test_the_header_line_stays_in_its_own_section():
r = ch.chunk_document(RULING, **DEFAULTS)
moti = [c for c in r["chunks"] if c["section"] == "motivazione"]
assert moti and "MOTIVI DELLA DECISIONE" in moti[0]["text"]
# --------------------------------------- 2. a filing, conclusions twice #
FILING = """ECC.MO TRIBUNALE DI ROMA
PREMESSO CHE
la parte ricorrente ha subito un danno.
CONCLUSIONI
Nel merito, accertare la responsabilita.
Voglia il Tribunale condannare il resistente alle spese.
"""
def test_adjacent_duplicate_sections_merge():
"""Filings routinely open the conclusions twice. Left separate, one block
becomes two and the `parent` ids multiply, defeating the point of
`parent`."""
r = ch.chunk_document(FILING, **DEFAULTS)
labels = [s["section"] for s in r["sections"]]
assert labels.count("conclusioni") == 1, labels
conc = [s for s in r["sections"] if s["section"] == "conclusioni"][0]
assert conc["n_chars"] > 0
def test_conclusioni_is_never_folded_into_dispositivo():
"""In a ruling the dispositivo is what the judge DECIDED; in a filing the
conclusioni are what a party ASKED FOR. Merging them would let a lawyer's
demand surface as a holding -- the most damaging mistake this index can
make."""
labels = [s["label"] for s in ch.IT_LEGAL["sections"]]
assert "conclusioni" in labels and "dispositivo" in labels
assert labels.index("conclusioni") != labels.index("dispositivo")
r = ch.chunk_document(FILING, **DEFAULTS)
assert "dispositivo" not in [s["section"] for s in r["sections"]]
# -------------------------------------------------- 3. the flat fallback #
def test_a_document_with_no_header_falls_back_flat_and_says_so():
r = ch.chunk_document("Una lettera qualunque.\n\nSenza intestazioni note.",
**DEFAULTS)
assert r["structural_applied"] is False
assert all(c["section"] is None for c in r["chunks"])
assert len(r["sections"]) == 1
def test_structural_false_skips_sectioning_entirely():
r = ch.chunk_document(RULING, structural=False, **DEFAULTS)
assert r["structural_applied"] is False
assert [s["section"] for s in r["sections"]] == [None]
# ------------------------------- 4. one paragraph far over the target #
#: The real ceiling is target + overlap + 2, not target + overlap. The carried
#: tail is exactly `overlap` characters, the next piece can be `target`, and they
#: are joined with "\n\n" -- so 3200 + 480 + 2 = 3682 with the defaults.
#: docs/chunking.md says "3680" in prose and "observed max on the live corpus is
#: 3682" two lines later; the arithmetic and the measurement agree with each
#: other, and the prose is the odd one out. Asserting 3680 would fail a correct
#: port.
CEILING = TARGET + OVERLAP + 2
def test_an_over_long_paragraph_splits_on_sentences():
para = " ".join(f"Frase numero {i} del paragrafo." for i in range(400))
r = ch.chunk_document(para, structural=False, **DEFAULTS)
assert r["n_chunks"] > 1
assert all(c["n_chars"] <= CEILING for c in r["chunks"])
assert max(c["n_chars"] for c in r["chunks"]) > TARGET, (
"the carried tail should push at least one chunk past the target")
def test_a_single_sentence_over_the_target_is_sliced():
"""Nothing is ever discarded -- a chunker that drops text is worse than one
that chunks badly, because the loss is invisible in the index."""
one = "x" * (TARGET * 3) # no sentence boundary at all
r = ch.chunk_document(one, structural=False, refs=False, **DEFAULTS)
assert r["n_chunks"] >= 3
assert sum(1 for c in r["chunks"] for _ in c["text"]) >= len(one)
def test_no_text_is_lost_across_the_pieces():
para = " ".join(f"Periodo {i}." for i in range(2000))
r = ch.chunk_document(para, structural=False, refs=False,
overlap_tokens=0, chunk_tokens=50, chars_per_token=4)
joined = "".join(c["text"] for c in r["chunks"])
for i in (0, 999, 1999):
assert f"Periodo {i}." in joined
# --------------------------------------------- 5/6. overlap and markdown #
def test_the_overlap_is_a_raw_character_tail():
"""Not sentence- or token-aligned, and it may begin mid-word. That is the
production behaviour; 'improving' it moves every existing boundary."""
pieces = ["A" * 100, "B" * 100, "C" * 100]
out = ch.pack(pieces, target=150, overlap=20)
assert len(out) >= 2
assert out[1].startswith("A" * 20), out[1][:30]
def test_zero_overlap_starts_the_next_chunk_clean():
out = ch.pack(["A" * 100, "B" * 100], target=150, overlap=0)
assert out[1] == "B" * 100
def test_a_chunk_may_reach_target_plus_overlap_but_no_more():
pieces = ["Z" * 300 for _ in range(20)]
out = ch.pack(pieces, target=1000, overlap=200)
assert out, "nothing packed"
assert max(len(c) for c in out) <= 1000 + 200
def test_a_quoted_or_prefixed_header_is_recognised():
"""LEADING decoration is tolerated: `line_prefix` ([ \\t>*_#]*) plus the
`[*#>_ ]*` some patterns carry themselves."""
doc = "Intestazione\n\n> motivazione\nLe ragioni.\n\n* premesso che\nI fatti.\n"
r = ch.chunk_document(doc, **DEFAULTS)
labels = [s["section"] for s in r["sections"]]
assert "motivazione" in labels and "premessa" in labels, labels
@pytest.mark.parametrize("header,label", [
("**FATTO**", "fatto"),
("*FATTO*", "fatto"),
("## FATTO", "fatto"),
("> fatto", "fatto"),
("FATTO", "fatto"),
(" fatto ", "fatto"),
("__DISPOSITIVO__", "dispositivo"),
("# Dispositivo #", "dispositivo"),
("**PREMESSA**", "premessa"),
("> motivazione", "motivazione"),
("* premesso che", "premessa"),
])
def test_a_markdown_decorated_header_is_recognised(header, label):
"""Decoration on EITHER side of a header is tolerated.
`**FATTO**` used to be missed twice over: `^\\s*fatto\\s*$` tolerates only
whitespace before the word and requires the line to end right after it, so
both the leading and the trailing asterisks defeated it. The patterns that
already coped (`premesso che`, `conclusioni`) spell it `^\\s*[*#>_ ]*` for
exactly this reason; these three never did.
"""
doc = f"Testa\n\n{header}\nCorpo del testo.\n"
got = [s["section"] for s in ch.chunk_document(doc, **DEFAULTS)["sections"]]
assert label in got, got
def test_widening_the_anchors_is_purely_additive():
"""Every line that matched the old `^\\s*X\\s*$` still matches the new
`^[\\s*#>_]*X[\\s*#>_]*$`, because \\s is a subset of that class. So no
EXISTING match moves; what changes is which documents get sectioned at all.
That distinction is the whole risk of the change: a document whose headers
were markdown-emphasised used to fall back to flat chunking and now gets
labels, so its chunks differ from an index built before this and it needs
re-chunking. A document already sectioned is untouched.
"""
import re
for word in ("fatto", "dispositivo", "premessa"):
old = re.compile(rf"^\s*{word}\s*$", re.I | re.M)
new = re.compile(rf"^{ch._LEAD}{word}{ch._TRAIL}$", re.I | re.M)
for line in (word, f" {word}", f"{word} ", f"\t{word}\t"):
if old.search(line):
assert new.search(line), f"{word}: {line!r} no longer matches"
def test_an_already_sectioned_document_is_unchanged_by_the_widening():
"""The ruling has plain headers, so its chunking must be identical to what
it was before the anchors were widened -- asserted by boundary, not by
eyeballing."""
r = ch.chunk_document(RULING, **DEFAULTS)
assert [s["section"] for s in r["sections"]] == [
"intestazione", "fatto", "motivazione", "dispositivo"]
# ----------------------------------------------------- 7. binary input #
@pytest.mark.parametrize("raw,why", [
(b"PK\x03\x04rest of a zip", "zip/OOXML header"),
(b"%PDF-1.7\nstuff", "pdf header"),
("[Content_Types].xml and more".encode(), "ooxml marker"),
])
def test_container_bytes_are_refused(raw, why):
"""A chunker that accepts those bytes embeds garbage and the index cannot
tell afterwards. Measured on a production corpus: 0.59% of chunks with
control characters, 0.04% raw OOXML."""
assert ch.looks_binary(raw), why
with pytest.raises(ch.ChunkError) as e:
ch.chunk_document(raw.decode("utf-8", "ignore"), **DEFAULTS)
assert e.value.code == "not_text"
def test_control_characters_above_a_small_ratio_are_refused():
assert ch.looks_binary("text" + "\x00" * 50)
def test_ordinary_text_with_tabs_and_newlines_is_fine():
assert ch.looks_binary("riga uno\n\triga due\n\fpagina\n") is None
def test_require_text_false_opts_out_deliberately():
r = ch.chunk_document("%PDF-1.7 pretend", require_text=False,
structural=False, refs=False, **DEFAULTS)
assert r["n_chunks"] == 1
# ------------------------------------------------- 8. degenerate paths #
def test_empty_input_is_an_error_not_an_empty_answer():
for bad in ("", " \n\n "):
with pytest.raises(ch.ChunkError) as e:
ch.chunk_document(bad, **DEFAULTS)
assert e.value.code == "empty_input"
def test_carriage_returns_are_normalised_away():
r = ch.chunk_document("riga\r\nuno\r\n\r\nriga due", structural=False,
refs=False, **DEFAULTS)
assert "\r" not in r["chunks"][0]["text"]
# ------------------------------------------------------------ templates #
def test_an_unknown_template_name_is_an_error_never_a_fallback():
"""A caller who asked for their own template and silently got Italian legal
sectioning would index the wrong boundaries and never know."""
with pytest.raises(ch.ChunkError) as e:
ch.chunk_document(RULING, structure="no-such-template", **DEFAULTS)
assert e.value.code == "unknown_template"
assert e.value.status == 404
def test_an_inline_template_is_accepted():
tpl = {"name": "mine", "head_label": "head", "merge_adjacent": True,
"sections": [{"label": "uno", "pattern": r"^CAPO\s+I\b"},
{"label": "due", "pattern": r"^CAPO\s+II\b"}]}
doc = "prefazione\n\nCAPO I\nprimo testo\n\nCAPO II\nsecondo testo\n"
r = ch.chunk_document(doc, structure=tpl, **DEFAULTS)
assert [s["section"] for s in r["sections"]] == ["head", "uno", "due"]
assert r["structure"] == "mine"
def test_an_inline_template_beats_a_registered_one():
reg = {"mine": {"name": "registered", "sections":
[{"label": "r", "pattern": "^NOPE"}]}}
tpl = {"name": "inline", "sections": [{"label": "i", "pattern": r"^CAPO"}]}
r = ch.chunk_document("x\n\nCAPO\ny\n", structure=tpl, registered=reg,
**DEFAULTS)
assert r["structure"] == "inline"
def test_a_registered_template_is_selectable_by_name():
reg = {"mio": {"name": "mio", "head_label": None,
"sections": [{"label": "a", "pattern": r"^SEZIONE"}]}}
r = ch.chunk_document("testa\n\nSEZIONE\ncorpo\n", structure="mio",
registered=reg, **DEFAULTS)
assert r["structure"] == "mio"
# head_label None drops the text before the first match
assert [s["section"] for s in r["sections"]] == ["a"]
@pytest.mark.parametrize("tpl", [
{"sections": []},
{"sections": [{"label": "x"}]},
{"sections": [{"label": "x", "pattern": "([unclosed"}]},
{"sections": [{"label": "x", "pattern": "a" * 3000}]},
"not-an-object-and-not-registered",
])
def test_a_malformed_template_is_bad_template_not_a_crash(tpl):
with pytest.raises(ch.ChunkError) as e:
ch.chunk_document(RULING, structure=tpl, **DEFAULTS)
assert e.value.code in ("bad_template", "unknown_template")
def test_too_many_sections_is_refused():
tpl = {"sections": [{"label": f"s{i}", "pattern": f"^X{i}"}
for i in range(ch.MAX_TEMPLATE_SECTIONS + 1)]}
with pytest.raises(ch.ChunkError) as e:
ch.chunk_document(RULING, structure=tpl, **DEFAULTS)
assert e.value.code == "bad_template"
def test_one_alternation_is_scanned_not_one_pass_per_pattern():
"""Scanning per pattern and merging changes WHICH header claims an
overlapping position -- a different segmentation for the same document."""
src = (ROOT / "codai/api/chunking.py").read_text()
fn = src.split("def compile_template")[1].split("\ndef ")[0]
assert '"|".join(parts)' in fn
assert "finditer" in src.split("def split_sections")[1].split("\ndef ")[0]
# ----------------------------------------------------------- references #
def test_articles_normalise_with_their_code():
assert ch.extract_refs("ai sensi dell'art. 2043 c.c.") == ["art.2043 c.c."]
assert ch.extract_refs("art.112 cpc") == ["art.112 c.p.c."]
def test_an_article_suffix_keeps_its_form_without_spaces():
assert ch.extract_refs("art. 2055 bis c.c.") == ["art.2055bis c.c."]
def test_acts_normalise_to_kind_number_year():
assert ch.extract_refs("il d.lgs. 209/2005") == ["d.lgs.209/2005"]
assert ch.extract_refs("legge 24/2017") == ["l.24/2017"]
def test_references_are_deduplicated_in_first_seen_order():
txt = "art. 2054 c.c., poi art. 2043 c.c., di nuovo art. 2054 c.c."
assert ch.extract_refs(txt) == ["art.2054 c.c.", "art.2043 c.c."]
def test_extract_refs_false_skips_the_work():
r = ch.chunk_document(RULING, refs=False, **DEFAULTS)
assert all("norm_refs" not in c for c in r["chunks"])
def test_refs_are_per_chunk_not_per_document():
r = ch.chunk_document(RULING, **DEFAULTS)
assert any(c.get("norm_refs") for c in r["chunks"])
assert any(c.get("norm_refs") == [] for c in r["chunks"])
# -------------------------------------------------------------- offsets #
def test_offsets_are_omitted_by_default():
r = ch.chunk_document(RULING, **DEFAULTS)
assert all("start" not in c for c in r["chunks"])
def test_offsets_point_into_the_normalised_text():
r = ch.chunk_document(RULING, include_offsets=True, **DEFAULTS)
norm = RULING.replace("\r", "")
for c in r["chunks"]:
if "start" in c:
assert norm[c["start"]:c["end"]] == c["text"]
# ------------------------------------------------- it is not the windower #
def test_chunking_never_pools_and_never_drops():
"""The distinction from embed_long_input, asserted: many chunks out, and
the text accounted for."""
r = ch.chunk_document(RULING, overlap_tokens=0, **{
k: v for k, v in DEFAULTS.items() if k != "overlap_tokens"})
assert r["n_chunks"] > 1
for c in r["chunks"]:
assert c["text"].strip()
def test_the_endpoint_needs_no_model_or_tokenizer():
"""Token counts are approximated as characters / chars_per_token so this
stays a pure CPU operation -- which is what lets a light-core image serve
it with no GPU and no pod."""
src = (ROOT / "codai/api/chunking.py").read_text()
for forbidden in ("AutoTokenizer", "import torch", "from transformers"):
assert forbidden not in src, forbidden
def test_the_params_are_echoed_so_a_caller_can_tell_what_applied():
r = ch.chunk_document(RULING, chunk_tokens=100, overlap_tokens=10,
chars_per_token=5)
assert r["params"] == {"chunk_tokens": 100, "overlap_tokens": 10,
"chars_per_token": 5}
# ------------------------------------------- forwarding it to a remote #
#
# /v1/chunk rides the EMBEDDINGS capability: chunks are on their way to be
# embedded, the embeddings image is full-core so it carries the chunker, and a
# `chunk` capability of its own would be one nothing advertises -- or would rent
# a GPU pod to run a regex split.
def test_the_chunk_endpoints_are_forwardable():
from codai.api.remote_gateway import capability_for, _SKIP_PATHS
for p in ("/v1/chunk", "/v1/chunk/file", "/v1/chunk/templates"):
assert capability_for(p) == "embeddings", p
assert p not in _SKIP_PATHS, p
def test_a_named_registered_template_is_inlined_before_it_leaves():
"""A template registered HERE does not exist on the remote. Forwarded by
name it comes back 404 -- or, worse, matches a DIFFERENT template with the
same name and returns different boundaries with nothing to say so."""
import json
from codai.api import remote_gateway as rg
import codai.api.chunking as chk
mine = {"name": "mine", "sections": [{"label": "x", "pattern": "^CAPO"}]}
orig = chk._registered_templates
chk._registered_templates = lambda: {"mine": mine}
try:
body = json.dumps({"text": "t", "structure": "mine"}).encode()
out = rg._inline_chunk_template("/v1/chunk",
{"content-type": "application/json"}, body)
assert json.loads(out)["structure"] == mine
finally:
chk._registered_templates = orig
def test_a_builtin_name_is_left_alone():
"""Every coderai has the built-ins; sending the object would only make the
request bigger."""
import json
from codai.api import remote_gateway as rg
body = json.dumps({"text": "t", "structure": "it-legal"}).encode()
out = rg._inline_chunk_template("/v1/chunk",
{"content-type": "application/json"}, body)
assert json.loads(out)["structure"] == "it-legal"
def test_an_inline_template_is_passed_through_untouched():
import json
from codai.api import remote_gateway as rg
tpl = {"name": "inline", "sections": [{"label": "a", "pattern": "^X"}]}
body = json.dumps({"text": "t", "structure": tpl}).encode()
out = rg._inline_chunk_template("/v1/chunk",
{"content-type": "application/json"}, body)
assert json.loads(out)["structure"] == tpl
@pytest.mark.parametrize("path,ctype,body", [
("/v1/embeddings", "application/json", b'{"structure":"mine"}'),
("/v1/chunk", "multipart/form-data; boundary=x", b"--x--"),
("/v1/chunk", "application/json", b"not json at all"),
("/v1/chunk", "application/json", b'[]'),
("/v1/chunk", "application/json", b'{}'),
])
def test_inlining_never_raises_and_leaves_what_it_cannot_handle(path, ctype, body):
"""A body it cannot parse is forwarded untouched and the remote answers for
itself -- a rewrite hook that can throw would take down every forwarded
request, not just a chunk one."""
from codai.api import remote_gateway as rg
assert rg._inline_chunk_template(path, {"content-type": ctype}, body) == body
tests/test_ocr_dual_mode.py
View file @ 29233812
......@@ -344,9 +344,9 @@ def test_pages_come_back_in_document_order(monkeypatch):
seen = []
def _post(base, key, blob, name, dpi, opts, detect):
def _post(base, key, blob, name, dpi, opts, detect, engine="surya"):
seen.append(name)
return {"engine": "surya", "pages": [{"text": name, "width": 1, "height": 2}]}
return {"engine": engine, "pages": [{"text": name, "width": 1, "height": 2}]}
monkeypatch.setattr(om, "_post_page_to_pod", _post)
name, pages, blank, stamps, sigs = asyncio.run(mgr._read_on_pods(
_pdf(5), "d.pdf", "application/pdf", "m",
......@@ -468,3 +468,109 @@ def test_the_fallback_message_names_the_cause_not_just_the_effect():
fn = src.split("async def _read_on_pods")[1].split("\n def ")[0]
assert "placement answered 'local'" in fn
assert "nothing rented" in fn
# ------------------------------- all three VLM engines, not just surya #
#
# surya needed a 1.6 GB client venv baked into the image. chandra and olmOCR in
# served mode speak only HTTP (PIL + requests, already in the main venv), so for
# them the pod side costs nothing beyond the rasteriser. Leaving them on the
# orchestrator while surya moved was an asymmetry, not a design.
class _MultiCfg:
surya_serve = "model"
surya_model_id = "datalab-to/surya-ocr-2"
chandra_serve = "model"
chandra_model_id = "datalab-to/chandra-ocr-2"
olmocr_serve = "model"
olmocr_model_id = "allenai/olmOCR-2-7B-1025-FP8"
default_engine = "surya"
max_concurrency = 8
page_window = 4
blank_page_max_chars = 0
fallback_engines = "off"
@pytest.mark.parametrize("engine,expect", [
("surya", "datalab-to/surya-ocr-2"),
("chandra", "datalab-to/chandra-ocr-2"),
("olmocr", "allenai/olmOCR-2-7B-1025-FP8"),
])
def test_every_vlm_engine_can_take_the_pod_path(monkeypatch, engine, expect):
monkeypatch.setattr(rw, "_model_entry", lambda mid: {
"runpod": {"image": "ghcr.io/nextime/coderai-vllm-ocr:latest"}})
assert om.pod_pipeline_model(_MultiCfg(), engine) == expect
@pytest.mark.parametrize("engine", ["surya", "chandra", "olmocr"])
def test_a_plain_vllm_image_keeps_every_engine_local(monkeypatch, engine):
monkeypatch.setattr(rw, "_model_entry", lambda mid: {
"runpod": {"image": "ghcr.io/nextime/coderai-vllm:latest"}})
assert om.pod_pipeline_model(_MultiCfg(), engine) == ""
def test_a_non_vlm_engine_never_takes_the_pod_path(monkeypatch):
"""paddle and docTR run in-process; there is no model server to drive."""
monkeypatch.setattr(rw, "_model_entry", lambda mid: {
"runpod": {"image": "ghcr.io/nextime/coderai-vllm-ocr:latest"}})
for eng in ("paddle", "doctr", "nonsense"):
assert om.pod_pipeline_model(_MultiCfg(), eng) == ""
def test_each_engine_reads_its_own_serve_mode(monkeypatch):
"""chandra on `server` mode must not be dragged onto the pod path by
surya's `model` mode."""
class _C(_MultiCfg):
chandra_serve = "server"
monkeypatch.setattr(rw, "_model_entry", lambda mid: {
"runpod": {"image": "ghcr.io/nextime/coderai-vllm-ocr:latest"}})
assert om.pod_pipeline_model(_C(), "chandra") == ""
assert om.pod_pipeline_model(_C(), "surya") == "datalab-to/surya-ocr-2"
def test_the_dispatch_tells_the_pod_WHICH_engine(monkeypatch):
"""Hardcoded "surya" sent a chandra page to surya, which the pod would have
served with the wrong model -- and the text would have come back looking
plausible."""
from codai.api import pod_http
seen = {}
class _R:
status_code = 200
def json(self):
return {"engine": "chandra", "pages": []}
monkeypatch.setattr(pod_http, "post",
lambda url, **kw: seen.update(kw) or _R())
om._post_page_to_pod("https://1.2.3.4:8000/v1", "t", b"%PDF-x", "0.pdf",
250, {}, None, "chandra")
assert seen["data"]["engine"] == "chandra"
def test_the_pod_env_is_shaped_by_the_served_model_not_hardcoded():
src = (ROOT / "codai/api/runpod_worker.py").read_text()
blk = src.split("if pod_runs_ocr_pipeline(mcfg, image):")[1] \
.split('if cap == "ocr":')[0]
assert 'env[f"CODERAI_OCR_{_eng.upper()}_ENABLED"]' in blk
assert 'env[f"CODERAI_OCR_{_eng.upper()}_SERVE"]' in blk
assert '"CODERAI_OCR_SURYA_ENABLED"' not in blk, "still hardcoded to surya"
def test_only_surya_is_pointed_at_a_baked_venv():
"""chandra and olmOCR have none -- they speak HTTP from the main venv, which
is why the pod side is nearly free for them."""
src = (ROOT / "codai/api/runpod_worker.py").read_text()
blk = src.split("if pod_runs_ocr_pipeline(mcfg, image):")[1] \
.split('if cap == "ocr":')[0]
assert 'if _eng == "surya":' in blk
venv = blk.split('CODERAI_OCR_SURYA_VENV')[0]
assert venv.rstrip().endswith('if _eng == "surya":') or 'if _eng == "surya"' in venv
def test_each_engine_inherits_its_own_licence_decision():
"""surya is GPL, chandra's weights are a modified OpenRAIL-M. A pod cannot
accept either for itself."""
src = (ROOT / "codai/api/runpod_worker.py").read_text()
blk = src.split("if pod_runs_ocr_pipeline(mcfg, image):")[1] \
.split('if cap == "ocr":')[0]
assert "CODERAI_OCR_SURYA_ACCEPT_LICENSE" in blk
assert "CODERAI_OCR_CHANDRA_ACCEPT_LICENSE" in blk
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment