DriptoBhattacharyya commited on
Commit
6e8672b
·
verified ·
1 Parent(s): 1cb8456

Deploy AstraNexus (Docker: fine-tuned encoder + custom FastAPI frontend)

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. Dockerfile +15 -0
  2. README.md +20 -13
  3. astranexus/__init__.py +0 -0
  4. astranexus/assets/vis_canvas.html +157 -0
  5. astranexus/cluster/__init__.py +0 -0
  6. astranexus/cluster/cluster.py +36 -0
  7. astranexus/cluster/embed.py +99 -0
  8. astranexus/cluster/enrich.py +30 -0
  9. astranexus/cluster/eval.py +8 -0
  10. astranexus/cluster/ft_encoder.py +138 -0
  11. astranexus/config.py +26 -0
  12. astranexus/core/__init__.py +0 -0
  13. astranexus/core/schema.py +54 -0
  14. astranexus/graph.py +232 -0
  15. astranexus/labeler/__init__.py +0 -0
  16. astranexus/labeler/fallback.py +47 -0
  17. astranexus/labeler/hf_serverless.py +60 -0
  18. astranexus/labeler/nemotron.py +191 -0
  19. astranexus/lg_studio.py +31 -0
  20. astranexus/pipeline.py +60 -0
  21. astranexus/serialize/__init__.py +0 -0
  22. astranexus/serialize/canvas.py +30 -0
  23. astranexus/serialize/render.py +30 -0
  24. astranexus/sources/__init__.py +0 -0
  25. astranexus/sources/base.py +16 -0
  26. astranexus/sources/gmail_api.py +183 -0
  27. astranexus/sources/imap_gmail.py +204 -0
  28. astranexus/sources/mock_json.py +22 -0
  29. astranexus/ui/__init__.py +0 -0
  30. astranexus/ui/app.py +247 -0
  31. astranexus/web/__init__.py +0 -0
  32. astranexus/web/google_oauth.py +122 -0
  33. astranexus/web/modal_client.py +52 -0
  34. data/images/msg0000.png +0 -0
  35. data/images/msg0002.png +0 -0
  36. data/images/msg0004.png +0 -0
  37. data/images/msg0005.png +0 -0
  38. data/images/msg0007.png +0 -0
  39. data/images/msg0012.png +0 -0
  40. data/images/msg0014.png +0 -0
  41. data/images/msg0015.png +0 -0
  42. data/images/msg0016.png +0 -0
  43. data/images/msg0017.png +0 -0
  44. data/images/msg0021.png +0 -0
  45. data/images/msg0022.png +0 -0
  46. data/images/msg0024.png +0 -0
  47. data/images/msg0034.png +0 -0
  48. data/images/msg0035.png +0 -0
  49. data/images/msg0038.png +0 -0
  50. data/images/msg0043.png +0 -0
Dockerfile ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.11-slim
2
+ RUN apt-get update && apt-get install -y --no-install-recommends build-essential \
3
+ && rm -rf /var/lib/apt/lists/*
4
+ RUN useradd -m -u 1000 user
5
+ USER user
6
+ ENV HOME=/home/user PATH=/home/user/.local/bin:$PATH \
7
+ ASTRANEXUS_FINETUNED=1 HF_HOME=/home/user/.cache/huggingface
8
+ WORKDIR /home/user/app
9
+ COPY --chown=user requirements.txt .
10
+ # CPU torch first (avoid the multi-GB CUDA wheel), then the rest.
11
+ RUN pip install --no-cache-dir --user torch torchvision --index-url https://download.pytorch.org/whl/cpu
12
+ RUN pip install --no-cache-dir --user -r requirements.txt
13
+ COPY --chown=user . .
14
+ EXPOSE 7860
15
+ CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "7860"]
README.md CHANGED
@@ -1,13 +1,20 @@
1
- ---
2
- title: AstraNexus
3
- emoji: 🐨
4
- colorFrom: indigo
5
- colorTo: indigo
6
- sdk: gradio
7
- sdk_version: 6.18.0
8
- python_version: '3.13'
9
- app_file: app.py
10
- pinned: false
11
- ---
12
-
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: AstraNexus
3
+ emoji: "🌌"
4
+ colorFrom: indigo
5
+ colorTo: purple
6
+ sdk: docker
7
+ app_port: 7860
8
+ pinned: false
9
+ short_description: Inbox constellations via a fine-tuned multimodal encoder
10
+ ---
11
+
12
+ # ✦ AstraNexus
13
+
14
+ Unsupervised email organizer. A **fine-tuned** multimodal encoder (MiniLM text +
15
+ SigLIP vision) embeds each email, HDBSCAN clusters them into **constellations**,
16
+ cross-account threads show as **wormholes**. Custom FastAPI frontend (landing at
17
+ `/`, Gradio app mounted at `/app`); LangGraph orchestration emits `trace.json`;
18
+ runs fully local (no cloud LLM APIs on the default path).
19
+
20
+ Fine-tuned encoder: [DriptoBhattacharyya/astranexus-mm-encoder](https://huggingface.co/DriptoBhattacharyya/astranexus-mm-encoder).
astranexus/__init__.py ADDED
File without changes
astranexus/assets/vis_canvas.html ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <!DOCTYPE html>
2
+ <html>
3
+ <head>
4
+ <meta charset="utf-8"/>
5
+ <script src="https://unpkg.com/vis-network/standalone/umd/vis-network.min.js"></script>
6
+ <style>
7
+ html,body{margin:0;background:#0a0e17;font-family:system-ui,Segoe UI,Roboto,sans-serif;color:#c0caf5;}
8
+ #wrap{position:relative;width:100%;height:620px;}
9
+ #net{width:100%;height:100%;}
10
+ .panel{position:absolute;background:rgba(13,19,32,.92);border:1px solid #29304a;
11
+ border-radius:10px;padding:12px 14px;font-size:12px;backdrop-filter:blur(4px);}
12
+ #legend{top:10px;left:10px;max-width:230px;max-height:90%;overflow:auto;}
13
+ #legend h4{margin:0 0 6px;font-size:12px;color:#7aa2f7;letter-spacing:.04em;}
14
+ .leg-row{display:flex;align-items:center;gap:7px;margin:3px 0;cursor:pointer;}
15
+ .leg-row:hover{color:#fff;}
16
+ .sw{width:11px;height:11px;border-radius:50%;flex:0 0 auto;}
17
+ .leg-name{flex:1;white-space:nowrap;overflow:hidden;text-overflow:ellipsis;}
18
+ .leg-cnt{opacity:.6;}
19
+ .key{margin-top:8px;padding-top:8px;border-top:1px solid #29304a;font-size:11px;opacity:.85;line-height:1.6;}
20
+ #detail{top:10px;right:10px;width:280px;display:none;}
21
+ #detail .x{float:right;cursor:pointer;opacity:.6;font-size:14px;}
22
+ #detail .x:hover{opacity:1;}
23
+ #detail .subj{font-size:13px;font-weight:600;color:#fff;margin:0 16px 8px 0;line-height:1.35;}
24
+ #detail .row{margin:3px 0;}
25
+ #detail .lbl{color:#7aa2f7;}
26
+ #detail .snip{margin-top:8px;padding-top:8px;border-top:1px solid #29304a;color:#a9b1d6;line-height:1.45;}
27
+ .chip{display:inline-block;padding:1px 7px;border-radius:8px;font-size:10px;margin-left:4px;}
28
+ </style>
29
+ </head>
30
+ <body>
31
+ <div id="wrap">
32
+ <div id="net"></div>
33
+ <div id="legend" class="panel"></div>
34
+ <div id="detail" class="panel"></div>
35
+ </div>
36
+ <script>
37
+ const G = __GRAPH_JSON__;
38
+ const HUE = g => (g * 47) % 360;
39
+ const colorFor = g => `hsl(${HUE(g)},70%,55%)`;
40
+ const clusterById = {}; (G.clusters||[]).forEach(c => clusterById[c.id] = c);
41
+ const detailById = {};
42
+
43
+ const nodes = [];
44
+ G.nodes.forEach(n => {
45
+ detailById[n.id] = n;
46
+ const light = 35 + Math.round((n.urgency || 0) * 45);
47
+ nodes.push({
48
+ id: n.id, label: "", title: n.label || "(no subject)",
49
+ // shape MUST be "dot" for size to apply — the default "ellipse" sizes to the
50
+ // (empty) label and ignores size/value, which made all stars look equal.
51
+ shape: "dot", size: n.unread ? 24 : 11,
52
+ color: {background: `hsl(${HUE(n.group)},70%,${light}%)`,
53
+ border: n.has_image ? "#e0af68" : "#1a1f2e"},
54
+ borderWidth: n.has_image ? 4 : 1,
55
+ shadow: n.has_image ? {enabled: true, color: "#e0af68", size: 12} : false,
56
+ font: {color: "#c0caf5", size: 10}
57
+ });
58
+ });
59
+
60
+ const edges = [];
61
+ (G.clusters || []).forEach(c => {
62
+ if (c.id < 0) return; // noise: no hub
63
+ nodes.push({
64
+ id: "hub_" + c.id, label: c.name || ("Cluster " + c.id),
65
+ title: `${c.name || "Cluster " + c.id} — ${c.count} emails`,
66
+ shape: "star", size: 22, value: 22,
67
+ color: {background: colorFor(c.id), border: "#c0caf5"},
68
+ font: {color: "#ffffff", size: 16, strokeWidth: 3, strokeColor: "#0a0e17"}
69
+ });
70
+ });
71
+ // faint spokes pull each constellation's members around its named hub
72
+ G.nodes.forEach(n => {
73
+ if (n.group >= 0)
74
+ edges.push({from: "hub_" + n.group, to: n.id, width: 1, smooth: false,
75
+ color: {color: `hsla(${HUE(n.group)},60%,50%,0.22)`}});
76
+ });
77
+ // wormholes = same thread across different accounts
78
+ (G.edges || []).forEach(e => {
79
+ if (e.kind === "wormhole")
80
+ edges.push({from: e.src, to: e.dst, color: "#f7768e", dashes: true,
81
+ width: 2, smooth: {type: "curvedCW", roundness: 0.2}});
82
+ });
83
+
84
+ const data = {nodes: new vis.DataSet(nodes), edges: new vis.DataSet(edges)};
85
+ const opts = {
86
+ physics: {stabilization: true,
87
+ barnesHut: {gravitationalConstant: -4500, springLength: 110, springConstant: 0.03}},
88
+ interaction: {hover: true, tooltipDelay: 120}
89
+ };
90
+ const network = new vis.Network(document.getElementById("net"), data, opts);
91
+
92
+ // ---- legend ----
93
+ (function buildLegend() {
94
+ const el = document.getElementById("legend");
95
+ const real = (G.clusters || []).filter(c => c.id >= 0)
96
+ .sort((a, b) => b.count - a.count);
97
+ let html = "<h4>✦ CONSTELLATIONS</h4>";
98
+ real.forEach(c => {
99
+ html += `<div class="leg-row" data-hub="hub_${c.id}">
100
+ <span class="sw" style="background:${colorFor(c.id)}"></span>
101
+ <span class="leg-name">${esc(c.name || "Cluster " + c.id)}</span>
102
+ <span class="leg-cnt">${c.count}</span></div>`;
103
+ });
104
+ html += `<div class="key">
105
+ ● big = unread &nbsp; ● small = read<br>
106
+ <span style="color:#e0af68">◉ gold halo</span> = has image<br>
107
+ <span style="color:#f7768e">- - -</span> = cross-account thread<br>
108
+ brightness ∝ urgency &nbsp;·&nbsp; <b>double-click</b> a star</div>`;
109
+ el.innerHTML = html;
110
+ el.querySelectorAll(".leg-row").forEach(r =>
111
+ r.onclick = () => network.focus(r.dataset.hub, {scale: 1.1, animation: true}));
112
+ })();
113
+
114
+ // ---- detail panel ----
115
+ function esc(s) { return (s || "").replace(/[&<>"]/g, c =>
116
+ ({"&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;"}[c])); }
117
+ function fmtDate(iso) { if (!iso) return ""; const d = new Date(iso);
118
+ return isNaN(d) ? iso : d.toLocaleString(); }
119
+
120
+ function showEmail(n) {
121
+ const c = clusterById[n.group];
122
+ const cname = c ? (c.name || "Cluster " + n.group) : "unclustered";
123
+ const urg = Math.round((n.urgency || 0) * 100);
124
+ document.getElementById("detail").innerHTML = `
125
+ <span class="x" onclick="hide()">✕</span>
126
+ <div class="subj">${esc(n.label) || "(no subject)"}</div>
127
+ <div class="row"><span class="lbl">From:</span> ${esc(n.sender) || "—"}</div>
128
+ <div class="row"><span class="lbl">Account:</span> ${esc(n.account)}</div>
129
+ <div class="row"><span class="lbl">Date:</span> ${esc(fmtDate(n.ts)) || "—"}</div>
130
+ <div class="row"><span class="lbl">Constellation:</span>
131
+ <span class="chip" style="background:${colorFor(n.group)};color:#0a0e17">${esc(cname)}</span></div>
132
+ <div class="row"><span class="lbl">Urgency:</span> ${urg}%
133
+ ${n.unread ? '<span class="chip" style="background:#7aa2f7;color:#0a0e17">unread</span>' : ""}
134
+ ${n.has_image ? '<span class="chip" style="background:#e0af68;color:#0a0e17">image</span>' : ""}</div>
135
+ <div class="snip">${esc(n.snippet) || "<i>no preview</i>"}</div>`;
136
+ document.getElementById("detail").style.display = "block";
137
+ }
138
+ function showCluster(c) {
139
+ document.getElementById("detail").innerHTML = `
140
+ <span class="x" onclick="hide()">✕</span>
141
+ <div class="subj">✦ ${esc(c.name || "Cluster " + c.id)}</div>
142
+ <div class="row"><span class="lbl">Emails:</span> ${c.count}</div>
143
+ <div class="row"><span class="lbl">Avg urgency:</span> ${Math.round((c.urgency || 0) * 100)}%</div>
144
+ <div class="snip">${esc(c.rationale) || "<i>grouped by the multimodal encoder</i>"}</div>`;
145
+ document.getElementById("detail").style.display = "block";
146
+ }
147
+ window.hide = () => { document.getElementById("detail").style.display = "none"; };
148
+
149
+ network.on("doubleClick", p => {
150
+ if (!p.nodes.length) return hide();
151
+ const id = p.nodes[0];
152
+ if (String(id).startsWith("hub_")) return showCluster(clusterById[id.slice(4)]);
153
+ showEmail(detailById[id]);
154
+ });
155
+ </script>
156
+ </body>
157
+ </html>
astranexus/cluster/__init__.py ADDED
File without changes
astranexus/cluster/cluster.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from astranexus.core.schema import Cluster
3
+
4
+
5
+ def _reduce(X: np.ndarray, seed: int) -> np.ndarray:
6
+ if X.shape[0] < 10:
7
+ return X
8
+ import umap
9
+ n_comp = min(10, X.shape[0] - 2)
10
+ # smaller n_neighbors → more local structure → finer, more numerous clusters
11
+ reducer = umap.UMAP(n_components=n_comp, metric="cosine",
12
+ n_neighbors=min(10, X.shape[0] - 1), random_state=seed)
13
+ return reducer.fit_transform(X)
14
+
15
+
16
+ def cluster_embeddings(X: np.ndarray, seed: int = 7) -> list[Cluster]:
17
+ import hdbscan
18
+ n = X.shape[0]
19
+ reduced = _reduce(X, seed)
20
+ # Granularity: the old min_cluster_size=n//20 forced a few huge blobs on big
21
+ # inboxes (everything collapsed into one constellation). A small floor +
22
+ # cluster_selection_method='leaf' surfaces many fine-grained topics instead.
23
+ min_cluster_size = max(3, n // 40)
24
+ labels = hdbscan.HDBSCAN(min_cluster_size=min_cluster_size,
25
+ min_samples=1,
26
+ cluster_selection_method="leaf",
27
+ metric="euclidean").fit_predict(reduced)
28
+ real = {l for l in labels if l != -1}
29
+ if len(real) < 2: # graceful collapse: each point its own star
30
+ return [Cluster(id=i, rationale="", name="", email_ids=[str(i)], urgency=0.0)
31
+ for i in range(n)]
32
+ groups: dict[int, list[str]] = {}
33
+ for idx, lab in enumerate(labels):
34
+ groups.setdefault(int(lab), []).append(str(idx))
35
+ return [Cluster(id=cid, rationale="", name="", email_ids=members, urgency=0.0)
36
+ for cid, members in sorted(groups.items())]
astranexus/cluster/embed.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from astranexus.core.schema import Email
3
+ from astranexus.cluster.enrich import load_images
4
+
5
+
6
+ def _l2(x: np.ndarray) -> np.ndarray:
7
+ n = np.linalg.norm(x, axis=1, keepdims=True)
8
+ n[n == 0] = 1.0
9
+ return x / n
10
+
11
+
12
+ class Encoder:
13
+ """Multimodal late-fusion encoder.
14
+
15
+ Text branch = MiniLM. Image branch = SigLIP (filled in SP2).
16
+ Fusion = concat(text_weight * text_unit, (1-text_weight) * image_unit), renormalized.
17
+ Emails without images get a zero image block (so they cluster on text alone).
18
+ SP3 swaps fine-tuned weights behind this same encode() signature.
19
+ """
20
+
21
+ def __init__(self, text_model: str = "all-MiniLM-L6-v2",
22
+ image_model: str = "google/siglip-base-patch16-224",
23
+ text_weight: float = 0.6, use_images: bool = True, root: str = "."):
24
+ from sentence_transformers import SentenceTransformer
25
+ self._text = SentenceTransformer(text_model)
26
+ self.image_model = image_model
27
+ self.text_weight = text_weight
28
+ self.use_images = use_images
29
+ self.root = root
30
+ self._siglip = None
31
+ self._proc = None
32
+ self._cache: dict[str, np.ndarray] = {}
33
+
34
+ # ---- text branch ----
35
+ def _text_vecs(self, emails: list[Email]) -> np.ndarray:
36
+ docs = [f"{e.subject} {e.snippet}" for e in emails]
37
+ v = self._text.encode(docs, normalize_embeddings=True, show_progress_bar=False)
38
+ return np.asarray(v, dtype=np.float32)
39
+
40
+ # ---- image branch ----
41
+ def _ensure_siglip(self):
42
+ if self._siglip is None:
43
+ import torch
44
+ from transformers import SiglipVisionModel, SiglipImageProcessor
45
+ # Vision tower + the concrete (PIL-based) image processor only.
46
+ # Avoids the text tokenizer's sentencepiece dep AND AutoImageProcessor's
47
+ # torchvision requirement, since we only embed images.
48
+ self._siglip = SiglipVisionModel.from_pretrained(self.image_model).eval()
49
+ self._proc = SiglipImageProcessor.from_pretrained(self.image_model)
50
+ self._torch = torch
51
+ return self._siglip
52
+
53
+ def _embed_images(self, images) -> np.ndarray:
54
+ """Mean-pooled, L2-normalized SigLIP vector for one email's images."""
55
+ self._ensure_siglip()
56
+ inputs = self._proc(images=images, return_tensors="pt")
57
+ with self._torch.no_grad():
58
+ feats = self._siglip(**inputs).pooler_output
59
+ v = feats.cpu().numpy().astype(np.float32)
60
+ v = _l2(v).mean(axis=0)
61
+ nrm = np.linalg.norm(v)
62
+ return v / nrm if nrm else v
63
+
64
+ def _image_vecs(self, emails: list[Email]):
65
+ """Return (matrix (n, D_img), any_images: bool). Zero rows for no-image emails."""
66
+ per_email, dim = {}, None
67
+ for e in emails:
68
+ if not e.has_image:
69
+ continue
70
+ key = "|".join(e.attachments)
71
+ if key in self._cache:
72
+ vec = self._cache[key]
73
+ else:
74
+ imgs = load_images(e, root=self.root)
75
+ if not imgs:
76
+ continue
77
+ vec = self._embed_images(imgs)
78
+ self._cache[key] = vec
79
+ per_email[e.id] = vec
80
+ dim = vec.shape[0]
81
+ if dim is None:
82
+ return None, False
83
+ mat = np.zeros((len(emails), dim), dtype=np.float32)
84
+ for idx, e in enumerate(emails):
85
+ if e.id in per_email:
86
+ mat[idx] = per_email[e.id]
87
+ return mat, True
88
+
89
+ # ---- fusion ----
90
+ def encode(self, emails: list[Email]) -> np.ndarray:
91
+ text = self._text_vecs(emails) # (n, 384), unit rows
92
+ if not self.use_images:
93
+ return text
94
+ img, any_img = self._image_vecs(emails)
95
+ if not any_img:
96
+ return text # nothing multimodal → text-only, unchanged contract
97
+ a = self.text_weight
98
+ fused = np.concatenate([a * text, (1.0 - a) * img], axis=1)
99
+ return _l2(fused)
astranexus/cluster/enrich.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from astranexus.core.schema import Email
3
+
4
+
5
+ def load_images(email: Email, root: str = "."):
6
+ """Return a list of PIL images for an email's attachments.
7
+
8
+ PDF → first page rendered via pymupdf; image files opened via Pillow.
9
+ Missing/corrupt attachments are skipped (warn, never raise).
10
+ """
11
+ from PIL import Image
12
+ out = []
13
+ for rel in email.attachments:
14
+ path = os.path.join(root, rel) if not os.path.isabs(rel) else rel
15
+ if not os.path.exists(path):
16
+ continue
17
+ try:
18
+ if path.lower().endswith(".pdf"):
19
+ import fitz # pymupdf
20
+ doc = fitz.open(path)
21
+ page = doc.load_page(0)
22
+ pix = page.get_pixmap(dpi=150)
23
+ img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
24
+ out.append(img)
25
+ doc.close()
26
+ else:
27
+ out.append(Image.open(path).convert("RGB"))
28
+ except Exception as ex: # corrupt file → skip
29
+ print(f"[enrich] skipped {path}: {ex}")
30
+ return out
astranexus/cluster/eval.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score
2
+
3
+
4
+ def score(true_labels: list[str], pred_ids: list[int]) -> dict:
5
+ return {
6
+ "ari": float(adjusted_rand_score(true_labels, pred_ids)),
7
+ "nmi": float(normalized_mutual_info_score(true_labels, pred_ids)),
8
+ }
astranexus/cluster/ft_encoder.py ADDED
@@ -0,0 +1,138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SP3 fine-tuned encoder — inference mirror of data/sp3_kaggle.py:MMEncoder.
2
+
3
+ Loads the LoRA adapters (text + vision) and the learned fusion head saved to
4
+ sp3_out/, and exposes the SAME .encode(emails) -> np.ndarray contract as
5
+ cluster.embed.Encoder, so it drops straight into the pipeline.
6
+
7
+ The text LoRA was trained on the AutoModel (BertModel) module paths, NOT the
8
+ SentenceTransformer wrapper — so this must use AutoModel + mean-pool to match.
9
+ """
10
+ import os
11
+ import numpy as np
12
+
13
+ from astranexus.core.schema import Email
14
+ from astranexus.cluster.enrich import load_images
15
+
16
+ TEXT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
17
+ # Must match what the adapters were trained against (Kaggle used the large/384).
18
+ IMAGE_MODEL = os.environ.get("SP3_IMAGE_MODEL", "google/siglip-large-patch16-384")
19
+ PROJ_DIM = 256
20
+
21
+
22
+ def _mean_pool(last_hidden, mask):
23
+ import torch # noqa: F401 — ops below are torch
24
+ m = mask.unsqueeze(-1).float()
25
+ return (last_hidden * m).sum(1) / m.sum(1).clamp(min=1e-9)
26
+
27
+
28
+ def _load_adapter_robust(base, adapter_dir, device):
29
+ """PeftModel.from_pretrained, then REPAIR key mismatches caused by
30
+ transformers-version drift in SigLIP's module layout (Kaggle saved keys with
31
+ a `vision_model.` prefix that newer transformers flattens away). Without this
32
+ the vision LoRA loads as a silent no-op and the encoder emits garbage.
33
+ """
34
+ import os
35
+ from peft import PeftModel, get_peft_model_state_dict, set_peft_model_state_dict
36
+ from safetensors.torch import load_file
37
+
38
+ model = PeftModel.from_pretrained(base, adapter_dir)
39
+ want = get_peft_model_state_dict(model) # canonical save-format keys
40
+ file_sd = load_file(os.path.join(adapter_dir, "adapter_model.safetensors"))
41
+
42
+ want_vm = any("vision_model" in k for k in want)
43
+ have_vm = any("vision_model" in k for k in file_sd)
44
+ if have_vm and not want_vm:
45
+ file_sd = {k.replace("base_model.model.vision_model.", "base_model.model."): v
46
+ for k, v in file_sd.items()}
47
+ elif want_vm and not have_vm:
48
+ file_sd = {k.replace("base_model.model.", "base_model.model.vision_model."): v
49
+ for k, v in file_sd.items()}
50
+
51
+ res = set_peft_model_state_dict(model, file_sd)
52
+ missing = getattr(res, "unexpected_keys", None)
53
+ # sanity: after repair, the loaded keys must cover what the model expects
54
+ still_missing = [k for k in want if k not in file_sd]
55
+ if still_missing:
56
+ raise RuntimeError(f"adapter repair failed for {adapter_dir}: "
57
+ f"{len(still_missing)}/{len(want)} keys unmatched "
58
+ f"(e.g. {still_missing[0]})")
59
+ return model.to(device).eval()
60
+
61
+
62
+ class FineTunedEncoder:
63
+ """Drop-in replacement for Encoder using the SP3 LoRA adapters + fusion head."""
64
+
65
+ def __init__(self, out_dir: str = "sp3_out", device: str = "cpu", root: str = "."):
66
+ import torch
67
+ import torch.nn as nn
68
+ import torch.nn.functional as F
69
+ from transformers import AutoModel, AutoTokenizer, SiglipVisionModel, SiglipImageProcessor
70
+
71
+ self._torch, self._F = torch, F
72
+ self.device = device
73
+ self.root = root
74
+ self._cache: dict[str, np.ndarray] = {}
75
+
76
+ # --- text branch: MiniLM (AutoModel) + LoRA ---
77
+ self.tok = AutoTokenizer.from_pretrained(TEXT_MODEL)
78
+ text = AutoModel.from_pretrained(TEXT_MODEL)
79
+ self.text = _load_adapter_robust(text, os.path.join(out_dir, "text_lora"), device)
80
+
81
+ # --- vision branch: SigLIP-large + LoRA ---
82
+ self.vproc = SiglipImageProcessor.from_pretrained(IMAGE_MODEL)
83
+ vis = SiglipVisionModel.from_pretrained(IMAGE_MODEL)
84
+ self.vis = _load_adapter_robust(vis, os.path.join(out_dir, "vision_lora"), device)
85
+
86
+ t_dim = self.text.config.hidden_size
87
+ self.v_dim = self.vis.config.hidden_size
88
+
89
+ # --- learned fusion head (proj.pt) ---
90
+ self.proj = nn.Sequential(nn.Linear(t_dim + self.v_dim, PROJ_DIM), nn.GELU(),
91
+ nn.Linear(PROJ_DIM, PROJ_DIM)).to(device)
92
+ self.proj.load_state_dict(torch.load(os.path.join(out_dir, "proj.pt"), map_location=device))
93
+ self.proj.eval()
94
+
95
+ # ---- branches (mirror MMEncoder) ----
96
+ def _encode_text(self, subjects, snippets):
97
+ torch, F = self._torch, self._F
98
+ docs = [f"{s} {b}" for s, b in zip(subjects, snippets)]
99
+ enc = self.tok(docs, padding=True, truncation=True, max_length=64, return_tensors="pt").to(self.device)
100
+ out = self.text(**enc).last_hidden_state
101
+ return F.normalize(_mean_pool(out, enc["attention_mask"]), dim=1)
102
+
103
+ def _encode_images(self, pil_per_email):
104
+ torch, F = self._torch, self._F
105
+ flat, owner = [], []
106
+ for i, imgs in enumerate(pil_per_email):
107
+ for im in imgs:
108
+ flat.append(im); owner.append(i)
109
+ vecs = torch.zeros(len(pil_per_email), self.v_dim, device=self.device)
110
+ if flat:
111
+ px = self.vproc(images=flat, return_tensors="pt").to(self.device)
112
+ feats = F.normalize(self.vis(**px).pooler_output, dim=1)
113
+ agg = torch.zeros(len(pil_per_email), self.v_dim, device=self.device)
114
+ cnt = torch.zeros(len(pil_per_email), 1, device=self.device)
115
+ for k, o in enumerate(owner):
116
+ agg[o] += feats[k]; cnt[o] += 1
117
+ mask = (cnt.squeeze(1) > 0)
118
+ agg = F.normalize(agg / cnt.clamp(min=1), dim=1) * mask.unsqueeze(1)
119
+ vecs = agg
120
+ return vecs
121
+
122
+ def _load_email_images(self, e: Email):
123
+ if not e.has_image:
124
+ return []
125
+ return load_images(e, root=self.root)
126
+
127
+ # ---- fusion / public contract ----
128
+ def encode(self, emails: list[Email], batch: int = 16) -> np.ndarray:
129
+ torch = self._torch
130
+ out = []
131
+ with torch.no_grad():
132
+ for i in range(0, len(emails), batch):
133
+ b = emails[i:i + batch]
134
+ t = self._encode_text([e.subject for e in b], [e.snippet for e in b])
135
+ v = self._encode_images([self._load_email_images(e) for e in b])
136
+ fused = self._F.normalize(self.proj(torch.cat([t, v], dim=1)), dim=1)
137
+ out.append(fused.cpu().numpy().astype(np.float32))
138
+ return np.concatenate(out)
astranexus/config.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Load a local .env into os.environ if present. No external dependency.
2
+
3
+ Import this early (it runs on import). Real secrets live in .env (gitignored);
4
+ .env.example documents the keys.
5
+ """
6
+ import os
7
+ import pathlib
8
+
9
+ _ROOT = pathlib.Path(__file__).resolve().parents[1]
10
+
11
+
12
+ def load_env(path: str | os.PathLike | None = None) -> None:
13
+ p = pathlib.Path(path) if path else _ROOT / ".env"
14
+ if not p.exists():
15
+ return
16
+ for line in p.read_text(encoding="utf-8").splitlines():
17
+ line = line.strip()
18
+ if not line or line.startswith("#") or "=" not in line:
19
+ continue
20
+ key, _, val = line.partition("=")
21
+ key, val = key.strip(), val.strip().strip('"').strip("'")
22
+ if key and val and key not in os.environ: # don't clobber real env
23
+ os.environ[key] = val
24
+
25
+
26
+ load_env()
astranexus/core/__init__.py ADDED
File without changes
astranexus/core/schema.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass
2
+ from datetime import datetime
3
+
4
+
5
+ @dataclass(frozen=True)
6
+ class Email:
7
+ id: str
8
+ account: str
9
+ sender: str
10
+ recipient: str
11
+ ts: datetime
12
+ subject: str
13
+ snippet: str
14
+ unread: bool
15
+ thread_key: str
16
+ attachments: list[str]
17
+ has_image: bool
18
+
19
+
20
+ @dataclass(frozen=True)
21
+ class Cluster:
22
+ id: int
23
+ rationale: str
24
+ name: str
25
+ email_ids: list[str]
26
+ urgency: float
27
+
28
+
29
+ @dataclass(frozen=True)
30
+ class CanvasNode:
31
+ id: str
32
+ label: str
33
+ group: int
34
+ urgency: float
35
+ unread: bool
36
+ account: str
37
+ has_image: bool = False
38
+ sender: str = ""
39
+ snippet: str = ""
40
+ ts: str = "" # ISO 8601, for the detail panel
41
+
42
+
43
+ @dataclass(frozen=True)
44
+ class CanvasEdge:
45
+ src: str
46
+ dst: str
47
+ kind: str
48
+
49
+
50
+ @dataclass(frozen=True)
51
+ class CanvasGraph:
52
+ nodes: list[CanvasNode]
53
+ edges: list[CanvasEdge]
54
+ clusters: list[Cluster]
astranexus/graph.py ADDED
@@ -0,0 +1,232 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SP4 — LangGraph orchestration over the SP1/SP2 units.
2
+
3
+ Wires fetch → embed → cluster → label → serialize as a LangGraph StateGraph,
4
+ sharing a single Encoder instance across nodes, and serializes the execution
5
+ history to trace.json (the hackathon Agent-Trace requirement).
6
+
7
+ Node interfaces are the SAME functions the plain pipeline uses — this is a
8
+ re-wiring, not a rewrite.
9
+ """
10
+ import json, operator, time
11
+ from datetime import datetime, timezone
12
+ from dataclasses import replace
13
+ from typing import Annotated, Any, Callable, Optional, TypedDict
14
+
15
+ from langgraph.graph import StateGraph, END
16
+ from astranexus.cluster.cluster import cluster_embeddings
17
+ from astranexus.serialize.canvas import to_canvas
18
+ from astranexus.core.schema import CanvasGraph
19
+
20
+
21
+ class GraphState(TypedDict, total=False):
22
+ source: Any
23
+ seed: int
24
+ emails: list
25
+ X: Any
26
+ clusters: list
27
+ canvas: CanvasGraph
28
+ trace: Annotated[list, operator.add] # reducer: nodes append entries
29
+ # SP6 agentic loop:
30
+ attempts: int # number of label passes done
31
+ critique: str # judge feedback fed into the next label pass
32
+ judge_score: float # latest judge verdict
33
+
34
+
35
+ def _entry(node: str, t0: float, summary: dict) -> dict:
36
+ return {
37
+ "node": node,
38
+ "ts": datetime.now(timezone.utc).isoformat(),
39
+ "duration_ms": round((time.perf_counter() - t0) * 1000, 1),
40
+ "summary": summary,
41
+ }
42
+
43
+
44
+ def build_graph(encoder, labeler: Callable):
45
+ """Compile a StateGraph closing over a shared encoder + a labeler fn."""
46
+
47
+ def n_fetch(state: GraphState):
48
+ t0 = time.perf_counter()
49
+ emails = state["source"].fetch()
50
+ return {"emails": emails,
51
+ "trace": [_entry("fetch", t0, {"n_emails": len(emails),
52
+ "accounts": sorted({e.account for e in emails})})]}
53
+
54
+ def n_embed(state: GraphState):
55
+ t0 = time.perf_counter()
56
+ X = encoder.encode(state["emails"])
57
+ return {"X": X, "trace": [_entry("embed", t0, {"shape": list(X.shape)})]}
58
+
59
+ def n_cluster(state: GraphState):
60
+ t0 = time.perf_counter()
61
+ emails = state["emails"]
62
+ clusters = cluster_embeddings(state["X"], state.get("seed", 7))
63
+ clusters = [replace(c, email_ids=[emails[int(i)].id for i in c.email_ids]) for c in clusters]
64
+ noise = sum(len(c.email_ids) for c in clusters if c.id == -1)
65
+ return {"clusters": clusters,
66
+ "trace": [_entry("cluster", t0, {"n_clusters": len([c for c in clusters if c.id != -1]),
67
+ "noise": noise})]}
68
+
69
+ def n_label(state: GraphState):
70
+ t0 = time.perf_counter()
71
+ clusters = labeler(state["clusters"], state["emails"])
72
+ return {"clusters": clusters,
73
+ "trace": [_entry("label", t0, {"names": [c.name for c in clusters][:12]})]}
74
+
75
+ def n_serialize(state: GraphState):
76
+ t0 = time.perf_counter()
77
+ canvas = to_canvas(state["clusters"], state["emails"])
78
+ worm = sum(1 for e in canvas.edges if e.kind == "wormhole")
79
+ return {"canvas": canvas,
80
+ "trace": [_entry("serialize", t0, {"nodes": len(canvas.nodes),
81
+ "wormholes": worm})]}
82
+
83
+ g = StateGraph(GraphState)
84
+ g.add_node("fetch", n_fetch)
85
+ g.add_node("embed", n_embed)
86
+ g.add_node("cluster", n_cluster)
87
+ g.add_node("label", n_label)
88
+ g.add_node("serialize", n_serialize)
89
+ g.set_entry_point("fetch")
90
+ g.add_edge("fetch", "embed")
91
+ g.add_edge("embed", "cluster")
92
+ g.add_edge("cluster", "label")
93
+ g.add_edge("label", "serialize")
94
+ g.add_edge("serialize", END)
95
+ return g.compile()
96
+
97
+
98
+ def run_graph(source, labeler: Optional[Callable] = None, encoder=None, seed: int = 7,
99
+ trace_path: str = "trace.json", frozen: bool = False) -> CanvasGraph:
100
+ """Execute the orchestration graph and write trace.json. Returns the CanvasGraph.
101
+
102
+ frozen=True forces the off-the-shelf encoder (the live A/B comparison)."""
103
+ if encoder is None:
104
+ from astranexus.pipeline import _get_encoder
105
+ encoder = _get_encoder(frozen)
106
+ if labeler is None:
107
+ from astranexus.labeler.hf_serverless import LLMLabeler
108
+ labeler = LLMLabeler().label
109
+
110
+ app = build_graph(encoder, labeler)
111
+ started = datetime.now(timezone.utc).isoformat()
112
+ final = app.invoke({"source": source, "seed": seed, "trace": []})
113
+
114
+ if trace_path:
115
+ doc = {
116
+ "agent": "AstraNexus",
117
+ "started": started,
118
+ "finished": datetime.now(timezone.utc).isoformat(),
119
+ "graph": ["fetch", "embed", "cluster", "label", "serialize"],
120
+ "steps": final["trace"],
121
+ }
122
+ with open(trace_path, "w", encoding="utf-8") as f:
123
+ json.dump(doc, f, indent=2)
124
+ return final["canvas"]
125
+
126
+
127
+ # ── SP6: agentic graph (Nemotron label + LLM-as-judge self-correction) ────────
128
+ def build_agentic_graph(encoder, nemotron, tau: float = 0.7, max_attempts: int = 3):
129
+ """Compile the agentic StateGraph:
130
+
131
+ fetch → embed → cluster → label → judge ─┐
132
+ ▲ │ score<tau & attempts<max
133
+ └─────────────┘
134
+ └→ serialize (otherwise)
135
+
136
+ `nemotron` is a NemotronClient (Modal-vLLM primary, NIM fallback). The judge
137
+ scores the labels; a low score loops back to label with the critique injected,
138
+ up to `max_attempts` passes. Every pass + verdict is recorded in the trace.
139
+ """
140
+ from astranexus.labeler.nemotron import NemotronLabeler, judge_labels
141
+
142
+ def n_fetch(state: GraphState):
143
+ t0 = time.perf_counter()
144
+ emails = state["source"].fetch()
145
+ return {"emails": emails,
146
+ "trace": [_entry("fetch", t0, {"n_emails": len(emails),
147
+ "accounts": sorted({e.account for e in emails})})]}
148
+
149
+ def n_embed(state: GraphState):
150
+ t0 = time.perf_counter()
151
+ X = encoder.encode(state["emails"])
152
+ return {"X": X, "trace": [_entry("embed", t0, {"shape": list(X.shape)})]}
153
+
154
+ def n_cluster(state: GraphState):
155
+ t0 = time.perf_counter()
156
+ emails = state["emails"]
157
+ clusters = cluster_embeddings(state["X"], state.get("seed", 7))
158
+ clusters = [replace(c, email_ids=[emails[int(i)].id for i in c.email_ids]) for c in clusters]
159
+ noise = sum(len(c.email_ids) for c in clusters if c.id == -1)
160
+ return {"clusters": clusters,
161
+ "trace": [_entry("cluster", t0, {"n_clusters": len([c for c in clusters if c.id != -1]),
162
+ "noise": noise})]}
163
+
164
+ def n_label(state: GraphState):
165
+ t0 = time.perf_counter()
166
+ attempt = state.get("attempts", 0) + 1
167
+ labeler = NemotronLabeler(nemotron, critique=state.get("critique", ""))
168
+ clusters = labeler.label(state["clusters"], state["emails"])
169
+ return {"clusters": clusters, "attempts": attempt,
170
+ "trace": [_entry("label", t0, {"attempt": attempt,
171
+ "names": [c.name for c in clusters if c.id != -1][:12]})]}
172
+
173
+ def n_judge(state: GraphState):
174
+ t0 = time.perf_counter()
175
+ v = judge_labels(state["clusters"], state["emails"], nemotron)
176
+ return {"judge_score": v["score"], "critique": v["critique"],
177
+ "trace": [_entry("judge", t0, {"score": v["score"], "attempt": state.get("attempts", 0),
178
+ "critique": v["critique"][:160]})]}
179
+
180
+ def n_serialize(state: GraphState):
181
+ t0 = time.perf_counter()
182
+ canvas = to_canvas(state["clusters"], state["emails"])
183
+ worm = sum(1 for e in canvas.edges if e.kind == "wormhole")
184
+ return {"canvas": canvas,
185
+ "trace": [_entry("serialize", t0, {"nodes": len(canvas.nodes), "wormholes": worm})]}
186
+
187
+ def route(state: GraphState) -> str:
188
+ """Retry labeling while the judge is unhappy and we have attempts left."""
189
+ if state.get("judge_score", 1.0) < tau and state.get("attempts", 0) < max_attempts:
190
+ return "label"
191
+ return "serialize"
192
+
193
+ g = StateGraph(GraphState)
194
+ for name, fn in [("fetch", n_fetch), ("embed", n_embed), ("cluster", n_cluster),
195
+ ("label", n_label), ("judge", n_judge), ("serialize", n_serialize)]:
196
+ g.add_node(name, fn)
197
+ g.set_entry_point("fetch")
198
+ g.add_edge("fetch", "embed")
199
+ g.add_edge("embed", "cluster")
200
+ g.add_edge("cluster", "label")
201
+ g.add_edge("label", "judge")
202
+ g.add_conditional_edges("judge", route, {"label": "label", "serialize": "serialize"})
203
+ g.add_edge("serialize", END)
204
+ return g.compile()
205
+
206
+
207
+ def run_agentic_graph(source, nemotron, encoder=None, seed: int = 7,
208
+ tau: float = 0.7, max_attempts: int = 3,
209
+ trace_path: str = "trace.json", frozen: bool = False) -> CanvasGraph:
210
+ """Execute the agentic graph and write trace.json. Returns the CanvasGraph."""
211
+ if encoder is None:
212
+ from astranexus.pipeline import _get_encoder
213
+ encoder = _get_encoder(frozen)
214
+ app = build_agentic_graph(encoder, nemotron, tau=tau, max_attempts=max_attempts)
215
+ started = datetime.now(timezone.utc).isoformat()
216
+ final = app.invoke({"source": source, "seed": seed, "trace": [], "attempts": 0, "critique": ""},
217
+ config={"recursion_limit": 50})
218
+ if trace_path:
219
+ doc = {
220
+ "agent": "AstraNexus",
221
+ "mode": "agentic",
222
+ "started": started,
223
+ "finished": datetime.now(timezone.utc).isoformat(),
224
+ "graph": ["fetch", "embed", "cluster", "label", "judge", "serialize"],
225
+ "tau": tau, "max_attempts": max_attempts,
226
+ "final_score": final.get("judge_score"),
227
+ "attempts": final.get("attempts"),
228
+ "steps": final["trace"],
229
+ }
230
+ with open(trace_path, "w", encoding="utf-8") as f:
231
+ json.dump(doc, f, indent=2)
232
+ return final["canvas"]
astranexus/labeler/__init__.py ADDED
File without changes
astranexus/labeler/fallback.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ from collections import Counter
3
+ from dataclasses import replace
4
+ from astranexus.core.schema import Email, Cluster
5
+
6
+ _STOP = set("the a an of to and for your you re fwd is are on in at with this that".split())
7
+ _URGENT = re.compile(r"\b(urgent|asap|overdue|deadline|due|immediately|important)\b", re.I)
8
+ _LOWPRI = re.compile(r"\b(newsletter|digest|unsubscribe|no-?reply|promo)\b", re.I)
9
+
10
+
11
+ def _keywords(texts: list[str], k: int = 2) -> str:
12
+ words = re.findall(r"[a-zA-Z]{3,}", " ".join(texts).lower())
13
+ freq = Counter(w for w in words if w not in _STOP)
14
+ top = [w.capitalize() for w, _ in freq.most_common(k)]
15
+ return " ".join(top) or "Cluster"
16
+
17
+
18
+ def email_urgency(e: Email) -> float:
19
+ """Per-email urgency in [0,1]. Keyword cues + unread, around a 0.4 baseline."""
20
+ text = f"{e.subject} {e.snippet}"
21
+ score = 0.4
22
+ if _URGENT.search(text):
23
+ score += 0.35
24
+ if _LOWPRI.search(text):
25
+ score -= 0.3
26
+ if e.unread:
27
+ score += 0.15
28
+ return max(0.0, min(1.0, score))
29
+
30
+
31
+ def _urgency(emails: list[Email]) -> float:
32
+ """Cluster-level urgency = mean of its members' per-email urgency."""
33
+ if not emails:
34
+ return 0.0
35
+ return sum(email_urgency(e) for e in emails) / len(emails)
36
+
37
+
38
+ def label_fallback(clusters: list[Cluster], emails: list[Email]) -> list[Cluster]:
39
+ by_id = {e.id: e for e in emails}
40
+ out = []
41
+ for c in clusters:
42
+ members = [by_id[i] for i in c.email_ids if i in by_id]
43
+ texts = [f"{e.subject} {e.snippet}" for e in members]
44
+ name = "Cosmic Dust" if c.id == -1 else _keywords(texts)
45
+ out.append(replace(c, name=name, rationale="grouped by shared keywords",
46
+ urgency=_urgency(members)))
47
+ return out
astranexus/labeler/hf_serverless.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json, re
2
+ import astranexus.config # noqa: F401 — ensures HF_TOKEN from .env is loaded
3
+ from dataclasses import replace
4
+ from astranexus.core.schema import Email, Cluster
5
+ from astranexus.labeler.fallback import label_fallback
6
+
7
+ _GENERIC = {"misc", "other", "updates", "general", "stuff", "cluster", ""}
8
+
9
+ _PROMPT = (
10
+ "You name an email cluster. Given these emails, reply ONLY with JSON "
11
+ '{{"name": "<1-3 word specific name>", "urgency": <0..1>, "rationale": "<one short reason these group>"}}. '
12
+ "Forbidden names: Misc, Other, Updates, General.\nEmails:\n{body}"
13
+ )
14
+
15
+
16
+ class LLMLabeler:
17
+ def __init__(self, client=None, model: str = "Qwen/Qwen2.5-7B-Instruct"):
18
+ self.client = client
19
+ self.model = model
20
+
21
+ def _ensure_client(self):
22
+ if self.client is None:
23
+ import os
24
+ from huggingface_hub import InferenceClient
25
+ self.client = InferenceClient(self.model, token=os.environ.get("HF_TOKEN"))
26
+ return self.client
27
+
28
+ def _complete(self, prompt: str) -> str:
29
+ """Chat-completion call (current HF serverless API for instruct models)."""
30
+ resp = self._ensure_client().chat_completion(
31
+ messages=[{"role": "user", "content": prompt}],
32
+ max_tokens=120, temperature=0.2)
33
+ return resp.choices[0].message.content
34
+
35
+ def _label_one(self, c: Cluster, members: list[Email]) -> Cluster | None:
36
+ body = "\n".join(f"- {e.subject}: {e.snippet}" for e in members[:5])
37
+ try:
38
+ raw = self._complete(_PROMPT.format(body=body))
39
+ m = re.search(r"\{.*\}", raw, re.S)
40
+ data = json.loads(m.group(0))
41
+ name = str(data["name"]).strip()
42
+ if name.lower() in _GENERIC:
43
+ return None
44
+ urg = max(0.0, min(1.0, float(data.get("urgency", 0.5))))
45
+ return replace(c, name=name, urgency=urg,
46
+ rationale=str(data.get("rationale", "")).strip() or "semantically related")
47
+ except Exception:
48
+ return None
49
+
50
+ def label(self, clusters: list[Cluster], emails: list[Email]) -> list[Cluster]:
51
+ by_id = {e.id: e for e in emails}
52
+ out = []
53
+ for c in clusters:
54
+ members = [by_id[i] for i in c.email_ids if i in by_id]
55
+ if c.id == -1:
56
+ out.append(replace(c, name="Cosmic Dust", rationale="low-signal bulk mail", urgency=0.1))
57
+ continue
58
+ labeled = self._label_one(c, members)
59
+ out.append(labeled if labeled is not None else label_fallback([c], emails)[0])
60
+ return out
astranexus/labeler/nemotron.py ADDED
@@ -0,0 +1,191 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Nemotron-driven labeler + LLM-as-a-Judge with Pydantic structured output (SP6).
2
+
3
+ Labeler and judge share one `NemotronClient` (an OpenAI-compatible chat client +
4
+ model name). Structured output is enforced with Pydantic schemas via
5
+ `beta.chat.completions.parse`, so the SAME code runs against either:
6
+ * Modal-self-hosted Nemotron-3-Nano-4B served by vLLM (primary — modal_app.py), or
7
+ * NVIDIA NIM hosted Nemotron (fallback / local dev — `nim_client`).
8
+
9
+ Nemotron Nano are hybrid *reasoning* models; a `/no_think` system turn disables
10
+ the chain-of-thought so the structured answer isn't truncated by the token budget.
11
+
12
+ The judge scores the labeling; the graph loops back to the labeler with the
13
+ judge's critique when the score is below threshold (self-correction).
14
+ """
15
+ import os
16
+ from dataclasses import replace
17
+
18
+ from pydantic import BaseModel, Field
19
+
20
+ from astranexus.core.schema import Cluster, Email
21
+ from astranexus.labeler.fallback import label_fallback
22
+
23
+ # The 4B-Nano runs self-hosted on Modal (modal_app.py) — it is NOT on the NIM
24
+ # hosted endpoint. NIM is the *fallback* path, so it uses the nearest hosted
25
+ # Nemotron (9B-v2). Both satisfy "built with Nemotron".
26
+ NIM_MODEL = os.environ.get("NEMOTRON_NIM_MODEL", "nvidia/nvidia-nemotron-nano-9b-v2")
27
+ NIM_BASE_URL = os.environ.get("NIM_BASE_URL", "https://integrate.api.nvidia.com/v1")
28
+ MODAL_MODEL = os.environ.get("NEMOTRON_MODAL_MODEL", "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
29
+
30
+ _GENERIC = {"misc", "other", "updates", "general", "stuff", "cluster", ""}
31
+
32
+
33
+ class Label(BaseModel):
34
+ """Structured cluster label the model must return."""
35
+ name: str = Field(description="1-3 word specific name; never Misc/Other/Updates/General",
36
+ max_length=40)
37
+ urgency: float = Field(ge=0.0, le=1.0, description="0=ignorable, 1=act now")
38
+ rationale: str = Field(description="one short reason these emails group", max_length=200)
39
+
40
+
41
+ class Verdict(BaseModel):
42
+ """Structured judgement of a whole labeling pass."""
43
+ score: float = Field(ge=0.0, le=1.0, description="1.0 = all names specific & accurate")
44
+ critique: str = Field(description="concrete actionable fixes naming bad clusters",
45
+ max_length=400)
46
+
47
+
48
+ _LABEL_PROMPT = (
49
+ "These emails were grouped automatically by a clustering model, so a few may "
50
+ "be outliers. Name the DOMINANT theme shared by the MAJORITY — ignore the odd "
51
+ "straggler. Give a specific 1-3 word name (never Misc, Other, Updates, or "
52
+ "General), an urgency in [0,1], and one short rationale.{critique}\nEmails:\n{body}"
53
+ )
54
+
55
+ _JUDGE_PROMPT = (
56
+ "You judge an email-clustering app's constellation NAMES. The clusters are "
57
+ "produced by a separate model and may contain a few unrelated emails; that is "
58
+ "NOT the namer's fault. Judge ONLY whether each name is the best concise, "
59
+ "specific summary of the MAJORITY of its emails. Score 0..1 (1.0 = every name "
60
+ "accurately captures its majority theme and is non-generic). Penalize generic "
61
+ "names (Misc/Updates), names that miss the majority theme, and duplicates — "
62
+ "but do NOT penalize the clustering for including a few outliers. Give a "
63
+ "one-sentence critique naming only names that are genuinely wrong and a better "
64
+ "name for each.\nLabels:\n{body}"
65
+ )
66
+
67
+
68
+ class NemotronClient:
69
+ """OpenAI-compatible chat client + model, with Pydantic structured output."""
70
+
71
+ def __init__(self, client, model: str, no_think: bool = True):
72
+ self.client, self.model, self.no_think = client, model, no_think
73
+
74
+ def parse(self, prompt: str, schema: type[BaseModel]) -> BaseModel | None:
75
+ msgs = []
76
+ if self.no_think:
77
+ msgs.append({"role": "system", "content": "/no_think"})
78
+ msgs.append({"role": "user", "content": prompt})
79
+ resp = self.client.beta.chat.completions.parse(
80
+ model=self.model, messages=msgs, response_format=schema,
81
+ max_tokens=400, temperature=0.2)
82
+ return resp.choices[0].message.parsed
83
+
84
+
85
+ def nim_client(model: str = NIM_MODEL, api_key: str | None = None,
86
+ base_url: str = NIM_BASE_URL) -> NemotronClient:
87
+ """A NemotronClient backed by the NVIDIA NIM hosted endpoint (fallback path)."""
88
+ from openai import OpenAI
89
+ return NemotronClient(
90
+ OpenAI(api_key=api_key or os.environ.get("NVIDIA_API_KEY"), base_url=base_url),
91
+ model)
92
+
93
+
94
+ def modal_client(base_url: str, model: str = MODAL_MODEL,
95
+ api_key: str = "EMPTY") -> NemotronClient:
96
+ """A NemotronClient backed by an OpenAI-compatible Modal server (unused now —
97
+ kept for a future vLLM server path)."""
98
+ from openai import OpenAI
99
+ return NemotronClient(OpenAI(api_key=api_key, base_url=base_url), model)
100
+
101
+
102
+ class LocalNemotronClient:
103
+ """NemotronClient-compatible client for an in-process HuggingFace transformers
104
+ model (the Modal self-hosted path). Nemotron-3-Nano-4B is a hybrid
105
+ (NemotronHForCausalLM) model, so we generate with `/no_think` and a tight
106
+ prompt, then extract+validate JSON with Pydantic (no guided decoding needed).
107
+ Same `.parse(prompt, schema)` contract as NemotronClient."""
108
+
109
+ def __init__(self, model, tokenizer, no_think: bool = True, max_new_tokens: int = 256):
110
+ self.model, self.tok = model, tokenizer
111
+ self.no_think, self.max_new_tokens = no_think, max_new_tokens
112
+
113
+ def _generate(self, prompt: str) -> str:
114
+ import torch
115
+ msgs = []
116
+ if self.no_think:
117
+ msgs.append({"role": "system", "content": "/no_think"})
118
+ msgs.append({"role": "user", "content": prompt})
119
+ ids = self.tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt")
120
+ ids = ids.to(self.model.device)
121
+ with torch.no_grad():
122
+ out = self.model.generate(ids, max_new_tokens=self.max_new_tokens, do_sample=False,
123
+ pad_token_id=self.tok.eos_token_id)
124
+ return self.tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True)
125
+
126
+ def parse(self, prompt: str, schema: type[BaseModel]) -> BaseModel | None:
127
+ import re
128
+ text = self._generate(prompt)
129
+ m = re.search(r"\{.*\}", text, re.S)
130
+ if not m:
131
+ return None
132
+ return schema.model_validate_json(m.group(0))
133
+
134
+
135
+ class NemotronLabeler:
136
+ """Names clusters via a NemotronClient. Pass `critique` on a retry so the
137
+ model sees the judge's feedback. Falls back to heuristics per-cluster."""
138
+
139
+ def __init__(self, client: NemotronClient, critique: str = ""):
140
+ self.client, self.critique = client, critique
141
+
142
+ def _label_one(self, c: Cluster, members: list[Email]) -> Cluster | None:
143
+ body = "\n".join(f"- {e.subject}: {e.snippet}" for e in members[:5])
144
+ crit = (f"\nA reviewer flagged the previous attempt: {self.critique}"
145
+ if self.critique else "")
146
+ try:
147
+ out = self.client.parse(_LABEL_PROMPT.format(body=body, critique=crit), Label)
148
+ if out is None:
149
+ return None
150
+ name = out.name.strip()
151
+ if name.lower() in _GENERIC:
152
+ return None
153
+ urg = max(0.0, min(1.0, out.urgency)) # clamp: model may exceed [0,1]
154
+ return replace(c, name=name, urgency=urg,
155
+ rationale=out.rationale.strip() or "semantically related")
156
+ except Exception:
157
+ return None
158
+
159
+ def label(self, clusters: list[Cluster], emails: list[Email]) -> list[Cluster]:
160
+ by_id = {e.id: e for e in emails}
161
+ out = []
162
+ for c in clusters:
163
+ members = [by_id[i] for i in c.email_ids if i in by_id]
164
+ if c.id == -1:
165
+ out.append(replace(c, name="Cosmic Dust", rationale="low-signal bulk mail", urgency=0.1))
166
+ continue
167
+ labeled = self._label_one(c, members)
168
+ out.append(labeled if labeled is not None else label_fallback([c], emails)[0])
169
+ return out
170
+
171
+
172
+ def judge_labels(clusters: list[Cluster], emails: list[Email],
173
+ client: NemotronClient) -> dict:
174
+ """Score the current labeling. Returns {"score": float, "critique": str}.
175
+ On any failure returns a passing score so the loop never hard-blocks."""
176
+ by_id = {e.id: e for e in emails}
177
+ lines = []
178
+ for c in clusters:
179
+ if c.id == -1:
180
+ continue
181
+ sample = [by_id[i].subject for i in c.email_ids[:3] if i in by_id]
182
+ lines.append(f'NAME="{c.name}" RATIONALE="{c.rationale}" emails={sample}')
183
+ if not lines:
184
+ return {"score": 1.0, "critique": ""}
185
+ try:
186
+ out = client.parse(_JUDGE_PROMPT.format(body="\n".join(lines)), Verdict)
187
+ if out is None:
188
+ return {"score": 1.0, "critique": ""}
189
+ return {"score": max(0.0, min(1.0, out.score)), "critique": out.critique.strip()}
190
+ except Exception:
191
+ return {"score": 1.0, "critique": ""}
astranexus/lg_studio.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Zero-arg graph factory for `langgraph dev` / LangGraph Studio.
2
+
3
+ `build_agentic_graph` takes (encoder, nemotron), which the LangGraph CLI can't
4
+ supply, so this exposes a no-arg `make_graph()` it can import. Both deps are
5
+ wrapped lazily so importing/compiling the graph (to draw the topology in Studio)
6
+ never triggers the heavy SigLIP load or an LLM call — those only happen if a node
7
+ actually runs. Studio thus shows the full SP6 agentic loop:
8
+ fetch→embed→cluster→label→judge→(retry|serialize).
9
+ """
10
+ from astranexus.graph import build_agentic_graph
11
+
12
+
13
+ class _LazyEncoder:
14
+ """Defers the real encoder (and its model download) until first encode()."""
15
+
16
+ def encode(self, emails):
17
+ from astranexus.pipeline import _get_encoder
18
+ return _get_encoder().encode(emails)
19
+
20
+
21
+ class _LazyNemotron:
22
+ """Defers building the NIM Nemotron client until the first .parse() call, so
23
+ Studio can render topology without NVIDIA_API_KEY set."""
24
+
25
+ def parse(self, prompt, schema):
26
+ from astranexus.labeler.nemotron import nim_client
27
+ return nim_client().parse(prompt, schema)
28
+
29
+
30
+ def make_graph():
31
+ return build_agentic_graph(_LazyEncoder(), _LazyNemotron())
astranexus/pipeline.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from dataclasses import replace
3
+ from astranexus.sources.base import Source
4
+ from astranexus.cluster.embed import Encoder
5
+ from astranexus.cluster.cluster import cluster_embeddings
6
+ from astranexus.serialize.canvas import to_canvas
7
+ from astranexus.core.schema import CanvasGraph, Cluster
8
+
9
+ _encoders: dict = {} # cache keyed by variant: "auto" | "frozen"
10
+
11
+
12
+ def _get_encoder(frozen: bool = False):
13
+ """Encoder for the active variant, cached.
14
+
15
+ frozen=True forces the off-the-shelf baseline (the A/B comparison + honest
16
+ before/after eval). frozen=False prefers the SP3 fine-tuned encoder when its
17
+ adapters are present (sp3_out/) and ASTRANEXUS_FINETUNED=1, else off-the-shelf.
18
+ """
19
+ key = "frozen" if frozen else "auto"
20
+ if key in _encoders:
21
+ return _encoders[key]
22
+ if (not frozen and os.environ.get("ASTRANEXUS_FINETUNED") == "1"
23
+ and os.path.isdir("sp3_out/text_lora")):
24
+ try:
25
+ from astranexus.cluster.ft_encoder import FineTunedEncoder
26
+ _encoders[key] = FineTunedEncoder()
27
+ print("[encoder] using SP3 fine-tuned adapters (sp3_out/)")
28
+ return _encoders[key]
29
+ except Exception as ex: # adapters missing/corrupt → safe fallback
30
+ print(f"[encoder] fine-tuned load failed ({ex}); falling back to off-the-shelf")
31
+ _encoders[key] = Encoder()
32
+ return _encoders[key]
33
+
34
+
35
+ def encoder_label(frozen: bool = False) -> str:
36
+ """Human-readable name of the active encoder (for the UI / proof it's the
37
+ fine-tuned model running, not the off-the-shelf fallback)."""
38
+ enc = _get_encoder(frozen)
39
+ if type(enc).__name__ == "FineTunedEncoder":
40
+ return "SP3 fine-tuned · MiniLM+SigLIP-large-384 (LoRA)"
41
+ return "off-the-shelf · MiniLM+SigLIP-base-224"
42
+
43
+
44
+ def run(source: Source, labeler=None, seed: int = 7) -> CanvasGraph:
45
+ emails = source.fetch()
46
+ if len(emails) < 5: # too few to cluster: flat star field
47
+ clusters = [Cluster(id=i, rationale="", name=e.subject[:20], email_ids=[e.id], urgency=0.0)
48
+ for i, e in enumerate(emails)]
49
+ return to_canvas(clusters, emails)
50
+
51
+ X = _get_encoder().encode(emails)
52
+ clusters = cluster_embeddings(X, seed=seed)
53
+ # remap row-index ids -> real email ids
54
+ clusters = [replace(c, email_ids=[emails[int(i)].id for i in c.email_ids]) for c in clusters]
55
+
56
+ if labeler is None:
57
+ from astranexus.labeler.hf_serverless import LLMLabeler
58
+ labeler = LLMLabeler().label
59
+ clusters = labeler(clusters, emails)
60
+ return to_canvas(clusters, emails)
astranexus/serialize/__init__.py ADDED
File without changes
astranexus/serialize/canvas.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from collections import defaultdict
2
+ from itertools import combinations
3
+ from astranexus.core.schema import Email, Cluster, CanvasNode, CanvasEdge, CanvasGraph
4
+ from astranexus.labeler.fallback import email_urgency
5
+
6
+
7
+ def to_canvas(clusters: list[Cluster], emails: list[Email]) -> CanvasGraph:
8
+ group_of = {eid: c.id for c in clusters for eid in c.email_ids}
9
+
10
+ nodes = []
11
+ for e in emails:
12
+ gid = group_of.get(e.id, -1)
13
+ # per-email urgency (so brightness varies star-to-star, not per-cluster)
14
+ nodes.append(CanvasNode(id=e.id, label=e.subject, group=gid,
15
+ urgency=email_urgency(e), unread=e.unread,
16
+ account=e.account, has_image=e.has_image,
17
+ sender=e.sender, snippet=e.snippet,
18
+ ts=e.ts.isoformat() if e.ts else ""))
19
+
20
+ threads: dict[str, list[Email]] = defaultdict(list)
21
+ for e in emails:
22
+ threads[e.thread_key].append(e)
23
+
24
+ edges = []
25
+ for members in threads.values():
26
+ for x, y in combinations(members, 2):
27
+ if x.account != y.account: # cross-account only = wormhole
28
+ edges.append(CanvasEdge(src=x.id, dst=y.id, kind="wormhole"))
29
+
30
+ return CanvasGraph(nodes=nodes, edges=edges, clusters=clusters)
astranexus/serialize/render.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import html as _html
2
+ import json, pathlib
3
+ from astranexus.core.schema import CanvasGraph
4
+
5
+ _TEMPLATE = pathlib.Path(__file__).resolve().parents[1] / "assets" / "vis_canvas.html"
6
+
7
+
8
+ def render_html(g: CanvasGraph) -> str:
9
+ payload = {
10
+ "nodes": [n.__dict__ for n in g.nodes],
11
+ "edges": [{"src": e.src, "dst": e.dst, "kind": e.kind} for e in g.edges],
12
+ "clusters": [{"id": c.id, "name": c.name, "count": len(c.email_ids),
13
+ "rationale": c.rationale, "urgency": c.urgency}
14
+ for c in g.clusters],
15
+ }
16
+ template = _TEMPLATE.read_text(encoding="utf-8")
17
+ return template.replace("__GRAPH_JSON__", json.dumps(payload))
18
+
19
+
20
+ def render_iframe(g: CanvasGraph, height: int = 620) -> str:
21
+ """Wrap the canvas doc in an <iframe srcdoc> so its <script> runs.
22
+
23
+ gr.HTML (Gradio 6) injects via innerHTML and does NOT execute inline
24
+ <script>; an iframe is a separate document that does. The vis-network
25
+ template already targets window.parent, so it was built for this.
26
+ """
27
+ srcdoc = _html.escape(render_html(g), quote=True)
28
+ return (f'<iframe srcdoc="{srcdoc}" '
29
+ f'style="width:100%;height:{height}px;border:0;background:#0a0e17" '
30
+ f'sandbox="allow-scripts"></iframe>')
astranexus/sources/__init__.py ADDED
File without changes
astranexus/sources/base.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ from typing import Protocol
3
+ from astranexus.core.schema import Email
4
+
5
+ _PREFIX = re.compile(r"^\s*((re|fwd|fw)\s*:\s*)+", re.IGNORECASE)
6
+ _WS = re.compile(r"\s+")
7
+
8
+
9
+ def normalize_thread_key(subject: str) -> str:
10
+ s = _PREFIX.sub("", subject)
11
+ s = _WS.sub(" ", s).strip().lower()
12
+ return s
13
+
14
+
15
+ class Source(Protocol):
16
+ def fetch(self) -> list[Email]: ...
astranexus/sources/gmail_api.py ADDED
@@ -0,0 +1,183 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Gmail REST API source (HTTPS / port 443).
2
+
3
+ The HF Space firewalls outbound IMAP (993/143), so live mode can't use
4
+ imap_gmail there. Gmail's REST API rides 443 (wide open — it's how the Space
5
+ pulls models), so we read the inbox over it with OAuth user credentials.
6
+
7
+ Emits the SAME Email shape as the IMAP source, so embed→cluster→label→serialize
8
+ is unchanged. google-* libs are imported lazily so the rest of the app (and the
9
+ test suite) runs without them installed.
10
+ """
11
+ import base64
12
+ import hashlib
13
+ import os
14
+ import sys
15
+ from dataclasses import replace
16
+ from datetime import datetime, timezone
17
+ from email.utils import parsedate_to_datetime
18
+ from astranexus.core.schema import Email
19
+ from astranexus.sources.base import normalize_thread_key
20
+
21
+ SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
22
+ GET_CHUNK = 50 # Gmail batch endpoint recommends <=50 sub-requests
23
+ IMG_DIR = "data/_live_images" # transient cache for fetched image attachments
24
+ MIN_IMG_BYTES = 8000 # skip logos / tracking pixels / signature icons
25
+ MAX_IMG_PER_EMAIL = 1 # one meaningful image per email keeps the vision pass cheap
26
+ _EXT = {"image/png": ".png", "image/jpeg": ".jpg", "image/jpg": ".jpg",
27
+ "image/gif": ".gif", "image/webp": ".webp"}
28
+
29
+
30
+ def build_credentials(client_id: str, client_secret: str, refresh_token: str,
31
+ token_uri: str = "https://oauth2.googleapis.com/token"):
32
+ """An OAuth user Credentials that auto-mints access tokens from a refresh
33
+ token. Lazy google import keeps this off the default code path."""
34
+ from google.oauth2.credentials import Credentials
35
+ return Credentials(None, refresh_token=refresh_token, client_id=client_id,
36
+ client_secret=client_secret, token_uri=token_uri, scopes=SCOPES)
37
+
38
+
39
+ def _build_service(credentials):
40
+ from googleapiclient.discovery import build
41
+ return build("gmail", "v1", credentials=credentials, cache_discovery=False)
42
+
43
+
44
+ def _header(headers: list, name: str) -> str:
45
+ for h in headers:
46
+ if h.get("name", "").lower() == name.lower():
47
+ return h.get("value", "")
48
+ return ""
49
+
50
+
51
+ def _has_image(payload: dict) -> bool:
52
+ if (payload.get("mimeType") or "").startswith("image/"):
53
+ return True
54
+ return any(_has_image(p) for p in (payload.get("parts") or []))
55
+
56
+
57
+ def _iter_parts(payload: dict):
58
+ yield payload
59
+ for p in payload.get("parts") or []:
60
+ yield from _iter_parts(p)
61
+
62
+
63
+ def _image_parts(payload: dict) -> list[dict]:
64
+ """Flatten the MIME tree to its image parts (needs format='full')."""
65
+ out = []
66
+ for p in _iter_parts(payload):
67
+ if (p.get("mimeType") or "").startswith("image/"):
68
+ body = p.get("body") or {}
69
+ out.append({"mime": p["mimeType"], "data": body.get("data"),
70
+ "attachmentId": body.get("attachmentId"),
71
+ "size": body.get("size") or 0})
72
+ return out
73
+
74
+
75
+ def message_to_email(msg: dict, account: str) -> Email:
76
+ """Map a Gmail API message resource -> our Email. Uses the API's own
77
+ `snippet` and the UNREAD label, so no body download is needed."""
78
+ payload = msg.get("payload", {}) or {}
79
+ headers = payload.get("headers", []) or []
80
+ subject = _header(headers, "Subject")
81
+ try:
82
+ ts = parsedate_to_datetime(_header(headers, "Date"))
83
+ except Exception:
84
+ ts = datetime.now(timezone.utc)
85
+ return Email(
86
+ id=hashlib.sha1(f"{account}:{msg.get('id', '')}".encode()).hexdigest()[:12],
87
+ account=account, sender=_header(headers, "From"),
88
+ recipient=_header(headers, "To"), ts=ts, subject=subject,
89
+ snippet=(msg.get("snippet") or "")[:200],
90
+ unread=("UNREAD" in (msg.get("labelIds") or [])),
91
+ thread_key=normalize_thread_key(subject), attachments=[],
92
+ has_image=_has_image(payload),
93
+ )
94
+
95
+
96
+ class GmailAPISource:
97
+ """Source backed by the Gmail REST API. Pass `credentials` (built from a
98
+ refresh token) for production, or inject a ready `service` in tests."""
99
+
100
+ def __init__(self, account: str, credentials=None, service=None, limit: int = 100):
101
+ self.account, self.credentials, self.limit = account, credentials, limit
102
+ self._service = service
103
+
104
+ def _svc(self):
105
+ if self._service is None:
106
+ self._service = _build_service(self.credentials)
107
+ return self._service
108
+
109
+ def fetch(self) -> list[Email]:
110
+ svc = self._svc()
111
+ users = svc.users()
112
+ listing = users.messages().list(
113
+ userId="me", maxResults=self.limit, labelIds=["INBOX"]).execute()
114
+ ids = [m["id"] for m in (listing.get("messages") or [])][:self.limit]
115
+
116
+ raw, dropped = [], 0
117
+
118
+ def _collect(req_id, response, exc):
119
+ nonlocal dropped
120
+ if exc is not None: # don't silently lose mail
121
+ dropped += 1
122
+ print(f"[gmail_api] {self.account} drop {req_id}: {exc}",
123
+ file=sys.stderr, flush=True)
124
+ elif response:
125
+ raw.append(response)
126
+
127
+ # format='full' (not 'metadata') so the MIME tree carries image parts we
128
+ # can download for the glow + the multimodal vision branch.
129
+ for i in range(0, len(ids), GET_CHUNK):
130
+ batch = svc.new_batch_http_request(callback=_collect)
131
+ for mid in ids[i:i + GET_CHUNK]:
132
+ batch.add(users.messages().get(userId="me", id=mid, format="full"))
133
+ batch.execute()
134
+
135
+ out = []
136
+ for msg in raw:
137
+ e = message_to_email(msg, self.account)
138
+ paths = self._save_images(svc, msg) # download real images, if any
139
+ out.append(replace(e, attachments=paths, has_image=bool(paths)))
140
+
141
+ n_img = sum(1 for e in out if e.has_image)
142
+ print(f"[gmail_api] {self.account}: {len(out)}/{len(ids)} fetched, "
143
+ f"{dropped} dropped, {n_img} with image", file=sys.stderr, flush=True)
144
+ return out
145
+
146
+ def _save_images(self, svc, msg: dict) -> list[str]:
147
+ """Download up to MAX_IMG_PER_EMAIL meaningful images for one message to
148
+ IMG_DIR; return their relative paths. Logos/pixels (< MIN_IMG_BYTES) are
149
+ skipped so the glow + vision pass key on real content, not signatures."""
150
+ parts = [p for p in _image_parts(msg.get("payload", {}) or {})
151
+ if (p["size"] or 0) >= MIN_IMG_BYTES or p["data"]]
152
+ if not parts:
153
+ return []
154
+ parts.sort(key=lambda p: p["size"] or 0, reverse=True) # biggest first
155
+ os.makedirs(IMG_DIR, exist_ok=True)
156
+ eid = hashlib.sha1(f"{self.account}:{msg.get('id', '')}".encode()).hexdigest()[:12]
157
+ saved = []
158
+ for j, p in enumerate(parts):
159
+ if len(saved) >= MAX_IMG_PER_EMAIL:
160
+ break
161
+ data = p["data"]
162
+ if not data and p["attachmentId"]:
163
+ try:
164
+ att = svc.users().messages().attachments().get(
165
+ userId="me", messageId=msg["id"], id=p["attachmentId"]).execute()
166
+ data = att.get("data")
167
+ except Exception as ex:
168
+ print(f"[gmail_api] {self.account} attachment fail: {ex}",
169
+ file=sys.stderr, flush=True)
170
+ continue
171
+ if not data:
172
+ continue
173
+ try:
174
+ rawb = base64.urlsafe_b64decode(data + "===") # tolerant padding
175
+ except Exception:
176
+ continue
177
+ if len(rawb) < MIN_IMG_BYTES: # gate after decode
178
+ continue
179
+ path = os.path.join(IMG_DIR, f"{eid}_{j}{_EXT.get(p['mime'], '.img')}")
180
+ with open(path, "wb") as f:
181
+ f.write(rawb)
182
+ saved.append(path)
183
+ return saved
astranexus/sources/imap_gmail.py ADDED
@@ -0,0 +1,204 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import email, hashlib, imaplib, re, socket, ssl, sys, time
2
+ from email.header import decode_header, make_header
3
+ from datetime import datetime, timezone
4
+ from email.utils import parsedate_to_datetime
5
+ from astranexus.core.schema import Email
6
+ from astranexus.sources.base import normalize_thread_key
7
+
8
+ # Per-socket op timeout. Kept tight so a blocked/dropped outbound port fails
9
+ # fast instead of hanging (a blocked egress to 993 just times out silently).
10
+ SOCKET_TIMEOUT = 15
11
+
12
+
13
+ def probe_connectivity() -> str:
14
+ """TCP-connect to the ports the live path needs, so we can tell a code bug
15
+ from an HF Space egress block. 443 should pass (HF pulls models over it);
16
+ if 993/143 time out while 443 works, outbound IMAP is firewalled."""
17
+ targets = [("Google HTTPS :443", "gmail.googleapis.com", 443),
18
+ ("Gmail IMAP-SSL:993", "imap.gmail.com", 993),
19
+ ("Gmail IMAP :143", "imap.gmail.com", 143)]
20
+ lines = []
21
+ for label, host, port in targets:
22
+ t0 = time.perf_counter()
23
+ try:
24
+ ai = socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM)[0]
25
+ s = socket.socket(ai[0], ai[1], ai[2])
26
+ s.settimeout(8)
27
+ s.connect(ai[4])
28
+ s.close()
29
+ lines.append(f"✅ {label} OK ({(time.perf_counter() - t0) * 1000:.0f} ms)")
30
+ except Exception as ex:
31
+ lines.append(f"❌ {label} {type(ex).__name__}: {ex} "
32
+ f"({(time.perf_counter() - t0) * 1000:.0f} ms)")
33
+ return "\n".join(lines)
34
+ # UIDs per FETCH command. One command pulls a whole chunk in a single
35
+ # round-trip, so 200 mails = ~4 round-trips instead of 200 (the long-lived
36
+ # single socket over 200 sequential fetches is what Gmail/egress was dropping).
37
+ FETCH_CHUNK = 50
38
+ # Per-chunk reconnect attempts before we give up on that chunk and move on
39
+ # (graceful degrade: a partial map beats a crashed demo).
40
+ MAX_RETRIES = 2
41
+
42
+
43
+ def _hdr(msg, key: str) -> str:
44
+ v = msg.get(key, "")
45
+ try:
46
+ return str(make_header(decode_header(v)))
47
+ except Exception:
48
+ return v
49
+
50
+
51
+ def _first_sentence(msg) -> str:
52
+ body = ""
53
+ if msg.is_multipart():
54
+ for part in msg.walk():
55
+ if part.get_content_type() == "text/plain":
56
+ body = part.get_payload(decode=True).decode(errors="ignore")
57
+ break
58
+ else:
59
+ body = (msg.get_payload(decode=True) or b"").decode(errors="ignore")
60
+ body = " ".join(body.split())
61
+ return (body.split(". ")[0][:200]) if body else ""
62
+
63
+
64
+ def parse_message(raw: bytes, account: str, uid: str, unread: bool) -> Email:
65
+ msg = email.message_from_bytes(raw)
66
+ subject = _hdr(msg, "Subject")
67
+ try:
68
+ ts = parsedate_to_datetime(msg.get("Date"))
69
+ except Exception:
70
+ ts = datetime.now(timezone.utc)
71
+ return Email(
72
+ id=hashlib.sha1(f"{account}:{uid}".encode()).hexdigest()[:12],
73
+ account=account, sender=_hdr(msg, "From"), recipient=_hdr(msg, "To"),
74
+ ts=ts, subject=subject, snippet=_first_sentence(msg), unread=unread,
75
+ thread_key=normalize_thread_key(subject), attachments=[],
76
+ has_image=any(p.get_content_maintype() == "image" for p in msg.walk()) if msg.is_multipart() else False,
77
+ )
78
+
79
+
80
+ class _IPv4IMAP4SSL(imaplib.IMAP4_SSL):
81
+ """IMAP4_SSL that connects over IPv4 only.
82
+
83
+ HF Spaces (and many containers) have no IPv6 route, yet Gmail's DNS returns
84
+ AAAA records, so Python tries IPv6 first and fails with
85
+ `OSError [Errno 101] Network is unreachable` at connect time — before any
86
+ login. Resolving A-records only and connecting to the IPv4 address avoids
87
+ it; SNI + cert validation still use the hostname via server_hostname.
88
+ """
89
+
90
+ def _create_socket(self, timeout=None):
91
+ timeout = timeout or SOCKET_TIMEOUT
92
+ last = None
93
+ for family, socktype, proto, _canon, sockaddr in socket.getaddrinfo(
94
+ self.host, self.port, socket.AF_INET, socket.SOCK_STREAM):
95
+ try:
96
+ sock = socket.socket(family, socktype, proto)
97
+ sock.settimeout(timeout)
98
+ sock.connect(sockaddr)
99
+ return self.ssl_context.wrap_socket(sock, server_hostname=self.host)
100
+ except OSError as ex:
101
+ last = ex
102
+ raise last or OSError(f"no IPv4 route to {self.host}:{self.port}")
103
+
104
+
105
+ class GmailSource:
106
+ def __init__(self, account: str, app_password: str, limit: int = 100):
107
+ self.account, self.app_password, self.limit = account, app_password, limit
108
+
109
+ def _connect(self) -> imaplib.IMAP4_SSL:
110
+ ctx = ssl.create_default_context()
111
+ M = _IPv4IMAP4SSL("imap.gmail.com", 993, ssl_context=ctx, timeout=SOCKET_TIMEOUT)
112
+ M.login(self.account, self.app_password)
113
+ # readonly => EXAMINE, so fetching never sets the \Seen flag on the
114
+ # user's real mail (the old (RFC822) fetch was silently marking
115
+ # everything as read).
116
+ M.select("INBOX", readonly=True)
117
+ return M
118
+
119
+ def _connect_with_retry(self) -> imaplib.IMAP4_SSL:
120
+ """The initial connect must survive a transient network blip too —
121
+ not just the per-chunk fetches."""
122
+ last = None
123
+ for attempt in range(MAX_RETRIES + 1):
124
+ try:
125
+ return self._connect()
126
+ except TimeoutError as ex:
127
+ # a silent timeout means the port is firewalled/dropped, not a
128
+ # transient blip — retrying just burns the user's time.
129
+ print(f"[imap] {self.account}: connect timed out — outbound 993 "
130
+ f"likely blocked (see the Diagnose button).", file=sys.stderr, flush=True)
131
+ raise
132
+ except (OSError, imaplib.IMAP4.abort, ssl.SSLError) as ex:
133
+ last = ex
134
+ print(f"[imap] {self.account}: connect attempt {attempt + 1} failed "
135
+ f"({type(ex).__name__}: {ex})", file=sys.stderr, flush=True)
136
+ raise last
137
+
138
+ def _reconnect(self, M) -> imaplib.IMAP4_SSL:
139
+ try:
140
+ M.logout()
141
+ except Exception:
142
+ pass
143
+ return self._connect()
144
+
145
+ def _parse_chunk(self, data, unseen: set) -> list[Email]:
146
+ out = []
147
+ for item in data:
148
+ if not isinstance(item, tuple) or item[1] is None:
149
+ continue # the b')' separators imaplib interleaves between messages
150
+ meta, raw = item[0], item[1]
151
+ m = re.match(rb"\s*(\d+)", meta or b"")
152
+ seqno = m.group(1) if m else b"?"
153
+ out.append(parse_message(raw, self.account, seqno.decode(), seqno in unseen))
154
+ return out
155
+
156
+ def fetch(self) -> list[Email]:
157
+ M = self._connect_with_retry()
158
+ try:
159
+ uids = M.search(None, "ALL")[1][0].split()[-self.limit:]
160
+ unseen = set(M.search(None, "UNSEEN")[1][0].split())
161
+ out, i = [], 0
162
+ while i < len(uids):
163
+ group = uids[i:i + FETCH_CHUNK]
164
+ seq = b",".join(group).decode()
165
+ for attempt in range(MAX_RETRIES + 1):
166
+ try:
167
+ typ, data = M.fetch(seq, "(BODY.PEEK[])")
168
+ out.extend(self._parse_chunk(data, unseen))
169
+ break
170
+ except (OSError, imaplib.IMAP4.abort, ssl.SSLError) as ex:
171
+ # socket died mid-run (this is the Errno 101 path).
172
+ if attempt >= MAX_RETRIES:
173
+ print(f"[imap] {self.account}: skipping chunk after "
174
+ f"{MAX_RETRIES} retries ({type(ex).__name__}: {ex})",
175
+ file=sys.stderr, flush=True)
176
+ break
177
+ print(f"[imap] {self.account}: reconnecting after "
178
+ f"{type(ex).__name__} (attempt {attempt + 1})",
179
+ file=sys.stderr, flush=True)
180
+ M = self._reconnect(M)
181
+ i += FETCH_CHUNK
182
+ return out
183
+ finally:
184
+ try:
185
+ M.logout()
186
+ except Exception:
187
+ pass
188
+
189
+
190
+ class MultiSource:
191
+ def __init__(self, sources: list[GmailSource]):
192
+ self.sources = sources
193
+
194
+ def fetch(self) -> list[Email]:
195
+ out, errors = [], []
196
+ for s in self.sources:
197
+ try:
198
+ out.extend(s.fetch())
199
+ except Exception as ex: # one bad mailbox must not kill the rest
200
+ errors.append(f"{getattr(s, 'account', '?')}: {type(ex).__name__}: {ex}")
201
+ print(f"[imap] mailbox failed -> {errors[-1]}", file=sys.stderr, flush=True)
202
+ if not out and errors: # every mailbox failed: surface it, don't return empty
203
+ raise RuntimeError("All mailboxes failed: " + " | ".join(errors))
204
+ return out
astranexus/sources/mock_json.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ from datetime import datetime
3
+ from astranexus.core.schema import Email
4
+ from astranexus.sources.base import normalize_thread_key
5
+
6
+
7
+ class MockSource:
8
+ def __init__(self, path: str):
9
+ self.path = path
10
+
11
+ def fetch(self) -> list[Email]:
12
+ rows = json.loads(open(self.path, encoding="utf-8").read())
13
+ out = []
14
+ for r in rows:
15
+ out.append(Email(
16
+ id=r["id"], account=r["account"], sender=r["sender"],
17
+ recipient=r["recipient"], ts=datetime.fromisoformat(r["ts"]),
18
+ subject=r["subject"], snippet=r["snippet"], unread=bool(r["unread"]),
19
+ thread_key=normalize_thread_key(r["subject"]),
20
+ attachments=r.get("attachments", []), has_image=bool(r.get("has_image", False)),
21
+ ))
22
+ return out
astranexus/ui/__init__.py ADDED
File without changes
astranexus/ui/app.py ADDED
@@ -0,0 +1,247 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os, sys, traceback
2
+ import gradio as gr
3
+ import astranexus.config # noqa: F401 — loads .env into os.environ on import
4
+ from astranexus.sources.mock_json import MockSource
5
+ from astranexus.sources.imap_gmail import GmailSource, MultiSource
6
+ from astranexus.labeler.fallback import label_fallback
7
+ from astranexus.graph import run_graph
8
+ from astranexus.serialize.render import render_iframe
9
+
10
+ # Newest-N emails fetched per connected account. Capped (not a UI slider) to keep
11
+ # a live multi-account map fast + reliable on the free CPU Space. Override with the
12
+ # FETCH_LIMIT env var if you want to push it.
13
+ FETCH_LIMIT = int(os.environ.get("FETCH_LIMIT", "200"))
14
+
15
+ # Shown in the canvas while a map runs. Gives the output column real height so the
16
+ # Gradio progress overlay is visible (an empty gr.HTML collapses to ~0px, which
17
+ # made the app look frozen during the ~2-3 min CPU encode).
18
+ _LOADING_HTML = (
19
+ '<div style="height:600px;display:flex;flex-direction:column;align-items:center;'
20
+ 'justify-content:center;gap:16px;color:#7aa2f7;font-family:system-ui,Segoe UI,sans-serif">'
21
+ '<div style="font-size:46px;animation:anpulse 1.4s ease-in-out infinite">&#10022;</div>'
22
+ '<div style="font-size:16px">Fetching, embedding &amp; clustering your inbox&hellip;</div>'
23
+ '<div style="opacity:.6;font-size:13px">~2&ndash;3 min on the free CPU Space &mdash; hang tight.</div>'
24
+ '<style>@keyframes anpulse{0%,100%{opacity:.35;transform:scale(.9)}50%{opacity:1;transform:scale(1.1)}}</style>'
25
+ '</div>')
26
+
27
+ DARK = gr.themes.Base(primary_hue="indigo", neutral_hue="slate").set(
28
+ body_background_fill="#0a0e17", block_background_fill="#0d1320",
29
+ body_text_color="#c0caf5")
30
+
31
+
32
+ def _oauth_enabled() -> bool:
33
+ from astranexus.web import google_oauth
34
+ return google_oauth.enabled()
35
+
36
+
37
+ # ── live sources ────────────────────────────────────────────────────────────
38
+ def _session_sources(request):
39
+ """Gmail API sources for every Google account signed in this session
40
+ (the Space path — IMAP is firewalled there)."""
41
+ if request is None:
42
+ return []
43
+ from astranexus.web.google_oauth import credentials_for, COOKIE
44
+ from astranexus.sources.gmail_api import GmailAPISource
45
+ sid = request.cookies.get(COOKIE)
46
+ return [GmailAPISource(email, creds, limit=FETCH_LIMIT)
47
+ for email, creds in credentials_for(sid).items()]
48
+
49
+
50
+ # ── connected-account views ─────────────────────────────────────────────────
51
+ def _connected_md(accounts):
52
+ if not accounts:
53
+ return "_No mailboxes connected. Add one above, or use the offline fixture._"
54
+ rows = "\n".join(f"- ✉️ `{a['email']}`" for a in accounts)
55
+ return f"**Connected mailboxes ({len(accounts)}):**\n{rows}"
56
+
57
+
58
+ def connected_google(request: gr.Request = None):
59
+ """Signed-in Google accounts for this session, refreshed on page load
60
+ (so the list updates after the OAuth redirect lands back on /app)."""
61
+ emails, note = [], ""
62
+ if request is not None:
63
+ from astranexus.web.google_oauth import accounts_for, COOKIE
64
+ emails = accounts_for(request.cookies.get(COOKIE))
65
+ if request.query_params.get("auth") == "error":
66
+ note = "\n\n⚠️ _Sign-in didn't complete — try again._"
67
+ if not emails:
68
+ return "_No Google accounts connected. Click **Sign in with Google** above._" + note
69
+ rows = "\n".join(f"- ✉️ `{e}`" for e in emails)
70
+ return f"**Connected Google accounts ({len(emails)}):**\n{rows}" + note
71
+
72
+
73
+ def add_mailbox(email, app_pw, accounts):
74
+ """Append a mailbox to the session list (local IMAP path). Passwords live
75
+ only in this in-memory state — never written to disk, never rendered."""
76
+ accounts = list(accounts or [])
77
+ email = (email or "").strip()
78
+ if not email or not app_pw:
79
+ return accounts, _connected_md(accounts), "", gr.update(), "Enter both an email and an App Password."
80
+ if any(a["email"].lower() == email.lower() for a in accounts):
81
+ return accounts, _connected_md(accounts), "", gr.update(), f"Mailbox already connected: {email}"
82
+ accounts.append({"email": email, "pw": app_pw})
83
+ # auto-uncheck the fixture so the next Map Inbox uses the live mailbox(es)
84
+ return (accounts, _connected_md(accounts), "", gr.update(value=False),
85
+ f"Added {email}. ({len(accounts)} connected) — fixture unchecked; click Map Inbox.")
86
+
87
+
88
+ def _run_pipeline(source, use_base, scope):
89
+ """Run the AI pipeline with a 3-tier fallback so the demo never hard-breaks:
90
+ 1. Modal GPU — whole agentic graph + self-hosted Nemotron-3-Nano-4B (SP6)
91
+ 2. local agentic — same graph, Nemotron via NIM hosted (needs NVIDIA_API_KEY)
92
+ 3. local heuristic — LangGraph + keyword labels (no external LLM)
93
+ Returns (CanvasGraph, status_markdown). Writes trace.json in every tier.
94
+ `use_base` forces the off-the-shelf encoder (live A/B vs the fine-tuned one)."""
95
+ import json as _json
96
+ from astranexus.web import modal_client
97
+ from astranexus.pipeline import encoder_label
98
+
99
+ emails = source.fetch()
100
+
101
+ class _ListSource: # reuse the single fetch across fallback tiers
102
+ def fetch(self):
103
+ return emails
104
+
105
+ # 1) Modal GPU — the whole pipeline runs remotely on an L4 (use_base drives
106
+ # the A/B encoder toggle there: fine-tuned multimodal vs off-the-shelf)
107
+ if modal_client.enabled():
108
+ try:
109
+ g, trace = modal_client.run_remote(emails, use_base=use_base)
110
+ with open("trace.json", "w", encoding="utf-8") as f:
111
+ _json.dump(trace, f, indent=2)
112
+ sc, at = trace.get("final_score"), trace.get("attempts")
113
+ verdict = f"judge {sc:.2f} in {at} pass(es)" if sc is not None else "judged"
114
+ enc = trace.get("encoder", "")
115
+ return g, (f"Mapped {len(g.nodes)} stars · {len(g.clusters)} constellations · {scope} · "
116
+ f"runtime: **Modal L4 GPU** · model: **Nemotron-3-Nano-4B** · "
117
+ f"encoder: {enc} · {verdict} · trace → trace.json")
118
+ except Exception as ex:
119
+ print(f"[modal] remote failed, falling back to local: {ex}", file=sys.stderr, flush=True)
120
+
121
+ # 2) Local agentic — Nemotron via NIM hosted endpoint
122
+ if os.environ.get("NVIDIA_API_KEY"):
123
+ try:
124
+ from astranexus.graph import run_agentic_graph
125
+ from astranexus.labeler.nemotron import NIM_MODEL, nim_client
126
+ g = run_agentic_graph(_ListSource(), nim_client(), trace_path="trace.json", frozen=use_base)
127
+ return g, (f"Mapped {len(g.nodes)} stars · {len(g.clusters)} constellations · {scope} · "
128
+ f"runtime: local · model: **{NIM_MODEL}** (NIM) · "
129
+ f"encoder: {encoder_label(use_base)} · trace → trace.json")
130
+ except Exception as ex:
131
+ print(f"[nim] agentic failed, falling back to heuristic: {ex}", file=sys.stderr, flush=True)
132
+
133
+ # 3) Heuristic fallback — keyword labels, no external LLM required
134
+ g = run_graph(_ListSource(), labeler=lambda cl, em: label_fallback(cl, em),
135
+ trace_path="trace.json", frozen=use_base)
136
+ return g, (f"Mapped {len(g.nodes)} stars · {len(g.clusters)} constellations · {scope} · "
137
+ f"runtime: local · encoder: {encoder_label(use_base)} · trace → trace.json")
138
+
139
+
140
+ def map_inbox(use_mock, accounts, use_base, request: gr.Request = None):
141
+ # Generator: yields a loading panel first (immediate visible feedback during the
142
+ # long encode), then the final map. Gradio streams each yield to the UI.
143
+ if use_mock:
144
+ source, scope = MockSource("data/mock_mailbox.json"), "offline fixture"
145
+ else:
146
+ srcs = _session_sources(request) # OAuth (Space)
147
+ if not srcs and accounts: # IMAP (local)
148
+ srcs = [GmailSource(a["email"], a["pw"], limit=FETCH_LIMIT) for a in accounts]
149
+ if not srcs:
150
+ yield None, "No accounts connected — sign in with Google (or tick the offline fixture)."
151
+ return
152
+ source, scope = MultiSource(srcs), f"most recent {FETCH_LIMIT}/account × {len(srcs)}"
153
+ yield _LOADING_HTML, "⏳ Mapping… fetch → embed → cluster → Nemotron label → judge. Hang tight."
154
+ try:
155
+ g, status = _run_pipeline(source, use_base, scope)
156
+ yield render_iframe(g), status
157
+ except Exception as ex:
158
+ tb = traceback.format_exc()
159
+ print(tb, file=sys.stderr, flush=True) # -> HF Space "Logs" tab
160
+ detail = f"\n\n<details><summary>Show traceback</summary>\n\n```\n{tb}\n```\n\n</details>"
161
+ yield None, f"❌ {type(ex).__name__}: {ex}{detail}"
162
+
163
+
164
+ # target="_blank": the Space runs inside HF's cross-origin iframe, so "_top"
165
+ # navigation to /oauth/login is silently blocked. Opening a new tab lands on the
166
+ # hf.space origin directly, where the session cookie + OAuth redirect work.
167
+ _SIGNIN_HTML = (
168
+ '<a href="/oauth/login" target="_blank" rel="noopener" style="display:inline-block;'
169
+ 'background:#7aa2f7;color:#0a0e17;font-weight:700;text-decoration:none;'
170
+ 'padding:10px 20px;border-radius:8px;margin:4px 0">Sign in with Google &#8599;</a>'
171
+ '&nbsp;<a href="/oauth/logout" target="_blank" rel="noopener" style="color:#7aa2f7;font-size:13px">'
172
+ 'sign out</a>'
173
+ '<div style="font-size:12px;opacity:.7;margin-top:6px">Opens a new tab — '
174
+ 'finish sign-in there, then return and reload, or just map from that tab.</div>'
175
+ )
176
+
177
+ # Link to the open agent trace (LangGraph execution log) the last map wrote.
178
+ # target="_blank" because the Space runs inside HF's cross-origin iframe.
179
+ _TRACE_HTML = (
180
+ '<a href="/trace.json" target="_blank" rel="noopener" style="display:inline-block;'
181
+ 'border:1px solid #7aa2f7;color:#7aa2f7;text-decoration:none;padding:6px 14px;'
182
+ 'border-radius:8px;font-size:13px;margin-top:8px">🛰 View agent trace (trace.json) &#8599;</a>'
183
+ '<div style="font-size:12px;opacity:.65;margin-top:4px">The LangGraph run log — '
184
+ 'nodes, timings, judge scores &amp; retries. Generated by your latest map.</div>'
185
+ )
186
+
187
+
188
+ def build_app(theme=None, oauth_enabled=None):
189
+ # Gradio 5 (Docker Space) honors theme on Blocks; Gradio 6 wants it on
190
+ # launch() — pass theme=None there and set it in launch().
191
+ oauth = _oauth_enabled() if oauth_enabled is None else oauth_enabled
192
+ with gr.Blocks(title="AstraNexus", theme=theme) as demo:
193
+ accounts = gr.State([])
194
+ gr.Markdown("## ✦ AstraNexus — inbox constellations")
195
+ gr.Markdown(
196
+ "Map your inboxes into visual **constellations** — clustered by a "
197
+ "**fine-tuned multimodal encoder** (MiniLM text + SigLIP vision, LoRA), "
198
+ "named & reasoned over by **NVIDIA Nemotron-3-Nano-4B**. Cross-account "
199
+ "threads appear as **wormholes**; emails with images get a **gold halo**.\n\n"
200
+ "**Under the hood** — the whole AI pipeline runs on a **Modal** GPU, "
201
+ "orchestrated as a **LangGraph** agentic graph: "
202
+ "`fetch → embed → cluster → label → judge → (retry if low) → serialize`. "
203
+ "An **LLM-as-judge** scores the labels and loops back to relabel until they pass. "
204
+ "Tick **A/B** to compare the fine-tuned encoder against the off-the-shelf SigLIP. "
205
+ "Layered fallbacks (Modal → NIM-hosted Nemotron → keyword heuristic) keep it alive. "
206
+ "Every run writes an open **agent trace** (button below).\n\n"
207
+ "Use the offline **fixture** (a synthetic dataset for demo/testing), or connect "
208
+ "Google. Credentials are session-only — never stored or logged.")
209
+ with gr.Row():
210
+ with gr.Column(scale=1):
211
+ use_mock = gr.Checkbox(value=True, label="Use fixture (offline)")
212
+ if oauth:
213
+ gr.HTML(_SIGNIN_HTML)
214
+ connected = gr.Markdown(connected_google())
215
+ account = app_pw = add_btn = None
216
+ else:
217
+ account = gr.Textbox(label="Gmail address", placeholder="you@gmail.com")
218
+ app_pw = gr.Textbox(label="App password", type="password")
219
+ add_btn = gr.Button("➕ Add Mailbox")
220
+ connected = gr.Markdown(_connected_md([]))
221
+ use_base = gr.Checkbox(value=False,
222
+ label="A/B: off-the-shelf SigLIP (vs fine-tuned)")
223
+ btn = gr.Button("Map Inbox", variant="primary")
224
+ gr.HTML(_TRACE_HTML)
225
+ gr.Markdown(f"<sub>Live mode maps your **most recent {FETCH_LIMIT}** "
226
+ f"emails per account (free-tier cap, read-only).</sub>")
227
+ gr.Markdown(
228
+ "<sub>**Why invite-only for live Gmail?** Google restricts reading "
229
+ "people's mail, so live access is limited to approved test accounts "
230
+ "(my own mailboxes). Everyone else uses the synthetic **fixture** — "
231
+ "a generated dataset built purely for demo/testing, no real email. "
232
+ "Want in? Email **dripto.215@gmail.com**. Data stays in-session, never stored.</sub>")
233
+ status = gr.Markdown()
234
+ with gr.Column(scale=3):
235
+ canvas = gr.HTML()
236
+ if oauth:
237
+ # refresh the signed-in list when the page (re)loads after redirect
238
+ demo.load(connected_google, None, connected)
239
+ else:
240
+ add_btn.click(add_mailbox, [account, app_pw, accounts],
241
+ [accounts, connected, app_pw, use_mock, status])
242
+ btn.click(map_inbox, [use_mock, accounts, use_base], [canvas, status])
243
+ return demo
244
+
245
+
246
+ if __name__ == "__main__":
247
+ build_app().launch(theme=DARK)
astranexus/web/__init__.py ADDED
File without changes
astranexus/web/google_oauth.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Server-side 'Sign in with Google' for the HF Space.
2
+
3
+ Live mode needs the Gmail REST API (outbound IMAP is firewalled), which needs
4
+ OAuth user credentials. This wires the Authorization-Code web flow onto the
5
+ FastAPI app: /oauth/login -> Google consent -> /oauth/callback -> store the
6
+ refresh token in an in-memory session keyed by a cookie.
7
+
8
+ Single uvicorn worker on the Space, so in-memory sessions are fine; nothing is
9
+ written to disk. Visitors can sign in multiple accounts (each appends), which
10
+ is what powers the cross-account 'wormhole' demo.
11
+ """
12
+ import os
13
+ import secrets as _secrets
14
+ import sys
15
+ from urllib.parse import urlencode
16
+
17
+ SCOPES = ["https://www.googleapis.com/auth/gmail.readonly",
18
+ "openid", "https://www.googleapis.com/auth/userinfo.email"]
19
+ _AUTH_URI = "https://accounts.google.com/o/oauth2/v2/auth"
20
+ _TOKEN_URI = "https://oauth2.googleapis.com/token"
21
+ _USERINFO = "https://openidconnect.googleapis.com/v1/userinfo"
22
+
23
+ COOKIE = "an_session"
24
+ # session_id -> {"accounts": {email: {client_id, client_secret, refresh_token}},
25
+ # "state": <csrf>}
26
+ SESSIONS: dict = {}
27
+
28
+
29
+ def enabled() -> bool:
30
+ """True when the Space has the OAuth client secrets configured."""
31
+ return bool(os.environ.get("GOOGLE_CLIENT_ID") and os.environ.get("GOOGLE_CLIENT_SECRET"))
32
+
33
+
34
+ def redirect_uri() -> str:
35
+ """The callback Google redirects to — must be registered verbatim on the
36
+ OAuth client. SPACE_HOST is set inside HF Space containers."""
37
+ host = os.environ.get("SPACE_HOST") or os.environ.get("OAUTH_HOST", "localhost:7860")
38
+ scheme = "http" if host.startswith("localhost") or host.startswith("127.") else "https"
39
+ return f"{scheme}://{host}/oauth/callback"
40
+
41
+
42
+ def accounts_for(session_id) -> list[str]:
43
+ s = SESSIONS.get(session_id or "")
44
+ return sorted(s["accounts"].keys()) if s else []
45
+
46
+
47
+ def credentials_for(session_id) -> dict:
48
+ """email -> google Credentials, for every account signed in this session."""
49
+ from astranexus.sources.gmail_api import build_credentials
50
+ s = SESSIONS.get(session_id or "")
51
+ if not s:
52
+ return {}
53
+ return {email: build_credentials(**creds) for email, creds in s["accounts"].items()}
54
+
55
+
56
+ def attach(app) -> None:
57
+ """Register the OAuth routes on a FastAPI app."""
58
+ import requests
59
+ from fastapi import Request
60
+ from fastapi.responses import RedirectResponse
61
+
62
+ print(f"[oauth] enabled={enabled()} redirect_uri={redirect_uri()}",
63
+ file=sys.stderr, flush=True)
64
+
65
+ @app.get("/oauth/login")
66
+ def login(request: Request):
67
+ sid = request.cookies.get(COOKIE) or _secrets.token_urlsafe(24)
68
+ state = _secrets.token_urlsafe(16)
69
+ SESSIONS.setdefault(sid, {"accounts": {}})["state"] = state
70
+ params = urlencode({
71
+ "client_id": os.environ["GOOGLE_CLIENT_ID"],
72
+ "redirect_uri": redirect_uri(),
73
+ "response_type": "code",
74
+ "scope": " ".join(SCOPES),
75
+ "access_type": "offline", # we want a refresh token
76
+ "prompt": "consent", # force it even on re-auth
77
+ "include_granted_scopes": "true",
78
+ "state": state,
79
+ })
80
+ resp = RedirectResponse(f"{_AUTH_URI}?{params}")
81
+ resp.set_cookie(COOKIE, sid, httponly=True, secure=True,
82
+ samesite="lax", max_age=86400)
83
+ return resp
84
+
85
+ @app.get("/oauth/callback")
86
+ def callback(request: Request):
87
+ sid = request.cookies.get(COOKIE)
88
+ sess = SESSIONS.get(sid or "")
89
+ code = request.query_params.get("code")
90
+ state = request.query_params.get("state")
91
+ if not sess or not code or state != sess.get("state"):
92
+ return RedirectResponse("/app?auth=error")
93
+ try:
94
+ tok = requests.post(_TOKEN_URI, timeout=15, data={
95
+ "code": code,
96
+ "client_id": os.environ["GOOGLE_CLIENT_ID"],
97
+ "client_secret": os.environ["GOOGLE_CLIENT_SECRET"],
98
+ "redirect_uri": redirect_uri(),
99
+ "grant_type": "authorization_code",
100
+ }).json()
101
+ refresh = tok.get("refresh_token")
102
+ access = tok.get("access_token")
103
+ email = requests.get(_USERINFO, timeout=15,
104
+ headers={"Authorization": f"Bearer {access}"}).json().get("email")
105
+ if refresh and email:
106
+ sess["accounts"][email] = {
107
+ "client_id": os.environ["GOOGLE_CLIENT_ID"],
108
+ "client_secret": os.environ["GOOGLE_CLIENT_SECRET"],
109
+ "refresh_token": refresh,
110
+ }
111
+ return RedirectResponse("/app?auth=ok")
112
+ print(f"[oauth] callback missing token/email: {tok.get('error', tok)}",
113
+ file=sys.stderr, flush=True)
114
+ except Exception as ex:
115
+ print(f"[oauth] callback failed: {type(ex).__name__}: {ex}",
116
+ file=sys.stderr, flush=True)
117
+ return RedirectResponse("/app?auth=error")
118
+
119
+ @app.get("/oauth/logout")
120
+ def logout(request: Request):
121
+ SESSIONS.pop(request.cookies.get(COOKIE) or "", None)
122
+ return RedirectResponse("/app")
astranexus/web/modal_client.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Space-side client for the Modal pipeline (SP6).
2
+
3
+ The HF Space fetches emails locally (OAuth lives here), then hands the whole AI
4
+ pipeline to the Modal GPU endpoint. On any failure the caller falls back to the
5
+ local pipeline, so the demo never hard-breaks.
6
+ """
7
+ import os
8
+
9
+ from astranexus.core.schema import CanvasEdge, CanvasGraph, CanvasNode, Cluster
10
+
11
+ TIMEOUT = int(os.environ.get("MODAL_TIMEOUT", "600"))
12
+
13
+
14
+ def endpoint() -> str:
15
+ return os.environ.get("MODAL_ENDPOINT_URL", "").strip()
16
+
17
+
18
+ def enabled() -> bool:
19
+ return bool(endpoint())
20
+
21
+
22
+ def _email_to_row(e) -> dict:
23
+ return {
24
+ "id": e.id, "account": e.account, "sender": e.sender,
25
+ "recipient": e.recipient, "ts": e.ts.isoformat(), "subject": e.subject,
26
+ "snippet": e.snippet, "unread": bool(e.unread), "thread_key": e.thread_key,
27
+ # attachments are fixture-relative paths shipped into the Modal image; live
28
+ # Gmail sends [] (metadata-only), so Modal degrades to text for those.
29
+ "attachments": list(e.attachments), "has_image": bool(e.has_image),
30
+ }
31
+
32
+
33
+ def _canvas_from_dict(d: dict) -> CanvasGraph:
34
+ return CanvasGraph(
35
+ nodes=[CanvasNode(**n) for n in d["nodes"]],
36
+ edges=[CanvasEdge(**e) for e in d["edges"]],
37
+ clusters=[Cluster(**c) for c in d["clusters"]],
38
+ )
39
+
40
+
41
+ def run_remote(emails, use_base: bool = False, tau: float = 0.7,
42
+ max_attempts: int = 3, seed: int = 7):
43
+ """POST emails to the Modal pipeline. Returns (CanvasGraph, trace_dict).
44
+ `use_base` drives the Modal A/B encoder toggle (off-the-shelf vs fine-tuned).
45
+ Raises on any transport/HTTP error so the caller can fall back."""
46
+ import requests
47
+ payload = {"emails": [_email_to_row(e) for e in emails], "use_base": bool(use_base),
48
+ "tau": tau, "max_attempts": max_attempts, "seed": seed}
49
+ resp = requests.post(endpoint(), json=payload, timeout=TIMEOUT)
50
+ resp.raise_for_status()
51
+ data = resp.json()
52
+ return _canvas_from_dict(data["canvas"]), data.get("trace", {})
data/images/msg0000.png ADDED
data/images/msg0002.png ADDED
data/images/msg0004.png ADDED
data/images/msg0005.png ADDED
data/images/msg0007.png ADDED
data/images/msg0012.png ADDED
data/images/msg0014.png ADDED
data/images/msg0015.png ADDED
data/images/msg0016.png ADDED
data/images/msg0017.png ADDED
data/images/msg0021.png ADDED
data/images/msg0022.png ADDED
data/images/msg0024.png ADDED
data/images/msg0034.png ADDED
data/images/msg0035.png ADDED
data/images/msg0038.png ADDED
data/images/msg0043.png ADDED