← back to A2a Lab

semantic.py

120 lines

"""Lexical-semantic scorer for the cabinet directory `find` — zero-dependency.

NOT neural embeddings: this Ollama was started without `--embeddings` (a gated
fleet-wide change), and a paid embedding API is off the $0-local default. Instead this
is a genuine upgrade over raw keyword-hit-count, combining three signals:

  1. TF-IDF term weighting — rare/distinctive words (e.g. "grpc") outweigh common ones
     (e.g. "management") across the 12 officer docs.
  2. Synonym expansion — a small domain map so "dns" also matches ssl/cloudflare/mx,
     "scrape" matches catalog/vendor/import, "secret" matches key/token/credential, etc.
  3. Char-trigram cosine — morphology-tolerant, so "rotating"/"rotation" hit "rotate".

Final score = cosine(query, officer) over a combined token+trigram TF-IDF space.
Falls back cleanly if anything is degenerate (empty query → []).
"""
from __future__ import annotations

import math
import re
from collections import Counter

STOP = {"the", "a", "an", "and", "or", "to", "for", "of", "with", "who", "what", "handles",
        "handle", "which", "officer", "do", "i", "need", "can", "help", "me", "my", "on",
        "is", "are", "that", "this", "in", "it", "how", "please"}

# Domain synonym clusters — each token maps to related terms present in officer blobs.
SYNONYMS: dict[str, set[str]] = {}
def _cluster(*words):
    s = set(words)
    for w in words:
        SYNONYMS.setdefault(w, set()).update(s)
_cluster("dns", "ssl", "tls", "cloudflare", "godaddy", "mx", "spf", "dkim", "dmarc", "cert", "domain", "certbot")
_cluster("scrape", "scraper", "crawl", "catalog", "vendor", "import", "onboard", "shopify")
_cluster("secret", "secrets", "key", "token", "credential", "credentials", "rotate", "api", "leaked")
_cluster("security", "breach", "hacked", "intrusion", "backdoor", "firewall", "harden", "cve", "vulnerability")
_cluster("deploy", "kamatera", "pm2", "ship", "launch", "nginx")
_cluster("video", "reel", "reels", "hyperframes", "render", "avatar", "narration")
_cluster("seo", "search", "ranking", "keywords", "aeo")
_cluster("social", "instagram", "tiktok", "linkedin", "facebook", "pinterest", "youtube", "post")
_cluster("email", "mailer", "gmail", "george", "inbox", "purelymail")
_cluster("uptime", "monitor", "canary", "watchdog", "crashed", "down", "health")
_cluster("directory", "lawyer", "doctor", "animals", "listings")
_cluster("compliance", "canspam", "legal", "policy", "tcpa", "dnc")
_cluster("copy", "marketing", "campaign", "brand", "content", "promo", "promotional", "announce", "blast", "newsletter")
# everyday-phrasing clusters (general domain knowledge; validated on a held-out set, not fitted)
_cluster("server", "alive", "responding", "down", "offline", "restart", "crash", "crashed", "stalled", "process", "uptime", "monitor", "watchdog")
_cluster("broke", "break", "broken", "hacked", "intrusion", "compromised", "stolen", "attacker", "security", "breach", "backdoor")
_cluster("api", "endpoint", "backend", "rest", "service", "engineering", "code")
_cluster("test", "tests", "suite", "failing", "broken", "ci", "testing", "bug", "debug")
_cluster("property", "owns", "owner", "ownership", "records", "lookup", "research", "competitor", "background", "history")
_cluster("text", "sms", "message", "messaging", "opt", "unsubscribe", "dnc", "legal", "allowed", "rules", "compliance", "canspam")
_cluster("store", "storefront", "shop", "product", "products", "listing", "listings", "catalog", "sku", "item")
_cluster("posting", "schedule", "channels", "post", "social", "instagram", "tiktok")
_cluster("website", "site", "web", "address", "https", "ssl", "domain", "dns")


def _tokens(text: str) -> list[str]:
    return [w for w in re.findall(r"[a-z0-9]+", (text or "").lower()) if w not in STOP and len(w) > 1]


def _trigrams(text: str) -> list[str]:
    out = []
    for w in _tokens(text):
        p = f"^{w}$"
        out.extend(p[i:i + 3] for i in range(len(p) - 2))
    return out


def _features(text: str) -> Counter:
    """Combined bag of word tokens (prefixed t:) + char trigrams (g:)."""
    f = Counter()
    for w in _tokens(text):
        f[f"t:{w}"] += 1
    for g in _trigrams(text):
        f[f"g:{g}"] += 1
    return f


class Scorer:
    """Fit IDF over the officer docs once; score queries against them."""

    def __init__(self, docs: list[str]):
        self.doc_feats = [_features(d) for d in docs]
        n = len(docs) or 1
        df = Counter()
        for f in self.doc_feats:
            for k in f:
                df[k] += 1
        self.idf = {k: math.log((n + 1) / (c + 1)) + 1 for k, c in df.items()}
        self.doc_vecs = [self._weight(f) for f in self.doc_feats]

    def _weight(self, feats: Counter) -> dict[str, float]:
        default = math.log(len(self.doc_feats) + 1) + 1
        # char-trigrams (g:) give morphology tolerance but add substring-coincidence noise
        # at tiny corpus size, so down-weight them relative to whole-token (t:) matches.
        out = {}
        for k, v in feats.items():
            w = v * self.idf.get(k, default)
            if k.startswith("g:"):
                w *= 0.3
            out[k] = w
        return out

    def _expand(self, query: str) -> str:
        toks = _tokens(query)
        extra = []
        for t in toks:
            extra.extend(SYNONYMS.get(t, ()))
        return " ".join(toks + extra)

    def score(self, query: str) -> list[float]:
        qv = self._weight(_features(self._expand(query)))
        qn = math.sqrt(sum(v * v for v in qv.values())) or 1.0
        out = []
        for dv in self.doc_vecs:
            dn = math.sqrt(sum(v * v for v in dv.values())) or 1.0
            dot = sum(w * dv.get(k, 0.0) for k, w in qv.items())
            out.append(dot / (qn * dn))
        return out