@pentatonic-ai/ai-agent-sdk 0.10.1 → 0.10.2

package/packages/memory-engine-v2/extractor-async/worker.py

@@ -40,7 +40,9 @@ import psycopg
 import psycopg.rows
 
 from confidence import corroborated_confidence
+from entity_id import entity_id, normalize_surface_form
 from noise_filter import is_noise_entity_name
+from sensitive_filter import SKIP_SENSITIVE_CONTENT, is_sensitive_event
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 log = logging.getLogger("extractor-async")
@@ -74,6 +76,13 @@ LLM_MAX_TOKENS_PER_EVENT = int(os.environ.get("LLM_MAX_TOKENS_PER_EVENT", "300")
 
 WORKER_ID = f"{socket.gethostname()}:{os.getpid()}"
 
+# Trace logging — captures raw teacher I/O per distilled event so we can
+# train a student model (BART/FLAN-T5) on the teacher's distribution.
+# Default off; opt-in per environment. See migration 003.
+DISTILL_TRACE_ENABLED = os.environ.get(
+    "DISTILL_TRACE_ENABLED", "false"
+).strip().lower() in ("true", "1", "yes", "on")
+
 
 # KV-text output format constants. We dropped JSON output (and the
 # `guided_json` schema enforcement that went with it) because a single
@@ -124,8 +133,16 @@ RULES:
 - Each event MUST start with a `=== event K ===` header (zero-indexed, \
 matching the input index). NEVER skip an event — if an event has \
 nothing to extract, emit ONLY the header.
-- ENT lines have exactly 3 fields: literal `ENT`, type, name.
+- ENT lines have 3 or 4 fields: literal `ENT`, type, name, [email].
   type ∈ {person, org, product, place, project, concept, topic, date, other}
+  email (OPTIONAL, person only): when the event body or attributes
+  show an email address that unambiguously identifies the person,
+  append it as the 4th field. This pairs the name+email forms so a
+  later event seeing only the email resolves to the same entity.
+  Examples:
+    ENT|person|Alex Wong|alex@example.com
+    ENT|person|Acme Corp           (org, no email)
+    ENT|person|Sam Patel           (person, email not visible)
 - FCT lines have exactly 6 fields: `FCT`, category, subject, \
 predicate, object, statement.
   category ∈ {decision, commitment, state, mention, observation, preference}
@@ -144,6 +161,12 @@ A whole file is one entity, not twenty.
 - Output ONLY the formatted records. No header, no footer, no prose."""
 
 
+# Teacher-prompt fingerprint for trace logging. If the prompt changes,
+# the hash changes — lets training-data exports filter by teacher
+# version so we never mix outputs from a retired prompt.
+SYSTEM_PROMPT_HASH = hashlib.sha256(BATCH_SYSTEM_PROMPT.encode()).hexdigest()[:16]
+
+
 def build_event_block(idx: int, event: dict[str, Any]) -> str:
     """Render one event as `[event K]\nheader\n---\ncontent` block."""
     src = event.get("source_kind", "unknown")
@@ -200,11 +223,24 @@ def _parse_kv_records(text: str, expected_n: int) -> list[dict[str, Any]]:
         # maxsplit so statement / name fields can contain colons or
         # other reserved-looking content without breaking parsing.
         if line.startswith("ENT|"):
-            parts = line.split("|", 2)
-            if len(parts) == 3 and parts[2].strip():
-                current["entities"].append(
-                    {"type": parts[1].strip().lower(), "name": parts[2].strip()}
-                )
+            # ENT|type|name|email? — email is optional, person-only.
+            # Use maxsplit=3 so a literal `|` in name (which the prompt
+            # forbids but the model might still emit) doesn't get
+            # parsed as an email field.
+            parts = line.split("|", 3)
+            if len(parts) >= 3 and parts[2].strip():
+                etype = parts[1].strip().lower()
+                name = parts[2].strip()
+                ent: dict[str, Any] = {"type": etype, "name": name}
+                # Promote 4th-field email into aliases when present
+                # and it actually looks like an email. Non-email
+                # 4th fields (random text the model added) are dropped
+                # — better to silently strip junk than poison aliases.
+                if len(parts) == 4:
+                    email = parts[3].strip()
+                    if email and "@" in email and " " not in email:
+                        ent["aliases"] = [email]
+                current["entities"].append(ent)
         elif line.startswith("FCT|"):
             parts = line.split("|", 5)
             if len(parts) == 6 and parts[5].strip():
@@ -232,6 +268,38 @@ def _parse_kv_records(text: str, expected_n: int) -> list[dict[str, Any]]:
     return results
 
 
+def _split_event_blocks(text: str, expected_n: int) -> list[str]:
+    """Slice raw LLM output by `=== event K ===` headers.
+
+    Returns expected_n slices in event-index order. Each slice is the
+    verbatim text between a header and the next header (or end-of-text),
+    stripped of trailing whitespace. Events the model omitted come back
+    as empty strings — same shape contract as _parse_kv_records.
+
+    Separate from the parser so trace logging stays decoupled from
+    extraction semantics: the parser drops malformed lines silently
+    (correctness), but the trace wants the raw output verbatim
+    (training fidelity)."""
+    slices: list[str] = [""] * expected_n
+    current_idx: int | None = None
+    current_lines: list[str] = []
+
+    def flush() -> None:
+        if current_idx is not None and 0 <= current_idx < expected_n:
+            slices[current_idx] = "\n".join(current_lines).rstrip()
+
+    for raw in text.splitlines():
+        m = EVENT_HEADER_RE.match(raw.strip()) if raw.strip() else None
+        if m:
+            flush()
+            current_idx = int(m.group(1))
+            current_lines = []
+        elif current_idx is not None:
+            current_lines.append(raw)
+    flush()
+    return slices
+
+
 async def call_llm_batch(
     client: httpx.AsyncClient, events: list[dict[str, Any]]
 ) -> list[dict[str, Any]]:
@@ -279,7 +347,15 @@ async def call_llm_batch(
         text = data.get("message", {}).get("content", "")
     if not text:
         raise RuntimeError(f"llm returned no content: {json.dumps(data)[:300]}")
-    return _parse_kv_records(text, n)
+    parsed = _parse_kv_records(text, n)
+    # Attach the per-event raw slice so downstream trace logging gets
+    # the model's verbatim output for THIS event without re-splitting
+    # the chunk-level text. Parser semantics are unaffected — the
+    # raw_slice key is ignored by upsert paths.
+    slices = _split_event_blocks(text, n)
+    for record, slice_text in zip(parsed, slices):
+        record["raw_slice"] = slice_text
+    return parsed
 
 
 # --------------------------------------------------------------------
@@ -288,6 +364,10 @@ async def call_llm_batch(
 
 
 def _content_id(*parts: str) -> str:
+    """Deterministic content-addressed id for facts / relationships.
+    Entity ids are minted via `entity_id()` from entity_id.py (see
+    `upsert_entities` below); this helper covers the non-entity
+    content-hash needs."""
     return hashlib.sha256("\x1f".join(parts).encode()).hexdigest()[:32]
 
 
@@ -299,12 +379,33 @@ def upsert_entities(
     disclosure_class: str,
     entities: list[dict],
 ) -> dict[str, str]:
-    """Insert (or merge) entities; return a name→id map so facts and
-    relationships can link to the inserted rows.
-
-    ID is sha256(arena:entity_type:canonical_name)[:32] so the same
-    entity in the same arena converges across events. Aliases and
-    provenance_event_ids array-append on conflict; never replace."""
+    """Alias-aware insert (or merge) of entities; returns a name→id
+    map so facts and relationships can link to the inserted rows.
+
+    Two concerns layered together:
+
+    1. **ID derivation** uses the shared `entity_id()` helper from
+       entity_id.py: `e_` + 24 hex of sha256("{arena}|{entity_type}|
+       {normalize_surface_form(name)}"). BYTE-IDENTICAL to extractor-
+       sync's id derivation, so the same person extracted by both
+       passes converges to the same row instead of fragmenting across
+       two id schemes. (RFC step 1.)
+
+    2. **Resolution at upsert** — before INSERT, check for existing
+       rows in the same (arena, entity_type) whose canonical_name or
+       aliases overlap any incoming surface form. If matched, merge
+       into the existing row. Per-form `pg_advisory_xact_lock`
+       serialises concurrent writers (sync + async on the same event)
+       on the same surface form. (RFC steps 2 + 2a.)
+
+    MIRROR of extractor-sync/server.py:_upsert_entities — same
+    resolution algorithm. Kept as separate Python because the sync
+    extractor uses async psycopg and the async worker uses sync
+    psycopg; the SQL is identical.
+
+    Returns name→id where `name` is the LLM-emitted surface form
+    (canonical) so facts/relationships using the same surface form
+    in the same LLM batch resolve to the right id."""
     name_to_id: dict[str, str] = {}
     if not entities:
         return name_to_id
@@ -314,39 +415,84 @@ def upsert_entities(
             name = (e.get("name") or "").strip()
             if not name:
                 continue
-            # Drop junk names before they enter the graph. See
-            # noise_filter.py — patterns are anchored to live-arena
-            # noise (pronouns, hostnames, paths, agent-worktree
-            # labels). Skipping here means name_to_id never carries
-            # the bad name, so any fact/relationship the LLM tried to
-            # attach to it gets dropped downstream (subj/obj resolve
-            # to None ⇒ filtered out by upsert_facts /
-            # upsert_relationships).
+            # Drop junk names before they enter the graph.
             if is_noise_entity_name(etype, name):
                 continue
             aliases = [a for a in (e.get("aliases") or []) if a]
-            eid = _content_id(arena, etype, name)
-            name_to_id[name] = eid
+
+            # Sort (don't `list(set(...))`) so lock acquisition order
+            # is deterministic across processes — set-iteration order
+            # depends on Python's per-process hash randomisation, so
+            # sync and async extractors processing the same person
+            # could otherwise acquire the same locks in opposite
+            # orders and deadlock.
+            forms_original = sorted({name, *aliases})
+            forms_normalized = sorted({normalize_surface_form(f) for f in forms_original})
+
+            # 1. Advisory lock per surface form. Serialises concurrent
+            # writers (sync + async on the same event) on the same
+            # person without blocking anything else. See RFC §2a.
+            for f in forms_normalized:
+                cur.execute(
+                    "SELECT pg_advisory_xact_lock(hashtext(%s))",
+                    (f"{arena}|{etype}|{f}",),
+                )
+
+            # 2. Resolve via canonical-name (normalised, case-insensitive)
+            # or aliases overlap.
             cur.execute(
                 """
-                INSERT INTO entities (
-                  id, arena, entity_type, canonical_name, aliases,
-                  provenance_event_ids, participant_set, disclosure_class
-                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::disclosure_class)
-                ON CONFLICT (id) DO UPDATE SET
-                  aliases = (
-                    SELECT ARRAY(SELECT DISTINCT UNNEST(entities.aliases || EXCLUDED.aliases))
-                  ),
-                  provenance_event_ids = (
-                    SELECT ARRAY(SELECT DISTINCT UNNEST(entities.provenance_event_ids || EXCLUDED.provenance_event_ids))
-                  ),
-                  last_seen = NOW()
+                SELECT id FROM entities
+                WHERE arena = %s
+                  AND entity_type = %s
+                  AND (
+                    lower(canonical_name) = ANY(%s::text[])
+                    OR aliases && %s::text[]
+                  )
+                LIMIT 1
                 """,
-                (
-                    eid, arena, etype, name, aliases,
-                    [event_id], participant_set, disclosure_class,
-                ),
+                (arena, etype, forms_normalized, forms_original),
             )
+            row = cur.fetchone()
+
+            if row is not None:
+                # 3a. Existing match — merge aliases + provenance.
+                # Canonical stays as-was (accrete-only per RFC §2b).
+                eid = row[0]
+                cur.execute(
+                    """
+                    UPDATE entities SET
+                      aliases = ARRAY(SELECT DISTINCT UNNEST(aliases || %s::text[])),
+                      provenance_event_ids = ARRAY(SELECT DISTINCT UNNEST(provenance_event_ids || %s::text[])),
+                      last_seen = NOW()
+                    WHERE id = %s
+                    """,
+                    (aliases, [event_id], eid),
+                )
+            else:
+                # 3b. No match — insert new.
+                eid = entity_id(arena, etype, name)
+                cur.execute(
+                    """
+                    INSERT INTO entities (
+                      id, arena, entity_type, canonical_name, aliases,
+                      provenance_event_ids, participant_set, disclosure_class
+                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::disclosure_class)
+                    ON CONFLICT (id) DO UPDATE SET
+                      aliases = (
+                        SELECT ARRAY(SELECT DISTINCT UNNEST(entities.aliases || EXCLUDED.aliases))
+                      ),
+                      provenance_event_ids = (
+                        SELECT ARRAY(SELECT DISTINCT UNNEST(entities.provenance_event_ids || EXCLUDED.provenance_event_ids))
+                      ),
+                      last_seen = NOW()
+                    """,
+                    (
+                        eid, arena, etype, name, aliases,
+                        [event_id], participant_set, disclosure_class,
+                    ),
+                )
+            name_to_id[name] = eid
     return name_to_id
 
 
@@ -486,6 +632,43 @@ def upsert_relationships(
     return inserted
 
 
+# --------------------------------------------------------------------
+# Distillation trace logging
+# --------------------------------------------------------------------
+
+
+def _insert_trace(
+    conn: psycopg.Connection,
+    *,
+    event_id: str,
+    user_prompt: str,
+    raw_response: str,
+    llm_chunk_ms: float | None,
+) -> None:
+    """Append a (user_prompt, raw_response) pair to distillation_traces.
+
+    Audit-only — not on the hot path of distillation semantics. The
+    caller wraps this in a try/except and never lets a trace-insert
+    failure poison the upsert path. Skip rows with empty raw_response
+    (no signal to train on); the model occasionally emits a header
+    with no body."""
+    if not raw_response.strip():
+        return
+    with conn.cursor() as cur:
+        cur.execute(
+            """
+            INSERT INTO distillation_traces (
+              event_id, user_prompt, raw_response,
+              llm_model, system_prompt_hash, llm_chunk_ms
+            ) VALUES (%s, %s, %s, %s, %s, %s)
+            """,
+            (
+                event_id, user_prompt, raw_response,
+                LLM_MODEL, SYSTEM_PROMPT_HASH, llm_chunk_ms,
+            ),
+        )
+
+
 # --------------------------------------------------------------------
 # Queue mechanics
 # --------------------------------------------------------------------
@@ -570,6 +753,29 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
                 """,
                 (list(SKIP_ATTRIBUTE_SOURCES),),
             )
+        # Pre-filter: content-guardrail sensitive events (interpersonal
+        # gossip about a colleague). Never distil gossip into the entity
+        # graph — the subject has no standing there. attributes is jsonb;
+        # cast defensively in case a producer wrote json.
+        if SKIP_SENSITIVE_CONTENT:
+            cur.execute(
+                """
+                UPDATE distillation_queue dq SET
+                  status = 'done',
+                  completed_at = NOW(),
+                  last_error = 'filtered: content-guardrail sensitive'
+                FROM events e
+                WHERE dq.event_id = e.id
+                  AND dq.status = 'pending'
+                  AND (
+                    (e.attributes::jsonb)->>'sensitivity_class' = 'interpersonal'
+                    OR (
+                      jsonb_typeof((e.attributes::jsonb)->'sensitive_about') = 'array'
+                      AND jsonb_array_length((e.attributes::jsonb)->'sensitive_about') > 0
+                    )
+                  )
+                """
+            )
         # Pre-filter: events older than the window.
         cur.execute(
             """
@@ -704,14 +910,24 @@ async def process_batch(
     for item in items:
         events_by_qid[item["id"]] = fetch_event(conn, item["event_id"])
 
-    # Drop items whose event is missing (mark done up-front, no LLM call).
+    # Drop items whose event is missing (mark done up-front, no LLM call),
+    # and — content guardrail — any sensitive event that slipped past the
+    # claim-time pre-filter (defense in depth; the pure predicate is the
+    # testable contract for the rule). Never distil interpersonal gossip.
     callable_items: list[dict[str, Any]] = []
     for item in items:
-        if events_by_qid[item["id"]] is None:
+        ev = events_by_qid[item["id"]]
+        if ev is None:
             log.warning(
                 f"event {item['event_id']} missing — marking queue {item['id']} done"
             )
             mark_done(conn, item["id"])
+        elif SKIP_SENSITIVE_CONTENT and is_sensitive_event(ev):
+            log.info(
+                f"content-guardrail: filtering sensitive event {item['event_id']} "
+                f"— not distilling gossip into the graph"
+            )
+            mark_done(conn, item["id"])
         else:
             callable_items.append(item)
 
@@ -740,7 +956,7 @@ async def process_batch(
 
     # Flatten chunk_outcomes back to per-item results, paired with items.
     for (chunk_items, _chunk_events), (per_item, llm_ms) in zip(chunks, chunk_outcomes):
-        for item, result in zip(chunk_items, per_item):
+        for local_idx, (item, result) in enumerate(zip(chunk_items, per_item)):
             queue_id = item["id"]
             event_id = item["event_id"]
             attempts = item["attempts"]
@@ -781,6 +997,24 @@ async def process_batch(
                     f"relationships={n_rels}"
                     + (f" llm_ms={llm_ms:.0f}/chunk" if not stub_mode else "")
                 )
+                # Trace logging — best-effort, never breaks the worker.
+                # Captures (input, output) so a student model can be
+                # trained on the teacher's distribution. Skipped in
+                # stub mode (no real LLM output to record).
+                if DISTILL_TRACE_ENABLED and not stub_mode:
+                    try:
+                        _insert_trace(
+                            conn,
+                            event_id=event_id,
+                            user_prompt=build_event_block(local_idx, event),
+                            raw_response=result.get("raw_slice", ""),
+                            llm_chunk_ms=llm_ms,
+                        )
+                    except Exception as trace_exc:
+                        log.warning(
+                            f"trace insert failed queue_id={queue_id} "
+                            f"event_id={event_id}: {trace_exc}"
+                        )
             except Exception as exc:
                 err = f"{type(exc).__name__}: {exc}"
                 log.warning(

package/packages/memory-engine-v2/extractor-sync/Dockerfile

@@ -5,7 +5,7 @@ WORKDIR /app
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-COPY server.py .
+COPY entity_id.py server.py .
 
 EXPOSE 8101
 CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8101", "--workers", "2"]

package/packages/memory-engine-v2/extractor-sync/entity_id.py

@@ -0,0 +1,57 @@
+"""Canonical entity-ID scheme — SHARED, byte-identical across extractor-sync and
+extractor-async.
+
+The two extractors run as separate Docker services with PER-SERVICE build contexts
+(docker-compose `context: ./extractor-sync` / `./extractor-async`), so a single
+importable module can't be COPY'd into both. This file is therefore DUPLICATED in
+each service dir, and tests/test_entity_id_parity.py fails if the copies ever drift.
+
+Why this exists: both passes must key an entity (person / org / …) by the SAME id so
+the same entity converges across the deterministic (sync) and LLM (async) passes.
+Before this, the two services keyed entities DIFFERENTLY — sync as
+`e_` + sha256("{arena}|{type}|{name.lower().strip()}")[:24]; async as
+sha256("\\x1f".join(parts))[:32] (no lowercasing, no prefix) — so even identical
+names produced different ids and never merged. We unify on the sync scheme: sync's
+existing rows are unaffected, and the async pass converges onto them.
+
+Step 1 of RFC-entity-reconciliation.md (the foundation for alias-aware resolution).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+import unicodedata
+
+_WHITESPACE = re.compile(r"\s+")
+
+
+def normalize_surface_form(value: str) -> str:
+    """Normalize a surface form (person name, email, org name, …) for identity
+    keying. Steps, in order:
+
+      1. None → "" (defensive; some producers can hand a missing field as None).
+      2. Unicode NFKC (compatibility decomposition + canonical composition) —
+         collapses width / ligature / decomposed-accent variants. Without this
+         "Café" (precomposed U+00E9) and "Cafe\\u0301" (decomposed e+combining
+         acute) — which render identically — would key as different entities.
+         Same for fullwidth Latin ("ＣＡＲＬＹ" ↔ "CARLY") and ligatures
+         ("fi" U+FB01 ↔ "fi"). Real-world relevant for vCard imports, Mac
+         pasteboard, IME inputs, internationalised name sources.
+      3. Trim outer whitespace, then collapse internal `\\s+` to a single space —
+         "Carly  Snider" (slack-autocomplete double-space) ↔ "Carly Snider".
+      4. Lowercase. Email casing is case-insensitive per spec; person-name
+         casing varies by producer (gmail header casing vs slack profile).
+    """
+    s = unicodedata.normalize("NFKC", value or "")
+    return _WHITESPACE.sub(" ", s.strip()).lower()
+
+
+def entity_id(arena: str, entity_type: str, canonical_name: str) -> str:
+    """Deterministic entity id. The same (arena, entity_type, normalized
+    canonical_name) yields the same id across BOTH extractor passes, so re-extraction
+    and cross-pass extraction converge. Format is preserved from extractor-sync
+    (`e_` + 24 hex of sha256) so its existing rows are unaffected.
+    """
+    key = f"{arena}|{entity_type}|{normalize_surface_form(canonical_name)}"
+    return "e_" + hashlib.sha256(key.encode()).hexdigest()[:24]