« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Token limit compression strategies."""
3from __future__ import annotations
5from datetime import UTC, datetime
7from lexigram.ai.rag.context_compression.base import AbstractCompressor
8from lexigram.ai.rag.context_compression.types import (
9 CompressionResult,
10 CompressionStrategy,
11)
14class TokenLimitCompressor(AbstractCompressor):
15 """Simple truncation to token limit.
17 Truncates context to fit within token limit, preserving the beginning
18 and optionally the end of the text.
20 Example:
21 >>> compressor = TokenLimitCompressor(
22 ... max_tokens=500,
23 ... preserve_end=True,
24 ... end_ratio=0.2
25 ... )
26 >>> result = await compressor.compress(very_long_context)
27 """
29 def __init__(
30 self,
31 max_tokens: int = 500,
32 preserve_end: bool = False,
33 end_ratio: float = 0.2,
34 ):
35 """Initialize token limit compressor.
37 Args:
38 max_tokens: Maximum tokens to keep.
39 preserve_end: Whether to preserve end of text.
40 end_ratio: Ratio of tokens to preserve from end (if preserve_end=True).
41 """
42 self.max_tokens = max_tokens
43 self.preserve_end = preserve_end
44 self.end_ratio = end_ratio
46 async def compress(
47 self,
48 context: str | list[str],
49 query: str | None = None,
50 **kwargs,
51 ) -> CompressionResult:
52 """Compress by truncating to token limit."""
53 original_text = self._normalize_context(context)
54 original_tokens = self._estimate_tokens(original_text)
56 if original_tokens <= self.max_tokens:
57 # Already within limit
58 compressed_text = original_text
59 elif self.preserve_end:
60 # Keep beginning and end
61 start_tokens = int(self.max_tokens * (1 - self.end_ratio))
62 end_tokens = int(self.max_tokens * self.end_ratio)
64 # Approximate character counts
65 start_chars = start_tokens * 4
66 end_chars = end_tokens * 4
68 start_text = original_text[:start_chars]
69 end_text = original_text[-end_chars:]
71 compressed_text = start_text + "\n[...]\n" + end_text
72 else:
73 # Keep only beginning
74 max_chars = self.max_tokens * 4
75 compressed_text = original_text[:max_chars]
77 compressed_tokens = self._estimate_tokens(compressed_text)
78 compression_ratio = (
79 compressed_tokens / original_tokens if original_tokens > 0 else 1.0
80 )
82 return CompressionResult(
83 original_text=original_text,
84 compressed_text=compressed_text,
85 original_tokens=original_tokens,
86 compressed_tokens=compressed_tokens,
87 compression_ratio=compression_ratio,
88 strategy=CompressionStrategy.TOKEN_LIMIT,
89 metadata={
90 "max_tokens": self.max_tokens,
91 "preserve_end": self.preserve_end,
92 "end_ratio": self.end_ratio if self.preserve_end else 0.0,
93 "timestamp": datetime.now(UTC).isoformat(),
94 },
95 )