@pentatonic-ai/ai-agent-sdk 0.9.5 → 0.10.0

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

@@ -0,0 +1,11 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY server.py .
+
+EXPOSE 8101
+CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8101", "--workers", "2"]

package/packages/memory-engine-v2/extractor-sync/requirements.txt

@@ -0,0 +1,4 @@
+fastapi==0.115.0
+uvicorn[standard]==0.32.0
+psycopg[binary,pool]==3.2.3
+pydantic==2.9.2

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

@@ -0,0 +1,424 @@
+"""extractor-sync — deterministic fast-path extraction.
+
+Called inline by compat on the write path. No LLM. Pure Python rules
+per source_kind. Writes events + provisional facts/entities into
+org-model, enqueues an entry on distillation_queue for the LLM
+distillation worker to upgrade later.
+
+Per-source rule registry is in `RULES` at the bottom of this file.
+Adding a new source = adding an entry there, not changing this file's
+core. Each rule is `(event) -> (entities[], facts[], relationships[])`
+where entities/facts/relationships are provisional — the async LLM
+distillation worker will upgrade them with high-confidence content.
+
+Why this exists (vs putting it in compat): keeping the extractor
+behind its own HTTP boundary lets the compat shim be the pure wire-
+contract layer + lets the extractor be replaced/re-implemented
+without touching compat. Also lets it scale horizontally — N extractor
+replicas behind one compat — if the deterministic path ever becomes a
+hot spot.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import os
+import re
+import time
+from contextlib import asynccontextmanager
+from typing import Any
+
+import psycopg
+import psycopg.rows
+from fastapi import FastAPI, HTTPException
+from psycopg_pool import AsyncConnectionPool
+from pydantic import BaseModel
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger("extractor-sync")
+
+PG_DSN = os.environ.get("PG_DSN", "postgresql://pme:local-dev-pw@org-model:5432/org_model")
+
+# Connection pool. Sized for high-concurrency bursts: with TES at
+# shardCount=8 and BATCH_SIZE=50 in the DO drain, compat's parallel
+# asyncio.gather in /store-batch can fan ~400 simultaneous _extract
+# calls at us. max_size=20 caused engine_500 cascades under that load
+# (DLQ losses observed 2026-05-17). 200 covers the 8×50 ceiling with
+# headroom; min_size=8 keeps a few warm so the first batch in a burst
+# doesn't pay connection-setup cost.
+_pool: AsyncConnectionPool | None = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global _pool
+    _pool = AsyncConnectionPool(
+        conninfo=PG_DSN,
+        min_size=8,
+        max_size=50,
+        kwargs={"row_factory": psycopg.rows.dict_row},
+        open=False,
+    )
+    await _pool.open()
+    log.info("extractor-sync: pool opened (min=8 max=50)")
+    yield
+    await _pool.close()
+
+
+app = FastAPI(title="pme2-extractor-sync", lifespan=lifespan)
+
+
+# ----------------------------------------------------------------------
+# Wire shape
+# ----------------------------------------------------------------------
+
+
+class ExtractRequest(BaseModel):
+    """One event for extraction.
+
+    Matches the relevant subset of TES's STORE_MEMORY envelope shape.
+    compat already validates the broader payload; we just need the
+    fields the rules read.
+    """
+
+    arena: str
+    clientId: str
+    userId: str | None = None
+    event_type: str = "STORE_MEMORY"
+    source_kind: str  # 'chat' | 'note' | 'doc' | 'event' | 'ticket' | 'commit' | 'system' | 'agent'
+    source_id: str | None = None
+    content: str
+    attributes: dict[str, Any] = {}
+    emitted_at: str | None = None  # ISO; defaults to now() in DB
+
+
+class ExtractResponse(BaseModel):
+    event_id: str
+    entities_extracted: int
+    facts_extracted: int
+    relationships_extracted: int
+    distillation_queued: bool
+
+
+# ----------------------------------------------------------------------
+# Hashing + ID helpers
+# ----------------------------------------------------------------------
+
+
+def _content_hash(arena: str, content: str) -> str:
+    """sha256(arena:content)[:32] — matches the v1 _arena_scoped_rid
+    convention so caller-supplied predicted IDs continue to work."""
+    h = hashlib.sha256(f"{arena}:{content}".encode()).hexdigest()
+    return h[:32]
+
+
+def _entity_id(arena: str, entity_type: str, canonical_name: str) -> str:
+    """Deterministic entity ID — same canonical name in the same arena
+    always produces the same entity, so re-extractions converge."""
+    key = f"{arena}|{entity_type}|{canonical_name.lower().strip()}"
+    return "e_" + hashlib.sha256(key.encode()).hexdigest()[:24]
+
+
+def _fact_id(arena: str, category: str, subject: str | None, predicate: str | None,
+             object_: str | None, statement: str) -> str:
+    """Deterministic fact ID — same statement-shape converges across
+    re-extractions. Statement-included so two facts about the same
+    subject with different statements get separate IDs."""
+    key = f"{arena}|{category}|{subject or ''}|{predicate or ''}|{object_ or ''}|{statement}"
+    return "f_" + hashlib.sha256(key.encode()).hexdigest()[:24]
+
+
+def _relationship_id(arena: str, from_id: str, to_id: str, rel_type: str) -> str:
+    key = f"{arena}|{from_id}|{to_id}|{rel_type}"
+    return "r_" + hashlib.sha256(key.encode()).hexdigest()[:24]
+
+
+# ----------------------------------------------------------------------
+# Per-source extraction rules (deterministic, no LLM).
+#
+# Each rule returns `(entities, facts, relationships)`. The compat shim
+# may have already pulled out emails/mentions from the message envelope
+# itself; we treat those as authoritative inputs in `req.attributes`
+# and only extract from `content` where unambiguous (e.g., regex
+# matches on email addresses, @-mentions).
+#
+# These are intentionally cautious. Better to extract less and let the
+# async LLM distillation fill in than to extract noisy junk that
+# pollutes the org model. The keystone spec (#285) will set the κ /
+# false-inclusion bars for each rule.
+# ----------------------------------------------------------------------
+
+
+EMAIL_RE = re.compile(r"\b([a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,})\b")
+MENTION_RE = re.compile(r"<@([A-Z0-9]+)>")  # slack @-mentions
+URL_RE = re.compile(r"https?://\S+")
+
+
+def _extract_chat(req: ExtractRequest, event_id: str) -> tuple[list, list, list]:
+    """slack / teams / DM: extract participant emails + @-mentions.
+
+    Skip body distillation — the LLM worker handles that. We just want
+    structural entities (who) here so peopleList / spreading-activation
+    queries work on day one.
+    """
+    entities, facts, relationships = [], [], []
+    # Emails in body.
+    for m in EMAIL_RE.findall(req.content):
+        eid = _entity_id(req.arena, "person", m)
+        entities.append({
+            "id": eid,
+            "arena": req.arena,
+            "entity_type": "person",
+            "canonical_name": m,
+            "aliases": [m],
+            "provenance_event_ids": [event_id],
+            "participant_set": req.attributes.get("participant_set", [req.arena]),
+            "disclosure_class": req.attributes.get("disclosure_class", "private"),
+        })
+    # Slack-style mentions (uppercase user IDs).
+    for sid in MENTION_RE.findall(req.content):
+        eid = _entity_id(req.arena, "person", f"slack:{sid}")
+        entities.append({
+            "id": eid,
+            "arena": req.arena,
+            "entity_type": "person",
+            "canonical_name": f"slack:{sid}",
+            "aliases": [f"slack:{sid}"],
+            "provenance_event_ids": [event_id],
+            "participant_set": req.attributes.get("participant_set", [req.arena]),
+            "disclosure_class": req.attributes.get("disclosure_class", "private"),
+        })
+    return entities, facts, relationships
+
+
+def _extract_doc(req: ExtractRequest, event_id: str) -> tuple[list, list, list]:
+    """drive / notion / confluence: leave content distillation entirely
+    to the LLM worker. Just extract source URL as a structural entity."""
+    entities, facts, relationships = [], [], []
+    src = req.attributes.get("source_url") or req.attributes.get("path")
+    if src:
+        eid = _entity_id(req.arena, "document", src)
+        entities.append({
+            "id": eid,
+            "arena": req.arena,
+            "entity_type": "document",
+            "canonical_name": src,
+            "aliases": [src],
+            "provenance_event_ids": [event_id],
+            "participant_set": req.attributes.get("participant_set", [req.arena]),
+            "disclosure_class": req.attributes.get("disclosure_class", "private"),
+        })
+    return entities, facts, relationships
+
+
+def _extract_note(req: ExtractRequest, event_id: str) -> tuple[list, list, list]:
+    """gmail / drafts: extract from + to emails, subject."""
+    entities, facts, relationships = [], [], []
+    for key in ("from_email", "to_email", "cc_email"):
+        val = req.attributes.get(key)
+        if isinstance(val, str) and "@" in val:
+            eid = _entity_id(req.arena, "person", val)
+            entities.append({
+                "id": eid,
+                "arena": req.arena,
+                "entity_type": "person",
+                "canonical_name": val,
+                "aliases": [val],
+                "provenance_event_ids": [event_id],
+                "participant_set": req.attributes.get("participant_set", [req.arena]),
+                "disclosure_class": req.attributes.get("disclosure_class", "private"),
+            })
+    return entities, facts, relationships
+
+
+def _extract_event(req: ExtractRequest, event_id: str) -> tuple[list, list, list]:
+    """calendar: extract organizer + attendees as people entities."""
+    entities, facts, relationships = [], [], []
+    organizer = req.attributes.get("organizer_email")
+    attendees = req.attributes.get("attendee_emails") or []
+    for email in [organizer, *attendees]:
+        if isinstance(email, str) and "@" in email:
+            eid = _entity_id(req.arena, "person", email)
+            entities.append({
+                "id": eid,
+                "arena": req.arena,
+                "entity_type": "person",
+                "canonical_name": email,
+                "aliases": [email],
+                "provenance_event_ids": [event_id],
+                "participant_set": req.attributes.get("participant_set", [req.arena]),
+                "disclosure_class": req.attributes.get("disclosure_class", "private"),
+            })
+    return entities, facts, relationships
+
+
+def _extract_default(req: ExtractRequest, event_id: str) -> tuple[list, list, list]:
+    """Unknown source kind: extract bare emails from content, nothing else.
+    Async distillation does the heavy lifting."""
+    entities, facts, relationships = [], [], []
+    for m in EMAIL_RE.findall(req.content):
+        eid = _entity_id(req.arena, "person", m)
+        entities.append({
+            "id": eid,
+            "arena": req.arena,
+            "entity_type": "person",
+            "canonical_name": m,
+            "aliases": [m],
+            "provenance_event_ids": [event_id],
+            "participant_set": req.attributes.get("participant_set", [req.arena]),
+            "disclosure_class": req.attributes.get("disclosure_class", "private"),
+        })
+    return entities, facts, relationships
+
+
+RULES = {
+    "chat": _extract_chat,
+    "doc": _extract_doc,
+    "note": _extract_note,
+    "event": _extract_event,
+}
+
+
+# ----------------------------------------------------------------------
+# DB writes (idempotent upsert)
+# ----------------------------------------------------------------------
+
+
+async def _upsert_event(cur: psycopg.AsyncCursor, req: ExtractRequest,
+                        event_id: str, content_hash: str) -> None:
+    """ON CONFLICT DO NOTHING — re-emitting the same event is a no-op."""
+    await cur.execute(
+        """
+        INSERT INTO events (
+          id, arena, client_id, user_id, event_type, source_kind,
+          source_id, content, content_hash, participant_set,
+          participant_kind, disclosure_class, attributes
+        ) VALUES (
+          %s, %s, %s, %s, %s, %s::source_kind,
+          %s, %s, %s, %s,
+          %s::participant_kind, %s::disclosure_class, %s::jsonb
+        )
+        ON CONFLICT (id) DO NOTHING
+        """,
+        (
+            event_id, req.arena, req.clientId, req.userId,
+            req.event_type, req.source_kind, req.source_id,
+            req.content, content_hash,
+            req.attributes.get("participant_set", [req.arena]),
+            req.attributes.get("participant_kind", "unknown"),
+            req.attributes.get("disclosure_class", "private"),
+            psycopg.types.json.Json(req.attributes),
+        ),
+    )
+
+
+async def _upsert_entities(cur: psycopg.AsyncCursor, entities: list[dict]) -> None:
+    """Idempotent entity upsert: same canonical_name in same arena
+    converges to the same row. Aliases + provenance_event_ids grow
+    via array_append; never replace."""
+    for e in entities:
+        await 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()
+            """,
+            (e["id"], e["arena"], e["entity_type"], e["canonical_name"],
+             e["aliases"], e["provenance_event_ids"],
+             e["participant_set"], e["disclosure_class"]),
+        )
+
+
+async def _enqueue_distillation(cur: psycopg.AsyncCursor, event_id: str) -> None:
+    """Queue an LLM distillation pass on this event. Idempotent: if
+    we already have a pending/claimed/done entry for this event, no-op."""
+    await cur.execute(
+        """
+        INSERT INTO distillation_queue (event_id, status)
+        SELECT %s, 'pending'
+        WHERE NOT EXISTS (
+          SELECT 1 FROM distillation_queue
+          WHERE event_id = %s AND status IN ('pending', 'claimed', 'done')
+        )
+        """,
+        (event_id, event_id),
+    )
+
+
+# ----------------------------------------------------------------------
+# HTTP endpoints
+# ----------------------------------------------------------------------
+
+
+@app.get("/health")
+async def health():
+    """Cheap liveness probe — does NOT touch DB so compose healthchecks
+    don't gate on the DB being up at the wrong moment."""
+    return {"status": "healthy", "service": "extractor-sync"}
+
+
+@app.get("/health/deep")
+async def health_deep():
+    """Includes a DB round-trip to confirm the pool is alive."""
+    if _pool is None:
+        raise HTTPException(503, "pool not initialised")
+    try:
+        async with _pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute("SELECT 1")
+                row = await cur.fetchone()
+                if not row:
+                    raise HTTPException(503, "db query returned empty")
+    except Exception as e:
+        raise HTTPException(503, f"db error: {e}")
+    return {"status": "healthy", "service": "extractor-sync", "db": "ok"}
+
+
+@app.post("/extract", response_model=ExtractResponse)
+async def extract(req: ExtractRequest):
+    """Process one event. Idempotent: re-submitting the same event
+    converges to the same state (caller-safe to retry)."""
+    if _pool is None:
+        raise HTTPException(503, "pool not initialised")
+
+    t0 = time.perf_counter()
+    content_hash = _content_hash(req.arena, req.content)
+    event_id = content_hash  # identity == content-hash, per the schema doc
+
+    # Dispatch to the per-source rule.
+    rule = RULES.get(req.source_kind, _extract_default)
+    entities, facts, relationships = rule(req, event_id)
+
+    async with _pool.connection() as conn:
+        async with conn.cursor() as cur:
+            await _upsert_event(cur, req, event_id, content_hash)
+            await _upsert_entities(cur, entities)
+            # Facts + relationships are deliberately left to the async
+            # distillation worker — the deterministic path can't
+            # reliably extract decisions/commitments without LLM context.
+            # Future: deterministic relationships from explicit
+            # participant_set (e.g., calendar: "organizer
+            # communicated_with attendee") — leaving as a follow-up.
+            await _enqueue_distillation(cur, event_id)
+
+    dur_ms = (time.perf_counter() - t0) * 1000
+    log.info(f"extract event_id={event_id} entities={len(entities)} ms={dur_ms:.1f}")
+    return ExtractResponse(
+        event_id=event_id,
+        entities_extracted=len(entities),
+        facts_extracted=len(facts),
+        relationships_extracted=len(relationships),
+        distillation_queued=True,
+    )