feat: add context-budget compression engine for historical summaries (#409)
Create scripts/context_budget.py with: - compress_to_budget(): truncates text to fit word budget - assemble_historical_context(): assembles rolling, yearly, prev-week, and monthly context within configurable budget allocations (rolling=500, prev_week=200, yearly=500, month=300 words) Pruning rules implemented: - Predictions unconfirmed after 8 weeks are dropped - Trends with 0 signal for 4 weeks compress to one sentence - Noise patterns keep only top 3 Output is a single markdown string suitable for LLM prompt injection. Tests cover word counting, budget enforcement, multi-file assembly, graceful handling of missing files, and pruning behavior. Closes #402 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Juan Manuel Servera committed
Jun 12, 2026 at 12:36 UTC
d7a273b763e9110f41f52e2aaeb421ffa133f7c2
2 files changed
+592
scripts/context_budget.py
new
+315
@@ -0,0 +1,315 @@
1
+#!/usr/bin/env python3
2
+"""Context-budget compression engine for historical summaries.
3
+
4
+Assembles rolling, yearly, previous-week, and monthly context into a single
5
+markdown string suitable for LLM prompt injection, respecting a total word
6
+budget with per-section allocations.
7
+
8
+Pruning rules:
9
+ - Predictions unconfirmed after 8 weeks → drop
10
+ - Trends with 0 signal for 4 weeks → compress to one sentence
11
+ - Noise patterns → keep only top 3
12
+
13
+CLI:
14
+ python scripts/context_budget.py --rolling path --yearly path [--prev-week path] [--max-words 1500]
15
+"""
16
+
17
+from __future__ import annotations
18
+
19
+import argparse
20
+import re
21
+import sys
22
+from datetime import datetime, timezone, timedelta
23
+from pathlib import Path
24
+from typing import Optional
25
+
26
+
27
+# Default budget allocations (words)
28
+BUDGET_ROLLING = 500
29
+BUDGET_PREV_WEEK = 200
30
+BUDGET_YEARLY = 500
31
+BUDGET_MONTH = 300
32
+
33
+STALE_PREDICTION_WEEKS = 8
34
+STALE_TREND_WEEKS = 4
35
+MAX_NOISE_PATTERNS = 3
36
+
37
+
38
+def word_count(text: str) -> int:
39
+ """Count words in text."""
40
+ return len(text.split())
41
+
42
+
43
+def compress_to_budget(text: str, max_words: int) -> str:
44
+ """Truncate/summarize text to fit within a word budget.
45
+
46
+ Preserves complete lines where possible, trimming from the end.
47
+ Returns the compressed text ending with '...' if truncated.
48
+ """
49
+ if not text or not text.strip():
50
+ return ""
51
+ if max_words <= 0:
52
+ return ""
53
+
54
+ words = text.split()
55
+ if len(words) <= max_words:
56
+ return text.strip()
57
+
58
+ # Preserve line structure: keep complete lines until budget
59
+ lines = text.strip().splitlines()
60
+ result_lines: list[str] = []
61
+ total_words = 0
62
+ for line in lines:
63
+ line_words = len(line.split())
64
+ if total_words + line_words > max_words:
65
+ # Partial last line if we have room
66
+ remaining = max_words - total_words
67
+ if remaining > 0:
68
+ partial = " ".join(line.split()[:remaining])
69
+ result_lines.append(partial + "...")
70
+ elif not result_lines:
71
+ # Edge case: first line exceeds budget
72
+ result_lines.append(" ".join(words[:max_words]) + "...")
73
+ break
74
+ result_lines.append(line)
75
+ total_words += line_words
76
+
77
+ return "\n".join(result_lines)
78
+
79
+
80
+def _parse_date_from_line(line: str) -> Optional[datetime]:
81
+ """Try to extract a date from a line like '- [2026-01-15] ...'."""
82
+ match = re.search(r"\[(\d{4}-\d{2}-\d{2})\]", line)
83
+ if match:
84
+ try:
85
+ return datetime.strptime(match.group(1), "%Y-%m-%d").replace(
86
+ tzinfo=timezone.utc
87
+ )
88
+ except ValueError:
89
+ return None
90
+ return None
91
+
92
+
93
+def prune_stale_predictions(text: str, now: Optional[datetime] = None) -> str:
94
+ """Drop predictions unconfirmed after 8 weeks."""
95
+ if not text:
96
+ return ""
97
+ if now is None:
98
+ now = datetime.now(timezone.utc)
99
+
100
+ cutoff = now - timedelta(weeks=STALE_PREDICTION_WEEKS)
101
+ lines = text.splitlines()
102
+ result: list[str] = []
103
+
104
+ for line in lines:
105
+ # Predictions are bullet lines with dates and no confirmation marker
106
+ if re.match(r"\s*[-*]\s*\[", line):
107
+ date = _parse_date_from_line(line)
108
+ if date and date < cutoff:
109
+ # Check for confirmation markers
110
+ lower = line.lower()
111
+ if "confirmed" not in lower and "✓" not in line and "✅" not in line:
112
+ continue # Drop stale unconfirmed prediction
113
+ result.append(line)
114
+
115
+ return "\n".join(result)
116
+
117
+
118
+def compress_stale_trends(text: str, now: Optional[datetime] = None) -> str:
119
+ """Compress trends with 0 signal for 4+ weeks to one sentence."""
120
+ if not text:
121
+ return ""
122
+ if now is None:
123
+ now = datetime.now(timezone.utc)
124
+
125
+ cutoff = now - timedelta(weeks=STALE_TREND_WEEKS)
126
+ lines = text.splitlines()
127
+ result: list[str] = []
128
+ i = 0
129
+
130
+ while i < len(lines):
131
+ line = lines[i]
132
+ # Detect trend blocks: "### Trend: <name>" or "## Trend: <name>"
133
+ trend_match = re.match(r"(#{2,3})\s+[Tt]rend:\s*(.+)", line)
134
+ if trend_match:
135
+ heading_level = trend_match.group(1)
136
+ trend_name = trend_match.group(2).strip()
137
+ # Collect the trend block
138
+ block_lines = [line]
139
+ i += 1
140
+ has_recent_signal = False
141
+ while i < len(lines) and not re.match(r"#{2,3}\s+", lines[i]):
142
+ block_lines.append(lines[i])
143
+ # Check for recent signals
144
+ date = _parse_date_from_line(lines[i])
145
+ if date and date >= cutoff:
146
+ has_recent_signal = True
147
+ # Non-dated content counts as signal
148
+ if lines[i].strip() and not lines[i].startswith("#"):
149
+ if "no signal" in lines[i].lower() or "0 signal" in lines[i].lower():
150
+ pass
151
+ elif date is None and lines[i].strip().startswith(("-", "*")):
152
+ has_recent_signal = True
153
+ i += 1
154
+
155
+ if has_recent_signal:
156
+ result.extend(block_lines)
157
+ else:
158
+ # Compress to one sentence
159
+ result.append(f"- {trend_name}: no recent signal (stale)")
160
+ else:
161
+ result.append(line)
162
+ i += 1
163
+
164
+ return "\n".join(result)
165
+
166
+
167
+def keep_top_noise_patterns(text: str, max_patterns: int = MAX_NOISE_PATTERNS) -> str:
168
+ """Keep only the top N noise patterns from a noise section."""
169
+ if not text:
170
+ return ""
171
+
172
+ lines = text.splitlines()
173
+ result: list[str] = []
174
+ in_noise_section = False
175
+ noise_count = 0
176
+
177
+ for line in lines:
178
+ # Detect noise section headers
179
+ if re.match(r"#{2,3}\s+[Nn]oise", line):
180
+ in_noise_section = True
181
+ noise_count = 0
182
+ result.append(line)
183
+ continue
184
+
185
+ if in_noise_section:
186
+ # New section starts
187
+ if re.match(r"#{2,3}\s+", line) and not re.match(r"#{2,3}\s+[Nn]oise", line):
188
+ in_noise_section = False
189
+ result.append(line)
190
+ continue
191
+
192
+ # Count bullet items
193
+ if re.match(r"\s*[-*]\s+", line):
194
+ noise_count += 1
195
+ if noise_count <= max_patterns:
196
+ result.append(line)
197
+ continue
198
+
199
+ result.append(line)
200
+ else:
201
+ result.append(line)
202
+
203
+ return "\n".join(result)
204
+
205
+
206
+def _read_file_safe(path: Optional[Path]) -> str:
207
+ """Read a file, returning empty string if missing or unreadable."""
208
+ if path is None:
209
+ return ""
210
+ try:
211
+ return Path(path).read_text(encoding="utf-8")
212
+ except (FileNotFoundError, PermissionError, OSError):
213
+ return ""
214
+
215
+
216
+def assemble_historical_context(
217
+ rolling_path: Optional[str] = None,
218
+ yearly_path: Optional[str] = None,
219
+ prev_week_path: Optional[str] = None,
220
+ month_path: Optional[str] = None,
221
+ max_total_words: int = 1500,
222
+ now: Optional[datetime] = None,
223
+) -> str:
224
+ """Assemble context from multiple sources within budget.
225
+
226
+ Budget allocation:
227
+ - rolling: 500 words
228
+ - prev_week: 200 words
229
+ - yearly: 500 words
230
+ - month: 300 words
231
+
232
+ Total is capped at max_total_words (default 1500).
233
+ Returns a single markdown string for LLM prompt injection.
234
+ """
235
+ # Scale budgets if max_total_words differs from default
236
+ default_total = BUDGET_ROLLING + BUDGET_PREV_WEEK + BUDGET_YEARLY + BUDGET_MONTH
237
+ scale = max_total_words / default_total if default_total > 0 else 1.0
238
+
239
+ budget_rolling = int(BUDGET_ROLLING * scale)
240
+ budget_prev_week = int(BUDGET_PREV_WEEK * scale)
241
+ budget_yearly = int(BUDGET_YEARLY * scale)
242
+ budget_month = int(BUDGET_MONTH * scale)
243
+
244
+ # Read sources
245
+ rolling_raw = _read_file_safe(Path(rolling_path) if rolling_path else None)
246
+ yearly_raw = _read_file_safe(Path(yearly_path) if yearly_path else None)
247
+ prev_week_raw = _read_file_safe(Path(prev_week_path) if prev_week_path else None)
248
+ month_raw = _read_file_safe(Path(month_path) if month_path else None)
249
+
250
+ # Apply pruning rules
251
+ rolling_pruned = keep_top_noise_patterns(
252
+ compress_stale_trends(prune_stale_predictions(rolling_raw, now=now), now=now)
253
+ )
254
+ yearly_pruned = prune_stale_predictions(yearly_raw, now=now)
255
+ prev_week_pruned = prev_week_raw
256
+ month_pruned = compress_stale_trends(
257
+ prune_stale_predictions(month_raw, now=now), now=now
258
+ )
259
+
260
+ # Compress each section to its budget
261
+ rolling_compressed = compress_to_budget(rolling_pruned, budget_rolling)
262
+ prev_week_compressed = compress_to_budget(prev_week_pruned, budget_prev_week)
263
+ yearly_compressed = compress_to_budget(yearly_pruned, budget_yearly)
264
+ month_compressed = compress_to_budget(month_pruned, budget_month)
265
+
266
+ # Assemble final markdown
267
+ sections: list[str] = []
268
+
269
+ if rolling_compressed:
270
+ sections.append(f"## Rolling Context\n\n{rolling_compressed}")
271
+ if prev_week_compressed:
272
+ sections.append(f"## Previous Week\n\n{prev_week_compressed}")
273
+ if yearly_compressed:
274
+ sections.append(f"## Yearly Context\n\n{yearly_compressed}")
275
+ if month_compressed:
276
+ sections.append(f"## Monthly Context\n\n{month_compressed}")
277
+
278
+ assembled = "\n\n".join(sections)
279
+
280
+ # Final enforcement: ensure total stays within max_total_words
281
+ if word_count(assembled) > max_total_words:
282
+ assembled = compress_to_budget(assembled, max_total_words)
283
+
284
+ return assembled
285
+
286
+
287
+def main(argv: Optional[list[str]] = None) -> int:
288
+ parser = argparse.ArgumentParser(
289
+ description="Assemble historical context within a word budget."
290
+ )
291
+ parser.add_argument("--rolling", type=str, help="Path to rolling context file")
292
+ parser.add_argument("--yearly", type=str, help="Path to yearly context file")
293
+ parser.add_argument("--prev-week", type=str, help="Path to previous week file")
294
+ parser.add_argument("--month", type=str, help="Path to monthly context file")
295
+ parser.add_argument(
296
+ "--max-words", type=int, default=1500, help="Maximum total words (default: 1500)"
297
+ )
298
+
299
+ args = parser.parse_args(argv)
300
+
301
+ result = assemble_historical_context(
302
+ rolling_path=args.rolling,
303
+ yearly_path=args.yearly,
304
+ prev_week_path=args.prev_week,
305
+ month_path=args.month,
306
+ max_total_words=args.max_words,
307
+ )
308
+
309
+ if result:
310
+ print(result)
311
+ return 0
312
+
313
+
314
+if __name__ == "__main__":
315
+ sys.exit(main())
tests/test_context_budget.py
new
+277
@@ -0,0 +1,277 @@
1
+"""Tests for scripts/context_budget.py."""
2
+
3
+import sys
4
+from datetime import datetime, timezone, timedelta
5
+from pathlib import Path
6
+
7
+import pytest
8
+
9
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
10
+
11
+from scripts.context_budget import (
12
+ word_count,
13
+ compress_to_budget,
14
+ assemble_historical_context,
15
+ prune_stale_predictions,
16
+ compress_stale_trends,
17
+ keep_top_noise_patterns,
18
+ BUDGET_ROLLING,
19
+ BUDGET_PREV_WEEK,
20
+ BUDGET_YEARLY,
21
+ BUDGET_MONTH,
22
+ STALE_PREDICTION_WEEKS,
23
+ MAX_NOISE_PATTERNS,
24
+)
25
+
26
+
27
+# --- word_count tests ---
28
+
29
+
30
+class TestWordCount:
31
+ def test_empty_string(self):
32
+ assert word_count("") == 0
33
+
34
+ def test_single_word(self):
35
+ assert word_count("hello") == 1
36
+
37
+ def test_multiple_words(self):
38
+ assert word_count("one two three four five") == 5
39
+
40
+ def test_multiline(self):
41
+ text = "line one\nline two\nline three"
42
+ assert word_count(text) == 6
43
+
44
+ def test_extra_whitespace(self):
45
+ assert word_count(" hello world ") == 2
46
+
47
+
48
+# --- compress_to_budget tests ---
49
+
50
+
51
+class TestCompressToBudget:
52
+ def test_empty_text(self):
53
+ assert compress_to_budget("", 100) == ""
54
+
55
+ def test_whitespace_only(self):
56
+ assert compress_to_budget(" \n ", 100) == ""
57
+
58
+ def test_within_budget(self):
59
+ text = "This is a short sentence."
60
+ result = compress_to_budget(text, 100)
61
+ assert result == text.strip()
62
+ assert word_count(result) <= 100
63
+
64
+ def test_exceeds_budget_truncates(self):
65
+ text = " ".join(f"word{i}" for i in range(100))
66
+ result = compress_to_budget(text, 10)
67
+ assert word_count(result) <= 11 # 10 + possible trailing "..."
68
+ assert result.endswith("...")
69
+
70
+ def test_zero_budget(self):
71
+ assert compress_to_budget("hello world", 0) == ""
72
+
73
+ def test_negative_budget(self):
74
+ assert compress_to_budget("hello world", -5) == ""
75
+
76
+ def test_preserves_line_structure(self):
77
+ text = "line one two\nline three four five\nline six seven eight nine ten"
78
+ result = compress_to_budget(text, 5)
79
+ # Should keep first line (3 words) and partial of second
80
+ assert word_count(result) <= 6 # budget + trailing word
81
+ assert "..." in result or word_count(result) <= 5
82
+
83
+ def test_budget_enforcement_strict(self):
84
+ text = " ".join(["word"] * 1000)
85
+ result = compress_to_budget(text, 50)
86
+ # Word count should not greatly exceed budget
87
+ assert word_count(result) <= 51
88
+
89
+
90
+# --- prune_stale_predictions tests ---
91
+
92
+
93
+class TestPruneStale:
94
+ def test_drops_old_unconfirmed(self):
95
+ now = datetime(2026, 6, 12, tzinfo=timezone.utc)
96
+ text = "- [2026-03-01] prediction about AI growth\n- [2026-06-01] recent prediction"
97
+ result = prune_stale_predictions(text, now=now)
98
+ assert "2026-03-01" not in result
99
+ assert "2026-06-01" in result
100
+
101
+ def test_keeps_confirmed(self):
102
+ now = datetime(2026, 6, 12, tzinfo=timezone.utc)
103
+ text = "- [2026-01-01] old but confirmed ✓"
104
+ result = prune_stale_predictions(text, now=now)
105
+ assert "2026-01-01" in result
106
+
107
+ def test_keeps_recent(self):
108
+ now = datetime(2026, 6, 12, tzinfo=timezone.utc)
109
+ recent = (now - timedelta(weeks=2)).strftime("%Y-%m-%d")
110
+ text = f"- [{recent}] recent prediction"
111
+ result = prune_stale_predictions(text, now=now)
112
+ assert recent in result
113
+
114
+ def test_empty_text(self):
115
+ assert prune_stale_predictions("") == ""
116
+
117
+
118
+# --- compress_stale_trends tests ---
119
+
120
+
121
+class TestCompressStaleTrends:
122
+ def test_compresses_stale_trend(self):
123
+ now = datetime(2026, 6, 12, tzinfo=timezone.utc)
124
+ text = (
125
+ "## Trend: Blockchain hype\n"
126
+ "- [2026-01-01] some old signal\n"
127
+ "- [2026-01-15] another old signal\n"
128
+ )
129
+ result = compress_stale_trends(text, now=now)
130
+ assert "Blockchain hype: no recent signal" in result
131
+ assert "some old signal" not in result
132
+
133
+ def test_keeps_active_trend(self):
134
+ now = datetime(2026, 6, 12, tzinfo=timezone.utc)
135
+ recent = (now - timedelta(days=5)).strftime("%Y-%m-%d")
136
+ text = f"## Trend: AI growth\n- [{recent}] strong signal this week\n"
137
+ result = compress_stale_trends(text, now=now)
138
+ assert "strong signal this week" in result
139
+
140
+ def test_empty_text(self):
141
+ assert compress_stale_trends("") == ""
142
+
143
+
144
+# --- keep_top_noise_patterns tests ---
145
+
146
+
147
+class TestNoisePatterns:
148
+ def test_keeps_only_top_n(self):
149
+ text = (
150
+ "## Noise\n"
151
+ "- noise pattern 1\n"
152
+ "- noise pattern 2\n"
153
+ "- noise pattern 3\n"
154
+ "- noise pattern 4\n"
155
+ "- noise pattern 5\n"
156
+ )
157
+ result = keep_top_noise_patterns(text, max_patterns=3)
158
+ assert "noise pattern 1" in result
159
+ assert "noise pattern 2" in result
160
+ assert "noise pattern 3" in result
161
+ assert "noise pattern 4" not in result
162
+ assert "noise pattern 5" not in result
163
+
164
+ def test_fewer_than_max(self):
165
+ text = "## Noise\n- pattern 1\n- pattern 2\n"
166
+ result = keep_top_noise_patterns(text, max_patterns=3)
167
+ assert "pattern 1" in result
168
+ assert "pattern 2" in result
169
+
170
+ def test_non_noise_sections_untouched(self):
171
+ text = "## Signals\n- sig 1\n- sig 2\n- sig 3\n- sig 4\n- sig 5\n"
172
+ result = keep_top_noise_patterns(text, max_patterns=3)
173
+ assert "sig 5" in result # Not filtered
174
+
175
+
176
+# --- assemble_historical_context tests ---
177
+
178
+
179
+class TestAssemble:
180
+ def test_assembles_from_files(self, tmp_path):
181
+ rolling = tmp_path / "rolling.md"
182
+ yearly = tmp_path / "yearly.md"
183
+ prev_week = tmp_path / "prev_week.md"
184
+
185
+ rolling.write_text("Rolling context with some words here today.")
186
+ yearly.write_text("Yearly summary of important events and trends.")
187
+ prev_week.write_text("Last week's highlights and observations.")
188
+
189
+ result = assemble_historical_context(
190
+ rolling_path=str(rolling),
191
+ yearly_path=str(yearly),
192
+ prev_week_path=str(prev_week),
193
+ max_total_words=1500,
194
+ )
195
+
196
+ assert "## Rolling Context" in result
197
+ assert "## Yearly Context" in result
198
+ assert "## Previous Week" in result
199
+
200
+ def test_handles_missing_files(self, tmp_path):
201
+ result = assemble_historical_context(
202
+ rolling_path=str(tmp_path / "nonexistent.md"),
203
+ yearly_path=None,
204
+ prev_week_path=None,
205
+ max_total_words=1500,
206
+ )
207
+ # Should not crash, returns empty or partial
208
+ assert isinstance(result, str)
209
+
210
+ def test_all_none_paths(self):
211
+ result = assemble_historical_context(
212
+ rolling_path=None,
213
+ yearly_path=None,
214
+ prev_week_path=None,
215
+ max_total_words=1500,
216
+ )
217
+ assert result == ""
218
+
219
+ def test_budget_enforcement(self, tmp_path):
220
+ # Create a large file that exceeds budget
221
+ large_text = " ".join(["word"] * 2000)
222
+ rolling = tmp_path / "rolling.md"
223
+ rolling.write_text(large_text)
224
+
225
+ yearly = tmp_path / "yearly.md"
226
+ yearly.write_text(large_text)
227
+
228
+ result = assemble_historical_context(
229
+ rolling_path=str(rolling),
230
+ yearly_path=str(yearly),
231
+ max_total_words=100,
232
+ )
233
+ # Total words should respect the budget (with small margin for headers)
234
+ assert word_count(result) <= 110
235
+
236
+ def test_output_is_markdown(self, tmp_path):
237
+ rolling = tmp_path / "rolling.md"
238
+ rolling.write_text("Some rolling context data points here.")
239
+
240
+ result = assemble_historical_context(
241
+ rolling_path=str(rolling),
242
+ max_total_words=1500,
243
+ )
244
+ # Should contain markdown headers
245
+ assert result.startswith("## ")
246
+
247
+ def test_budget_allocation_scaling(self, tmp_path):
248
+ """Verify budgets scale proportionally with max_total_words."""
249
+ text = " ".join(["word"] * 1000)
250
+ rolling = tmp_path / "rolling.md"
251
+ rolling.write_text(text)
252
+
253
+ # With max_total_words=750 (half of default 1500),
254
+ # rolling budget should be ~250 (half of 500)
255
+ result = assemble_historical_context(
256
+ rolling_path=str(rolling),
257
+ max_total_words=750,
258
+ )
259
+ # Section content should be roughly half the default rolling budget
260
+ content = result.replace("## Rolling Context\n\n", "")
261
+ assert word_count(content) <= 280 # ~250 + margin
262
+
263
+ def test_pruning_applied_during_assembly(self, tmp_path):
264
+ now = datetime(2026, 6, 12, tzinfo=timezone.utc)
265
+ rolling = tmp_path / "rolling.md"
266
+ rolling.write_text(
267
+ "- [2026-01-01] stale prediction from January\n"
268
+ "- [2026-06-10] fresh prediction from this week\n"
269
+ )
270
+
271
+ result = assemble_historical_context(
272
+ rolling_path=str(rolling),
273
+ max_total_words=1500,
274
+ now=now,
275
+ )
276
+ assert "2026-01-01" not in result
277
+ assert "2026-06-10" in result