mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 18:15:33 +00:00
fix(attachments,usage): validate content behind a BOM, judge by the live parser table
Review follow-ups. A BOM told the sniff which encoding to read, but was also taken as the verdict: three prepended bytes let any binary through, including as notes.txt. A BOM now only selects the test — UTF-8 falls through to the byte rules on the remainder, UTF-16/32 decode and judge the characters (NUL, unprintable, or replacement chars from bytes the decoder could not read). Real Notepad-Unicode text still passes, mp4-behind-a-BOM does not, in either language. The gate treated the full parser table as a given, but without docling the fallback extractor has no .tif/.tiff/.bmp/.webp/.vtt/.xml handler, so those suffixes skipped the content check and reached the plain-text fallthrough — the original bug, one install away. The worker now passes the keys of the extractor it actually built, making the second gate stricter than the route's static one rather than a copy of it. Cache bins: extraction coerced a missing value to 0 and only non-zero bins were recorded, so a provider reporting cached_tokens=0 persisted as NULL — indistinguishable from "not reported", and OpenAI reports exactly that on every uncached request. Bins are now carried as Optional and recorded when not None, which is what the nullable columns and the NULL-means-unknown comment already assumed. Anthropic's cache_read/cache_creation bins had the same shape and are fixed alongside; the int-or-None coercion is shared in llm/base.py.
This commit is contained in:
1 parent
f7cd94668f
commit
9db11f18f5
10 files changed
+341
-66
No files matched your search
@@ -6,7 +6,7 @@ from typing import Any, Dict, Generator, List, Optional, Tuple
|
||||
from anthropic import Anthropic
|
||||
|
||||
from application.core.settings import settings
|
||||
from application.llm.base import BaseLLM
|
||||
from application.llm.base import BaseLLM, optional_int
|
||||
from application.storage.storage_creator import StorageCreator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -387,8 +387,12 @@ class AnthropicLLM(BaseLLM):
|
||||
# message_start are reused below instead.
|
||||
if not output_only:
|
||||
base_input = int(getattr(usage, "input_tokens", 0) or 0)
|
||||
created = int(getattr(usage, "cache_creation_input_tokens", 0) or 0)
|
||||
cached = int(getattr(usage, "cache_read_input_tokens", 0) or 0)
|
||||
# None (unreported) and 0 (reported, no caching) must stay
|
||||
# distinguishable — see ``optional_int``.
|
||||
created = optional_int(
|
||||
getattr(usage, "cache_creation_input_tokens", None)
|
||||
)
|
||||
cached = optional_int(getattr(usage, "cache_read_input_tokens", None))
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
|
||||
@@ -398,11 +402,11 @@ class AnthropicLLM(BaseLLM):
|
||||
prompt = int(previous.get("prompt_tokens", 0) or 0)
|
||||
details = previous.get("prompt_tokens_details")
|
||||
else:
|
||||
prompt = base_input + created + cached
|
||||
prompt = base_input + (created or 0) + (cached or 0)
|
||||
details = {}
|
||||
if cached:
|
||||
if cached is not None:
|
||||
details["cached_tokens"] = cached
|
||||
if created:
|
||||
if created is not None:
|
||||
details["cache_creation_tokens"] = created
|
||||
details = details or None
|
||||
|
||||
|
||||
@@ -33,6 +33,29 @@ _STREAM_RETRYABLE_TRANSPORT_ERRORS = (
|
||||
)
|
||||
|
||||
|
||||
def optional_int(value) -> Optional[int]:
|
||||
"""Coerce a provider-reported count to int, keeping "unreported" as None.
|
||||
|
||||
Cache bins persist as NULL when a provider says nothing and as 0 when it
|
||||
reports zero, so the two must stay distinguishable all the way from the
|
||||
usage object: providers report ``cached_tokens: 0`` on every uncached
|
||||
request, and folding that into "unknown" would file the ordinary case as
|
||||
no-data.
|
||||
|
||||
Args:
|
||||
value: The raw attribute off a provider usage object.
|
||||
|
||||
Returns:
|
||||
The integer value, or None when absent or not a number.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
class BaseLLM(ABC):
|
||||
# Stamped onto the ``llm_stream_start`` event so dashboards can group
|
||||
# calls by vendor. Subclasses override.
|
||||
|
||||
+23
-19
@@ -9,7 +9,7 @@ from typing import Any, Callable
|
||||
from openai import BadRequestError, OpenAI
|
||||
|
||||
from application.core.settings import settings
|
||||
from application.llm.base import BaseLLM
|
||||
from application.llm.base import BaseLLM, optional_int
|
||||
from application.storage.storage_creator import StorageCreator
|
||||
|
||||
# Placeholder sent to OpenAI-compatible backends that require no credentials.
|
||||
@@ -1305,14 +1305,11 @@ class OpenAILLM(BaseLLM):
|
||||
}
|
||||
input_details = getattr(usage, "prompt_tokens_details", None)
|
||||
output_details = getattr(usage, "completion_tokens_details", None)
|
||||
try:
|
||||
cached = int(getattr(input_details, "cached_tokens", 0) or 0)
|
||||
written = int(getattr(input_details, "cache_write_tokens", 0) or 0)
|
||||
reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0)
|
||||
except (TypeError, ValueError):
|
||||
cached = 0
|
||||
written = 0
|
||||
reasoning = 0
|
||||
cached = optional_int(getattr(input_details, "cached_tokens", None))
|
||||
written = optional_int(getattr(input_details, "cache_write_tokens", None))
|
||||
reasoning = optional_int(
|
||||
getattr(output_details, "reasoning_tokens", None)
|
||||
)
|
||||
details = self._prompt_cache_details(cached, written)
|
||||
if details:
|
||||
result["prompt_tokens_details"] = details
|
||||
@@ -1322,18 +1319,21 @@ class OpenAILLM(BaseLLM):
|
||||
self._last_usage_claimed = False
|
||||
|
||||
@staticmethod
|
||||
def _prompt_cache_details(cached: int, written: int) -> dict:
|
||||
def _prompt_cache_details(cached: int | None, written: int | None) -> dict:
|
||||
"""Build the ``prompt_tokens_details`` breakdown from the two cache bins.
|
||||
|
||||
Only non-zero bins are included so a provider that never reports
|
||||
caching produces no details (and no misleading zero downstream).
|
||||
``cache_write_tokens`` is reported by newer OpenAI-family deployments
|
||||
that charge for cache writes.
|
||||
A bin is included whenever the provider reported it, zero included:
|
||||
downstream, an absent bin persists as NULL ("we don't know") and a
|
||||
reported zero as 0 ("no cache hits"), and OpenAI reports
|
||||
``cached_tokens: 0`` on every uncached request. Collapsing the two
|
||||
would file every ordinary request as unknown. ``cache_write_tokens``
|
||||
is reported by newer OpenAI-family deployments that charge for cache
|
||||
writes.
|
||||
"""
|
||||
details = {}
|
||||
if cached:
|
||||
if cached is not None:
|
||||
details["cached_tokens"] = cached
|
||||
if written:
|
||||
if written is not None:
|
||||
details["cache_write_tokens"] = written
|
||||
return details
|
||||
|
||||
@@ -1375,9 +1375,13 @@ class OpenAILLM(BaseLLM):
|
||||
"completion_tokens": completion,
|
||||
"total_tokens": int(getattr(usage, "total_tokens", 0) or prompt + completion),
|
||||
}
|
||||
cached = int(getattr(input_details, "cached_tokens", 0) or 0)
|
||||
written = int(getattr(input_details, "cache_write_tokens", 0) or 0)
|
||||
reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0)
|
||||
cached = optional_int(getattr(input_details, "cached_tokens", None))
|
||||
written = optional_int(
|
||||
getattr(input_details, "cache_write_tokens", None)
|
||||
)
|
||||
reasoning = optional_int(
|
||||
getattr(output_details, "reasoning_tokens", None)
|
||||
)
|
||||
details = self._prompt_cache_details(cached, written)
|
||||
if details:
|
||||
result["prompt_tokens_details"] = details
|
||||
|
||||
@@ -2,10 +2,11 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import codecs
|
||||
import io
|
||||
import os
|
||||
from contextlib import suppress
|
||||
from typing import BinaryIO
|
||||
from typing import BinaryIO, Container
|
||||
|
||||
from application.core.settings import settings
|
||||
from application.parser.file.constants import (
|
||||
@@ -64,16 +65,48 @@ _TEXT_SNIFF_BYTES = 8192
|
||||
_MAX_NONTEXT_RATIO = 0.10
|
||||
# Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC.
|
||||
_TEXT_CONTROL_BYTES = frozenset({0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x1B})
|
||||
# A UTF-16/32 file is half NUL bytes, so it has to be recognised by its BOM
|
||||
# before the NUL test — Notepad's "Unicode" .txt is ordinary text.
|
||||
_TEXT_CONTROL_CHARS = frozenset("\t\n\v\f\r\x1b")
|
||||
# A UTF-16/32 file is half NUL bytes, so the byte rules below would reject it.
|
||||
# Its BOM says which encoding to read it as; the decoded characters are then
|
||||
# judged instead. Longest BOM first — UTF-32-LE starts with the UTF-16-LE one.
|
||||
# A BOM is never a verdict on its own: three bytes must not buy a video a pass.
|
||||
_TEXT_BOMS = (
|
||||
b"\xef\xbb\xbf",
|
||||
b"\xff\xfe",
|
||||
b"\xfe\xff",
|
||||
b"\x00\x00\xfe\xff",
|
||||
(codecs.BOM_UTF32_LE, "utf-32-le"),
|
||||
(codecs.BOM_UTF32_BE, "utf-32-be"),
|
||||
(codecs.BOM_UTF8, None),
|
||||
(codecs.BOM_UTF16_LE, "utf-16-le"),
|
||||
(codecs.BOM_UTF16_BE, "utf-16-be"),
|
||||
)
|
||||
|
||||
|
||||
def _decoded_looks_like_text(text: str) -> bool:
|
||||
"""Return whether decoded characters read as text.
|
||||
|
||||
Args:
|
||||
text: Characters decoded from a BOM-marked sample.
|
||||
|
||||
Returns:
|
||||
False when a NUL character appears — the byte-level rule, one level
|
||||
up, since no text holds one — or when too many characters are
|
||||
unprintable: control, unassigned, private-use or surrogate, which is
|
||||
what binary decodes into. Replacement characters count as unprintable,
|
||||
so bytes the decoder could not read are evidence rather than silently
|
||||
dropped (a truncated character at the sample boundary is one of
|
||||
thousands and cannot swing the ratio).
|
||||
"""
|
||||
if not text:
|
||||
return True
|
||||
if "\x00" in text:
|
||||
return False
|
||||
nontext = sum(
|
||||
1
|
||||
for char in text
|
||||
if char == "�"
|
||||
or (not char.isprintable() and char not in _TEXT_CONTROL_CHARS)
|
||||
)
|
||||
return nontext / len(text) <= _MAX_NONTEXT_RATIO
|
||||
|
||||
|
||||
def looks_like_text(sample: bytes) -> bool:
|
||||
"""Return whether a leading byte sample reads as text rather than binary.
|
||||
|
||||
@@ -82,20 +115,30 @@ def looks_like_text(sample: bytes) -> bool:
|
||||
|
||||
Returns:
|
||||
False when the sample holds a NUL byte or too many other non-text
|
||||
control bytes, True otherwise. A Unicode BOM settles it as text.
|
||||
control bytes, True otherwise. A UTF-16/32 BOM switches the test to
|
||||
the decoded characters; the content is still what decides.
|
||||
"""
|
||||
if not sample:
|
||||
return True
|
||||
if sample.startswith(_TEXT_BOMS):
|
||||
body = sample
|
||||
for bom, encoding in _TEXT_BOMS:
|
||||
if not sample.startswith(bom):
|
||||
continue
|
||||
body = sample[len(bom) :]
|
||||
if encoding is not None:
|
||||
return _decoded_looks_like_text(body.decode(encoding, errors="replace"))
|
||||
# UTF-8 BOM: the byte rules still apply to everything after it.
|
||||
break
|
||||
if not body:
|
||||
return True
|
||||
if b"\x00" in sample:
|
||||
if b"\x00" in body:
|
||||
return False
|
||||
nontext = sum(
|
||||
1
|
||||
for byte in sample
|
||||
for byte in body
|
||||
if (byte < 0x20 and byte not in _TEXT_CONTROL_BYTES) or byte == 0x7F
|
||||
)
|
||||
return nontext / len(sample) <= _MAX_NONTEXT_RATIO
|
||||
return nontext / len(body) <= _MAX_NONTEXT_RATIO
|
||||
|
||||
|
||||
def file_looks_like_text(path: str | os.PathLike[str]) -> bool:
|
||||
@@ -116,25 +159,35 @@ def file_looks_like_text(path: str | os.PathLike[str]) -> bool:
|
||||
|
||||
|
||||
def enforce_parseable_attachment(
|
||||
path: str | os.PathLike[str], filename: str | None
|
||||
path: str | os.PathLike[str],
|
||||
filename: str | None,
|
||||
parser_extensions: Container[str] | None = None,
|
||||
) -> None:
|
||||
"""Reject an attachment that no parser handles and that is not plain text.
|
||||
|
||||
Suffixes with a parser (``ATTACHMENT_PARSER_EXTENSIONS``) are admitted
|
||||
unconditionally — a PDF is binary and parses fine. Everything else, .txt
|
||||
included, has to read as text: that is what keeps a video or an archive
|
||||
out of the plain-text fallthrough while leaving source, config and log
|
||||
files in, whatever the file happens to be named.
|
||||
Suffixes with a parser are admitted unconditionally — a PDF is binary and
|
||||
parses fine. Everything else, .txt included, has to read as text: that is
|
||||
what keeps a video or an archive out of the plain-text fallthrough while
|
||||
leaving source, config and log files in, whatever the file is named.
|
||||
|
||||
Args:
|
||||
path: Local path to the staged upload, readable before it is stored.
|
||||
filename: The upload's filename, used for the suffix and the message.
|
||||
parser_extensions: Suffixes to treat as parser-backed. Defaults to
|
||||
``ATTACHMENT_PARSER_EXTENSIONS``, which assumes the full parser
|
||||
table; callers holding the live extractor (the worker) pass its
|
||||
keys, so a docling-less install does not admit a .webp on trust
|
||||
and then read it as text.
|
||||
|
||||
Raises:
|
||||
UnsupportedUploadTypeError: When the file has no parser and its
|
||||
contents are binary.
|
||||
"""
|
||||
if has_attachment_parser(filename):
|
||||
if parser_extensions is None:
|
||||
has_parser = has_attachment_parser(filename)
|
||||
else:
|
||||
has_parser = attachment_extension(filename) in parser_extensions
|
||||
if has_parser:
|
||||
return
|
||||
if file_looks_like_text(path):
|
||||
return
|
||||
|
||||
+10
-6
@@ -1558,22 +1558,26 @@ class AttachmentRejectedError(Exception):
|
||||
"""
|
||||
|
||||
|
||||
def _reject_unparseable_attachment(local_path: str, filename: str) -> None:
|
||||
def _reject_unparseable_attachment(
|
||||
local_path: str, filename: str, parser_extensions
|
||||
) -> None:
|
||||
"""Reject an attachment with no parser whose contents are binary.
|
||||
|
||||
Defence in depth behind the route's gate: ``SimpleDirectoryReader`` opens
|
||||
any suffix it has no parser for as plain text, so a binary must fail here,
|
||||
visibly, rather than "succeed" as garbage.
|
||||
Defence in depth behind the route's gate, and stricter than it: this runs
|
||||
against the parser table actually loaded, so a suffix the route trusted
|
||||
(.webp with docling absent, say) is still refused rather than opened as
|
||||
plain text by ``SimpleDirectoryReader`` and stored as garbage.
|
||||
|
||||
Args:
|
||||
local_path: Filesystem path of the attachment about to be parsed.
|
||||
filename: The upload's original filename, which carries the suffix.
|
||||
parser_extensions: Suffixes the live ``file_extractor`` can parse.
|
||||
|
||||
Raises:
|
||||
AttachmentRejectedError: If the file has no parser and is not text.
|
||||
"""
|
||||
try:
|
||||
enforce_parseable_attachment(local_path, filename)
|
||||
enforce_parseable_attachment(local_path, filename, parser_extensions)
|
||||
except UnsupportedUploadTypeError as exc:
|
||||
raise AttachmentRejectedError(str(exc)) from exc
|
||||
|
||||
@@ -1775,7 +1779,7 @@ def attachment_worker(self, file_info, user):
|
||||
parser_name = type(_parser).__name__ if _parser is not None else "SimpleDirectoryReader"
|
||||
|
||||
def _parse_local_file(local_path: str, **kwargs) -> Document:
|
||||
_reject_unparseable_attachment(local_path, filename)
|
||||
_reject_unparseable_attachment(local_path, filename, set(file_extractor))
|
||||
_reject_attachment_zip_bomb(local_path)
|
||||
parse_path, is_temp_copy = _bounded_attachment_copy(local_path)
|
||||
try:
|
||||
|
||||
@@ -95,13 +95,28 @@ describe('looksLikeText', () => {
|
||||
});
|
||||
|
||||
it('accepts BOM-marked Unicode, which is half NUL bytes but still text', () => {
|
||||
// UTF-8, UTF-16 LE/BE, UTF-32 BE.
|
||||
// UTF-8, UTF-16 LE/BE, and UTF-32 BE (no TextDecoder — left to the server).
|
||||
expect(looksLikeText(bytes(0xef, 0xbb, 0xbf, 0x68, 0x69))).toBe(true);
|
||||
expect(looksLikeText(bytes(0xff, 0xfe, 0x68, 0x00, 0x69, 0x00))).toBe(true);
|
||||
expect(looksLikeText(bytes(0xfe, 0xff, 0x00, 0x68, 0x00, 0x69))).toBe(true);
|
||||
expect(looksLikeText(bytes(0x00, 0x00, 0xfe, 0xff, 0x00, 0x68))).toBe(true);
|
||||
});
|
||||
|
||||
it('still rejects binary behind a BOM — three bytes buy nothing', () => {
|
||||
// UTF-8 BOM + MP4 header: the byte rules apply to what follows it.
|
||||
expect(looksLikeText(bytes(0xef, 0xbb, 0xbf, ...MP4_HEADER))).toBe(false);
|
||||
// UTF-16 LE BOM + control-only code points.
|
||||
expect(
|
||||
looksLikeText(
|
||||
bytes(0xff, 0xfe, 0x01, 0x00, 0x02, 0x00, 0x03, 0x00, 0x04, 0x00),
|
||||
),
|
||||
).toBe(false);
|
||||
// UTF-16 BE BOM + bytes that decode to NUL characters.
|
||||
expect(looksLikeText(bytes(0xfe, 0xff, 0x00, 0x00, 0x00, 0x41))).toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
it('rejects a NUL byte and dense control bytes', () => {
|
||||
expect(looksLikeText(MP4_HEADER)).toBe(false);
|
||||
expect(looksLikeText(bytes(0x74, 0x78, 0x00, 0x74))).toBe(false);
|
||||
@@ -142,9 +157,11 @@ describe('partitionAttachmentFiles', () => {
|
||||
it('refuses a binary renamed to .txt — .txt is sniffed like any other suffix', async () => {
|
||||
const { supported, unsupported } = await partitionAttachmentFiles([
|
||||
file('notes.txt', MP4_HEADER),
|
||||
// A BOM in front of it changes nothing.
|
||||
file('bom.txt', bytes(0xef, 0xbb, 0xbf, ...MP4_HEADER)),
|
||||
]);
|
||||
expect(supported).toEqual([]);
|
||||
expect(unsupported.map((f) => f.name)).toEqual(['notes.txt']);
|
||||
expect(unsupported.map((f) => f.name)).toEqual(['notes.txt', 'bom.txt']);
|
||||
});
|
||||
|
||||
it('admits parser-backed types on their suffix, binary contents and all', async () => {
|
||||
|
||||
@@ -160,27 +160,82 @@ const TEXT_SNIFF_BYTES = 8192;
|
||||
const MAX_NONTEXT_RATIO = 0.1;
|
||||
// Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC.
|
||||
const TEXT_CONTROL_BYTES = new Set([0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x1b]);
|
||||
// A UTF-16/32 file is half NUL bytes, so it has to be recognised by its BOM
|
||||
// before the NUL test — Notepad's "Unicode" .txt is ordinary text.
|
||||
const TEXT_BOMS = [
|
||||
[0xef, 0xbb, 0xbf],
|
||||
[0xff, 0xfe],
|
||||
[0xfe, 0xff],
|
||||
[0x00, 0x00, 0xfe, 0xff],
|
||||
const TEXT_CONTROL_CHARS = new Set(['\t', '\n', '\v', '\f', '\r', '\u001b']);
|
||||
// \p{C}: control, format, private-use, surrogate and unassigned code points —
|
||||
// what binary decodes into.
|
||||
const NONTEXT_CHAR = /\p{C}/u;
|
||||
// A UTF-16/32 file is half NUL bytes, so the byte rules below would reject it.
|
||||
// Its BOM says which encoding to read it as; the decoded characters are then
|
||||
// judged instead. Longest BOM first — UTF-32-LE starts with the UTF-16-LE one.
|
||||
// A BOM is never a verdict on its own: three bytes must not buy a video a pass.
|
||||
const TEXT_BOMS: { bom: number[]; encoding: string | null }[] = [
|
||||
// TextDecoder has no UTF-32; leave those two to the server, which does.
|
||||
{ bom: [0xff, 0xfe, 0x00, 0x00], encoding: null },
|
||||
{ bom: [0x00, 0x00, 0xfe, 0xff], encoding: null },
|
||||
{ bom: [0xef, 0xbb, 0xbf], encoding: 'utf-8' },
|
||||
{ bom: [0xff, 0xfe], encoding: 'utf-16le' },
|
||||
{ bom: [0xfe, 0xff], encoding: 'utf-16be' },
|
||||
];
|
||||
|
||||
/** Whether a leading byte sample reads as text rather than binary. */
|
||||
/**
|
||||
* Whether decoded characters read as text rather than decoded binary.
|
||||
* Replacement characters count against it, so bytes the decoder could not
|
||||
* read are evidence rather than silently dropped.
|
||||
*/
|
||||
function decodedLooksLikeText(text: string): boolean {
|
||||
// The byte-level NUL rule, one level up: no text holds a NUL character.
|
||||
if (text.includes('\u0000')) return false;
|
||||
let total = 0;
|
||||
let nontext = 0;
|
||||
for (const char of text) {
|
||||
total += 1;
|
||||
if (
|
||||
char === '\ufffd' ||
|
||||
(!TEXT_CONTROL_CHARS.has(char) && NONTEXT_CHAR.test(char))
|
||||
)
|
||||
nontext += 1;
|
||||
}
|
||||
if (total === 0) return true;
|
||||
return nontext / total <= MAX_NONTEXT_RATIO;
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether a leading byte sample reads as text rather than binary. A UTF-16
|
||||
* BOM switches the test to the decoded characters; the content is still what
|
||||
* decides. Mirrors `looks_like_text` in application/upload_limits.py.
|
||||
*/
|
||||
export function looksLikeText(sample: Uint8Array): boolean {
|
||||
if (sample.length === 0) return true;
|
||||
if (TEXT_BOMS.some((bom) => bom.every((byte, i) => sample[i] === byte)))
|
||||
return true;
|
||||
|
||||
let body = sample;
|
||||
const marked = TEXT_BOMS.find(({ bom }) =>
|
||||
bom.every((byte, i) => sample[i] === byte),
|
||||
);
|
||||
if (marked) {
|
||||
// UTF-32, or a decoder this browser lacks: defer to the server rather
|
||||
// than refuse a file it would accept.
|
||||
if (marked.encoding === null) return true;
|
||||
body = sample.subarray(marked.bom.length);
|
||||
if (body.length === 0) return true;
|
||||
if (marked.encoding !== 'utf-8') {
|
||||
try {
|
||||
return decodedLooksLikeText(
|
||||
new TextDecoder(marked.encoding).decode(body),
|
||||
);
|
||||
} catch {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
// UTF-8 BOM: the byte rules still apply to everything after it.
|
||||
}
|
||||
|
||||
let nontext = 0;
|
||||
for (const byte of sample) {
|
||||
for (const byte of body) {
|
||||
if (byte === 0x00) return false;
|
||||
if ((byte < 0x20 && !TEXT_CONTROL_BYTES.has(byte)) || byte === 0x7f)
|
||||
nontext += 1;
|
||||
}
|
||||
return nontext / sample.length <= MAX_NONTEXT_RATIO;
|
||||
return nontext / body.length <= MAX_NONTEXT_RATIO;
|
||||
}
|
||||
|
||||
async function isSupportedAttachmentFile(file: File): Promise<boolean> {
|
||||
|
||||
@@ -1583,5 +1583,26 @@ def test_record_responses_metadata_captures_cache_write_bin(monkeypatch):
|
||||
),
|
||||
))
|
||||
# A pure cache-write turn (first request of a prefix) has reads=0 and
|
||||
# writes>0; the details must still be recorded.
|
||||
assert llm._last_usage["prompt_tokens_details"] == {"cache_write_tokens": 9}
|
||||
# writes>0; both bins were reported, so both are recorded — a reported
|
||||
# zero is "no cache hits", not "unknown".
|
||||
assert llm._last_usage["prompt_tokens_details"] == {
|
||||
"cached_tokens": 0,
|
||||
"cache_write_tokens": 9,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_record_responses_metadata_omits_unreported_cache_bins(monkeypatch):
|
||||
"""A provider that says nothing about caching must persist as NULL, not 0."""
|
||||
llm = _make_llm(monkeypatch, _responses_caps())
|
||||
llm._record_responses_metadata(_ns(
|
||||
id="resp_2",
|
||||
usage=_ns(
|
||||
input_tokens=10,
|
||||
output_tokens=7,
|
||||
total_tokens=17,
|
||||
input_tokens_details=None,
|
||||
output_tokens_details=None,
|
||||
),
|
||||
))
|
||||
assert "prompt_tokens_details" not in llm._last_usage
|
||||
@@ -121,6 +121,49 @@ def test_enforce_parseable_attachment_accepts_bom_marked_unicode_text(
|
||||
enforce_parseable_attachment(path, "notes.txt")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bom",
|
||||
[codecs.BOM_UTF8, codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE, codecs.BOM_UTF32_BE],
|
||||
)
|
||||
def test_enforce_parseable_attachment_rejects_binary_behind_a_bom(bom, tmp_path):
|
||||
"""A BOM says which encoding to read, not that the content is text.
|
||||
|
||||
Otherwise three prepended bytes buy any binary a pass.
|
||||
"""
|
||||
from application.upload_limits import (
|
||||
enforce_parseable_attachment,
|
||||
UnsupportedUploadTypeError,
|
||||
)
|
||||
|
||||
path = tmp_path / "notes.txt"
|
||||
path.write_bytes(bom + MP4_HEADER + bytes(range(256)) * 8)
|
||||
|
||||
with pytest.raises(UnsupportedUploadTypeError):
|
||||
enforce_parseable_attachment(path, "notes.txt")
|
||||
|
||||
|
||||
def test_enforce_parseable_attachment_uses_the_extractor_it_is_given(tmp_path):
|
||||
"""The worker holds the live parser table; a trimmed install must not admit on trust.
|
||||
|
||||
Without docling the fallback extractor has no .webp handler, so a .webp
|
||||
would otherwise skip the content check and be read as plain text.
|
||||
"""
|
||||
from application.upload_limits import (
|
||||
enforce_parseable_attachment,
|
||||
UnsupportedUploadTypeError,
|
||||
)
|
||||
|
||||
path = tmp_path / "scan.webp"
|
||||
path.write_bytes(b"RIFF\x00\x00\x00\x00WEBPVP8 " + bytes(range(256)))
|
||||
|
||||
# Default list: .webp is parser-backed, admitted on its name.
|
||||
enforce_parseable_attachment(path, "scan.webp")
|
||||
|
||||
# The extractor actually loaded has no .webp parser.
|
||||
with pytest.raises(UnsupportedUploadTypeError):
|
||||
enforce_parseable_attachment(path, "scan.webp", {".pdf", ".docx"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
[
|
||||
@@ -177,6 +220,10 @@ def test_enforce_parseable_attachment_accepts_text_without_a_parser(filename, tm
|
||||
(b"\x1b[31mred log line\x1b[0m\n", True),
|
||||
(codecs.BOM_UTF16_LE + "hi\n".encode("utf-16-le"), True),
|
||||
(codecs.BOM_UTF8 + b"hi\n", True),
|
||||
(codecs.BOM_UTF32_LE + "hi\n".encode("utf-32-le"), True),
|
||||
# A BOM in front of binary is still binary.
|
||||
(codecs.BOM_UTF8 + b"\x00\x01\x02", False),
|
||||
(codecs.BOM_UTF16_LE + MP4_HEADER, False),
|
||||
(b"text\x00with nul", False),
|
||||
(bytes(range(32)) * 4, False),
|
||||
(b"\x7f\x7f\x7f\x7f" + b"a" * 16, False),
|
||||
|
||||
@@ -379,6 +379,53 @@ class TestAttachmentTypeGuard:
|
||||
failed = [payload for name, payload in events if name == "attachment.failed"]
|
||||
assert failed and failed[0]["error"] == "Unsupported file type: .mp4"
|
||||
|
||||
def test_suffix_the_loaded_extractor_cannot_parse_is_refused(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch, tmp_path
|
||||
):
|
||||
"""The guard judges against the parser table actually loaded.
|
||||
|
||||
Without docling the fallback extractor has no .webp parser, so a
|
||||
.webp the route admitted on its name would otherwise be opened as
|
||||
plain text here.
|
||||
"""
|
||||
from application import worker
|
||||
|
||||
local_path = tmp_path / "scan.webp"
|
||||
local_path.write_bytes(b"RIFF\x00\x00\x00\x00WEBPVP8 " + bytes(range(256)))
|
||||
|
||||
events = []
|
||||
fake_storage = MagicMock(name="storage")
|
||||
fake_storage.process_file.side_effect = lambda path, callback: callback(
|
||||
str(local_path)
|
||||
)
|
||||
monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage)
|
||||
# The docling-less fallback table: images are handled, .webp is not.
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"get_default_file_extractor",
|
||||
lambda ocr_enabled=False, pdf_text_fast_path=False: {".png": object()},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"publish_user_event",
|
||||
lambda user, name, payload, **kwargs: events.append((name, payload)),
|
||||
)
|
||||
|
||||
file_info = {
|
||||
"filename": "scan.webp",
|
||||
"attachment_id": "507f1f77bcf86cd799439014",
|
||||
"path": "uploads/user1/attachments/scan.webp",
|
||||
"metadata": {"source": "chat"},
|
||||
}
|
||||
|
||||
with pytest.raises(
|
||||
worker.AttachmentRejectedError, match=r"Unsupported file type: \.webp"
|
||||
):
|
||||
worker.attachment_worker(task_self, file_info, "user1")
|
||||
|
||||
failed = [payload for name, payload in events if name == "attachment.failed"]
|
||||
assert failed and failed[0]["error"] == "Unsupported file type: .webp"
|
||||
|
||||
def test_text_without_a_parser_is_parsed(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch, tmp_path
|
||||
):
|
||||
|
||||
Reference in new issue
Block a user