zzstoatzz.io
search for standard sites
pub-search.waow.tech
search
zig
blog
atproto
31 kB
801 lines
1#!/usr/bin/env -S uv run --script --quiet
2# /// script
3# requires-python = ">=3.12,<3.14"
4# # numpy<2.2 + python<3.14: umap-learn's transitive numba/llvmlite have no wheel
5# # for numpy>=2.2 or py3.14, so an unpinned solve backtracks to ancient numba
6# # 0.53 and tries a source build that fails the py-version guard. keep both pins.
7# dependencies = ["httpx", "numpy<2.2", "scikit-learn", "umap-learn", "hdbscan", "pydantic-settings", "anthropic"]
8# ///
9"""
10Build atlas.json — 2D semantic map of the document index.
11
12Exports vectors from turbopuffer, projects to 2D via PCA+UMAP,
13clusters with HDBSCAN at two granularities, labels via c-TF-IDF.
14
15Usage:
16 ./scripts/build-atlas # writes site/atlas.json
17 ./scripts/build-atlas --output out.json
18"""
19
20import argparse
21import asyncio
22import gc
23import json
24import os
25import re
26import sys
27import time
28from collections import Counter
29from pathlib import Path
30
31import httpx
32import numpy as np
33from pydantic import ValidationError
34from pydantic_settings import BaseSettings, SettingsConfigDict
35from sklearn.decomposition import PCA
36from sklearn.feature_extraction.text import TfidfVectorizer
37
38
39# avatar cache: persists did → avatar URL across runs.
40# lives in site/ so it gets deployed alongside atlas.json — the deployed copy
41# is what the prefect flow (which clones fresh each run) reads on cold start.
42AVATAR_CACHE_PATH = Path(__file__).resolve().parent.parent / "site" / "atlas-avatar-cache.json"
43AVATAR_CACHE_REMOTE = "https://pub-search.waow.tech/atlas-avatar-cache.json"
44AVATAR_CACHE_TTL_SECONDS = 7 * 86400 # refresh entries weekly so changed avatars trickle in
45
46# theme cache: persists basePath → leaflet publication theme colors.
47# same deployment-roundtrip pattern as the avatar cache above.
48THEME_CACHE_PATH = Path(__file__).resolve().parent.parent / "site" / "atlas-theme-cache.json"
49THEME_CACHE_REMOTE = "https://pub-search.waow.tech/atlas-theme-cache.json"
50
51
52def log(msg: str) -> None:
53 print(msg, flush=True)
54
55
56class Settings(BaseSettings):
57 model_config = SettingsConfigDict(
58 env_file=os.environ.get("ENV_FILE", ".env"), extra="ignore"
59 )
60
61 turbopuffer_api_key: str
62 turbopuffer_namespace: str = "leaflet-search"
63 anthropic_api_key: str = "" # optional: enables LLM-refined cluster labels
64 turso_url: str = "" # optional: enables publication data (e.g. https://leaf-....turso.io)
65 turso_token: str = ""
66
67 @classmethod
68 def settings_customise_sources(cls, settings_cls, **kwargs):
69 """Dotenv file wins over environment variables."""
70 return (
71 kwargs["init_settings"],
72 kwargs["dotenv_settings"],
73 kwargs["env_settings"],
74 kwargs["file_secret_settings"],
75 )
76
77
78def tpuf_export(client: httpx.Client, settings: Settings) -> list[dict]:
79 """Export all vectors from turbopuffer via paginated query."""
80 url = f"https://api.turbopuffer.com/v2/namespaces/{settings.turbopuffer_namespace}/query"
81 headers = {
82 "Authorization": f"Bearer {settings.turbopuffer_api_key}",
83 "Content-Type": "application/json",
84 }
85 attrs = ["uri", "title", "platform", "base_path", "path", "did", "vector"]
86 all_rows = []
87 last_id = None
88
89 while True:
90 body: dict = {
91 "rank_by": ["id", "asc"],
92 "limit": 10000,
93 "include_attributes": attrs,
94 }
95 if last_id is not None:
96 body["filters"] = ["id", "Gt", last_id]
97
98 resp = client.post(url, headers=headers, json=body, timeout=120)
99 resp.raise_for_status()
100 rows = resp.json().get("rows", [])
101 if not rows:
102 break
103
104 all_rows.extend(rows)
105 last_id = rows[-1]["id"]
106 log(f" fetched {len(all_rows)} vectors so far...")
107
108 if len(rows) < 10000:
109 break
110
111 return all_rows
112
113
114def turso_query_publications(settings: Settings) -> dict[str, dict]:
115 """Fetch publication metadata from Turso. Returns basePath → {name, did, coverImage}."""
116 if not settings.turso_url or not settings.turso_token:
117 return {}
118
119 host = settings.turso_url.replace("https://", "").replace("http://", "").replace("libsql://", "")
120 url = f"https://{host}/v3/pipeline"
121 headers = {"Authorization": f"Bearer {settings.turso_token}"}
122
123 body = {
124 "requests": [
125 {
126 "type": "execute",
127 "stmt": {
128 "sql": (
129 "SELECT p.name, p.base_path, p.did, c.cover_image "
130 "FROM publications p "
131 "LEFT JOIN ("
132 " SELECT publication_uri, MIN(cover_image) as cover_image "
133 " FROM documents WHERE cover_image IS NOT NULL "
134 " GROUP BY publication_uri"
135 ") c ON c.publication_uri = p.uri"
136 )
137 },
138 },
139 {"type": "close"},
140 ]
141 }
142
143 try:
144 resp = httpx.post(url, headers=headers, json=body, timeout=30)
145 resp.raise_for_status()
146 results = resp.json().get("results", [])
147 if not results or "response" not in results[0]:
148 return {}
149 rows = results[0]["response"]["result"]["rows"]
150 cols = [c["name"] for c in results[0]["response"]["result"]["cols"]]
151 except httpx.HTTPError as e:
152 log(f" warning: turso request failed: {e}")
153 return {}
154 except (KeyError, IndexError) as e:
155 log(f" warning: unexpected turso response: {e}")
156 return {}
157 lookup = {}
158 for row in rows:
159 values = {cols[i]: (row[i]["value"] if row[i]["type"] != "null" else None) for i in range(len(cols))}
160 bp = values.get("base_path")
161 if bp:
162 lookup[bp] = {
163 "name": values.get("name") or bp,
164 "did": values.get("did") or "",
165 "coverImage": values.get("cover_image") or "",
166 }
167 return lookup
168
169
170def extract_terms(title: str) -> list[str]:
171 """Extract lowercase alphanumeric tokens from a title."""
172 if not title:
173 return []
174 return [t for t in re.findall(r"[a-z0-9]+", title.lower()) if len(t) > 1]
175
176
177def cluster_labels(titles_per_cluster: dict[int, list[str]], n_terms: int = 3) -> dict[int, str]:
178 """Compute c-TF-IDF labels for clusters from document titles."""
179 cluster_ids = sorted(titles_per_cluster.keys())
180 if not cluster_ids:
181 return {}
182
183 # build one "document" per cluster: concatenated titles
184 docs = []
185 for cid in cluster_ids:
186 docs.append(" ".join(titles_per_cluster[cid]))
187
188 # TF-IDF across cluster pseudo-documents
189 vectorizer = TfidfVectorizer(
190 max_features=5000,
191 stop_words="english",
192 token_pattern=r"[a-z0-9]{2,}",
193 lowercase=True,
194 )
195 try:
196 tfidf = vectorizer.fit_transform(docs)
197 except ValueError:
198 return {cid: f"cluster {cid}" for cid in cluster_ids}
199
200 feature_names = vectorizer.get_feature_names_out()
201 labels = {}
202 for i, cid in enumerate(cluster_ids):
203 row = tfidf[i].toarray().flatten()
204 top_idx = row.argsort()[-n_terms:][::-1]
205 top_terms = [feature_names[j] for j in top_idx if row[j] > 0]
206 labels[cid] = " ".join(top_terms) if top_terms else f"cluster {cid}"
207
208 return labels
209
210
211def llm_refine_labels(
212 tfidf_labels: dict[int, str],
213 titles_per_cluster: dict[int, list[str]],
214 cluster_counts: dict[int, int],
215 api_key: str,
216 tier: str = "coarse",
217 batch_size: int = 15,
218) -> dict[int, str]:
219 """Refine c-TF-IDF labels with Claude Haiku for human-readable topic names."""
220 from anthropic import AsyncAnthropic
221
222 client = AsyncAnthropic(api_key=api_key)
223 rng = np.random.default_rng(42)
224
225 cluster_ids = sorted(tfidf_labels.keys())
226 batches = [cluster_ids[i : i + batch_size] for i in range(0, len(cluster_ids), batch_size)]
227
228 # Split fixed instructions (system, cached) from variable per-batch payload
229 # (user). Each rebuild runs ~18 batches (5 coarse + 13 fine); the
230 # instructions are identical across all of them, so caching the system
231 # block drops input cost to ~10% on every batch after the first within
232 # the 5m TTL.
233 system_instructions = (
234 "You are labeling topic clusters for a document atlas — a 2D map where each cluster "
235 "represents a group of semantically similar blog posts and articles.\n\n"
236 "For each cluster you receive, write a short descriptive label (2-4 lowercase words) "
237 "that captures the theme. The label should read like a topic name on a map.\n\n"
238 "Good labels: \"ai language models\", \"bluesky development\", \"personal journals\", "
239 "\"spanish poetry\", \"game reviews\", \"home cooking\"\n"
240 "Bad labels: \"ai llms llm\", \"results april june\", \"39 let english\"\n\n"
241 "Respond with ONLY a JSON object mapping cluster ID (as string) to label. No other text."
242 )
243
244 async def label_batch(batch: list[int]) -> dict[int, str]:
245 parts = []
246 for cid in batch:
247 titles = titles_per_cluster.get(cid, [])
248 sample = list(rng.choice(titles, size=min(12, len(titles)), replace=False)) if titles else []
249 parts.append(
250 f"Cluster {cid} ({cluster_counts.get(cid, 0)} docs):\n"
251 f" Keywords: {tfidf_labels[cid]}\n"
252 f" Sample titles: {', '.join(repr(t) for t in sample)}"
253 )
254
255 resp = await client.messages.create(
256 model="claude-haiku-4-5-20251001",
257 max_tokens=2048,
258 system=[
259 {
260 "type": "text",
261 "text": system_instructions,
262 "cache_control": {"type": "ephemeral"},
263 }
264 ],
265 messages=[{"role": "user", "content": "\n\n".join(parts)}],
266 )
267 text = resp.content[0].text.strip()
268 # extract JSON from response (handle markdown code fences)
269 if text.startswith("```"):
270 text = text.split("\n", 1)[1].rsplit("```", 1)[0].strip()
271 try:
272 parsed = json.loads(text)
273 return {int(k): str(v) for k, v in parsed.items()}
274 except (json.JSONDecodeError, ValueError):
275 log(f" warning: failed to parse LLM response for batch, using TF-IDF fallback")
276 return {}
277
278 async def run_all():
279 results = await asyncio.gather(*[label_batch(b) for b in batches])
280 merged = {}
281 for r in results:
282 merged.update(r)
283 return merged
284
285 log(f" refining {tier} labels with LLM ({len(batches)} batches, {len(cluster_ids)} clusters)...")
286 refined = asyncio.run(run_all())
287
288 # merge: use LLM labels where available, fall back to c-TF-IDF
289 final = {}
290 for cid in cluster_ids:
291 final[cid] = refined.get(cid, tfidf_labels[cid])
292 log(f" {tier}: {len(refined)}/{len(cluster_ids)} labels refined by LLM")
293 return final
294
295
296def assign_outliers(coords_2d: np.ndarray, labels: np.ndarray, centroids: dict[int, np.ndarray]) -> np.ndarray:
297 """Assign outlier points (label == -1) to nearest cluster centroid."""
298 result = labels.copy()
299 outlier_mask = result == -1
300 if not outlier_mask.any() or not centroids:
301 return result
302
303 centroid_ids = sorted(centroids.keys())
304 centroid_arr = np.array([centroids[c] for c in centroid_ids])
305 outlier_coords = coords_2d[outlier_mask]
306
307 # nearest centroid for each outlier
308 dists = np.linalg.norm(outlier_coords[:, None] - centroid_arr[None, :], axis=2)
309 nearest = dists.argmin(axis=1)
310 result[outlier_mask] = np.array(centroid_ids)[nearest]
311
312 return result
313
314
315def load_avatar_cache(local_path: Path, remote_url: str) -> dict[str, dict]:
316 """Load avatar cache. Prefer local file (dev re-runs); fall back to the
317 previously-deployed copy on the CDN (prefect flow clones fresh each run).
318 Returns {} on any failure — caller treats that as a cold start."""
319 if local_path.exists():
320 try:
321 data = json.loads(local_path.read_text())
322 log(f" loaded {len(data)} avatar cache entries from {local_path.name}")
323 return data
324 except (json.JSONDecodeError, OSError) as e:
325 log(f" warning: local avatar cache unreadable ({e}), trying remote")
326
327 try:
328 resp = httpx.get(remote_url, timeout=15)
329 if resp.status_code == 200:
330 data = resp.json()
331 if isinstance(data, dict):
332 log(f" loaded {len(data)} avatar cache entries from {remote_url}")
333 return data
334 log(f" warning: remote avatar cache has unexpected shape, starting cold")
335 else:
336 log(f" remote avatar cache returned {resp.status_code}, starting cold")
337 except (httpx.HTTPError, json.JSONDecodeError) as e:
338 log(f" warning: failed to fetch remote avatar cache ({e}), starting cold")
339
340 return {}
341
342
343def save_avatar_cache(path: Path, entries: dict[str, dict]) -> None:
344 path.parent.mkdir(parents=True, exist_ok=True)
345 path.write_text(json.dumps(entries, separators=(",", ":"), sort_keys=True))
346
347
348async def resolve_avatars_async(
349 dids: list[str], cache: dict[str, dict], concurrency: int = 20
350) -> dict[str, dict]:
351 """Resolve bsky avatars for the given DIDs.
352
353 Uses cached entries within TTL, fetches the rest concurrently. Returns the
354 new cache state — exactly the DIDs passed in, so unused entries from prior
355 runs naturally drop out and the file stays roughly the size of the corpus.
356 Empty-string avatars are cached too, so we don't hammer the API every run
357 for authors who have no avatar set."""
358 now = time.time()
359 cutoff = now - AVATAR_CACHE_TTL_SECONDS
360
361 fresh: dict[str, dict] = {}
362 missing: list[str] = []
363 for did in dids:
364 entry = cache.get(did)
365 if isinstance(entry, dict) and float(entry.get("ts", 0)) > cutoff:
366 fresh[did] = entry
367 else:
368 missing.append(did)
369
370 log(f" cache: {len(fresh)} hits, {len(missing)} to fetch")
371 if not missing:
372 return fresh
373
374 sem = asyncio.Semaphore(concurrency)
375 async with httpx.AsyncClient(
376 timeout=httpx.Timeout(10.0),
377 limits=httpx.Limits(
378 max_connections=concurrency,
379 max_keepalive_connections=concurrency,
380 ),
381 ) as client:
382 async def fetch_one(did: str) -> tuple[str, str]:
383 async with sem:
384 try:
385 r = await client.get(
386 f"https://public.api.bsky.app/xrpc/app.bsky.actor.getProfile?actor={did}"
387 )
388 if r.status_code == 200:
389 return did, (r.json().get("avatar") or "")
390 except Exception:
391 pass
392 return did, ""
393
394 results = await asyncio.gather(*(fetch_one(d) for d in missing))
395
396 for did, avatar in results:
397 fresh[did] = {"avatar": avatar, "ts": now}
398
399 return fresh
400
401
402def _theme_color_hex(color: dict | None) -> str:
403 """pub.leaflet.theme.color#rgb{a} object → #rrggbb, or '' if absent."""
404 if not isinstance(color, dict):
405 return ""
406 try:
407 return "#%02x%02x%02x" % (int(color["r"]), int(color["g"]), int(color["b"]))
408 except (KeyError, ValueError, TypeError):
409 return ""
410
411
412async def resolve_themes_async(
413 pubs: list[dict], cache: dict[str, dict], concurrency: int = 16
414) -> dict[str, dict]:
415 """Resolve leaflet publication theme colors from their PDS records.
416
417 pub.leaflet.publication records carry the author's actual palette
418 (backgroundColor / accentBackground); the atlas uses it to paint that
419 pub's planets. One listRecords per DID covers all of its publications.
420 Empty results are cached too. Returns {basePath: {bg, accent, ts}} for
421 exactly the leaflet pubs passed in."""
422 now = time.time()
423 cutoff = now - AVATAR_CACHE_TTL_SECONDS
424
425 targets = [p for p in pubs if p.get("platform") == "leaflet" and str(p.get("did", "")).startswith("did:plc:")]
426 fresh: dict[str, dict] = {}
427 by_did: dict[str, list[str]] = {}
428 for p in targets:
429 entry = cache.get(p["basePath"])
430 if isinstance(entry, dict) and float(entry.get("ts", 0)) > cutoff:
431 fresh[p["basePath"]] = entry
432 else:
433 by_did.setdefault(p["did"], []).append(p["basePath"])
434
435 log(f" theme cache: {len(fresh)} hits, {len(by_did)} DIDs to fetch")
436 if not by_did:
437 return fresh
438
439 sem = asyncio.Semaphore(concurrency)
440 async with httpx.AsyncClient(
441 timeout=httpx.Timeout(10.0),
442 limits=httpx.Limits(
443 max_connections=concurrency,
444 max_keepalive_connections=concurrency,
445 ),
446 ) as client:
447 async def fetch_did(did: str, base_paths: list[str]) -> dict[str, dict]:
448 out = {bp: {"bg": "", "accent": "", "ts": now} for bp in base_paths}
449 async with sem:
450 try:
451 doc = (await client.get(f"https://plc.directory/{did}")).json()
452 pds = next(
453 s["serviceEndpoint"] for s in doc.get("service", [])
454 if s.get("id") == "#atproto_pds"
455 )
456 r = await client.get(
457 f"{pds}/xrpc/com.atproto.repo.listRecords",
458 params={"repo": did, "collection": "pub.leaflet.publication", "limit": 50},
459 )
460 if r.status_code != 200:
461 return out
462 for rec in r.json().get("records", []):
463 value = rec.get("value", {})
464 bp = value.get("base_path", "")
465 if bp not in out:
466 continue
467 theme = value.get("theme") or {}
468 out[bp] = {
469 "bg": _theme_color_hex(theme.get("backgroundColor") or theme.get("pageBackground")),
470 "accent": _theme_color_hex(theme.get("accentBackground") or theme.get("primary")),
471 "ts": now,
472 }
473 except Exception:
474 pass
475 return out
476
477 results = await asyncio.gather(*(fetch_did(d, bps) for d, bps in by_did.items()))
478
479 for chunk in results:
480 fresh.update(chunk)
481 return fresh
482
483
484def main():
485 parser = argparse.ArgumentParser(description="Build atlas.json")
486 parser.add_argument(
487 "--output", "-o",
488 default=str(Path(__file__).resolve().parent.parent / "site" / "atlas.json"),
489 help="Output path (default: site/atlas.json)",
490 )
491 args = parser.parse_args()
492
493 try:
494 settings = Settings() # type: ignore
495 except ValidationError as e:
496 # never print the exception directly: pydantic's str(ValidationError)
497 # embeds input_value — the full dict of ingested env/dotenv vars,
498 # including unrelated secrets — which would land in the prefect logs.
499 # surface only the offending field names, never their values.
500 fields = ", ".join(str(err["loc"][0]) for err in e.errors()) or "config"
501 print(f"error loading settings; missing/invalid: {fields}", file=sys.stderr)
502 print("required env vars: TURBOPUFFER_API_KEY", file=sys.stderr)
503 sys.exit(1)
504
505 t_start = time.monotonic()
506
507 # --- step 1: export vectors from turbopuffer ---
508 log("exporting vectors from turbopuffer...")
509 client = httpx.Client()
510 rows = tpuf_export(client, settings)
511 client.close()
512 log(f" got {len(rows)} documents")
513
514 if len(rows) < 10:
515 print("too few documents to build atlas", file=sys.stderr)
516 sys.exit(1)
517
518 # Parse into a dense float32 matrix without retaining a second Python list
519 # of vector objects. Turbopuffer returns vectors as JSON lists; holding
520 # rows + vectors + X at once was enough to OOM the Prefect pod as the
521 # corpus grew.
522 valid_rows = [row for row in rows if row.get("vector")]
523 first_vec = valid_rows[0]["vector"]
524 X = np.empty((len(valid_rows), len(first_vec)), dtype=np.float32)
525 metadata = []
526 for i, row in enumerate(valid_rows):
527 vec = row.pop("vector")
528 X[i] = vec
529 metadata.append({
530 "uri": row.get("uri", row.get("id", "")),
531 "title": row.get("title", ""),
532 "platform": row.get("platform", "other"),
533 "basePath": row.get("base_path", ""),
534 "path": row.get("path", ""),
535 "did": row.get("did", ""),
536 })
537 del rows, valid_rows
538 gc.collect()
539 log(f" {X.shape[0]} vectors, {X.shape[1]} dims")
540
541 # --- step 2: dimensionality reduction ---
542 log("PCA 1024 → 50...")
543 n_components_pca = min(50, X.shape[0] - 1, X.shape[1])
544 pca = PCA(n_components=n_components_pca, random_state=42)
545 X_pca = pca.fit_transform(X)
546 log(f" variance explained: {pca.explained_variance_ratio_.sum():.2%}")
547
548 log("UMAP 50 → 2...")
549 import umap
550 reducer = umap.UMAP(
551 n_components=2,
552 metric="cosine",
553 n_neighbors=15,
554 min_dist=0.1,
555 random_state=42,
556 )
557 X_2d = reducer.fit_transform(X_pca)
558
559 # normalize to [-1, 1] range
560 for dim in range(2):
561 lo, hi = X_2d[:, dim].min(), X_2d[:, dim].max()
562 if hi > lo:
563 X_2d[:, dim] = 2 * (X_2d[:, dim] - lo) / (hi - lo) - 1
564
565 log(f" 2D range: x=[{X_2d[:,0].min():.2f}, {X_2d[:,0].max():.2f}] y=[{X_2d[:,1].min():.2f}, {X_2d[:,1].max():.2f}]")
566
567 # --- step 3: clustering ---
568 import hdbscan
569
570 log("HDBSCAN coarse (min_cluster_size=50)...")
571 clusterer_coarse = hdbscan.HDBSCAN(min_cluster_size=50, min_samples=10)
572 labels_coarse_raw = clusterer_coarse.fit_predict(X_2d)
573 n_coarse = len(set(labels_coarse_raw)) - (1 if -1 in labels_coarse_raw else 0)
574 n_outliers_coarse = (labels_coarse_raw == -1).sum()
575 log(f" {n_coarse} clusters, {n_outliers_coarse} outliers")
576
577 # compute centroids for coarse
578 coarse_centroids = {}
579 for cid in set(labels_coarse_raw):
580 if cid == -1:
581 continue
582 mask = labels_coarse_raw == cid
583 coarse_centroids[cid] = X_2d[mask].mean(axis=0)
584
585 labels_coarse = assign_outliers(X_2d, labels_coarse_raw, coarse_centroids)
586 # recompute centroids after outlier assignment
587 for cid in set(labels_coarse):
588 mask = labels_coarse == cid
589 coarse_centroids[cid] = X_2d[mask].mean(axis=0)
590
591 log("HDBSCAN fine (min_cluster_size=20)...")
592 clusterer_fine = hdbscan.HDBSCAN(min_cluster_size=20, min_samples=3)
593 labels_fine_raw = clusterer_fine.fit_predict(X_2d)
594 n_fine = len(set(labels_fine_raw)) - (1 if -1 in labels_fine_raw else 0)
595 n_outliers_fine = (labels_fine_raw == -1).sum()
596 log(f" {n_fine} clusters, {n_outliers_fine} outliers")
597
598 fine_centroids = {}
599 for cid in set(labels_fine_raw):
600 if cid == -1:
601 continue
602 mask = labels_fine_raw == cid
603 fine_centroids[cid] = X_2d[mask].mean(axis=0)
604
605 labels_fine = assign_outliers(X_2d, labels_fine_raw, fine_centroids)
606 for cid in set(labels_fine):
607 mask = labels_fine == cid
608 fine_centroids[cid] = X_2d[mask].mean(axis=0)
609
610 # --- step 4: labels via c-TF-IDF + optional LLM refinement ---
611 log("computing cluster labels (c-TF-IDF on titles)...")
612
613 coarse_titles: dict[int, list[str]] = {}
614 fine_titles: dict[int, list[str]] = {}
615 for i, meta in enumerate(metadata):
616 title = meta.get("title", "")
617 if title:
618 coarse_titles.setdefault(int(labels_coarse[i]), []).append(title)
619 fine_titles.setdefault(int(labels_fine[i]), []).append(title)
620
621 coarse_labels = cluster_labels(coarse_titles)
622 fine_labels = cluster_labels(fine_titles)
623
624 if settings.anthropic_api_key:
625 log("refining labels with LLM...")
626 coarse_counts = {int(cid): int((labels_coarse == cid).sum()) for cid in set(labels_coarse)}
627 fine_counts = {int(cid): int((labels_fine == cid).sum()) for cid in set(labels_fine)}
628 try:
629 coarse_labels = llm_refine_labels(
630 coarse_labels, coarse_titles, coarse_counts,
631 settings.anthropic_api_key, tier="coarse",
632 )
633 fine_labels = llm_refine_labels(
634 fine_labels, fine_titles, fine_counts,
635 settings.anthropic_api_key, tier="fine",
636 )
637 except Exception as e:
638 log(f" LLM labeling failed, using c-TF-IDF fallback: {e}")
639 else:
640 log(" (no ANTHROPIC_API_KEY — using c-TF-IDF labels only)")
641
642 # map fine clusters to their parent coarse cluster (majority vote)
643 fine_to_coarse = {}
644 for fine_id in set(labels_fine):
645 fine_mask = labels_fine == fine_id
646 coarse_for_fine = labels_coarse[fine_mask]
647 counts = Counter(coarse_for_fine)
648 fine_to_coarse[int(fine_id)] = int(counts.most_common(1)[0][0])
649
650 log(f" coarse: {len(coarse_labels)} labels")
651 log(f" fine: {len(fine_labels)} labels")
652
653 # --- step 5: publication centroids ---
654 log("computing publication centroids...")
655 pub_lookup = turso_query_publications(settings)
656 if pub_lookup:
657 log(f" fetched {len(pub_lookup)} publications from turso")
658 else:
659 log(" (no TURSO_URL/TURSO_TOKEN — using basePath grouping only)")
660
661 # group documents by basePath → list of 2D positions
662 bp_groups: dict[str, list[int]] = {}
663 for i, meta in enumerate(metadata):
664 bp = meta.get("basePath", "")
665 if bp:
666 bp_groups.setdefault(bp, []).append(i)
667
668 publications = []
669 for bp, indices in bp_groups.items():
670 if len(indices) < 2:
671 continue
672 coords = X_2d[indices]
673 centroid = coords.mean(axis=0)
674 info = pub_lookup.get(bp, {})
675 # get DID from turso lookup first, fall back to documents
676 did = info.get("did", "")
677 if not did:
678 for idx in indices:
679 d = metadata[idx].get("did", "")
680 if d:
681 did = d
682 break
683 pub = {
684 "name": info.get("name") or bp,
685 "basePath": bp,
686 "did": did,
687 "cx": round(float(centroid[0]), 4),
688 "cy": round(float(centroid[1]), 4),
689 "count": len(indices),
690 "coverImage": info.get("coverImage", ""),
691 }
692 # determine dominant platform for this publication
693 plat_counts: dict[str, int] = {}
694 for idx in indices:
695 p = metadata[idx].get("platform", "other")
696 plat_counts[p] = plat_counts.get(p, 0) + 1
697 pub["platform"] = max(plat_counts, key=plat_counts.get) # type: ignore
698 publications.append(pub)
699
700 publications.sort(key=lambda p: p["count"], reverse=True)
701 log(f" {len(publications)} publications with 2+ documents")
702
703 # fetch author avatars for publications (fallback when no cover image).
704 # see resolve_avatars_async — concurrent fetch + persistent cache keep this
705 # step at ~5s steady-state instead of ~85s serial.
706 unique_dids = {p["did"] for p in publications if p["did"]}
707 log(f" resolving avatars for {len(unique_dids)} unique authors...")
708
709 cache = load_avatar_cache(AVATAR_CACHE_PATH, AVATAR_CACHE_REMOTE)
710 entries = asyncio.run(resolve_avatars_async(sorted(unique_dids), cache))
711 save_avatar_cache(AVATAR_CACHE_PATH, entries)
712
713 did_avatars = {did: e["avatar"] for did, e in entries.items() if e.get("avatar")}
714 log(f" resolved {len(did_avatars)} avatars ({len(entries) - len(did_avatars)} authors have none)")
715
716 # attach avatar URLs to publications
717 for pub in publications:
718 pub["avatar"] = did_avatars.get(pub["did"], "")
719
720 # fetch leaflet publication theme colors — the author's actual palette,
721 # used by the atlas to paint that publication's planets
722 log(" resolving leaflet publication themes...")
723 theme_cache = load_avatar_cache(THEME_CACHE_PATH, THEME_CACHE_REMOTE)
724 theme_entries = asyncio.run(resolve_themes_async(publications, theme_cache))
725 save_avatar_cache(THEME_CACHE_PATH, theme_entries)
726 themed = 0
727 for pub in publications:
728 t = theme_entries.get(pub["basePath"]) or {}
729 if t.get("accent"):
730 pub["themeAccent"] = t["accent"]
731 if t.get("bg"):
732 pub["themeBg"] = t["bg"]
733 themed += 1
734 log(f" {themed} publications have theme colors")
735
736 # --- step 6: build output ---
737 log("building output...")
738
739 points = []
740 for i, meta in enumerate(metadata):
741 points.append({
742 "x": round(float(X_2d[i, 0]), 4),
743 "y": round(float(X_2d[i, 1]), 4),
744 "uri": meta["uri"],
745 "title": meta["title"],
746 "platform": meta["platform"],
747 "basePath": meta["basePath"],
748 "path": meta["path"],
749 "clusterCoarse": int(labels_coarse[i]),
750 "clusterFine": int(labels_fine[i]),
751 })
752
753 coarse_clusters = []
754 for cid in sorted(coarse_centroids.keys()):
755 count = int((labels_coarse == cid).sum())
756 coarse_clusters.append({
757 "id": int(cid),
758 "label": coarse_labels.get(int(cid), f"cluster {cid}"),
759 "cx": round(float(coarse_centroids[cid][0]), 4),
760 "cy": round(float(coarse_centroids[cid][1]), 4),
761 "count": count,
762 })
763
764 fine_clusters = []
765 for cid in sorted(fine_centroids.keys()):
766 count = int((labels_fine == cid).sum())
767 fine_clusters.append({
768 "id": int(cid),
769 "label": fine_labels.get(int(cid), f"cluster {cid}"),
770 "cx": round(float(fine_centroids[cid][0]), 4),
771 "cy": round(float(fine_centroids[cid][1]), 4),
772 "count": count,
773 "parent": fine_to_coarse.get(int(cid), 0),
774 })
775
776 output = {
777 "points": points,
778 "clusters": {
779 "coarse": coarse_clusters,
780 "fine": fine_clusters,
781 },
782 "publications": publications,
783 "meta": {
784 "generatedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
785 "nDocuments": len(points),
786 "nPublications": len(publications),
787 },
788 }
789
790 out_path = Path(args.output)
791 out_path.parent.mkdir(parents=True, exist_ok=True)
792 out_path.write_text(json.dumps(output, separators=(",", ":")))
793
794 size_kb = out_path.stat().st_size / 1024
795 elapsed = time.monotonic() - t_start
796 log(f"\nwrote {out_path} ({size_kb:.0f} KB, {len(points)} points)")
797 log(f"done in {elapsed:.1f}s")
798
799
800if __name__ == "__main__":
801 main()