GalaxyGate Track winner: 2,746 messages in, 94 tokens out

"""Dedupe and supersession (PRD section 6.2's supersession rule).

Given a validated `ExtractionResult` and an open db connection: each fact's
entity is resolved/inserted via `db.upsert_entity`, then compared against
the current fact (if any) for that `(entity_id, predicate)`. Same value
widens confidence in place; a different value inserts the new fact and
marks the old one superseded; no prior fact just inserts. The episode (if
present) is written via `db.insert_episode`.

Entities carry no `category` in the extraction schema (section 6.1), but
the `entities` table requires one (section 6.2). Judgment call: an entity
referenced by a fact takes that fact's category (facts are the source of
truth for "what kind of thing is this"); an entity that appears only in
`extraction.entities` with no fact falls back to a category guessed from
its `kind` (a "person"/"people" kind maps to the `people` category), else
`projects`.
"""

from __future__ import annotations

import sqlite3

from vault import db
from vault.models import ExtractedEntity, ExtractionResult


def _guess_category(kind: str) -> str:
    k = (kind or "").lower()
    if "person" in k or "people" in k:
        return "people"
    return "projects"


def apply_extraction(
    conn: sqlite3.Connection,
    extraction: ExtractionResult,
    source_id: int,
) -> dict[str, int | None]:
    entity_meta: dict[str, ExtractedEntity] = {e.name: e for e in extraction.entities}
    fact_category_by_entity: dict[str, str] = {}

    added = 0
    superseded = 0

    for fact in extraction.facts:
        fact_category_by_entity.setdefault(fact.entity, fact.category)
        meta = entity_meta.get(fact.entity)
        kind = meta.kind if meta else "topic"
        description = meta.description if meta else ""

        entity_id = db.upsert_entity(conn, fact.entity, kind, description, fact.category)
        current = db.get_current_fact(conn, entity_id, fact.predicate)

        if current is None:
            db.insert_fact(
                conn,
                entity_id=entity_id,
                subject=fact.subject,
                predicate=fact.predicate,
                value=fact.value,
                category=fact.category,
                sensitive=fact.sensitive,
                confidence=fact.confidence,
                source_id=source_id,
            )
            added += 1
        elif current["value"].strip() == fact.value.strip():
            db.bump_confidence(conn, current["id"], fact.confidence)
        else:
            new_id = db.insert_fact(
                conn,
                entity_id=entity_id,
                subject=fact.subject,
                predicate=fact.predicate,
                value=fact.value,
                category=fact.category,
                sensitive=fact.sensitive,
                confidence=fact.confidence,
                source_id=source_id,
            )
            db.supersede_fact(conn, current["id"], new_id)
            added += 1
            superseded += 1

    # Entities mentioned only in extraction.entities (no fact referenced
    # them) still belong in the catalog.
    for name, meta in entity_meta.items():
        if name in fact_category_by_entity:
            continue
        category = _guess_category(meta.kind)
        db.upsert_entity(conn, name, meta.kind, meta.description, category)

    episode_id: int | None = None
    if extraction.episode is not None:
        episode_id = db.insert_episode(
            conn,
            source_id=source_id,
            summary=extraction.episode.summary,
            category=extraction.episode.category,
            tags=extraction.episode.tags,
            entities=extraction.episode.entities,
        )

    return {"added": added, "superseded": superseded, "episode_id": episode_id}

"""Gate: decide whether a prompt needs memory at all (PRD section 8).

``needs_memory`` returns False for prompts that don't reference anything
durable -- those get the profile block only, nothing else.
"""

from __future__ import annotations

import re

POSSESSIVES = re.compile(r"\b(my|our|mine)\b", re.IGNORECASE)

TRIGGER_PHRASES = ("remember", "recall", "continue", "last time", "we decided")


def _edit_distance(a: str, b: str) -> int:
    """Classic Levenshtein distance, no dependency required."""
    if a == b:
        return 0
    if not a:
        return len(b)
    if not b:
        return len(a)
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, start=1):
        cur = [i] + [0] * len(b)
        for j, cb in enumerate(b, start=1):
            cost = 0 if ca == cb else 1
            cur[j] = min(
                prev[j] + 1,  # deletion
                cur[j - 1] + 1,  # insertion
                prev[j - 1] + cost,  # substitution
            )
        prev = cur
    return prev[-1]


def _fuzzy_contains(prompt_lower: str, name_lower: str) -> bool:
    """True if name_lower appears in prompt_lower exactly, or within one
    edit distance over a same-length window of words."""
    if not name_lower:
        return False
    if name_lower in prompt_lower:
        return True
    words = re.findall(r"[a-z0-9']+", prompt_lower)
    name_words = re.findall(r"[a-z0-9']+", name_lower)
    n = len(name_words)
    if n == 0 or n > len(words):
        return False
    joined_name = " ".join(name_words)
    for i in range(len(words) - n + 1):
        window = " ".join(words[i : i + n])
        if _edit_distance(window, joined_name) <= 1:
            return True
    return False


def mentions_entity(prompt: str, catalog_names: list[str]) -> str | None:
    """Return the first catalog entity name fuzzy-mentioned in prompt, else
    None. Longer names are checked first so e.g. "acme-platform" wins over
    a shorter coincidental overlap."""
    lower = prompt.lower()
    for name in sorted(catalog_names, key=len, reverse=True):
        if _fuzzy_contains(lower, name.lower()):
            return name
    return None


def needs_memory(prompt: str, catalog_names: list[str], scope: str | None) -> bool:
    """True when the prompt has a possessive, names a catalog entity
    (fuzzy within one edit), contains a recall trigger phrase, or scope is
    explicitly given. Otherwise the caller should return the profile block
    only (PRD section 8)."""
    if scope:
        return True
    lower = prompt.lower()
    if POSSESSIVES.search(lower):
        return True
    if any(phrase in lower for phrase in TRIGGER_PHRASES):
        return True
    if mentions_entity(prompt, catalog_names) is not None:
        return True
    return False

"""Secret scrubbing (PRD section 9). Runs before any write to the db or
memory/: `scrub(text)` redacts, in order, `sk-`/`sk-ant-` API keys, AWS
`AKIA...` keys, GitHub `ghp_...` tokens, RunPod `rpa_...` keys, three-part
base64 JWTs, 13-19 digit runs that pass Luhn, `###-##-####` SSNs, and any
32+ character token with Shannon entropy above 4.0 bits/char. Each match is
replaced with `[REDACTED:<type>]` and counted by type.

Order matters: earlier, more specific patterns are redacted first, so by the
time the generic high-entropy pattern runs, already-redacted placeholders
(`[REDACTED:api_key]`, ...) no longer look like bare tokens and are not
double-counted.
"""

from __future__ import annotations

import math
import re

# --- specific patterns, most specific first -------------------------------

_API_KEY_RE = re.compile(r"\bsk-(?:ant-)?[A-Za-z0-9_-]{16,}\b")
_AWS_KEY_RE = re.compile(r"\bAKIA[0-9A-Z]{16}\b")
_GITHUB_TOKEN_RE = re.compile(r"\bghp_[A-Za-z0-9]{20,}\b")
_RUNPOD_KEY_RE = re.compile(r"\brpa_[A-Za-z0-9]{20,}\b")
_JWT_RE = re.compile(r"\b[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b")
_SSN_RE = re.compile(r"\b\d{3}-\d{2}-\d{4}\b")
_DIGIT_RUN_RE = re.compile(r"\b\d{13,19}\b")
_GENERIC_TOKEN_RE = re.compile(r"\b[A-Za-z0-9+/_=-]{32,}\b")
# Checked against a token the generic pattern already matched, so it only has
# to recognise the alphabet, not re-find the boundaries.
_HEX_RUN_RE = re.compile(r"[0-9a-fA-F]{32,}")
_HEX_ENTROPY_FLOOR = 3.0

_ENTROPY_THRESHOLD = 4.0

# Order in which the specific (non-entropy) patterns are applied. Each is a
# (label, pattern) pair; the label becomes `[REDACTED:<label>]`.
_SPECIFIC_PATTERNS: list[tuple[str, re.Pattern[str]]] = [
    ("api_key", _API_KEY_RE),
    ("aws_key", _AWS_KEY_RE),
    ("github_token", _GITHUB_TOKEN_RE),
    ("runpod_key", _RUNPOD_KEY_RE),
    ("jwt", _JWT_RE),
    ("ssn", _SSN_RE),
]


def _luhn_valid(digits: str) -> bool:
    total = 0
    for i, ch in enumerate(reversed(digits)):
        d = int(ch)
        if i % 2 == 1:
            d *= 2
            if d > 9:
                d -= 9
        total += d
    return total % 10 == 0


def _shannon_entropy(s: str) -> float:
    if not s:
        return 0.0
    counts: dict[str, int] = {}
    for ch in s:
        counts[ch] = counts.get(ch, 0) + 1
    length = len(s)
    return -sum((n / length) * math.log2(n / length) for n in counts.values())


def scrub(text: str) -> tuple[str, dict[str, int]]:
    """Redact secrets in `text`. Returns (clean_text, counts_by_type)."""
    counts: dict[str, int] = {}
    clean = text

    for label, pattern in _SPECIFIC_PATTERNS: