main
py 477 lines 15.4 KB
Raw
1 #!/usr/bin/env python3
2 """Assemble bounded historical context for weekly analysis prompts."""
3
4 # NOTE: This module coexists with scripts/context_budget.py (the older CLI-oriented
5 # budget engine). This module is canonical for pipeline-integrated historical context
6 # assembly used by analyze_fallback.py. The older module remains for standalone CLI use.
7
8 from __future__ import annotations
9
10 import argparse
11 import re
12 import sys
13 from dataclasses import dataclass
14 from datetime import datetime
15 from pathlib import Path
16 from typing import Iterable
17
18 ROOT = Path(__file__).resolve().parent.parent
19 DEFAULT_CONTENT_ROOT = ROOT / "content"
20 DEFAULT_MAX_WORDS = 1_500
21 DEFAULT_PROMPT_BUDGET_FRACTION = 0.15
22
23 SECTION_PRIORITY = ("rolling", "previous_week", "monthly", "yearly")
24
25 SECTION_SPECS = {
26 "rolling": {"label": "Rolling Last 4 Weeks", "target_words": 500, "min_words": 250},
27 "previous_week": {"label": "Previous Week Takeaways", "target_words": 200, "min_words": 100},
28 "monthly": {"label": "Month In Progress", "target_words": 200, "min_words": 100},
29 "yearly": {"label": "Yearly Narrative", "target_words": 500, "min_words": 150},
30 }
31
32 FRONTMATTER_PATTERN = re.compile(r"^---\n(?P<frontmatter>.*?)\n---\n?(?P<body>.*)\Z", re.DOTALL)
33
34
35 @dataclass(frozen=True)
36 class HistoricalContextSection:
37 key: str
38 label: str
39 source_path: str | None
40 words: int
41 token_estimate: int
42 content: str
43
44
45 @dataclass(frozen=True)
46 class HistoricalContextResult:
47 markdown: str
48 sections: tuple[HistoricalContextSection, ...]
49 max_words: int
50 max_tokens: int
51 word_count: int
52 token_estimate: int
53
54
55 @dataclass
56 class _SectionPlan:
57 key: str
58 label: str
59 source_path: Path | None
60 raw_content: str
61 target_words: int
62 min_words: int
63 current_words: int
64
65
66 def estimate_tokens(text: str) -> int:
67 return (len(text.encode("utf-8")) + 3) // 4
68
69
70 def word_count(text: str) -> int:
71 return len(re.findall(r"\S+", text))
72
73
74 def compress_to_budget(text: str, max_words: int) -> str:
75 if not text or max_words <= 0:
76 return ""
77 words = text.split()
78 if len(words) <= max_words:
79 return text.strip()
80
81 lines = text.strip().splitlines()
82 result_lines: list[str] = []
83 total_words = 0
84 for line in lines:
85 line_words = len(line.split())
86 if total_words + line_words > max_words:
87 remaining = max_words - total_words
88 if remaining > 0:
89 partial = " ".join(line.split()[:remaining])
90 result_lines.append(partial + "")
91 elif not result_lines:
92 result_lines.append(" ".join(words[:max_words]) + "")
93 break
94 result_lines.append(line)
95 total_words += line_words
96 return "\n".join(result_lines)
97
98
99 def _read_text(path: Path | None) -> str:
100 if path is None or not path.exists() or not path.is_file():
101 return ""
102 return path.read_text(encoding="utf-8").strip()
103
104
105 def _strip_frontmatter(markdown: str) -> str:
106 match = FRONTMATTER_PATTERN.match(markdown.strip())
107 return match.group("body").strip() if match else markdown.strip()
108
109
110 def _frontmatter_value(markdown: str, key: str) -> str:
111 match = FRONTMATTER_PATTERN.match(markdown.strip())
112 if not match:
113 return ""
114 for line in match.group("frontmatter").splitlines():
115 if ":" not in line or line.startswith((" ", "\t")):
116 continue
117 raw_key, value = line.split(":", 1)
118 if raw_key.strip() == key:
119 return value.strip().strip('"').strip("'")
120 return ""
121
122
123 def _extract_markdown_section(markdown: str, heading: str) -> str:
124 content = _strip_frontmatter(markdown)
125 lines = content.splitlines()
126 capture = False
127 result: list[str] = []
128 heading_pattern = re.compile(rf"^##\s+{re.escape(heading)}\s*$", re.IGNORECASE)
129 next_section_pattern = re.compile(r"^##\s+")
130
131 for line in lines:
132 if heading_pattern.match(line):
133 capture = True
134 result.append(line)
135 continue
136 if capture and next_section_pattern.match(line):
137 break
138 if capture:
139 result.append(line)
140 return "\n".join(result).strip()
141
142
143 def _join_nonempty(parts: Iterable[str]) -> str:
144 return "\n\n".join(part.strip() for part in parts if part and part.strip())
145
146
147 def _extract_previous_week_takeaways(markdown: str) -> str:
148 summary = _frontmatter_value(markdown, "summary")
149 parts: list[str] = []
150 if summary:
151 parts.append(f"- Prior weekly thesis: {summary}")
152 parts.extend(
153 section
154 for section in (
155 _extract_markdown_section(markdown, "Signal & Noise"),
156 _extract_markdown_section(markdown, "Blind Spots"),
157 _extract_markdown_section(markdown, "The Week Ahead"),
158 )
159 if section
160 )
161 return _join_nonempty(parts) or _strip_frontmatter(markdown)
162
163
164 def _extract_month_notes(markdown: str) -> str:
165 return _join_nonempty(
166 section
167 for section in (
168 _extract_markdown_section(markdown, "Month Overview"),
169 _extract_markdown_section(markdown, "Trends Observed"),
170 _extract_markdown_section(markdown, "Key Takeaways"),
171 )
172 if section
173 ) or _strip_frontmatter(markdown)
174
175
176 def _extract_yearly_narrative(markdown: str) -> str:
177 narrative = _extract_markdown_section(markdown, "Narrative")
178 if narrative:
179 arc = _extract_markdown_section(markdown, "Arc")
180 return _join_nonempty(section for section in (narrative, arc) if section)
181 return _join_nonempty(
182 section
183 for section in (
184 _extract_markdown_section(markdown, "Year in Review"),
185 _extract_markdown_section(markdown, "Biggest Trends"),
186 _extract_markdown_section(markdown, "Predictions Review"),
187 )
188 if section
189 ) or _strip_frontmatter(markdown)
190
191
192 def _parse_current_datetime(value: str) -> datetime | None:
193 normalized = value.strip()
194 if normalized.endswith("Z"):
195 normalized = normalized[:-1] + "+00:00"
196 try:
197 return datetime.fromisoformat(normalized)
198 except ValueError:
199 return None
200
201
202 def _resolve_month_path(content_root: Path, current_datetime: str) -> Path | None:
203 monthly_dir = content_root / "monthly"
204 if not monthly_dir.exists():
205 return None
206
207 dt = _parse_current_datetime(current_datetime)
208 target_year = dt.year if dt else None
209 target_month = dt.month if dt else None
210
211 candidates: list[tuple[int, int, Path]] = []
212 for path in sorted(monthly_dir.glob("*/*.md")):
213 try:
214 year = int(path.parent.name)
215 month = int(path.stem)
216 except ValueError:
217 continue
218 candidates.append((year, month, path))
219
220 if not candidates:
221 return None
222
223 if target_year is not None and target_month is not None:
224 eligible = [
225 item for item in candidates if (item[0], item[1]) <= (target_year, target_month)
226 ]
227 if eligible:
228 return eligible[-1][2]
229 return candidates[-1][2]
230
231
232 def _resolve_year_path(content_root: Path, current_datetime: str) -> Path | None:
233 yearly_dir = content_root / "yearly"
234 if not yearly_dir.exists():
235 return None
236 dt = _parse_current_datetime(current_datetime)
237 if dt:
238 exact = yearly_dir / f"{dt.year}.md"
239 if exact.exists():
240 return exact
241 candidates = sorted(path for path in yearly_dir.glob("*.md") if path.stem.isdigit())
242 return candidates[-1] if candidates else None
243
244
245 def extract_month_notes(markdown: str) -> str:
246 return _extract_month_notes(markdown)
247
248
249 def extract_yearly_narrative(markdown: str) -> str:
250 return _extract_yearly_narrative(markdown)
251
252
253 def resolve_latest_monthly_path(content_root: Path, current_datetime: str) -> Path | None:
254 return _resolve_month_path(content_root, current_datetime)
255
256
257 def resolve_latest_yearly_path(content_root: Path, current_datetime: str) -> Path | None:
258 return _resolve_year_path(content_root, current_datetime)
259
260
261 def _build_plans(
262 *,
263 current_datetime: str,
264 previous_summary_path: Path | None,
265 content_root: Path,
266 ) -> list[_SectionPlan]:
267 rolling_path = content_root / "rolling" / "last-month.md"
268 monthly_path = _resolve_month_path(content_root, current_datetime)
269 yearly_path = _resolve_year_path(content_root, current_datetime)
270
271 extracted = {
272 "rolling": _strip_frontmatter(_read_text(rolling_path)),
273 "previous_week": _extract_previous_week_takeaways(_read_text(previous_summary_path)),
274 "monthly": _extract_month_notes(_read_text(monthly_path)),
275 "yearly": _extract_yearly_narrative(_read_text(yearly_path)),
276 }
277 source_paths = {
278 "rolling": rolling_path if rolling_path.exists() else None,
279 "previous_week": previous_summary_path
280 if previous_summary_path and previous_summary_path.exists()
281 else None,
282 "monthly": monthly_path if monthly_path and monthly_path.exists() else None,
283 "yearly": yearly_path if yearly_path and yearly_path.exists() else None,
284 }
285
286 plans: list[_SectionPlan] = []
287 for key in SECTION_PRIORITY:
288 raw_content = extracted[key].strip()
289 if not raw_content:
290 continue
291 spec = SECTION_SPECS[key]
292 plans.append(
293 _SectionPlan(
294 key=key,
295 label=str(spec["label"]),
296 source_path=source_paths[key],
297 raw_content=raw_content,
298 target_words=int(spec["target_words"]),
299 min_words=int(spec["min_words"]),
300 current_words=min(word_count(raw_content), int(spec["target_words"])),
301 )
302 )
303 return plans
304
305
306 def _escape_boundaries(text: str) -> str:
307 """Defense-in-depth: escape untrusted-content boundary markers in assembled text."""
308 try:
309 from scripts.sanitize_repo_content import _escape_untrusted_boundaries
310 except ModuleNotFoundError: # pragma: no cover - script execution path
311 sys.path.insert(0, str(ROOT))
312 from scripts.sanitize_repo_content import _escape_untrusted_boundaries
313
314 return _escape_untrusted_boundaries(text)
315
316
317 def _render_sections(
318 plans: Iterable[_SectionPlan],
319 ) -> tuple[str, tuple[HistoricalContextSection, ...]]:
320 rendered_sections: list[str] = []
321 metadata: list[HistoricalContextSection] = []
322 for plan in plans:
323 if plan.current_words <= 0:
324 continue
325 content = compress_to_budget(plan.raw_content, plan.current_words)
326 if not content:
327 continue
328 # Defense-in-depth: escape boundary markers even though the caller
329 # (analyze_fallback.py) also escapes the final assembled string.
330 content = _escape_boundaries(content)
331 rendered_sections.append(f"### {plan.label}\n\n{content}")
332 metadata.append(
333 HistoricalContextSection(
334 key=plan.key,
335 label=plan.label,
336 source_path=plan.source_path.as_posix() if plan.source_path else None,
337 words=word_count(content),
338 token_estimate=estimate_tokens(content),
339 content=content,
340 )
341 )
342 return "\n\n".join(rendered_sections).strip(), tuple(metadata)
343
344
345 def _reduce_plans(
346 plans: list[_SectionPlan],
347 *,
348 max_words: int,
349 max_tokens: int,
350 ) -> tuple[str, tuple[HistoricalContextSection, ...]]:
351 if not plans:
352 return "", ()
353
354 for _ in range(400):
355 rendered, metadata = _render_sections(plans)
356 if not rendered:
357 return "", ()
358 if word_count(rendered) <= max_words and estimate_tokens(rendered) <= max_tokens:
359 return rendered, metadata
360
361 reduced = False
362 for key in reversed(SECTION_PRIORITY):
363 for plan in plans:
364 if plan.key != key or plan.current_words <= 0:
365 continue
366 floor = 0 if plan.current_words <= plan.min_words else plan.min_words
367 step = 50 if plan.current_words - floor > 100 else 25
368 next_words = max(floor, plan.current_words - step)
369 if next_words == plan.current_words and floor == 0:
370 next_words = 0
371 if next_words < plan.current_words:
372 plan.current_words = next_words
373 reduced = True
374 break
375 if reduced:
376 break
377 if not reduced:
378 break
379
380 return _render_sections(plans)
381
382
383 def build_historical_context(
384 *,
385 current_datetime: str,
386 previous_summary_path: Path | None,
387 content_root: Path = DEFAULT_CONTENT_ROOT,
388 max_words: int = DEFAULT_MAX_WORDS,
389 prompt_token_budget: int = 90_000,
390 prompt_budget_fraction: float = DEFAULT_PROMPT_BUDGET_FRACTION,
391 ) -> HistoricalContextResult:
392 max_tokens = max(0, int(prompt_token_budget * prompt_budget_fraction))
393 if max_words <= 0 or max_tokens <= 0:
394 return HistoricalContextResult("", (), max_words, max_tokens, 0, 0)
395
396 plans = _build_plans(
397 current_datetime=current_datetime,
398 previous_summary_path=previous_summary_path,
399 content_root=content_root,
400 )
401 rendered, sections = _reduce_plans(plans, max_words=max_words, max_tokens=max_tokens)
402 return HistoricalContextResult(
403 markdown=rendered,
404 sections=sections,
405 max_words=max_words,
406 max_tokens=max_tokens,
407 word_count=word_count(rendered),
408 token_estimate=estimate_tokens(rendered),
409 )
410
411
412 def assemble_historical_context(
413 *,
414 current_datetime: str,
415 previous_summary_path: Path | None,
416 content_root: Path = DEFAULT_CONTENT_ROOT,
417 max_words: int = DEFAULT_MAX_WORDS,
418 prompt_token_budget: int = 90_000,
419 prompt_budget_fraction: float = DEFAULT_PROMPT_BUDGET_FRACTION,
420 ) -> str:
421 return build_historical_context(
422 current_datetime=current_datetime,
423 previous_summary_path=previous_summary_path,
424 content_root=content_root,
425 max_words=max_words,
426 prompt_token_budget=prompt_token_budget,
427 prompt_budget_fraction=prompt_budget_fraction,
428 ).markdown
429
430
431 def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
432 parser = argparse.ArgumentParser(
433 description="Assemble bounded historical context for weekly analysis prompts."
434 )
435 parser.add_argument(
436 "--current-datetime", required=True, help="ISO-8601 timestamp for the current analysis run."
437 )
438 parser.add_argument(
439 "--previous-summary",
440 type=Path,
441 default=None,
442 help="Path to the previous week's markdown summary.",
443 )
444 parser.add_argument(
445 "--content-root", type=Path, default=DEFAULT_CONTENT_ROOT, help="Path to the content/ root."
446 )
447 parser.add_argument(
448 "--max-words",
449 type=int,
450 default=DEFAULT_MAX_WORDS,
451 help="Maximum total historical-context words.",
452 )
453 parser.add_argument(
454 "--prompt-token-budget",
455 type=int,
456 default=90_000,
457 help="Total prompt-token budget used to derive the 15%% historical-context ceiling.",
458 )
459 return parser.parse_args(argv)
460
461
462 def main(argv: list[str] | None = None) -> int:
463 args = parse_args(argv)
464 result = build_historical_context(
465 current_datetime=args.current_datetime,
466 previous_summary_path=args.previous_summary,
467 content_root=args.content_root,
468 max_words=args.max_words,
469 prompt_token_budget=args.prompt_token_budget,
470 )
471 if result.markdown:
472 print(result.markdown)
473 return 0
474
475
476 if __name__ == "__main__":
477 raise SystemExit(main())