102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
"""Pure unit coverage for the point_wire_bytes byte-quoting estimator."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import replace
|
|
from uuid import uuid4
|
|
|
|
from ragcore.application.indexes.models import point_wire_bytes
|
|
from ragcore.infrastructure.qdrant.models import (
|
|
CollectionKind,
|
|
CollectionProfile,
|
|
IndexPayload,
|
|
IndexPoint,
|
|
)
|
|
|
|
|
|
def _profile() -> CollectionProfile:
|
|
return CollectionProfile(
|
|
kind=CollectionKind.DENSE,
|
|
profile="quota-test",
|
|
version="v1",
|
|
dimensions=4,
|
|
model_digest="a" * 64,
|
|
)
|
|
|
|
|
|
def _point(profile: CollectionProfile, *, text: str = "safe passage") -> IndexPoint:
|
|
chunk_id = uuid4()
|
|
return IndexPoint(
|
|
point_id=chunk_id,
|
|
dense_vector=(0.1, 0.2, 0.3, 0.4),
|
|
sparse_indices=(1, 2),
|
|
sparse_values=(0.5, 0.5),
|
|
payload=IndexPayload(
|
|
workspace_id=uuid4(),
|
|
space_id=uuid4(),
|
|
document_id=uuid4(),
|
|
document_version_id=uuid4(),
|
|
chunk_id=chunk_id,
|
|
parent_chunk_id=None,
|
|
source_id=uuid4(),
|
|
status="published",
|
|
content_type="text/markdown",
|
|
language="en",
|
|
trust_level="source",
|
|
location={"page": 1, "block": "heading"},
|
|
parser_version="text-v1",
|
|
chunk_profile_version=profile.version,
|
|
embedding_model_digest=profile.model_digest,
|
|
text=text,
|
|
excerpt=text,
|
|
),
|
|
)
|
|
|
|
|
|
def test_wire_bytes_grows_with_non_text_metadata_fields() -> None:
|
|
profile = _profile()
|
|
small = _point(profile)
|
|
large_location = replace(
|
|
small.payload, location={f"key_{index}": f"value_{index}" for index in range(200)}
|
|
)
|
|
large = replace(small, payload=large_location)
|
|
|
|
assert point_wire_bytes(large) > point_wire_bytes(small)
|
|
|
|
|
|
def test_wire_bytes_grows_with_sparse_vector_size() -> None:
|
|
profile = _profile()
|
|
small = _point(profile)
|
|
large = replace(
|
|
small,
|
|
sparse_indices=tuple(range(500)),
|
|
sparse_values=tuple(0.1 for _ in range(500)),
|
|
)
|
|
|
|
assert point_wire_bytes(large) > point_wire_bytes(small)
|
|
|
|
|
|
def test_wire_bytes_counts_real_utf8_bytes_for_multibyte_text() -> None:
|
|
profile = _profile()
|
|
ascii_point = _point(profile, text="a" * 100)
|
|
multibyte_point = _point(profile, text="漢" * 100) # each character is 3 UTF-8 bytes
|
|
|
|
assert point_wire_bytes(multibyte_point) > point_wire_bytes(ascii_point) + 150
|
|
|
|
|
|
def test_wire_bytes_includes_fingerprint_field_overhead() -> None:
|
|
import json
|
|
|
|
profile = _profile()
|
|
point = _point(profile)
|
|
raw_payload_bytes = len(
|
|
json.dumps(dict(point.payload.as_dict()), sort_keys=True, separators=(",", ":")).encode(
|
|
"utf-8"
|
|
)
|
|
)
|
|
|
|
# The adapter attaches three 64-character fingerprint fields after the payload is
|
|
# built; the estimate must already include their size so quota cannot be evaded
|
|
# by omitting them.
|
|
assert point_wire_bytes(point) > raw_payload_bytes + 3 * 64
|