diff --git a/headroom/compression/handlers/json_handler.py b/headroom/compression/handlers/json_handler.py index bd9796d5c..941ef2b87 100644 --- a/headroom/compression/handlers/json_handler.py +++ b/headroom/compression/handlers/json_handler.py @@ -15,7 +15,7 @@ from __future__ import annotations import json from dataclasses import dataclass from enum import Enum -from typing import Any +from typing import Any, cast from headroom.compression.handlers.base import BaseStructureHandler, HandlerResult from headroom.compression.masks import EntropyScore, StructureMask @@ -370,7 +370,7 @@ class JSONStructureHandler(BaseStructureHandler): return tokens -def extract_json_schema(content: str) -> dict[str, Any]: +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. @@ -408,6 +408,11 @@ def extract_json_schema(content: str) -> dict[str, Any]: try: parsed = json.loads(content) - return _extract(parsed) + 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 {} diff --git a/headroom/compression/universal.py b/headroom/compression/universal.py index 05219691a..2025b25bc 100644 --- a/headroom/compression/universal.py +++ b/headroom/compression/universal.py @@ -187,11 +187,10 @@ class UniversalCompressor: try: from headroom.transforms.llmlingua_compressor import compress_with_llmlingua - result = compress_with_llmlingua( + return compress_with_llmlingua( text, - target_ratio=self.config.compression_ratio_target, + compression_rate=self.config.compression_ratio_target, ) - return result.compressed except ImportError: return self._simple_compress(text) except Exception as e: @@ -371,8 +370,8 @@ class UniversalCompressor: self._ccr_store = CompressionStore() key = self._ccr_store.store( - original_content=original, - compressed_content=compressed, + original, + compressed, original_tokens=self._estimate_tokens(original), compressed_tokens=self._estimate_tokens(compressed), )