Files
Artemis-Eapil/artemis/llm/argument_fingerprint.py
T

54 lines
1.8 KiB
Python

"""Bounded fingerprints of received argument deltas, without retaining text."""
import hashlib
PREFIX_BYTES = 4096
MAX_PREFIXES = 16
MAX_TRACKED_TOOLS = 64
class ArgumentFingerprint:
def __init__(self):
self.characters = 0
self.bytes = 0
self.deltas = 0
self.encoding_valid = True
self.digest = hashlib.sha256()
self.prefixes: list[dict] = []
def observe(self, arguments: str) -> None:
self.characters += len(arguments)
if not arguments:
return
self.deltas += 1
if not self.encoding_valid:
return
try:
encoded = arguments.encode("utf-8")
except UnicodeEncodeError:
# Diagnostics must not repair, reject or otherwise alter the stream.
self.encoding_valid = False
return
offset = 0
while offset < len(encoded) and len(self.prefixes) < MAX_PREFIXES:
size = min(PREFIX_BYTES - self.bytes % PREFIX_BYTES, len(encoded) - offset)
self.digest.update(encoded[offset : offset + size])
self.bytes += size
offset += size
if self.bytes % PREFIX_BYTES == 0:
self.prefixes.append({"bytes": self.bytes, "sha256": self.digest.hexdigest()})
if offset < len(encoded):
self.digest.update(encoded[offset:])
self.bytes += len(encoded) - offset
def snapshot(self) -> dict:
return {
"argument_chars": self.characters,
"argument_bytes": self.bytes if self.encoding_valid else None,
"argument_sha256": self.digest.hexdigest() if self.encoding_valid else None,
"argument_deltas": self.deltas,
"argument_prefixes": [dict(prefix) for prefix in self.prefixes],
"argument_encoding_valid": self.encoding_valid,
}