contextpress 0.2.0__tar.gz → 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {contextpress-0.2.0 → contextpress-0.3.0}/.gitignore +3 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/CHANGELOG.md +7 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/PKG-INFO +25 -2
- {contextpress-0.2.0 → contextpress-0.3.0}/README.md +24 -1
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/__init__.py +1 -1
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/compression.py +19 -3
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/core.py +39 -4
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/pipeline.py +78 -21
- contextpress-0.3.0/contextpress/registry.py +85 -0
- contextpress-0.3.0/examples/estimate_and_stats.py +21 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/pyproject.toml +1 -1
- contextpress-0.3.0/tests/test_v03.py +128 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/CITATION.cff +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/CONTRIBUTING.md +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/LICENSE +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/NOTICE +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/_bootstrap.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/__init__.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/_helpers.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/adapters.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/base.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/models.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/normalizer.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/profiles.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/py.typed +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/stats.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/__init__.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/base.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/budget.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/filler.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/recency.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/repetition.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/resolution.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/examples/llm_tier_ollama.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/examples/llm_tier_openai.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/__init__.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_budget.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_filler.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_llm_helpers.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_models.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_normalizer.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_pipeline.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_recency.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_repetition.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_resolution.py +0 -0
- {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_stats.py +0 -0
|
@@ -4,6 +4,13 @@ All notable changes to `contextpress` are recorded here.
|
|
|
4
4
|
The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
|
|
5
5
|
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
|
6
6
|
|
|
7
|
+
## [0.3.0] - 2026-07-14
|
|
8
|
+
|
|
9
|
+
- **`estimate_tokens(messages)`** — count tokens before compressing (same tiktoken encoding as budget).
|
|
10
|
+
- **Stage registry** — `ContextManager.register_stage()` for custom `BaseStrategy` stages in `stages=[...]`.
|
|
11
|
+
- **Tier 2 `llm_mode`** — `replace_all` (default), `dedupe_only`, or `summarize_only`.
|
|
12
|
+
- Example: `examples/estimate_and_stats.py`.
|
|
13
|
+
|
|
7
14
|
## [0.2.0] - 2026-05-24
|
|
8
15
|
|
|
9
16
|
- **`return_stats=True`** on `ContextManager.compress()` returns a `CompressionResult` with token/turn counts, stages run, and per-stage turn deltas.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: contextpress
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: Deterministic context compression for LLM chat, RAG, and agent pipelines
|
|
5
5
|
Project-URL: Homepage, https://github.com/Taha-azizi/contextpress
|
|
6
6
|
Project-URL: Documentation, https://github.com/Taha-azizi/contextpress#readme
|
|
@@ -290,6 +290,27 @@ print(result.stats.tokens_saved, result.stats.stages_run)
|
|
|
290
290
|
compressed = result.messages
|
|
291
291
|
```
|
|
292
292
|
|
|
293
|
+
Check token count **before** compressing:
|
|
294
|
+
|
|
295
|
+
```python
|
|
296
|
+
before = cm.estimate_tokens(messages)
|
|
297
|
+
```
|
|
298
|
+
|
|
299
|
+
### Custom stages (0.3+)
|
|
300
|
+
|
|
301
|
+
Register a `BaseStrategy` and include it in `stages=`:
|
|
302
|
+
|
|
303
|
+
```python
|
|
304
|
+
from contextpress.strategies.base import BaseStrategy
|
|
305
|
+
|
|
306
|
+
class MyStage(BaseStrategy):
|
|
307
|
+
def process(self, conversation):
|
|
308
|
+
...
|
|
309
|
+
|
|
310
|
+
cm.register_stage("my_stage", MyStage)
|
|
311
|
+
out = cm.compress(messages, stages=["filler", "my_stage", "budget"], token_budget=500)
|
|
312
|
+
```
|
|
313
|
+
|
|
293
314
|
### Minimal examples
|
|
294
315
|
|
|
295
316
|
```python
|
|
@@ -405,7 +426,9 @@ After **Tier 1** finishes, you can attach an **`LLMBackend`** for semantic compr
|
|
|
405
426
|
2. If the combined transcript is long enough (default **1500** characters; set **`llm_min_input_chars=0`** to always run), calls **`summarize(transcript, max_tokens)`**.
|
|
406
427
|
3. **System turns are unchanged** in order and content. **All other turns are replaced** by a **single assistant** message whose content is the LLM summary (metadata includes `source: contextpress_llm_tier`). If the LLM call fails, the Tier 1 conversation is returned and a **warning** is emitted.
|
|
407
428
|
|
|
408
|
-
Optional constructor knobs: **`llm_min_input_chars`**, **`llm_max_summary_tokens`**.
|
|
429
|
+
Optional constructor knobs: **`llm_min_input_chars`**, **`llm_max_summary_tokens`**, **`llm_mode`**.
|
|
430
|
+
|
|
431
|
+
**`llm_mode`** (0.3+): `replace_all` (default — dedupe then one summary turn), `dedupe_only` (dedupe, keep turns), `summarize_only` (append summary, keep turns).
|
|
409
432
|
|
|
410
433
|
**Install SDKs** (not bundled): `pip install openai`, `anthropic`, and/or **`ollama`** (for local Ollama), or `pip install "contextpress[llm]"` from this repo’s `pyproject.toml` to pull all optional LLM clients.
|
|
411
434
|
|
|
@@ -50,6 +50,27 @@ print(result.stats.tokens_saved, result.stats.stages_run)
|
|
|
50
50
|
compressed = result.messages
|
|
51
51
|
```
|
|
52
52
|
|
|
53
|
+
Check token count **before** compressing:
|
|
54
|
+
|
|
55
|
+
```python
|
|
56
|
+
before = cm.estimate_tokens(messages)
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
### Custom stages (0.3+)
|
|
60
|
+
|
|
61
|
+
Register a `BaseStrategy` and include it in `stages=`:
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
from contextpress.strategies.base import BaseStrategy
|
|
65
|
+
|
|
66
|
+
class MyStage(BaseStrategy):
|
|
67
|
+
def process(self, conversation):
|
|
68
|
+
...
|
|
69
|
+
|
|
70
|
+
cm.register_stage("my_stage", MyStage)
|
|
71
|
+
out = cm.compress(messages, stages=["filler", "my_stage", "budget"], token_budget=500)
|
|
72
|
+
```
|
|
73
|
+
|
|
53
74
|
### Minimal examples
|
|
54
75
|
|
|
55
76
|
```python
|
|
@@ -165,7 +186,9 @@ After **Tier 1** finishes, you can attach an **`LLMBackend`** for semantic compr
|
|
|
165
186
|
2. If the combined transcript is long enough (default **1500** characters; set **`llm_min_input_chars=0`** to always run), calls **`summarize(transcript, max_tokens)`**.
|
|
166
187
|
3. **System turns are unchanged** in order and content. **All other turns are replaced** by a **single assistant** message whose content is the LLM summary (metadata includes `source: contextpress_llm_tier`). If the LLM call fails, the Tier 1 conversation is returned and a **warning** is emitted.
|
|
167
188
|
|
|
168
|
-
Optional constructor knobs: **`llm_min_input_chars`**, **`llm_max_summary_tokens`**.
|
|
189
|
+
Optional constructor knobs: **`llm_min_input_chars`**, **`llm_max_summary_tokens`**, **`llm_mode`**.
|
|
190
|
+
|
|
191
|
+
**`llm_mode`** (0.3+): `replace_all` (default — dedupe then one summary turn), `dedupe_only` (dedupe, keep turns), `summarize_only` (append summary, keep turns).
|
|
169
192
|
|
|
170
193
|
**Install SDKs** (not bundled): `pip install openai`, `anthropic`, and/or **`ollama`** (for local Ollama), or `pip install "contextpress[llm]"` from this repo’s `pyproject.toml` to pull all optional LLM clients.
|
|
171
194
|
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
|
|
3
3
|
from __future__ import annotations
|
|
4
4
|
|
|
5
|
-
from contextpress.profiles import Profile
|
|
5
|
+
from contextpress.profiles import Profile, StageConfig
|
|
6
6
|
|
|
7
7
|
STAGE_ORDER: tuple[str, ...] = (
|
|
8
8
|
"filler",
|
|
@@ -14,6 +14,14 @@ STAGE_ORDER: tuple[str, ...] = (
|
|
|
14
14
|
|
|
15
15
|
VALID_STAGES = frozenset(STAGE_ORDER)
|
|
16
16
|
|
|
17
|
+
|
|
18
|
+
def known_stages() -> frozenset[str]:
|
|
19
|
+
"""Built-in plus registered custom stage names."""
|
|
20
|
+
from contextpress.registry import registered_stage_names
|
|
21
|
+
|
|
22
|
+
return VALID_STAGES | registered_stage_names()
|
|
23
|
+
|
|
24
|
+
|
|
17
25
|
_NON_BUDGET_ORDER: tuple[str, ...] = tuple(s for s in STAGE_ORDER if s != "budget")
|
|
18
26
|
|
|
19
27
|
# NLP stages only; budget is toggled from token_budget (see apply_stage_selection)
|
|
@@ -38,6 +46,7 @@ __all__ = [
|
|
|
38
46
|
"STAGE_ORDER",
|
|
39
47
|
"VALID_STAGES",
|
|
40
48
|
"apply_stage_selection",
|
|
49
|
+
"known_stages",
|
|
41
50
|
"normalize_compression_level",
|
|
42
51
|
]
|
|
43
52
|
|
|
@@ -60,6 +69,7 @@ def apply_stage_selection(
|
|
|
60
69
|
stages: list[str] | None,
|
|
61
70
|
disable: list[str] | None,
|
|
62
71
|
token_budget: int | None,
|
|
72
|
+
custom_stages: dict[str, StageConfig] | None = None,
|
|
63
73
|
) -> None:
|
|
64
74
|
"""
|
|
65
75
|
Mutates ``profile`` in place: sets each stage's ``enabled`` from explicit
|
|
@@ -70,12 +80,15 @@ def apply_stage_selection(
|
|
|
70
80
|
if stages is not None:
|
|
71
81
|
if not stages:
|
|
72
82
|
raise ValueError("stages= must list at least one stage name when provided")
|
|
73
|
-
unknown = [s for s in stages if s not in
|
|
83
|
+
unknown = [s for s in stages if s not in known_stages()]
|
|
74
84
|
if unknown:
|
|
75
|
-
raise ValueError(f"unknown stage name(s): {unknown}; valid: {sorted(
|
|
85
|
+
raise ValueError(f"unknown stage name(s): {unknown}; valid: {sorted(known_stages())}")
|
|
76
86
|
want = frozenset(stages)
|
|
77
87
|
for name in _NON_BUDGET_ORDER:
|
|
78
88
|
getattr(profile, name).enabled = name in want
|
|
89
|
+
if custom_stages:
|
|
90
|
+
for name, cfg in custom_stages.items():
|
|
91
|
+
cfg.enabled = name in want
|
|
79
92
|
else:
|
|
80
93
|
level = normalize_compression_level(compression)
|
|
81
94
|
preset = _COMPRESSION_PRESETS[level]
|
|
@@ -85,6 +98,9 @@ def apply_stage_selection(
|
|
|
85
98
|
|
|
86
99
|
if disable:
|
|
87
100
|
for name in disable:
|
|
101
|
+
if custom_stages and name in custom_stages:
|
|
102
|
+
custom_stages[name].enabled = False
|
|
103
|
+
continue
|
|
88
104
|
if not hasattr(profile, name):
|
|
89
105
|
continue
|
|
90
106
|
getattr(profile, name).enabled = False
|
|
@@ -6,12 +6,16 @@ from typing import TYPE_CHECKING, Any
|
|
|
6
6
|
|
|
7
7
|
from contextpress.compression import apply_stage_selection, normalize_compression_level
|
|
8
8
|
from contextpress.normalizer import denormalize_output, normalize_messages
|
|
9
|
-
from contextpress.pipeline import Pipeline
|
|
9
|
+
from contextpress.pipeline import VALID_LLM_MODES, Pipeline
|
|
10
10
|
from contextpress.profiles import PROFILES, Profile, StageConfig
|
|
11
|
-
from contextpress.
|
|
11
|
+
from contextpress.registry import register_stage as _register_stage
|
|
12
|
+
from contextpress.stats import CompressionResult, CompressionStats, count_conversation_tokens
|
|
12
13
|
|
|
13
14
|
if TYPE_CHECKING:
|
|
15
|
+
from collections.abc import Callable
|
|
16
|
+
|
|
14
17
|
from contextpress.llm.base import LLMBackend
|
|
18
|
+
from contextpress.strategies.base import BaseStrategy
|
|
15
19
|
|
|
16
20
|
|
|
17
21
|
def _validate_token_budget(token_budget: int | None) -> None:
|
|
@@ -39,9 +43,14 @@ class ContextManager:
|
|
|
39
43
|
compression: str = "medium",
|
|
40
44
|
llm_min_input_chars: int = 1500,
|
|
41
45
|
llm_max_summary_tokens: int = 2048,
|
|
46
|
+
llm_mode: str = "replace_all",
|
|
42
47
|
):
|
|
43
48
|
if type not in PROFILES:
|
|
44
49
|
raise ValueError(f"unknown context type {type!r}")
|
|
50
|
+
if llm_mode not in VALID_LLM_MODES:
|
|
51
|
+
raise ValueError(
|
|
52
|
+
f"unknown llm_mode {llm_mode!r}; use one of: {sorted(VALID_LLM_MODES)}"
|
|
53
|
+
)
|
|
45
54
|
self._type = type
|
|
46
55
|
self._profile: Profile = copy.deepcopy(PROFILES[type])
|
|
47
56
|
self._compression: str = normalize_compression_level(compression)
|
|
@@ -49,6 +58,25 @@ class ContextManager:
|
|
|
49
58
|
self.llm_backend = llm_backend
|
|
50
59
|
self.llm_min_input_chars = int(llm_min_input_chars)
|
|
51
60
|
self.llm_max_summary_tokens = int(llm_max_summary_tokens)
|
|
61
|
+
self.llm_mode = llm_mode
|
|
62
|
+
self._custom_stages: dict[str, StageConfig] = {}
|
|
63
|
+
|
|
64
|
+
def estimate_tokens(self, messages: Any, *, model: str | None = None) -> int:
|
|
65
|
+
"""Count tokens for ``messages`` using the same encoding as the budget stage."""
|
|
66
|
+
conv, _ = normalize_messages(messages, context_type=self._type)
|
|
67
|
+
return count_conversation_tokens(conv, model if model is not None else self.model)
|
|
68
|
+
|
|
69
|
+
def register_stage(
|
|
70
|
+
self,
|
|
71
|
+
name: str,
|
|
72
|
+
factory: Callable[..., BaseStrategy],
|
|
73
|
+
*,
|
|
74
|
+
before: str = "budget",
|
|
75
|
+
aggressiveness: float = 0.5,
|
|
76
|
+
) -> None:
|
|
77
|
+
"""Register a custom pipeline stage (see ``contextpress.registry``)."""
|
|
78
|
+
_register_stage(name, factory, before=before)
|
|
79
|
+
self._custom_stages[name] = StageConfig(enabled=False, aggressiveness=aggressiveness)
|
|
52
80
|
|
|
53
81
|
def compress(
|
|
54
82
|
self,
|
|
@@ -68,6 +96,7 @@ class ContextManager:
|
|
|
68
96
|
_validate_token_budget(token_budget)
|
|
69
97
|
level = compression if compression is not None else self._compression
|
|
70
98
|
profile = copy.deepcopy(self._profile)
|
|
99
|
+
custom_stages = copy.deepcopy(self._custom_stages)
|
|
71
100
|
apply_stage_selection(
|
|
72
101
|
profile,
|
|
73
102
|
base_profile=self._profile,
|
|
@@ -75,6 +104,7 @@ class ContextManager:
|
|
|
75
104
|
stages=stages,
|
|
76
105
|
disable=disable,
|
|
77
106
|
token_budget=token_budget,
|
|
107
|
+
custom_stages=custom_stages,
|
|
78
108
|
)
|
|
79
109
|
|
|
80
110
|
conv, ctx = normalize_messages(messages, context_type=self._type)
|
|
@@ -86,6 +116,8 @@ class ContextManager:
|
|
|
86
116
|
llm_backend=self.llm_backend,
|
|
87
117
|
llm_min_input_chars=self.llm_min_input_chars,
|
|
88
118
|
llm_max_summary_tokens=self.llm_max_summary_tokens,
|
|
119
|
+
llm_mode=self.llm_mode,
|
|
120
|
+
custom_stages=custom_stages,
|
|
89
121
|
)
|
|
90
122
|
out = pipeline.run(conv, stats=stats)
|
|
91
123
|
messages_out = denormalize_output(out, ctx)
|
|
@@ -100,9 +132,12 @@ class ContextManager:
|
|
|
100
132
|
|
|
101
133
|
def configure(self, stage: str, **kwargs: Any) -> None:
|
|
102
134
|
"""Patch ``StageConfig`` fields on the live profile (e.g. aggressiveness, enabled)."""
|
|
103
|
-
if
|
|
135
|
+
if stage in self._custom_stages:
|
|
136
|
+
sc = self._custom_stages[stage]
|
|
137
|
+
elif hasattr(self._profile, stage):
|
|
138
|
+
sc = getattr(self._profile, stage)
|
|
139
|
+
else:
|
|
104
140
|
raise ValueError(f"unknown stage {stage!r}")
|
|
105
|
-
sc: StageConfig = getattr(self._profile, stage)
|
|
106
141
|
for k, v in kwargs.items():
|
|
107
142
|
if hasattr(sc, k):
|
|
108
143
|
setattr(sc, k, v)
|
|
@@ -27,6 +27,11 @@ from contextpress.compression import STAGE_ORDER
|
|
|
27
27
|
from contextpress.models import Conversation, Turn
|
|
28
28
|
from contextpress.normalizer import extract_text_for_processing
|
|
29
29
|
from contextpress.profiles import Profile, StageConfig
|
|
30
|
+
from contextpress.registry import (
|
|
31
|
+
build_custom_strategy,
|
|
32
|
+
effective_stage_order,
|
|
33
|
+
registered_stage_names,
|
|
34
|
+
)
|
|
30
35
|
from contextpress.stats import CompressionStats, count_conversation_tokens
|
|
31
36
|
from contextpress.strategies.base import BaseStrategy
|
|
32
37
|
from contextpress.strategies.budget import BudgetStrategy
|
|
@@ -38,6 +43,8 @@ from contextpress.strategies.resolution import ResolutionStrategy
|
|
|
38
43
|
if TYPE_CHECKING:
|
|
39
44
|
from contextpress.llm.base import LLMBackend
|
|
40
45
|
|
|
46
|
+
VALID_LLM_MODES = frozenset({"replace_all", "dedupe_only", "summarize_only"})
|
|
47
|
+
|
|
41
48
|
|
|
42
49
|
def clone_turn(t: Turn) -> Turn:
|
|
43
50
|
if isinstance(t.content, str):
|
|
@@ -80,6 +87,8 @@ class Pipeline:
|
|
|
80
87
|
*,
|
|
81
88
|
llm_min_input_chars: int = 1500,
|
|
82
89
|
llm_max_summary_tokens: int = 2048,
|
|
90
|
+
llm_mode: str = "replace_all",
|
|
91
|
+
custom_stages: dict[str, StageConfig] | None = None,
|
|
83
92
|
):
|
|
84
93
|
self.profile = profile
|
|
85
94
|
self.token_budget = token_budget
|
|
@@ -87,6 +96,12 @@ class Pipeline:
|
|
|
87
96
|
self.llm_backend = llm_backend # None = Tier 1 only
|
|
88
97
|
self.llm_min_input_chars = max(0, int(llm_min_input_chars))
|
|
89
98
|
self.llm_max_summary_tokens = max(64, int(llm_max_summary_tokens))
|
|
99
|
+
if llm_mode not in VALID_LLM_MODES:
|
|
100
|
+
raise ValueError(
|
|
101
|
+
f"unknown llm_mode {llm_mode!r}; use one of: {sorted(VALID_LLM_MODES)}"
|
|
102
|
+
)
|
|
103
|
+
self.llm_mode = llm_mode
|
|
104
|
+
self.custom_stages = custom_stages or {}
|
|
90
105
|
|
|
91
106
|
def run(
|
|
92
107
|
self,
|
|
@@ -100,11 +115,12 @@ class Pipeline:
|
|
|
100
115
|
stats.token_budget = self.token_budget
|
|
101
116
|
|
|
102
117
|
result = clone_conversation(conversation)
|
|
103
|
-
|
|
118
|
+
stage_order = effective_stage_order()
|
|
119
|
+
for stage_name in stage_order:
|
|
104
120
|
if stage_name == "budget" and self.token_budget is None:
|
|
105
121
|
continue
|
|
106
|
-
stage_config =
|
|
107
|
-
if not stage_config.enabled:
|
|
122
|
+
stage_config = self._stage_config(stage_name)
|
|
123
|
+
if stage_config is None or not stage_config.enabled:
|
|
108
124
|
continue
|
|
109
125
|
before_turns = len(result.turns)
|
|
110
126
|
strategy = self._build_strategy(stage_name, stage_config)
|
|
@@ -124,12 +140,26 @@ class Pipeline:
|
|
|
124
140
|
|
|
125
141
|
return result
|
|
126
142
|
|
|
143
|
+
def _stage_config(self, name: str) -> StageConfig | None:
|
|
144
|
+
if name in self.custom_stages:
|
|
145
|
+
return self.custom_stages[name]
|
|
146
|
+
if hasattr(self.profile, name):
|
|
147
|
+
return getattr(self.profile, name)
|
|
148
|
+
return None
|
|
149
|
+
|
|
127
150
|
def _build_strategy(self, name: str, config: StageConfig) -> BaseStrategy:
|
|
128
151
|
kwargs: dict[str, Any] = {
|
|
129
152
|
"aggressiveness": config.aggressiveness,
|
|
130
153
|
"conv_type": self.profile.name,
|
|
131
154
|
"role_aware": self.profile.role_aware,
|
|
132
155
|
}
|
|
156
|
+
if name in registered_stage_names():
|
|
157
|
+
return build_custom_strategy(
|
|
158
|
+
name,
|
|
159
|
+
config,
|
|
160
|
+
conv_type=self.profile.name,
|
|
161
|
+
role_aware=self.profile.role_aware,
|
|
162
|
+
)
|
|
133
163
|
if name == "filler":
|
|
134
164
|
return FillerStrategy(**kwargs)
|
|
135
165
|
if name == "repetition":
|
|
@@ -166,25 +196,36 @@ class Pipeline:
|
|
|
166
196
|
stats.llm_dedup_turns_before = len(ns_turns)
|
|
167
197
|
|
|
168
198
|
texts = [extract_text_for_processing(t) for t in ns_turns]
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
199
|
+
if self.llm_mode in ("replace_all", "dedupe_only"):
|
|
200
|
+
try:
|
|
201
|
+
keep_idx = self.llm_backend.deduplicate(texts)
|
|
202
|
+
except Exception:
|
|
203
|
+
keep_idx = list(range(len(ns_turns)))
|
|
204
|
+
valid = sorted(
|
|
205
|
+
{
|
|
206
|
+
i
|
|
207
|
+
for i in keep_idx
|
|
208
|
+
if type(i) is int and not isinstance(i, bool) and 0 <= i < len(ns_turns)
|
|
209
|
+
}
|
|
210
|
+
)
|
|
211
|
+
if not valid:
|
|
212
|
+
valid = list(range(len(ns_turns)))
|
|
213
|
+
if len(valid) < len(ns_turns):
|
|
214
|
+
ns_turns = [ns_turns[i] for i in valid]
|
|
215
|
+
texts = [texts[i] for i in valid]
|
|
185
216
|
|
|
186
|
-
|
|
187
|
-
|
|
217
|
+
if stats is not None:
|
|
218
|
+
stats.llm_dedup_turns_after = len(ns_turns)
|
|
219
|
+
|
|
220
|
+
if self.llm_mode == "dedupe_only":
|
|
221
|
+
new_turns = list(system_turns) + [clone_turn(t) for t in ns_turns]
|
|
222
|
+
if stats is not None:
|
|
223
|
+
stats.llm_tier_applied = True
|
|
224
|
+
return Conversation(
|
|
225
|
+
turns=new_turns,
|
|
226
|
+
type=conversation.type,
|
|
227
|
+
metadata=copy.deepcopy(conversation.metadata),
|
|
228
|
+
)
|
|
188
229
|
|
|
189
230
|
lines = [f"{t.role}: {txt}" for t, txt in zip(ns_turns, texts, strict=True)]
|
|
190
231
|
combined = "\n\n".join(lines)
|
|
@@ -207,6 +248,22 @@ class Pipeline:
|
|
|
207
248
|
if stats is not None:
|
|
208
249
|
stats.llm_tier_applied = True
|
|
209
250
|
|
|
251
|
+
if self.llm_mode == "summarize_only":
|
|
252
|
+
new_turns = list(system_turns) + [clone_turn(t) for t in ns_turns]
|
|
253
|
+
new_turns.append(
|
|
254
|
+
Turn(
|
|
255
|
+
role="assistant",
|
|
256
|
+
content=summary,
|
|
257
|
+
metadata={"source": "contextpress_llm_tier", "mode": "summarize_only"},
|
|
258
|
+
compressed=True,
|
|
259
|
+
)
|
|
260
|
+
)
|
|
261
|
+
return Conversation(
|
|
262
|
+
turns=new_turns,
|
|
263
|
+
type=conversation.type,
|
|
264
|
+
metadata=copy.deepcopy(conversation.metadata),
|
|
265
|
+
)
|
|
266
|
+
|
|
210
267
|
new_turns: list[Turn] = list(system_turns)
|
|
211
268
|
new_turns.append(
|
|
212
269
|
Turn(
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
"""Register custom pipeline stages without subclassing ``Pipeline``."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from collections.abc import Callable
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
from contextpress.profiles import StageConfig
|
|
9
|
+
from contextpress.strategies.base import BaseStrategy
|
|
10
|
+
|
|
11
|
+
# (name, factory, insert_before)
|
|
12
|
+
_CustomEntry = tuple[str, Callable[..., BaseStrategy], str]
|
|
13
|
+
_REGISTRY: list[_CustomEntry] = []
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def _builtin_stages() -> tuple[str, ...]:
|
|
17
|
+
from contextpress.compression import STAGE_ORDER
|
|
18
|
+
|
|
19
|
+
return STAGE_ORDER
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def register_stage(
|
|
23
|
+
name: str,
|
|
24
|
+
factory: Callable[..., BaseStrategy],
|
|
25
|
+
*,
|
|
26
|
+
before: str = "budget",
|
|
27
|
+
) -> None:
|
|
28
|
+
"""Register a custom stage factory.
|
|
29
|
+
|
|
30
|
+
``factory`` receives ``aggressiveness``, ``conv_type``, ``role_aware``, and
|
|
31
|
+
optional ``**kwargs`` — same as built-in strategies. Custom stages run when
|
|
32
|
+
named in ``stages=`` on ``compress()``.
|
|
33
|
+
|
|
34
|
+
Parameters
|
|
35
|
+
----------
|
|
36
|
+
name:
|
|
37
|
+
Unique stage name (must not collide with a built-in stage).
|
|
38
|
+
before:
|
|
39
|
+
Built-in or custom stage name to insert before (default: ``budget``).
|
|
40
|
+
"""
|
|
41
|
+
if not name or name in _builtin_stages():
|
|
42
|
+
raise ValueError(f"cannot register stage {name!r}: name invalid or reserved")
|
|
43
|
+
if any(e[0] == name for e in _REGISTRY):
|
|
44
|
+
raise ValueError(f"stage {name!r} is already registered")
|
|
45
|
+
_REGISTRY.append((name, factory, before))
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def registered_stage_names() -> frozenset[str]:
|
|
49
|
+
return frozenset(name for name, _, _ in _REGISTRY)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def clear_registry() -> None:
|
|
53
|
+
"""Remove all custom stages (for tests)."""
|
|
54
|
+
_REGISTRY.clear()
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def effective_stage_order() -> list[str]:
|
|
58
|
+
"""Built-in ``STAGE_ORDER`` with registered custom stages inserted."""
|
|
59
|
+
order = list(_builtin_stages())
|
|
60
|
+
for name, _, before in _REGISTRY:
|
|
61
|
+
if before not in order:
|
|
62
|
+
raise RuntimeError(f"custom stage {name!r}: anchor {before!r} not in stage order")
|
|
63
|
+
order.insert(order.index(before), name)
|
|
64
|
+
return order
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def build_custom_strategy(
|
|
68
|
+
name: str,
|
|
69
|
+
config: StageConfig,
|
|
70
|
+
*,
|
|
71
|
+
conv_type: str,
|
|
72
|
+
role_aware: bool,
|
|
73
|
+
extra_kwargs: dict[str, Any] | None = None,
|
|
74
|
+
) -> BaseStrategy:
|
|
75
|
+
for reg_name, factory, _ in _REGISTRY:
|
|
76
|
+
if reg_name == name:
|
|
77
|
+
kwargs: dict[str, Any] = {
|
|
78
|
+
"aggressiveness": config.aggressiveness,
|
|
79
|
+
"conv_type": conv_type,
|
|
80
|
+
"role_aware": role_aware,
|
|
81
|
+
}
|
|
82
|
+
if extra_kwargs:
|
|
83
|
+
kwargs.update(extra_kwargs)
|
|
84
|
+
return factory(**kwargs)
|
|
85
|
+
raise ValueError(f"no registered factory for custom stage {name!r}")
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
"""Estimate tokens before compressing, then inspect stats after."""
|
|
2
|
+
|
|
3
|
+
from contextpress import ContextManager
|
|
4
|
+
|
|
5
|
+
messages = [
|
|
6
|
+
{"role": "system", "content": "You are helpful."},
|
|
7
|
+
{"role": "user", "content": "Old topic: " + "blah " * 50},
|
|
8
|
+
{"role": "assistant", "content": "Basically, honestly, " + "yak " * 50},
|
|
9
|
+
{"role": "user", "content": "What is 2+2?"},
|
|
10
|
+
{"role": "assistant", "content": "Four."},
|
|
11
|
+
]
|
|
12
|
+
|
|
13
|
+
cm = ContextManager(type="chat", compression="medium")
|
|
14
|
+
before = cm.estimate_tokens(messages)
|
|
15
|
+
print("tokens before:", before)
|
|
16
|
+
|
|
17
|
+
result = cm.compress(messages, token_budget=200, return_stats=True)
|
|
18
|
+
print("tokens after:", result.stats.tokens_after)
|
|
19
|
+
print("saved:", result.stats.tokens_saved)
|
|
20
|
+
print("stages:", result.stats.stages_run)
|
|
21
|
+
print("turns:", result.stats.turns_before, "->", result.stats.turns_after)
|
|
@@ -0,0 +1,128 @@
|
|
|
1
|
+
import copy
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
|
|
5
|
+
from contextpress import ContextManager
|
|
6
|
+
from contextpress.compression import known_stages
|
|
7
|
+
from contextpress.models import Conversation, Turn
|
|
8
|
+
from contextpress.pipeline import Pipeline
|
|
9
|
+
from contextpress.profiles import PROFILES
|
|
10
|
+
from contextpress.registry import clear_registry, register_stage
|
|
11
|
+
from contextpress.strategies.base import BaseStrategy
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class _DropMarkedStrategy(BaseStrategy):
|
|
15
|
+
def process(self, conversation: Conversation) -> Conversation:
|
|
16
|
+
kept = [
|
|
17
|
+
copy.deepcopy(t)
|
|
18
|
+
for t in conversation.turns
|
|
19
|
+
if t.role == "system"
|
|
20
|
+
or "DROP" not in (t.content if isinstance(t.content, str) else str(t.content))
|
|
21
|
+
]
|
|
22
|
+
return Conversation(
|
|
23
|
+
turns=kept,
|
|
24
|
+
type=conversation.type,
|
|
25
|
+
metadata=copy.deepcopy(conversation.metadata),
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@pytest.fixture(autouse=True)
|
|
30
|
+
def _reset_registry():
|
|
31
|
+
clear_registry()
|
|
32
|
+
yield
|
|
33
|
+
clear_registry()
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def test_estimate_tokens_dict_list():
|
|
37
|
+
cm = ContextManager(type="chat")
|
|
38
|
+
messages = [
|
|
39
|
+
{"role": "user", "content": "Hello world"},
|
|
40
|
+
{"role": "assistant", "content": "Hi there"},
|
|
41
|
+
]
|
|
42
|
+
n = cm.estimate_tokens(messages)
|
|
43
|
+
assert isinstance(n, int) and n > 0
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def test_register_stage_and_run():
|
|
47
|
+
cm = ContextManager(type="chat")
|
|
48
|
+
cm.register_stage("drop_marked", _DropMarkedStrategy)
|
|
49
|
+
assert "drop_marked" in known_stages()
|
|
50
|
+
messages = [
|
|
51
|
+
{"role": "user", "content": "keep me"},
|
|
52
|
+
{"role": "assistant", "content": "DROP this turn"},
|
|
53
|
+
]
|
|
54
|
+
out = cm.compress(messages, stages=["drop_marked"])
|
|
55
|
+
assert len(out) == 1
|
|
56
|
+
assert out[0]["content"] == "keep me"
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def test_register_stage_reserved_name_raises():
|
|
60
|
+
with pytest.raises(ValueError):
|
|
61
|
+
register_stage("filler", _DropMarkedStrategy)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def test_configure_custom_stage():
|
|
65
|
+
cm = ContextManager(type="chat")
|
|
66
|
+
cm.register_stage("drop_marked", _DropMarkedStrategy, aggressiveness=0.2)
|
|
67
|
+
cm.configure("drop_marked", aggressiveness=0.9)
|
|
68
|
+
assert cm._custom_stages["drop_marked"].aggressiveness == 0.9
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
class _DedupOnlyBackend:
|
|
72
|
+
def summarize(self, text: str, max_tokens: int) -> str:
|
|
73
|
+
raise AssertionError("summarize should not run in dedupe_only mode")
|
|
74
|
+
|
|
75
|
+
def deduplicate(self, turns: list[str]) -> list[int]:
|
|
76
|
+
return [0]
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
class _SummarizeOnlyBackend:
|
|
80
|
+
def summarize(self, text: str, max_tokens: int) -> str:
|
|
81
|
+
return "summary appended"
|
|
82
|
+
|
|
83
|
+
def deduplicate(self, turns: list[str]) -> list[int]:
|
|
84
|
+
raise AssertionError("deduplicate should not run in summarize_only mode")
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def test_llm_mode_dedupe_only():
|
|
88
|
+
conv = Conversation(
|
|
89
|
+
turns=[
|
|
90
|
+
Turn(role="user", content="hello"),
|
|
91
|
+
Turn(role="user", content="hello again"),
|
|
92
|
+
],
|
|
93
|
+
type="chat",
|
|
94
|
+
)
|
|
95
|
+
p = Pipeline(
|
|
96
|
+
PROFILES["chat"],
|
|
97
|
+
token_budget=None,
|
|
98
|
+
llm_backend=_DedupOnlyBackend(),
|
|
99
|
+
llm_mode="dedupe_only",
|
|
100
|
+
llm_min_input_chars=0,
|
|
101
|
+
)
|
|
102
|
+
out = p.run(conv)
|
|
103
|
+
assert len(out.turns) == 1
|
|
104
|
+
assert out.turns[0].content == "hello"
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def test_llm_mode_summarize_only():
|
|
108
|
+
conv = Conversation(
|
|
109
|
+
turns=[Turn(role="user", content="hello " * 200)],
|
|
110
|
+
type="chat",
|
|
111
|
+
)
|
|
112
|
+
p = Pipeline(
|
|
113
|
+
PROFILES["chat"],
|
|
114
|
+
token_budget=None,
|
|
115
|
+
llm_backend=_SummarizeOnlyBackend(),
|
|
116
|
+
llm_mode="summarize_only",
|
|
117
|
+
llm_min_input_chars=0,
|
|
118
|
+
)
|
|
119
|
+
out = p.run(conv)
|
|
120
|
+
assert len(out.turns) == 2
|
|
121
|
+
assert out.turns[0].role == "user"
|
|
122
|
+
assert out.turns[1].content == "summary appended"
|
|
123
|
+
assert out.turns[1].metadata.get("mode") == "summarize_only"
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def test_invalid_llm_mode_raises():
|
|
127
|
+
with pytest.raises(ValueError):
|
|
128
|
+
ContextManager(llm_mode="invalid")
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|