git.iko.soy/diagonal
on master help sign in
diagonal/host/validate.py
320 lines 13 kB history raw
1"""Post-validation of model replies (sections 6, 8 and 12). Repairs what it can, rejects the rest."""
2import os
3import re
4import zlib
5from urllib.parse import urlsplit
7HERE = os.path.dirname(os.path.realpath(__file__))
8COLORS = ["grey", "blue", "red", "yellow", "green", "pink", "purple", "cyan", "orange"]
9BANNED_WORDS = {"tab", "tabs", "group", "groups", "misc", "various", "stuff"}
10TRAILING_PUNCT = re.compile(r"[\s.,;:!?…\-–—。、，！？：；]+$")
13class ValidationError(Exception):
14 pass
17def _load_emoji():
18 categories = {}
19 with open(os.path.join(HERE, "emoji.txt"), encoding="utf-8") as f:
20 for line in f:
21 line = line.strip()
22 if not line or line.startswith("#"):
23 continue
24 name, rest = line.split(":", 1)
25 categories[name.strip()] = rest.split()
26 return categories
29CATEGORIES = _load_emoji()
32def _norm(e):
33 return str(e).replace("️", "").strip()
36CANONICAL = {_norm(e): e for emojis in CATEGORIES.values() for e in emojis}
39def all_emoji():
40 return list(CANONICAL.values())
43def canonical_emoji(e):
44 return CANONICAL.get(_norm(e)) if isinstance(e, str) else None
47# Majority-host fallback for an emoji outside the list: a small map of well-known sites.
48DOMAIN_CATEGORY = {
49 "technology": ["github.com", "gitlab.com", "stackoverflow.com", "stackexchange.com", "npmjs.com", "pypi.org", "crates.io",
50 "docs.rs", "rust-lang.org", "python.org", "developer.mozilla.org", "developer.apple.com", "developer.chrome.com",
51 "news.ycombinator.com", "vercel.com", "cloudflare.com", "aws.amazon.com", "nixos.org", "kernel.org"],
52 "reading": ["wikipedia.org", "medium.com", "substack.com", "nytimes.com", "theguardian.com", "bbc.co.uk", "bbc.com",
53 "reuters.com", "arstechnica.com", "theverge.com", "goodreads.com"],
54 "shopping": ["amazon.com", "amazon.co.uk", "amazon.de", "ebay.com", "etsy.com", "aliexpress.com", "bestbuy.com", "ikea.com"],
55 "travel": ["booking.com", "airbnb.com", "expedia.com", "skyscanner.net", "tripadvisor.com", "maps.google.com", "kayak.com"],
56 "media": ["youtube.com", "youtu.be", "netflix.com", "spotify.com", "twitch.tv", "vimeo.com", "soundcloud.com", "imdb.com"],
57 "work": ["docs.google.com", "sheets.google.com", "notion.so", "figma.com", "slack.com", "linear.app", "atlassian.net",
58 "trello.com", "asana.com", "calendar.google.com"],
59 "finance": ["paypal.com", "stripe.com", "coinbase.com", "bloomberg.com", "finance.yahoo.com"],
60 "science": ["arxiv.org", "nature.com", "sciencedirect.com", "pubmed.ncbi.nlm.nih.gov", "scholar.google.com"],
61 "communication": ["mail.google.com", "outlook.live.com", "outlook.office.com", "web.whatsapp.com", "discord.com", "messenger.com"],
63HOST_TO_CATEGORY = {h: c for c, hosts in DOMAIN_CATEGORY.items() for h in hosts}
66def _host(url):
67 try:
68 return (urlsplit(str(url)).hostname or "").removeprefix("www.")
69 except ValueError:
70 return ""
73def category_for_host(host):
74 parts = host.split(".")
75 for n in range(len(parts)):
76 c = HOST_TO_CATEGORY.get(".".join(parts[n:]))
77 if c:
78 return c
79 return "misc"
82def default_emoji(items):
83 hosts = [_host(it.get("url", "")) for it in items]
84 hosts = [h for h in hosts if h]
85 if not hosts:
86 return CATEGORIES["misc"][0]
87 majority = max(set(hosts), key=hosts.count)
88 return CATEGORIES[category_for_host(majority)][0]
91# Quotes go; an apostrophe between two letters ("Children's", "Rock'n'roll") is part of the word and stays.
92_QUOTE = re.compile(r"(?<![^\W\d_])['’]|['’](?![^\W\d_])|[\"“”‘«»`]")
93# A label cut short should not end mid-phrase: "Planning a trip to" reads as "Planning a trip".
94DANGLING = {"a", "an", "the", "to", "of", "for", "in", "on", "at", "by", "and", "or", "with", "from"}
97def repair_label(raw):
98 """Trim, drop quotes, banned words and trailing punctuation, cap at 4 words / 28 chars (16 wide characters
99 for no-space scripts), drop a dangling "to"/"of", capitalise. Kept in step with src/shared/label.ts."""
100 if not isinstance(raw, str):
101 return None
102 s = _QUOTE.sub("", raw)
103 s = TRAILING_PUNCT.sub("", " ".join(s.split()))
104 words = [w for w in s.split(" ") if w and re.sub(r"[^\w]", "", w.lower()) not in BANNED_WORDS]
105 words = words[:4]
106 if _no_space(" ".join(words)):
107 while len(words) > 1 and _width(" ".join(words)) > 16:
108 words.pop()
109 if _width(" ".join(words)) > 16:
110 words = [_cut_to_width(" ".join(words), 16)]
111 else:
112 while len(words) > 2 and len(" ".join(words)) > 28:
113 words.pop()
114 while len(words) > 1 and words[-1].lower() in DANGLING:
115 words.pop()
116 label = TRAILING_PUNCT.sub("", " ".join(words))
117 if not label:
118 return None
119 # "rust async" becomes "Rust async"; "iPhone tips" and "macOS updates" keep the brand's own casing.
120 if not any(c.isupper() for c in words[0][1:]):
121 label = label[0].upper() + label[1:]
122 return label if is_valid_label(label) else None
125# Chinese, Japanese, Korean, Thai, Lao, Khmer and Myanmar titles don't put spaces between words, so a label
126# with any of their characters ("Python教程") is checked by width instead of word count. Each such character
127# is about as wide as two Latin letters on the tab strip.
128NO_SPACE = re.compile(r"[\u0e00-\u0eff\u1000-\u109f\u1780-\u17ff\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uac00-\ud7af\uf900-\ufaff]")
129# Chinese, Japanese and Korean, for the host's token estimate.
130CJK = re.compile(r"[\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uac00-\ud7af\uf900-\ufaff]")
133def _no_space(label):
134 return bool(NO_SPACE.search(label))
137def _width(label):
138 wide = len(NO_SPACE.findall(label))
139 return wide + (len(label) - wide) / 2
142def _cut_to_width(label, most):
143 out = ""
144 for ch in label:
145 if _width(out + ch) > most:
146 break
147 out += ch
148 return out.strip()
151def is_valid_label(label):
152 words = label.split()
153 wide = _no_space(label)
154 return (
155 (2 <= len(label) <= 28 and _width(label) <= 16 if wide else 3 <= len(label) <= 28)
156 and (1 <= len(words) <= 4 if wide else 2 <= len(words) <= 4)
157 and not _QUOTE.search(label)
158 and not TRAILING_PUNCT.search(label)
159 and not any(w.lower() in BANNED_WORDS for w in words)
160 )
163def same_label(a, b):
164 return bool(a) and bool(b) and " ".join(a.split()).casefold() == " ".join(b.split()).casefold()
167def validate_name(out, payload):
168 if not isinstance(out, dict):
169 raise ValidationError("reply is not an object")
170 label = repair_label(out.get("title"))
171 if not label:
172 raise ValidationError(f"title {out.get('title')!r} breaks the label rules")
173 emoji = canonical_emoji(out.get("emoji")) or default_emoji(payload.get("items", []))
174 return {"title": label, "emoji": emoji}
177def _color_for(label):
178 choices = [c for c in COLORS if c != "grey"]
179 return choices[zlib.crc32(label.encode("utf-8")) % len(choices)]
182def validate_organize(out, payload):
183 """Every index in range and used once; `existing` in range; colours from the fixed set;
184 new groups under the minimum size or past maxGroups folded into leftovers; missing indexes appended."""
185 if not isinstance(out, dict):
186 raise ValidationError("reply is not an object")
187 items = payload.get("items", [])
188 n_items = len(items)
189 n_existing = len(payload.get("existingGroups") or [])
190 min_size = max(2, int(payload.get("minGroupSize", 2) or 2))
191 max_groups = max(1, int(payload.get("maxGroups", 8) or 8))
192 raw_groups = out.get("groups")
193 if not isinstance(raw_groups, list):
194 raise ValidationError("groups is not a list")
196 used = set()
197 joins, new = [], []
198 for g in raw_groups:
199 if not isinstance(g, dict):
200 continue
201 members = []
202 for i in g.get("members") or []:
203 if isinstance(i, bool) or not isinstance(i, int):
204 continue
205 if 0 <= i < n_items and i not in used:
206 used.add(i)
207 members.append(i)
208 if not members:
209 continue
210 ex = g.get("existing")
211 if isinstance(ex, list): # fm has returned the optional integer as [] or [n]
212 ex = ex[0] if len(ex) == 1 else None
213 if isinstance(ex, int) and not isinstance(ex, bool) and 0 <= ex < n_existing:
214 # Two topics can both point at one group (its example tabs got different topics): one join each.
215 join = next((j for j in joins if j["existing"] == ex), None)
216 if join:
217 join["members"].extend(members)
218 else:
219 joins.append({"title": "", "emoji": "", "color": "", "existing": ex, "members": members})
220 continue
221 new.append((members, g))
223 leftovers = []
224 new.sort(key=lambda mg: -len(mg[0]))
225 groups = list(joins)
226 for k, (members, g) in enumerate(new):
227 if len(members) < min_size or k >= max_groups:
228 leftovers.extend(members)
229 continue
230 label = repair_label(g.get("title")) or "" # "" = the extension names it through the naming loop
231 member_items = [items[i] for i in members]
232 emoji = canonical_emoji(g.get("emoji")) or default_emoji(member_items)
233 color = g.get("color") if g.get("color") in COLORS else _color_for(label or str(members))
234 groups.append({"title": label, "emoji": emoji, "color": color, "members": members})
236 for i in out.get("leftovers") or []:
237 if isinstance(i, int) and not isinstance(i, bool) and 0 <= i < n_items and i not in used:
238 used.add(i)
239 leftovers.append(i)
240 leftovers.extend(i for i in range(n_items) if i not in used and i not in leftovers)
241 return {"groups": groups, "leftovers": sorted(set(leftovers))}
244VAGUE_TOPICS = {"other", "others", "misc", "miscellaneous", "general", "various", "unknown", "none", "n/a", "-"}
247def _norm_topic(t):
248 t = " ".join(str(t or "").split()).strip(" .,:;!?\"'").casefold()
249 return "" if t in VAGUE_TOPICS else t
252def organize_from_topics(out, payload, owner):
253 """Tabs that share a topic form a group. A topic that an existing group's example tabs have joins that
254 group; others become new groups (titled later by the naming loop) or leftovers when too small.
255 `owner` maps the example tabs' indexes to their group number (prompts.topic_items)."""
256 if not isinstance(out, dict) or not isinstance(out.get("tabs"), list):
257 raise ValidationError("tabs is not a list")
258 n_items = len(payload.get("items", []))
259 tab_topic, group_votes = {}, {}
260 for t in out["tabs"]:
261 if not isinstance(t, dict):
262 continue
263 i, topic = t.get("index"), _norm_topic(t.get("topic"))
264 if isinstance(i, bool) or not isinstance(i, int) or not topic:
265 continue
266 if 0 <= i < n_items:
267 tab_topic.setdefault(i, topic)
268 elif i in owner:
269 votes = group_votes.setdefault(topic, {})
270 votes[owner[i]] = votes.get(owner[i], 0) + 1
271 if not tab_topic:
272 if not any(isinstance(t, dict) and _in_range(t.get("index"), n_items) for t in out["tabs"]):
273 raise ValidationError("no tab got a topic")
274 # Every tab got only a vague topic ("Other", "Misc"): the model found nothing in common.
275 return validate_organize({"groups": [], "leftovers": []}, payload)
276 clusters = {}
277 for i in sorted(tab_topic):
278 clusters.setdefault(tab_topic[i], []).append(i)
279 groups = []
280 for topic, members in clusters.items():
281 g = _winner(group_votes.get(topic) or {}, members, payload)
282 if g is not None:
283 groups.append({"existing": g, "members": members})
284 else:
285 groups.append({"title": "", "emoji": "", "color": "", "members": members})
286 return validate_organize({"groups": groups, "leftovers": []}, payload)
289def _in_range(i, n):
290 return isinstance(i, int) and not isinstance(i, bool) and 0 <= i < n
293_WORD = re.compile(r"\w{3,}")
296def _words(text):
297 return {w.casefold() for w in _WORD.findall(text or "")}
300def _winner(votes, members, payload):
301 """The existing group a topic's tabs join: the one whose example tabs share it most. When two groups tie,
302 the one whose title and examples share more words with these tabs; when that ties too, none (the tabs
303 form a group of their own, or stay loose), rather than whichever group happens to be listed first."""
304 if not votes:
305 return None
306 top = max(votes.values())
307 tied = sorted(k for k, v in votes.items() if v == top)
308 if len(tied) == 1:
309 return tied[0]
310 tab_words = set().union(*(_words(payload["items"][i].get("title")) for i in members))
311 groups = {g["g"]: g for g in payload.get("existingGroups") or []}
313 def overlap(k):
314 g = groups.get(k) or {}
315 return len(tab_words & _words(" ".join([g.get("title") or "", *(g.get("samples") or [])])))
317 scores = {k: overlap(k) for k in tied}
318 best = max(scores.values())
319 leaders = [k for k in tied if scores[k] == best]
320 return leaders[0] if len(leaders) == 1 and best > 0 else None