main
py 144 lines 4.67 KB
Raw
1 #!/usr/bin/env python3
2 """Pre-process raw crawl JSON to reduce token count for analysis prompts.
3
4 Extracts only fields needed by the analysis prompt, truncates descriptions,
5 and computes basic signals to produce a compact JSON suitable for LLM input.
6
7 CLI:
8 python scripts/preprocess_for_analysis.py \
9 --input data/raw/2026-W21.json \
10 --output data/raw/2026-W21-compact.json \
11 --max-desc-length 200
12 """
13
14 from __future__ import annotations
15
16 import argparse
17 import json
18 import sys
19 from datetime import datetime, timezone
20 from pathlib import Path
21
22 from scripts.sanitize_repo_content import sanitize_description
23
24
25 def estimate_tokens(text: str) -> int:
26 """Rough token estimate: characters / 4."""
27 return len(text) // 4
28
29
30 def compute_age_days(created_at: str | None, reference: datetime | None = None) -> int | None:
31 """Compute age in days from created_at ISO timestamp."""
32 if not created_at:
33 return None
34 try:
35 created = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
36 ref = reference or datetime.now(timezone.utc)
37 return max(0, (ref - created).days)
38 except (ValueError, TypeError):
39 return None
40
41
42 def compact_repo(repo: dict, max_desc: int, reference_date: datetime | None = None) -> dict:
43 """Extract and compact a single repo entry."""
44 raw_desc = repo.get("description") or ""
45 desc = sanitize_description(raw_desc, repo=repo, max_length=max_desc)
46 if not isinstance(desc, str):
47 desc = ""
48 return {
49 "name": repo.get("name", ""),
50 "desc": desc,
51 "stars": repo.get("stars", 0),
52 "gained": repo.get("stars_gained", repo.get("gained", 0)),
53 "topics": repo.get("topics", []),
54 "lang": repo.get("language"),
55 "age_days": compute_age_days(repo.get("created_at"), reference_date),
56 }
57
58
59 def compute_signals(repos: list[dict]) -> dict:
60 """Compute aggregate signals from the compacted repo list."""
61 topic_counts: dict[str, int] = {}
62 for r in repos:
63 for t in r.get("topics", []):
64 topic_counts[t] = topic_counts.get(t, 0) + 1
65 top_topics = sorted(topic_counts.items(), key=lambda x: -x[1])[:10]
66 return {"top_topics": [t for t, _ in top_topics]}
67
68
69 def preprocess(data: dict, max_desc: int = 200, reference_date: datetime | None = None) -> dict:
70 """Transform raw crawl JSON into compact analysis format."""
71 original_text = json.dumps(data)
72 original_tokens = estimate_tokens(original_text)
73
74 # Combine new_repos and trending_repos
75 all_repos = data.get("new_repos", []) + data.get("trending_repos", [])
76
77 # Deduplicate by name
78 seen = set()
79 unique_repos = []
80 for r in all_repos:
81 name = r.get("name", "")
82 if name not in seen:
83 seen.add(name)
84 unique_repos.append(r)
85
86 compacted = [compact_repo(r, max_desc, reference_date) for r in unique_repos]
87 signals = compute_signals(compacted)
88
89 result = {
90 "week": data.get("week", ""),
91 "repos": compacted,
92 "signals": signals,
93 }
94
95 compact_text = json.dumps(result)
96 compact_tokens = estimate_tokens(compact_text)
97 reduction_pct = (
98 round((1 - compact_tokens / original_tokens) * 100) if original_tokens > 0 else 0
99 )
100
101 result["stats"] = {
102 "original_tokens_est": original_tokens,
103 "compact_tokens_est": compact_tokens,
104 "reduction_pct": reduction_pct,
105 }
106
107 return result
108
109
110 def main(argv: list[str] | None = None) -> int:
111 parser = argparse.ArgumentParser(description="Pre-process raw JSON for analysis")
112 parser.add_argument("--input", required=True, help="Path to raw crawl JSON")
113 parser.add_argument("--output", help="Output path (default: input with -compact suffix)")
114 parser.add_argument("--max-desc-length", type=int, default=200, help="Max description length")
115 args = parser.parse_args(argv)
116
117 input_path = Path(args.input)
118 if not input_path.exists():
119 print(f"Error: input file not found: {input_path}", file=sys.stderr)
120 return 1
121
122 output_path = (
123 Path(args.output) if args.output else input_path.with_stem(input_path.stem + "-compact")
124 )
125
126 with open(input_path, encoding="utf-8") as f:
127 data = json.load(f)
128
129 result = preprocess(data, max_desc=args.max_desc_length)
130
131 output_path.parent.mkdir(parents=True, exist_ok=True)
132 with open(output_path, "w", encoding="utf-8") as f:
133 json.dump(result, f, indent=2)
134
135 stats = result["stats"]
136 print(f"Preprocessed: {input_path} -> {output_path}")
137 print(
138 f" Tokens: {stats['original_tokens_est']} -> {stats['compact_tokens_est']} ({stats['reduction_pct']}% reduction)"
139 )
140 return 0
141
142
143 if __name__ == "__main__":
144 sys.exit(main())