chopratejas commited on
Commit
ac3b21b
·
1 Parent(s): 02fd228

Add CCR marker support to all compressors and fix retrieval

Browse files

- Fix CCRToolInjector detection for LLMLingua (use standard marker format)
- Add CCR marker output to CodeAwareCompressor
- Add SmartCrusher.crush() method for ContentRouter integration
- Fix ContentRouter to pass CCRConfig to SmartCrusher
- Fix CCR response handler decompression error (remove encoding headers)
- Add content block processing for Anthropic tool_result format

headroom/proxy/server.py CHANGED
@@ -28,6 +28,7 @@ import asyncio
28
  import hashlib
29
  import json
30
  import logging
 
31
  import random
32
  import sys
33
  import time
@@ -63,7 +64,7 @@ from headroom.ccr import (
63
  ResponseHandlerConfig,
64
  parse_tool_call,
65
  )
66
- from headroom.config import CacheAlignerConfig, RollingWindowConfig, SmartCrusherConfig
67
  from headroom.providers import AnthropicProvider, OpenAIProvider
68
  from headroom.telemetry import get_telemetry_collector
69
  from headroom.tokenizers import get_tokenizer
@@ -72,6 +73,8 @@ from headroom.transforms import (
72
  CacheAligner,
73
  CodeAwareCompressor,
74
  CodeCompressorConfig,
 
 
75
  RollingWindow,
76
  SmartCrusher,
77
  TransformPipeline,
@@ -183,6 +186,9 @@ class ProxyConfig:
183
  # Code-aware compression (ON by default if installed)
184
  code_aware_enabled: bool = True # Enable AST-based code compression
185
 
 
 
 
186
  # Caching
187
  cache_enabled: bool = True
188
  cache_ttl_seconds: int = 3600 # 1 hour
@@ -675,30 +681,54 @@ class HeadroomProxy:
675
  self.anthropic_provider = AnthropicProvider()
676
  self.openai_provider = OpenAIProvider()
677
 
678
- # Initialize transforms
679
- transforms = [
680
- CacheAligner(CacheAlignerConfig(enabled=True)),
681
- SmartCrusher(
682
- SmartCrusherConfig( # type: ignore[arg-type]
683
- enabled=True,
684
- min_tokens_to_crush=config.min_tokens_to_crush,
685
- max_items_after_crush=config.max_items_after_crush,
686
- )
687
- ),
688
- RollingWindow(
689
- RollingWindowConfig(
690
- enabled=True,
691
- keep_system=True,
692
- keep_last_turns=config.keep_last_turns,
693
- )
694
- ),
695
- ]
696
-
697
- # Add LLMLingua if enabled and available
698
- self._llmlingua_status = self._setup_llmlingua(config, transforms)
699
-
700
- # Add CodeAware if enabled and available
701
- self._code_aware_status = self._setup_code_aware(config, transforms)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
702
 
703
  self.anthropic_pipeline = TransformPipeline(
704
  transforms=transforms,
@@ -873,24 +903,38 @@ class HeadroomProxy:
873
  logger.info(f"Caching: {'ENABLED' if self.config.cache_enabled else 'DISABLED'}")
874
  logger.info(f"Rate Limiting: {'ENABLED' if self.config.rate_limit_enabled else 'DISABLED'}")
875
 
 
 
 
 
 
 
876
  # LLMLingua status with helpful hint
877
  if self._llmlingua_status == "enabled":
878
  logger.info(
879
  f"LLMLingua: ENABLED (device={self.config.llmlingua_device}, "
880
  f"rate={self.config.llmlingua_target_rate})"
881
  )
 
 
882
  elif self._llmlingua_status == "available":
883
- logger.info("LLMLingua: available but disabled (use without --no-llmlingua)")
884
  elif self._llmlingua_status == "unavailable":
885
  logger.info("LLMLingua: not installed (pip install headroom-ai[llmlingua])")
 
 
886
 
887
  # Code-aware status
888
  if self._code_aware_status == "enabled":
889
  logger.info("Code-Aware: ENABLED (AST-based compression)")
 
 
890
  elif self._code_aware_status == "available":
891
- logger.info("Code-Aware: available but disabled (use without --no-code-aware)")
892
  elif self._code_aware_status == "unavailable":
893
  logger.info("Code-Aware: not installed (pip install headroom-ai[code])")
 
 
894
 
895
  # CCR status
896
  ccr_features = []
@@ -1233,6 +1277,7 @@ class HeadroomProxy:
1233
  logger.info(f"[{request_id}] CCR: Detected retrieval tool call, handling...")
1234
 
1235
  # Create API call function for continuation
 
1236
  async def api_call_fn(
1237
  msgs: list[dict], tls: list[dict] | None
1238
  ) -> dict[str, Any]:
@@ -1242,11 +1287,44 @@ class HeadroomProxy:
1242
  }
1243
  if tls is not None:
1244
  continuation_body["tools"] = tls
1245
- cont_response = await self._retry_request(
1246
- "POST", url, headers, continuation_body
1247
- )
1248
- result: dict[str, Any] = cont_response.json()
1249
- return result
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1250
 
1251
  # Handle CCR tool calls
1252
  try:
@@ -1259,14 +1337,25 @@ class HeadroomProxy:
1259
  )
1260
  # Update response content with final response
1261
  resp_json = final_resp_json
 
 
 
 
 
 
1262
  response = httpx.Response(
1263
  status_code=200,
1264
  content=json.dumps(final_resp_json).encode(),
1265
- headers=dict(response.headers),
1266
  )
1267
  logger.info(f"[{request_id}] CCR: Retrieval handled successfully")
1268
  except Exception as e:
1269
- logger.warning(f"[{request_id}] CCR: Response handling failed: {e}")
 
 
 
 
 
1270
  # Continue with original response
1271
 
1272
  total_latency = (time.time() - start_time) * 1000
@@ -2203,6 +2292,41 @@ def run_server(config: ProxyConfig | None = None):
2203
  uvicorn.run(app, host=config.host, port=config.port, log_level="warning")
2204
 
2205
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2206
  if __name__ == "__main__":
2207
  parser = argparse.ArgumentParser(description="Headroom Proxy Server")
2208
 
@@ -2232,12 +2356,24 @@ if __name__ == "__main__":
2232
  parser.add_argument("--log-file", help="Log file path")
2233
  parser.add_argument("--log-messages", action="store_true", help="Log full messages")
2234
 
 
 
 
 
 
 
 
2235
  # LLMLingua ML-based compression
2236
  parser.add_argument(
2237
  "--llmlingua",
2238
  action="store_true",
2239
  help="Enable LLMLingua-2 ML-based compression (requires: pip install headroom-ai[llmlingua])",
2240
  )
 
 
 
 
 
2241
  parser.add_argument(
2242
  "--llmlingua-device",
2243
  choices=["auto", "cuda", "cpu", "mps"],
@@ -2251,26 +2387,68 @@ if __name__ == "__main__":
2251
  help="LLMLingua target compression rate, 0.0-1.0 (default: 0.3 = keep 30%%)",
2252
  )
2253
 
 
 
 
 
 
 
 
 
 
 
 
 
2254
  args = parser.parse_args()
2255
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2256
  config = ProxyConfig(
2257
- host=args.host,
2258
- port=args.port,
2259
- optimize=not args.no_optimize,
2260
- min_tokens_to_crush=args.min_tokens,
2261
- max_items_after_crush=args.max_items,
2262
- cache_enabled=not args.no_cache,
2263
- cache_ttl_seconds=args.cache_ttl,
2264
- rate_limit_enabled=not args.no_rate_limit,
2265
- rate_limit_requests_per_minute=args.rpm,
2266
- rate_limit_tokens_per_minute=args.tpm,
2267
  budget_limit_usd=args.budget,
2268
  budget_period=args.budget_period,
2269
- log_file=args.log_file,
2270
- log_full_messages=args.log_messages,
2271
- llmlingua_enabled=args.llmlingua,
2272
- llmlingua_device=args.llmlingua_device,
2273
- llmlingua_target_rate=args.llmlingua_rate,
 
 
 
 
2274
  )
2275
 
2276
  run_server(config)
 
28
  import hashlib
29
  import json
30
  import logging
31
+ import os
32
  import random
33
  import sys
34
  import time
 
64
  ResponseHandlerConfig,
65
  parse_tool_call,
66
  )
67
+ from headroom.config import CacheAlignerConfig, CCRConfig, RollingWindowConfig, SmartCrusherConfig
68
  from headroom.providers import AnthropicProvider, OpenAIProvider
69
  from headroom.telemetry import get_telemetry_collector
70
  from headroom.tokenizers import get_tokenizer
 
73
  CacheAligner,
74
  CodeAwareCompressor,
75
  CodeCompressorConfig,
76
+ ContentRouter,
77
+ ContentRouterConfig,
78
  RollingWindow,
79
  SmartCrusher,
80
  TransformPipeline,
 
186
  # Code-aware compression (ON by default if installed)
187
  code_aware_enabled: bool = True # Enable AST-based code compression
188
 
189
+ # Smart content routing (routes each message to optimal compressor)
190
+ smart_routing: bool = True # Use ContentRouter for intelligent compression
191
+
192
  # Caching
193
  cache_enabled: bool = True
194
  cache_ttl_seconds: int = 3600 # 1 hour
 
681
  self.anthropic_provider = AnthropicProvider()
682
  self.openai_provider = OpenAIProvider()
683
 
684
+ # Initialize transforms based on routing mode
685
+ if config.smart_routing:
686
+ # Smart routing: ContentRouter handles all content types intelligently
687
+ # It lazy-loads compressors (including LLMLingua) only when needed
688
+ router_config = ContentRouterConfig(
689
+ enable_llmlingua=config.llmlingua_enabled,
690
+ enable_code_aware=config.code_aware_enabled,
691
+ )
692
+ transforms = [
693
+ CacheAligner(CacheAlignerConfig(enabled=True)),
694
+ ContentRouter(router_config),
695
+ RollingWindow(
696
+ RollingWindowConfig(
697
+ enabled=True,
698
+ keep_system=True,
699
+ keep_last_turns=config.keep_last_turns,
700
+ )
701
+ ),
702
+ ]
703
+ self._llmlingua_status = "lazy" if config.llmlingua_enabled else "disabled"
704
+ self._code_aware_status = "lazy" if config.code_aware_enabled else "disabled"
705
+ else:
706
+ # Legacy mode: sequential pipeline
707
+ transforms = [
708
+ CacheAligner(CacheAlignerConfig(enabled=True)),
709
+ SmartCrusher(
710
+ SmartCrusherConfig( # type: ignore[arg-type]
711
+ enabled=True,
712
+ min_tokens_to_crush=config.min_tokens_to_crush,
713
+ max_items_after_crush=config.max_items_after_crush,
714
+ ),
715
+ ccr_config=CCRConfig(
716
+ enabled=config.ccr_inject_tool,
717
+ inject_retrieval_marker=config.ccr_inject_tool, # Add CCR markers
718
+ ),
719
+ ),
720
+ RollingWindow(
721
+ RollingWindowConfig(
722
+ enabled=True,
723
+ keep_system=True,
724
+ keep_last_turns=config.keep_last_turns,
725
+ )
726
+ ),
727
+ ]
728
+ # Add LLMLingua if enabled and available
729
+ self._llmlingua_status = self._setup_llmlingua(config, transforms)
730
+ # Add CodeAware if enabled and available
731
+ self._code_aware_status = self._setup_code_aware(config, transforms)
732
 
733
  self.anthropic_pipeline = TransformPipeline(
734
  transforms=transforms,
 
903
  logger.info(f"Caching: {'ENABLED' if self.config.cache_enabled else 'DISABLED'}")
904
  logger.info(f"Rate Limiting: {'ENABLED' if self.config.rate_limit_enabled else 'DISABLED'}")
905
 
906
+ # Smart routing status
907
+ if self.config.smart_routing:
908
+ logger.info("Smart Routing: ENABLED (intelligent content detection)")
909
+ else:
910
+ logger.info("Smart Routing: DISABLED (legacy sequential mode)")
911
+
912
  # LLMLingua status with helpful hint
913
  if self._llmlingua_status == "enabled":
914
  logger.info(
915
  f"LLMLingua: ENABLED (device={self.config.llmlingua_device}, "
916
  f"rate={self.config.llmlingua_target_rate})"
917
  )
918
+ elif self._llmlingua_status == "lazy":
919
+ logger.info("LLMLingua: LAZY (will load when prose content detected)")
920
  elif self._llmlingua_status == "available":
921
+ logger.info("LLMLingua: available but disabled (use --llmlingua)")
922
  elif self._llmlingua_status == "unavailable":
923
  logger.info("LLMLingua: not installed (pip install headroom-ai[llmlingua])")
924
+ elif self._llmlingua_status == "disabled":
925
+ logger.info("LLMLingua: DISABLED")
926
 
927
  # Code-aware status
928
  if self._code_aware_status == "enabled":
929
  logger.info("Code-Aware: ENABLED (AST-based compression)")
930
+ elif self._code_aware_status == "lazy":
931
+ logger.info("Code-Aware: LAZY (will load when code content detected)")
932
  elif self._code_aware_status == "available":
933
+ logger.info("Code-Aware: available but disabled (use --code-aware)")
934
  elif self._code_aware_status == "unavailable":
935
  logger.info("Code-Aware: not installed (pip install headroom-ai[code])")
936
+ elif self._code_aware_status == "disabled":
937
+ logger.info("Code-Aware: DISABLED")
938
 
939
  # CCR status
940
  ccr_features = []
 
1277
  logger.info(f"[{request_id}] CCR: Detected retrieval tool call, handling...")
1278
 
1279
  # Create API call function for continuation
1280
+ # Use a fresh client to avoid potential decompression state issues
1281
  async def api_call_fn(
1282
  msgs: list[dict], tls: list[dict] | None
1283
  ) -> dict[str, Any]:
 
1287
  }
1288
  if tls is not None:
1289
  continuation_body["tools"] = tls
1290
+
1291
+ # Use clean headers for continuation
1292
+ continuation_headers = {
1293
+ k: v
1294
+ for k, v in headers.items()
1295
+ if k.lower()
1296
+ not in (
1297
+ "content-encoding",
1298
+ "transfer-encoding",
1299
+ "accept-encoding",
1300
+ "content-length",
1301
+ )
1302
+ }
1303
+
1304
+ # Use a fresh client for CCR continuations
1305
+ logger.info(f"CCR: Making continuation request with {len(msgs)} messages")
1306
+ async with httpx.AsyncClient(
1307
+ timeout=httpx.Timeout(120.0),
1308
+ ) as ccr_client:
1309
+ try:
1310
+ cont_response = await ccr_client.post(
1311
+ url,
1312
+ json=continuation_body,
1313
+ headers=continuation_headers,
1314
+ )
1315
+ logger.info(
1316
+ f"CCR: Got response status={cont_response.status_code}, "
1317
+ f"content-encoding={cont_response.headers.get('content-encoding')}"
1318
+ )
1319
+ result: dict[str, Any] = cont_response.json()
1320
+ logger.info("CCR: Parsed JSON successfully")
1321
+ return result
1322
+ except Exception as e:
1323
+ logger.error(
1324
+ f"CCR: API call failed: {e}, "
1325
+ f"response headers: {dict(cont_response.headers) if 'cont_response' in dir() else 'N/A'}"
1326
+ )
1327
+ raise
1328
 
1329
  # Handle CCR tool calls
1330
  try:
 
1337
  )
1338
  # Update response content with final response
1339
  resp_json = final_resp_json
1340
+ # Remove encoding headers since content is now uncompressed JSON
1341
+ ccr_response_headers = {
1342
+ k: v
1343
+ for k, v in response.headers.items()
1344
+ if k.lower() not in ("content-encoding", "content-length")
1345
+ }
1346
  response = httpx.Response(
1347
  status_code=200,
1348
  content=json.dumps(final_resp_json).encode(),
1349
+ headers=ccr_response_headers,
1350
  )
1351
  logger.info(f"[{request_id}] CCR: Retrieval handled successfully")
1352
  except Exception as e:
1353
+ import traceback
1354
+
1355
+ logger.warning(
1356
+ f"[{request_id}] CCR: Response handling failed: {e}\n"
1357
+ f"Traceback: {traceback.format_exc()}"
1358
+ )
1359
  # Continue with original response
1360
 
1361
  total_latency = (time.time() - start_time) * 1000
 
2292
  uvicorn.run(app, host=config.host, port=config.port, log_level="warning")
2293
 
2294
 
2295
+ def _get_env_bool(name: str, default: bool) -> bool:
2296
+ """Get boolean from environment variable."""
2297
+ val = os.environ.get(name)
2298
+ if val is None:
2299
+ return default
2300
+ return val.lower() in ("true", "1", "yes", "on")
2301
+
2302
+
2303
+ def _get_env_int(name: str, default: int) -> int:
2304
+ """Get integer from environment variable."""
2305
+ val = os.environ.get(name)
2306
+ if val is None:
2307
+ return default
2308
+ try:
2309
+ return int(val)
2310
+ except ValueError:
2311
+ return default
2312
+
2313
+
2314
+ def _get_env_float(name: str, default: float) -> float:
2315
+ """Get float from environment variable."""
2316
+ val = os.environ.get(name)
2317
+ if val is None:
2318
+ return default
2319
+ try:
2320
+ return float(val)
2321
+ except ValueError:
2322
+ return default
2323
+
2324
+
2325
+ def _get_env_str(name: str, default: str) -> str:
2326
+ """Get string from environment variable."""
2327
+ return os.environ.get(name, default)
2328
+
2329
+
2330
  if __name__ == "__main__":
2331
  parser = argparse.ArgumentParser(description="Headroom Proxy Server")
2332
 
 
2356
  parser.add_argument("--log-file", help="Log file path")
2357
  parser.add_argument("--log-messages", action="store_true", help="Log full messages")
2358
 
2359
+ # Smart routing (content-aware compression)
2360
+ parser.add_argument(
2361
+ "--no-smart-routing",
2362
+ action="store_true",
2363
+ help="Disable smart routing (use legacy sequential pipeline)",
2364
+ )
2365
+
2366
  # LLMLingua ML-based compression
2367
  parser.add_argument(
2368
  "--llmlingua",
2369
  action="store_true",
2370
  help="Enable LLMLingua-2 ML-based compression (requires: pip install headroom-ai[llmlingua])",
2371
  )
2372
+ parser.add_argument(
2373
+ "--no-llmlingua",
2374
+ action="store_true",
2375
+ help="Disable LLMLingua compression",
2376
+ )
2377
  parser.add_argument(
2378
  "--llmlingua-device",
2379
  choices=["auto", "cuda", "cpu", "mps"],
 
2387
  help="LLMLingua target compression rate, 0.0-1.0 (default: 0.3 = keep 30%%)",
2388
  )
2389
 
2390
+ # Code-aware compression
2391
+ parser.add_argument(
2392
+ "--code-aware",
2393
+ action="store_true",
2394
+ help="Enable AST-based code compression (requires: pip install headroom-ai[code])",
2395
+ )
2396
+ parser.add_argument(
2397
+ "--no-code-aware",
2398
+ action="store_true",
2399
+ help="Disable code-aware compression",
2400
+ )
2401
+
2402
  args = parser.parse_args()
2403
 
2404
+ # Environment variable defaults (HEADROOM_* prefix)
2405
+ # CLI args override env vars, env vars override ProxyConfig defaults
2406
+ env_smart_routing = _get_env_bool("HEADROOM_SMART_ROUTING", True)
2407
+ env_llmlingua = _get_env_bool("HEADROOM_LLMLINGUA_ENABLED", True)
2408
+ env_code_aware = _get_env_bool("HEADROOM_CODE_AWARE_ENABLED", True)
2409
+ env_optimize = _get_env_bool("HEADROOM_OPTIMIZE", True)
2410
+ env_cache = _get_env_bool("HEADROOM_CACHE_ENABLED", True)
2411
+ env_rate_limit = _get_env_bool("HEADROOM_RATE_LIMIT_ENABLED", True)
2412
+
2413
+ # Determine settings: CLI flags override env vars
2414
+ # --no-X explicitly disables, --X explicitly enables, neither uses env var
2415
+ smart_routing = env_smart_routing if not args.no_smart_routing else False
2416
+ llmlingua_enabled = (
2417
+ env_llmlingua
2418
+ if not (args.llmlingua or args.no_llmlingua)
2419
+ else (args.llmlingua or not args.no_llmlingua)
2420
+ )
2421
+ code_aware_enabled = (
2422
+ env_code_aware
2423
+ if not (args.code_aware or args.no_code_aware)
2424
+ else (args.code_aware or not args.no_code_aware)
2425
+ )
2426
+ optimize = env_optimize if not args.no_optimize else False
2427
+ cache_enabled = env_cache if not args.no_cache else False
2428
+ rate_limit_enabled = env_rate_limit if not args.no_rate_limit else False
2429
+
2430
  config = ProxyConfig(
2431
+ host=_get_env_str("HEADROOM_HOST", args.host),
2432
+ port=_get_env_int("HEADROOM_PORT", args.port),
2433
+ optimize=optimize,
2434
+ min_tokens_to_crush=_get_env_int("HEADROOM_MIN_TOKENS", args.min_tokens),
2435
+ max_items_after_crush=_get_env_int("HEADROOM_MAX_ITEMS", args.max_items),
2436
+ cache_enabled=cache_enabled,
2437
+ cache_ttl_seconds=_get_env_int("HEADROOM_CACHE_TTL", args.cache_ttl),
2438
+ rate_limit_enabled=rate_limit_enabled,
2439
+ rate_limit_requests_per_minute=_get_env_int("HEADROOM_RPM", args.rpm),
2440
+ rate_limit_tokens_per_minute=_get_env_int("HEADROOM_TPM", args.tpm),
2441
  budget_limit_usd=args.budget,
2442
  budget_period=args.budget_period,
2443
+ log_file=_get_env_str("HEADROOM_LOG_FILE", args.log_file)
2444
+ if args.log_file
2445
+ else os.environ.get("HEADROOM_LOG_FILE"),
2446
+ log_full_messages=args.log_messages or _get_env_bool("HEADROOM_LOG_MESSAGES", False),
2447
+ smart_routing=smart_routing,
2448
+ llmlingua_enabled=llmlingua_enabled,
2449
+ llmlingua_device=_get_env_str("HEADROOM_LLMLINGUA_DEVICE", args.llmlingua_device),
2450
+ llmlingua_target_rate=_get_env_float("HEADROOM_LLMLINGUA_RATE", args.llmlingua_rate),
2451
+ code_aware_enabled=code_aware_enabled,
2452
  )
2453
 
2454
  run_server(config)
headroom/transforms/code_compressor.py CHANGED
@@ -558,6 +558,12 @@ class CodeAwareCompressor(Transform):
558
  cache_key = None
559
  if self.config.enable_ccr and ratio < 0.8:
560
  cache_key = self._store_in_ccr(code, compressed, original_tokens)
 
 
 
 
 
 
561
 
562
  return CodeCompressionResult(
563
  compressed=compressed,
@@ -785,7 +791,8 @@ class CodeAwareCompressor(Transform):
785
 
786
  if total_body > keep_lines:
787
  omitted = total_body - keep_lines
788
- result_parts.append(f"{indent}# ... ({omitted} lines compressed)")
 
789
  result_parts.append(f"{indent}pass")
790
 
791
  return "\n".join(result_parts)
 
558
  cache_key = None
559
  if self.config.enable_ccr and ratio < 0.8:
560
  cache_key = self._store_in_ccr(code, compressed, original_tokens)
561
+ if cache_key:
562
+ # Add standard CCR marker format for CCRToolInjector detection
563
+ compressed += (
564
+ f"\n# [{original_tokens} items compressed to {compressed_tokens}. "
565
+ f"Retrieve more: hash={cache_key}]"
566
+ )
567
 
568
  return CodeCompressionResult(
569
  compressed=compressed,
 
791
 
792
  if total_body > keep_lines:
793
  omitted = total_body - keep_lines
794
+ # Simple, honest marker - no retrieval hints (causes hallucination)
795
+ result_parts.append(f"{indent}# [{omitted} lines omitted]")
796
  result_parts.append(f"{indent}pass")
797
 
798
  return "\n".join(result_parts)
headroom/transforms/content_router.py CHANGED
@@ -172,6 +172,9 @@ class ContentRouterConfig:
172
  mixed_content_threshold: Min distinct types to consider "mixed".
173
  min_section_tokens: Minimum tokens for a section to compress.
174
  fallback_strategy: Strategy when no compressor matches.
 
 
 
175
  """
176
 
177
  # Enable/disable specific compressors
@@ -189,6 +192,15 @@ class ContentRouterConfig:
189
  # Fallback
190
  fallback_strategy: CompressionStrategy = CompressionStrategy.PASSTHROUGH
191
 
 
 
 
 
 
 
 
 
 
192
 
193
  # Patterns for detecting mixed content
194
  _CODE_FENCE_PATTERN = re.compile(r"^```(\w*)\s*$", re.MULTILINE)
@@ -805,12 +817,18 @@ class ContentRouter(Transform):
805
  return self._code_compressor
806
 
807
  def _get_smart_crusher(self) -> Any:
808
- """Get SmartCrusher (lazy load)."""
809
  if self._smart_crusher is None:
810
  try:
 
811
  from .smart_crusher import SmartCrusher
812
 
813
- self._smart_crusher = SmartCrusher()
 
 
 
 
 
814
  except ImportError:
815
  logger.debug("SmartCrusher not available")
816
  return self._smart_crusher
@@ -889,9 +907,36 @@ class ContentRouter(Transform):
889
  transforms_applied = []
890
  warnings: list[str] = []
891
 
 
 
 
 
 
 
 
892
  for i, message in enumerate(messages):
 
893
  content = message.get("content", "")
894
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
895
  if not content or len(content.split()) < 50:
896
  # Skip small content
897
  transformed_messages.append(message)
@@ -900,6 +945,27 @@ class ContentRouter(Transform):
900
  # Get source hint if available
901
  source_hint = source_hints.get(i) or source_hints.get(str(i))
902
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
903
  # Route and compress
904
  result = self.compress(content, source_hint=source_hint, context=context)
905
 
@@ -923,6 +989,139 @@ class ContentRouter(Transform):
923
  warnings=warnings,
924
  )
925
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
926
  def should_apply(
927
  self,
928
  messages: list[dict[str, Any]],
 
172
  mixed_content_threshold: Min distinct types to consider "mixed".
173
  min_section_tokens: Minimum tokens for a section to compress.
174
  fallback_strategy: Strategy when no compressor matches.
175
+ skip_user_messages: Never compress user messages (they're the subject).
176
+ skip_recent_messages: Don't compress last N messages (likely the subject).
177
+ protect_analysis_context: Detect "analyze/review" intent, skip compression.
178
  """
179
 
180
  # Enable/disable specific compressors
 
192
  # Fallback
193
  fallback_strategy: CompressionStrategy = CompressionStrategy.PASSTHROUGH
194
 
195
+ # Protection: Don't compress content that's likely the subject of analysis
196
+ skip_user_messages: bool = True # User messages contain what they want analyzed
197
+ protect_recent_code: int = 4 # Don't compress CODE in last N messages (0 = disabled)
198
+ protect_analysis_context: bool = True # Detect "analyze/review" intent, protect code
199
+
200
+ # CCR (Compress-Cache-Retrieve) settings for SmartCrusher
201
+ ccr_enabled: bool = True # Enable CCR marker injection for reversible compression
202
+ ccr_inject_marker: bool = True # Add retrieval markers to compressed content
203
+
204
 
205
  # Patterns for detecting mixed content
206
  _CODE_FENCE_PATTERN = re.compile(r"^```(\w*)\s*$", re.MULTILINE)
 
817
  return self._code_compressor
818
 
819
  def _get_smart_crusher(self) -> Any:
820
+ """Get SmartCrusher (lazy load) with CCR config."""
821
  if self._smart_crusher is None:
822
  try:
823
+ from ..config import CCRConfig
824
  from .smart_crusher import SmartCrusher
825
 
826
+ # Pass CCR config for marker injection
827
+ ccr_config = CCRConfig(
828
+ enabled=self.config.ccr_enabled,
829
+ inject_retrieval_marker=self.config.ccr_inject_marker,
830
+ )
831
+ self._smart_crusher = SmartCrusher(ccr_config=ccr_config)
832
  except ImportError:
833
  logger.debug("SmartCrusher not available")
834
  return self._smart_crusher
 
907
  transforms_applied = []
908
  warnings: list[str] = []
909
 
910
+ # Check for analysis intent in the most recent user message
911
+ analysis_intent = False
912
+ if self.config.protect_analysis_context:
913
+ analysis_intent = self._detect_analysis_intent(messages)
914
+
915
+ num_messages = len(messages)
916
+
917
  for i, message in enumerate(messages):
918
+ role = message.get("role", "")
919
  content = message.get("content", "")
920
 
921
+ # Handle list content (Anthropic format with content blocks)
922
+ if isinstance(content, list):
923
+ transformed_message = self._process_content_blocks(
924
+ message, content, context, transforms_applied
925
+ )
926
+ transformed_messages.append(transformed_message)
927
+ continue
928
+
929
+ # Skip non-string content (other types)
930
+ if not isinstance(content, str):
931
+ transformed_messages.append(message)
932
+ continue
933
+
934
+ # Protection 1: Never compress user messages
935
+ if self.config.skip_user_messages and role == "user":
936
+ transformed_messages.append(message)
937
+ transforms_applied.append("router:protected:user_message")
938
+ continue
939
+
940
  if not content or len(content.split()) < 50:
941
  # Skip small content
942
  transformed_messages.append(message)
 
945
  # Get source hint if available
946
  source_hint = source_hints.get(i) or source_hints.get(str(i))
947
 
948
+ # Detect content type for protection decisions
949
+ detection = detect_content_type(content)
950
+ is_code = detection.content_type == ContentType.SOURCE_CODE
951
+
952
+ # Protection 2: Don't compress recent CODE
953
+ messages_from_end = num_messages - i
954
+ if (
955
+ self.config.protect_recent_code > 0
956
+ and messages_from_end <= self.config.protect_recent_code
957
+ and is_code
958
+ ):
959
+ transformed_messages.append(message)
960
+ transforms_applied.append("router:protected:recent_code")
961
+ continue
962
+
963
+ # Protection 3: Don't compress CODE when analysis intent detected
964
+ if analysis_intent and is_code:
965
+ transformed_messages.append(message)
966
+ transforms_applied.append("router:protected:analysis_context")
967
+ continue
968
+
969
  # Route and compress
970
  result = self.compress(content, source_hint=source_hint, context=context)
971
 
 
989
  warnings=warnings,
990
  )
991
 
992
+ def _process_content_blocks(
993
+ self,
994
+ message: dict[str, Any],
995
+ content_blocks: list[Any],
996
+ context: str,
997
+ transforms_applied: list[str],
998
+ ) -> dict[str, Any]:
999
+ """Process content blocks (Anthropic format) for tool_result compression.
1000
+
1001
+ Handles tool_result blocks by compressing their string content using
1002
+ the appropriate strategy (typically SmartCrusher for JSON arrays).
1003
+
1004
+ Args:
1005
+ message: The original message.
1006
+ content_blocks: List of content blocks.
1007
+ context: Context for compression.
1008
+ transforms_applied: List to append transform names to.
1009
+
1010
+ Returns:
1011
+ Transformed message with compressed content blocks.
1012
+ """
1013
+ import json
1014
+
1015
+ new_blocks = []
1016
+ any_compressed = False
1017
+
1018
+ for block in content_blocks:
1019
+ if not isinstance(block, dict):
1020
+ new_blocks.append(block)
1021
+ continue
1022
+
1023
+ block_type = block.get("type")
1024
+
1025
+ # Handle tool_result blocks
1026
+ if block_type == "tool_result":
1027
+ tool_content = block.get("content", "")
1028
+
1029
+ # Only process string content
1030
+ if isinstance(tool_content, str) and len(tool_content) > 500:
1031
+ # Try to detect if it's JSON array data (SmartCrusher target)
1032
+ try:
1033
+ parsed = json.loads(tool_content)
1034
+ if isinstance(parsed, list) and len(parsed) > 10:
1035
+ # Route to SmartCrusher for arrays
1036
+ result = self.compress(
1037
+ tool_content,
1038
+ source_hint="json_array",
1039
+ context=context,
1040
+ )
1041
+ if result.compression_ratio < 0.9:
1042
+ new_blocks.append(
1043
+ {
1044
+ **block,
1045
+ "content": result.compressed,
1046
+ }
1047
+ )
1048
+ transforms_applied.append(
1049
+ f"router:tool_result:{result.strategy_used.value}"
1050
+ )
1051
+ any_compressed = True
1052
+ continue
1053
+ except (json.JSONDecodeError, TypeError):
1054
+ # Not JSON, try general compression
1055
+ pass
1056
+
1057
+ # Try general compression for large non-JSON content
1058
+ result = self.compress(tool_content, context=context)
1059
+ if result.compression_ratio < 0.9:
1060
+ new_blocks.append({**block, "content": result.compressed})
1061
+ transforms_applied.append(
1062
+ f"router:tool_result:{result.strategy_used.value}"
1063
+ )
1064
+ any_compressed = True
1065
+ continue
1066
+
1067
+ # Keep block unchanged
1068
+ new_blocks.append(block)
1069
+
1070
+ if any_compressed:
1071
+ return {**message, "content": new_blocks}
1072
+ return message
1073
+
1074
+ def _detect_analysis_intent(self, messages: list[dict[str, Any]]) -> bool:
1075
+ """Detect if user wants to analyze/review code.
1076
+
1077
+ Looks at the most recent user message for analysis keywords.
1078
+
1079
+ Args:
1080
+ messages: Conversation messages.
1081
+
1082
+ Returns:
1083
+ True if analysis intent detected.
1084
+ """
1085
+ # Analysis keywords that suggest user wants full code details
1086
+ analysis_keywords = {
1087
+ "analyze",
1088
+ "analyse",
1089
+ "review",
1090
+ "audit",
1091
+ "inspect",
1092
+ "security",
1093
+ "vulnerability",
1094
+ "bug",
1095
+ "issue",
1096
+ "problem",
1097
+ "explain",
1098
+ "understand",
1099
+ "how does",
1100
+ "what does",
1101
+ "debug",
1102
+ "fix",
1103
+ "error",
1104
+ "wrong",
1105
+ "broken",
1106
+ "refactor",
1107
+ "improve",
1108
+ "optimize",
1109
+ "clean up",
1110
+ }
1111
+
1112
+ # Find most recent user message
1113
+ for message in reversed(messages):
1114
+ if message.get("role") == "user":
1115
+ content = message.get("content", "")
1116
+ if isinstance(content, str):
1117
+ content_lower = content.lower()
1118
+ for keyword in analysis_keywords:
1119
+ if keyword in content_lower:
1120
+ return True
1121
+ break
1122
+
1123
+ return False
1124
+
1125
  def should_apply(
1126
  self,
1127
  messages: list[dict[str, Any]],
headroom/transforms/llmlingua_compressor.py CHANGED
@@ -386,9 +386,8 @@ class LLMLinguaCompressor(Transform):
386
  if self.config.enable_ccr and ratio < 0.8:
387
  cache_key = self._store_in_ccr(content, compressed, original_tokens)
388
  if cache_key:
389
- compressed += (
390
- f"\n[LLMLingua: {original_tokens}→{compressed_tokens} tokens. hash={cache_key}]"
391
- )
392
 
393
  return LLMLinguaResult(
394
  compressed=compressed,
 
386
  if self.config.enable_ccr and ratio < 0.8:
387
  cache_key = self._store_in_ccr(content, compressed, original_tokens)
388
  if cache_key:
389
+ # Use standard CCR marker format for CCRToolInjector detection
390
+ compressed += f"\n[{original_tokens} items compressed to {compressed_tokens}. Retrieve more: hash={cache_key}]"
 
391
 
392
  return LLMLinguaResult(
393
  compressed=compressed,
headroom/transforms/smart_crusher.py CHANGED
@@ -697,6 +697,19 @@ class CompressionPlan:
697
  keep_count: int = 10
698
 
699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
700
  @dataclass
701
  class SmartCrusherConfig:
702
  """Configuration for smart crusher.
@@ -1325,6 +1338,29 @@ class SmartCrusher(Transform):
1325
  # NOTE: Error detection now uses structural outlier detection (_detect_structural_outliers)
1326
  # instead of hardcoded keywords. This scales to any data domain.
1327
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1328
  def _get_compression_store(self) -> CompressionStore:
1329
  """Get the compression store for CCR (lazy initialization).
1330
 
 
697
  keep_count: int = 10
698
 
699
 
700
+ @dataclass
701
+ class CrushResult:
702
+ """Result from SmartCrusher.crush() method.
703
+
704
+ Used by ContentRouter when routing JSON arrays to SmartCrusher.
705
+ """
706
+
707
+ compressed: str
708
+ original: str
709
+ was_modified: bool
710
+ strategy: str = "passthrough"
711
+
712
+
713
  @dataclass
714
  class SmartCrusherConfig:
715
  """Configuration for smart crusher.
 
1338
  # NOTE: Error detection now uses structural outlier detection (_detect_structural_outliers)
1339
  # instead of hardcoded keywords. This scales to any data domain.
1340
 
1341
+ def crush(self, content: str, query: str = "") -> CrushResult:
1342
+ """Crush content string directly (for use by ContentRouter).
1343
+
1344
+ This is a simplified interface for compressing a single content string,
1345
+ used by ContentRouter when routing JSON arrays to SmartCrusher.
1346
+
1347
+ Args:
1348
+ content: JSON string content to compress.
1349
+ query: Query context for relevance-based compression.
1350
+
1351
+ Returns:
1352
+ CrushResult with compressed content and metadata.
1353
+ """
1354
+ compressed, was_modified, analysis_info = self._smart_crush_content(
1355
+ content, query_context=query
1356
+ )
1357
+ return CrushResult(
1358
+ compressed=compressed,
1359
+ original=content,
1360
+ was_modified=was_modified,
1361
+ strategy=analysis_info or "passthrough",
1362
+ )
1363
+
1364
  def _get_compression_store(self) -> CompressionStore:
1365
  """Get the compression store for CCR (lazy initialization).
1366