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.
Files changed (46) hide show
  1. {contextpress-0.2.0 → contextpress-0.3.0}/.gitignore +3 -0
  2. {contextpress-0.2.0 → contextpress-0.3.0}/CHANGELOG.md +7 -0
  3. {contextpress-0.2.0 → contextpress-0.3.0}/PKG-INFO +25 -2
  4. {contextpress-0.2.0 → contextpress-0.3.0}/README.md +24 -1
  5. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/__init__.py +1 -1
  6. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/compression.py +19 -3
  7. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/core.py +39 -4
  8. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/pipeline.py +78 -21
  9. contextpress-0.3.0/contextpress/registry.py +85 -0
  10. contextpress-0.3.0/examples/estimate_and_stats.py +21 -0
  11. {contextpress-0.2.0 → contextpress-0.3.0}/pyproject.toml +1 -1
  12. contextpress-0.3.0/tests/test_v03.py +128 -0
  13. {contextpress-0.2.0 → contextpress-0.3.0}/CITATION.cff +0 -0
  14. {contextpress-0.2.0 → contextpress-0.3.0}/CONTRIBUTING.md +0 -0
  15. {contextpress-0.2.0 → contextpress-0.3.0}/LICENSE +0 -0
  16. {contextpress-0.2.0 → contextpress-0.3.0}/NOTICE +0 -0
  17. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/_bootstrap.py +0 -0
  18. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/__init__.py +0 -0
  19. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/_helpers.py +0 -0
  20. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/adapters.py +0 -0
  21. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/llm/base.py +0 -0
  22. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/models.py +0 -0
  23. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/normalizer.py +0 -0
  24. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/profiles.py +0 -0
  25. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/py.typed +0 -0
  26. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/stats.py +0 -0
  27. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/__init__.py +0 -0
  28. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/base.py +0 -0
  29. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/budget.py +0 -0
  30. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/filler.py +0 -0
  31. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/recency.py +0 -0
  32. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/repetition.py +0 -0
  33. {contextpress-0.2.0 → contextpress-0.3.0}/contextpress/strategies/resolution.py +0 -0
  34. {contextpress-0.2.0 → contextpress-0.3.0}/examples/llm_tier_ollama.py +0 -0
  35. {contextpress-0.2.0 → contextpress-0.3.0}/examples/llm_tier_openai.py +0 -0
  36. {contextpress-0.2.0 → contextpress-0.3.0}/tests/__init__.py +0 -0
  37. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_budget.py +0 -0
  38. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_filler.py +0 -0
  39. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_llm_helpers.py +0 -0
  40. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_models.py +0 -0
  41. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_normalizer.py +0 -0
  42. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_pipeline.py +0 -0
  43. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_recency.py +0 -0
  44. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_repetition.py +0 -0
  45. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_resolution.py +0 -0
  46. {contextpress-0.2.0 → contextpress-0.3.0}/tests/test_stats.py +0 -0
@@ -29,3 +29,6 @@ htmlcov/
29
29
  *.swp
30
30
  .DS_Store
31
31
  Thumbs.db
32
+
33
+ # Secrets (never commit)
34
+ .env
@@ -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.2.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
 
@@ -16,7 +16,7 @@ __all__ = [
16
16
  "CompressionResult",
17
17
  "CompressionStats",
18
18
  ]
19
- __version__ = "0.2.0"
19
+ __version__ = "0.3.0"
20
20
 
21
21
 
22
22
  def __getattr__(name: str) -> Any:
@@ -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 VALID_STAGES]
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(VALID_STAGES)}")
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.stats import CompressionResult, CompressionStats
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 not hasattr(self._profile, stage):
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
- for stage_name in self.STAGE_ORDER:
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 = getattr(self.profile, stage_name)
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
- try:
170
- keep_idx = self.llm_backend.deduplicate(texts)
171
- except Exception:
172
- keep_idx = list(range(len(ns_turns)))
173
- valid = sorted(
174
- {
175
- i
176
- for i in keep_idx
177
- if type(i) is int and not isinstance(i, bool) and 0 <= i < len(ns_turns)
178
- }
179
- )
180
- if not valid:
181
- valid = list(range(len(ns_turns)))
182
- if len(valid) < len(ns_turns):
183
- ns_turns = [ns_turns[i] for i in valid]
184
- texts = [texts[i] for i in valid]
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
- if stats is not None:
187
- stats.llm_dedup_turns_after = len(ns_turns)
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)
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "contextpress"
7
- version = "0.2.0"
7
+ version = "0.3.0"
8
8
  description = "Deterministic context compression for LLM chat, RAG, and agent pipelines"
9
9
  readme = "README.md"
10
10
  license = { file = "LICENSE" }
@@ -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