60 lines
2.2 KiB
Python
60 lines
2.2 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)
|
|
assert fingerprint.snapshot() == {
|
|
"argument_chars": 10,
|
|
"argument_bytes": None,
|
|
"argument_sha256": None,
|
|
"argument_deltas": 3,
|
|
"argument_prefixes": [],
|
|
"argument_encoding_valid": False,
|
|
}
|
|
|
|
|
|
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
|