headroom-ai 0.2.13__py3-none-any.whl

headroom/compression/handlers/json_handler.py

@@ -0,0 +1,418 @@
+"""JSON structure handler.
+
+Extracts structural elements from JSON content:
+- Keys (navigational - tells LLM what fields exist)
+- Brackets and colons (structural syntax)
+- Short values like booleans, nulls, small numbers
+
+Values (strings, long numbers, nested content) are marked as compressible.
+
+This enables the LLM to see the full schema while values are compressed.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, cast
+
+from headroom.compression.handlers.base import BaseStructureHandler, HandlerResult
+from headroom.compression.masks import EntropyScore, StructureMask
+
+
+class JSONTokenType(Enum):
+    """Types of JSON tokens for structure detection."""
+
+    KEY = "key"  # Object key (always structural)
+    STRING_VALUE = "string_value"  # String value (compressible)
+    NUMBER = "number"  # Numeric value (preserve if short)
+    BOOLEAN = "boolean"  # true/false (always structural)
+    NULL = "null"  # null (always structural)
+    BRACKET = "bracket"  # {, }, [, ] (always structural)
+    COLON = "colon"  # : (always structural)
+    COMMA = "comma"  # , (always structural)
+    WHITESPACE = "whitespace"  # spaces, newlines (compressible)
+
+
+@dataclass
+class JSONToken:
+    """A token in JSON content with its type and position."""
+
+    text: str
+    token_type: JSONTokenType
+    start: int
+    end: int
+
+    @property
+    def is_structural(self) -> bool:
+        """Whether this token should be preserved."""
+        return self.token_type in (
+            JSONTokenType.KEY,
+            JSONTokenType.BOOLEAN,
+            JSONTokenType.NULL,
+            JSONTokenType.BRACKET,
+            JSONTokenType.COLON,
+            JSONTokenType.COMMA,
+        )
+
+
+class JSONStructureHandler(BaseStructureHandler):
+    """Handler for JSON content.
+
+    Preserves:
+    - All keys (navigational - LLM sees what fields exist)
+    - Structural syntax ({, }, [, ], :, ,)
+    - Booleans and nulls (small, semantically important)
+    - High-entropy strings (UUIDs, hashes - identifiers)
+    - Short numbers (often IDs or important values)
+
+    Compresses:
+    - Long string values (descriptions, content)
+    - Whitespace
+    - Redundant array elements (after first few)
+
+    Example:
+        >>> handler = JSONStructureHandler()
+        >>> result = handler.get_mask('{"name": "Alice", "id": "usr_123"}')
+        >>> # Keys "name" and "id" preserved, values may be compressed
+    """
+
+    def __init__(
+        self,
+        preserve_short_values: bool = True,
+        short_value_threshold: int = 20,
+        preserve_high_entropy: bool = True,
+        entropy_threshold: float = 0.85,
+        max_array_items_full: int = 3,  # Keep first N items fully
+        max_number_digits: int = 10,  # Preserve numbers up to N digits
+    ):
+        """Initialize the JSON handler.
+
+        Args:
+            preserve_short_values: Preserve short string values.
+            short_value_threshold: Max length for "short" values.
+            preserve_high_entropy: Preserve high-entropy strings (UUIDs, etc.).
+            entropy_threshold: Entropy threshold for preservation.
+            max_array_items_full: Number of array items to keep in full.
+            max_number_digits: Max digits for numbers to preserve (often IDs).
+        """
+        super().__init__(name="json")
+        self.preserve_short_values = preserve_short_values
+        self.short_value_threshold = short_value_threshold
+        self.preserve_high_entropy = preserve_high_entropy
+        self.entropy_threshold = entropy_threshold
+        self.max_array_items_full = max_array_items_full
+        self.max_number_digits = max_number_digits
+
+    def can_handle(self, content: str) -> bool:
+        """Check if content is valid JSON."""
+        stripped = content.strip()
+        if not stripped.startswith(("{", "[")):
+            return False
+        try:
+            json.loads(stripped)
+            return True
+        except (json.JSONDecodeError, ValueError):
+            return False
+
+    def _extract_mask(
+        self,
+        content: str,
+        tokens: list[str],
+        **kwargs: Any,
+    ) -> HandlerResult:
+        """Extract structure mask from JSON content.
+
+        Args:
+            content: JSON content.
+            tokens: Character-level tokens.
+            **kwargs: Additional options.
+
+        Returns:
+            HandlerResult with mask marking structural elements.
+        """
+        # Tokenize JSON to identify structure
+        json_tokens = self._tokenize_json(content)
+
+        # Build character-level mask
+        mask = [False] * len(content)
+
+        # Track array depth for selective preservation
+        array_depth = 0
+        array_item_counts: dict[int, int] = {}  # depth -> count
+
+        for token in json_tokens:
+            # Track array items
+            if token.token_type == JSONTokenType.BRACKET:
+                if token.text == "[":
+                    array_depth += 1
+                    array_item_counts[array_depth] = 0
+                elif token.text == "]":
+                    if array_depth in array_item_counts:
+                        del array_item_counts[array_depth]
+                    array_depth = max(0, array_depth - 1)
+
+            # Count array items at commas
+            if token.token_type == JSONTokenType.COMMA and array_depth > 0:
+                array_item_counts[array_depth] = array_item_counts.get(array_depth, 0) + 1
+
+            # Determine if this token should be preserved
+            preserve = self._should_preserve_token(
+                token,
+                array_depth,
+                array_item_counts.get(array_depth, 0),
+            )
+
+            # Mark in mask
+            if preserve:
+                for i in range(token.start, min(token.end, len(mask))):
+                    mask[i] = True
+
+        # Convert to character tokens if needed
+        char_tokens = list(content) if tokens == list(content) else tokens
+
+        return HandlerResult(
+            mask=StructureMask(tokens=char_tokens, mask=mask),
+            handler_name=self.name,
+            confidence=1.0,
+            metadata={
+                "token_count": len(json_tokens),
+                "key_count": sum(1 for t in json_tokens if t.token_type == JSONTokenType.KEY),
+            },
+        )
+
+    def _should_preserve_token(
+        self,
+        token: JSONToken,
+        array_depth: int,
+        array_item_index: int,
+    ) -> bool:
+        """Determine if a token should be preserved.
+
+        Args:
+            token: The JSON token.
+            array_depth: Current array nesting depth.
+            array_item_index: Index of current item in array.
+
+        Returns:
+            True if token should be preserved.
+        """
+        # Always preserve structural tokens
+        if token.is_structural:
+            return True
+
+        # Whitespace is never preserved
+        if token.token_type == JSONTokenType.WHITESPACE:
+            return False
+
+        # Numbers: preserve short ones (often IDs)
+        if token.token_type == JSONTokenType.NUMBER:
+            return len(token.text) <= self.max_number_digits
+
+        # String values: selective preservation
+        if token.token_type == JSONTokenType.STRING_VALUE:
+            # Check if we're past the max array items threshold
+            if array_depth > 0 and array_item_index >= self.max_array_items_full:
+                # In deep array, be more aggressive
+                return False
+
+            # Preserve short values
+            if self.preserve_short_values and len(token.text) <= self.short_value_threshold:
+                return True
+
+            # Preserve high-entropy values (UUIDs, hashes)
+            if self.preserve_high_entropy:
+                # Strip quotes for entropy calculation
+                value = token.text.strip('"')
+                score = EntropyScore.compute(value, self.entropy_threshold)
+                if score.should_preserve:
+                    return True
+
+            return False
+
+        return False
+
+    def _tokenize_json(self, content: str) -> list[JSONToken]:
+        """Tokenize JSON content into typed tokens.
+
+        This is a simple tokenizer that identifies JSON structure.
+        It's not a full parser - just enough to identify keys vs values.
+
+        Args:
+            content: JSON content.
+
+        Returns:
+            List of JSONToken objects.
+        """
+        tokens: list[JSONToken] = []
+        i = 0
+        n = len(content)
+
+        # Track if we're expecting a key (after { or ,)
+        expect_key = False
+        brace_stack: list[str] = []
+
+        while i < n:
+            char = content[i]
+
+            # Whitespace
+            if char in " \t\n\r":
+                start = i
+                while i < n and content[i] in " \t\n\r":
+                    i += 1
+                tokens.append(JSONToken(content[start:i], JSONTokenType.WHITESPACE, start, i))
+                continue
+
+            # Brackets
+            if char in "{}[]":
+                tokens.append(JSONToken(char, JSONTokenType.BRACKET, i, i + 1))
+                if char == "{":
+                    brace_stack.append("{")
+                    expect_key = True
+                elif char == "}":
+                    if brace_stack and brace_stack[-1] == "{":
+                        brace_stack.pop()
+                    expect_key = False
+                elif char == "[":
+                    brace_stack.append("[")
+                    expect_key = False
+                elif char == "]":
+                    if brace_stack and brace_stack[-1] == "[":
+                        brace_stack.pop()
+                i += 1
+                continue
+
+            # Colon
+            if char == ":":
+                tokens.append(JSONToken(char, JSONTokenType.COLON, i, i + 1))
+                expect_key = False
+                i += 1
+                continue
+
+            # Comma
+            if char == ",":
+                tokens.append(JSONToken(char, JSONTokenType.COMMA, i, i + 1))
+                # After comma in object, expect key
+                if brace_stack and brace_stack[-1] == "{":
+                    expect_key = True
+                i += 1
+                continue
+
+            # String (key or value)
+            if char == '"':
+                start = i
+                i += 1
+                while i < n and content[i] != '"':
+                    if content[i] == "\\":
+                        i += 2  # Skip escaped character
+                    else:
+                        i += 1
+                i += 1  # Include closing quote
+
+                text = content[start:i]
+
+                # Determine if this is a key or value
+                # Look ahead for colon (skipping whitespace)
+                j = i
+                while j < n and content[j] in " \t\n\r":
+                    j += 1
+
+                is_key = j < n and content[j] == ":" and expect_key
+
+                if is_key:
+                    tokens.append(JSONToken(text, JSONTokenType.KEY, start, i))
+                    expect_key = False
+                else:
+                    tokens.append(JSONToken(text, JSONTokenType.STRING_VALUE, start, i))
+
+                continue
+
+            # Number
+            if char in "-0123456789":
+                start = i
+                # Match JSON number pattern
+                if char == "-":
+                    i += 1
+                while i < n and content[i] in "0123456789":
+                    i += 1
+                if i < n and content[i] == ".":
+                    i += 1
+                    while i < n and content[i] in "0123456789":
+                        i += 1
+                if i < n and content[i] in "eE":
+                    i += 1
+                    if i < n and content[i] in "+-":
+                        i += 1
+                    while i < n and content[i] in "0123456789":
+                        i += 1
+
+                tokens.append(JSONToken(content[start:i], JSONTokenType.NUMBER, start, i))
+                continue
+
+            # Boolean or null
+            if content[i : i + 4] == "true":
+                tokens.append(JSONToken("true", JSONTokenType.BOOLEAN, i, i + 4))
+                i += 4
+                continue
+            if content[i : i + 5] == "false":
+                tokens.append(JSONToken("false", JSONTokenType.BOOLEAN, i, i + 5))
+                i += 5
+                continue
+            if content[i : i + 4] == "null":
+                tokens.append(JSONToken("null", JSONTokenType.NULL, i, i + 4))
+                i += 4
+                continue
+
+            # Unknown character - skip
+            i += 1
+
+        return tokens
+
+
+def extract_json_schema(content: str) -> dict[str, Any] | list[Any]:
+    """Extract the schema (keys only) from JSON content.
+
+    Useful for understanding the structure without the values.
+
+    Args:
+        content: JSON content.
+
+    Returns:
+        Schema dictionary with keys and types (no values).
+
+    Example:
+        >>> extract_json_schema('{"name": "Alice", "age": 30}')
+        {'name': 'string', 'age': 'number'}
+    """
+
+    def _extract(obj: Any) -> Any:
+        if isinstance(obj, dict):
+            return {k: _extract(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            if obj:
+                return [_extract(obj[0])]  # Schema of first item
+            return []
+        elif isinstance(obj, str):
+            return "string"
+        elif isinstance(obj, bool):
+            return "boolean"
+        elif isinstance(obj, int):
+            return "integer"
+        elif isinstance(obj, float):
+            return "number"
+        elif obj is None:
+            return "null"
+        else:
+            return "unknown"
+
+    try:
+        parsed = json.loads(content)
+        result = _extract(parsed)
+        if isinstance(result, dict):
+            return cast(dict[str, Any], result)
+        elif isinstance(result, list):
+            return cast(list[Any], result)
+        return {}
+    except (json.JSONDecodeError, ValueError):
+        return {}