78 lines
3.1 KiB
Python
78 lines
3.1 KiB
Python
import hashlib
|
|
import json
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessageChunk
|
|
|
|
from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_PREFIXES, PREFIX_BYTES
|
|
from artemis.llm.diagnostics import RequestAttempt
|
|
|
|
|
|
@pytest.mark.parametrize("fragment_size", [1, 37, 4096, 100000])
|
|
def test_fingerprints_match_at_byte_offsets_despite_unicode_and_different_fragmentation(
|
|
fragment_size,
|
|
):
|
|
arguments = '{"target":"' + "私人内容🙂" * 8000 + '"}'
|
|
expected = arguments.encode("utf-8")
|
|
fingerprint = ArgumentFingerprint()
|
|
for offset in range(0, len(arguments), fragment_size):
|
|
fingerprint.observe(arguments[offset : offset + fragment_size])
|
|
result = fingerprint.snapshot()
|
|
assert result["argument_chars"] == len(arguments)
|
|
assert result["argument_bytes"] == len(expected)
|
|
assert result["argument_sha256"] == hashlib.sha256(expected).hexdigest()
|
|
assert result["argument_encoding_valid"] is True
|
|
assert result["argument_prefixes"] == [
|
|
{"bytes": size, "sha256": hashlib.sha256(expected[:size]).hexdigest()}
|
|
for size in range(PREFIX_BYTES, PREFIX_BYTES * MAX_PREFIXES + 1, PREFIX_BYTES)
|
|
]
|
|
assert "私人" not in json.dumps(result, ensure_ascii=False)
|
|
|
|
|
|
def test_invalid_utf8_is_marked_without_failing_or_forging_a_digest():
|
|
fingerprint = ArgumentFingerprint()
|
|
for value in ["safe", "\ud800", "later"]:
|
|
fingerprint.observe(value)
|
|
expected = {
|
|
"argument_chars": 10,
|
|
"argument_bytes": None,
|
|
"argument_sha256": None,
|
|
"argument_deltas": 3,
|
|
"argument_prefixes": [],
|
|
"argument_encoding_valid": False,
|
|
}
|
|
assert {key: fingerprint.snapshot()[key] for key in expected} == expected
|
|
|
|
|
|
def test_shape_counts_cross_delta_whitespace_and_repetition_without_retaining_text():
|
|
fingerprint = ArgumentFingerprint()
|
|
for fragment in ['{"text":"private', " ", " \t", "", " \t", " \t", '\u00a0"}']:
|
|
fingerprint.observe(fragment)
|
|
stats = fingerprint.snapshot()
|
|
assert stats["argument_json_whitespace_chars"] == 7
|
|
assert stats["argument_json_whitespace_run_max"] == 7
|
|
assert stats["argument_json_whitespace_run"] == 0
|
|
assert stats["argument_delta_chars_min"] == 1
|
|
assert stats["argument_delta_chars_max"] == len('{"text":"private')
|
|
assert stats["argument_same_delta_run_max"] == 3
|
|
assert stats["argument_same_delta_run"] == 1
|
|
assert stats["argument_deltas"] == 6
|
|
assert "private" not in repr(vars(fingerprint))
|
|
assert "private" not in json.dumps(stats)
|
|
|
|
|
|
def test_tool_tracking_is_bounded_but_total_count_keeps_all_received_arguments():
|
|
attempt = RequestAttempt(
|
|
None, request_id="bounded", attempt=1, timeout=300, max_retries=0, metadata={}
|
|
)
|
|
for index in range(65):
|
|
attempt.chunk(
|
|
AIMessageChunk(
|
|
content=[],
|
|
tool_call_chunks=[{"name": "click", "args": "{}", "index": index}],
|
|
)
|
|
)
|
|
assert len(attempt.data["tools"]) == len(attempt.tool_fingerprints) == 64
|
|
assert attempt.data["tool_stats_truncated"] is True
|
|
assert attempt.data["tool_argument_chars"] == 130
|