memd-engine 0.5.0__py3-none-any.whl

memd/harness/suites/longmemeval_synthetic.py

@@ -0,0 +1,221 @@
+"""LongMemEval-class synthetic suite.
+
+Deterministic generator (seeded): sessions of facts per user, then query
+classes mirroring LongMemEval's: single-hop, multi-hop, temporal, and the
+wedge classes - knowledge-update, knowledge-overwrite, abstention.
+
+The real LongMemEval is used when MEMD_LME_DATA points at a JSON file with
+either {"events": [...], "cases": [...]} in this suite's schema, or a list of
+rows. Two row shapes are accepted:
+  - the real LongMemEval release (longmemeval_s / _m / _oracle):
+      {"question_id", "question_type", "question", "question_date", "answer",
+       "answer_session_ids", "haystack_dates", "haystack_session_ids",
+       "haystack_sessions": [[{"role", "content", "has_answer"?}, ...], ...]}
+  - a flat legacy shape:
+      {"session": [{"content": ..., "role": ..., "user_id"?...}],
+       "question": ..., "answer": ...}
+Rows are converted to events+cases at load time. The synthetic suite exists
+so CI can gate every commit offline with zero network.
+"""
+from __future__ import annotations
+
+import datetime as _dt
+import json
+import os
+import random
+import re
+
+DAY = 86_400_000
+T0 = 1_700_000_000_000
+
+
+def generate_cases(seed: int = 42, users: int = 8) -> tuple[list[dict], list[dict]]:
+    rng = random.Random(seed)
+    names = ["Alice", "Bob", "Cara", "Dan", "Elena", "Frank", "Grace", "Hank"]
+    cities = ["Lisbon", "Oslo", "Kyoto", "Austin", "Nairobi", "Perth"]
+    tools = ["make ship", "cargo deploy", "npm run release", "just prod", "fab rollout"]
+    editors = ["Neovim", "Emacs", "VS Code", "Zed"]
+
+    events: list[dict] = []
+    cases: list[dict] = []
+
+    for u in range(users):
+        uid = f"u{u}"
+        # pools cycle so --users beyond len(names) works (the CLI exposes
+        # users as free-form; an IndexError here bricked the whole harness).
+        # Uniqueness comes from uid-scoped session ids and per-user queries.
+        name, city, tool, editor = (
+            names[u % len(names)], cities[u % len(cities)],
+            tools[u % len(tools)], editors[u % len(editors)],
+        )
+        base = T0 + u * DAY
+
+        def ev(day: int, content: str):
+            events.append({"content": content, "user_id": uid, "session_id": f"s{uid}-{day}",
+                           "role": "user", "t_event": base + day * DAY})
+
+        # stable facts
+        ev(0, f"My name is {name}.")
+        ev(0, f"I live in {city}.")
+        ev(1, f"We deploy with `{tool}`, never raw ftp.")
+        ev(2, f"I use {editor} as my editor.")
+        # temporal facts
+        ev(3, f"I moved to {city} in 2021.")
+        ev(4, f"My review meeting is scheduled on March 3rd.")
+
+        # single-hop
+        cases.append({
+            "id": f"{uid}-sh-deploy", "qclass": "single_hop", "user_id": uid,
+            "query": f"How do we deploy?",
+            "expected": tool,
+            "events_hint": None,
+        })
+        cases.append({
+            "id": f"{uid}-sh-editor", "qclass": "single_hop", "user_id": uid,
+            "query": f"What editor does {name} use?",
+            "expected": editor,
+        })
+        # multi-hop: name -> city via two facts
+        cases.append({
+            "id": f"{uid}-mh-city", "qclass": "multi_hop", "user_id": uid,
+            "query": f"Where does {name} live?",
+            "expected": city,
+        })
+        # temporal
+        cases.append({
+            "id": f"{uid}-tp-review", "qclass": "temporal", "user_id": uid,
+            "query": "When is my review meeting?",
+            "expected": "March 3rd",
+        })
+        # knowledge-update: old employer then new one; only latest correct
+        old_job, new_job = "Initech", "Initrode"
+        ev(10, f"I just started working at {old_job}.")
+        ev(40, f"I switched jobs, now I work at {new_job}.")
+        cases.append({
+            "id": f"{uid}-ku-employer", "qclass": "knowledge_update", "user_id": uid,
+            "query": f"Where does {name} work?",
+            "expected": new_job,
+            "stale": old_job,
+        })
+        return_value = f"ticket #{100+u}"
+        ev(20, f"My support ticket got resolved, reference {return_value}.")
+        cases.append({
+            "id": f"{uid}-sh-ticket", "qclass": "single_hop", "user_id": uid,
+            "query": "What was my support ticket reference?",
+            "expected": return_value,
+        })
+
+        # distractor volume: makes the full-context baseline pay real token
+        # cost and retrieval work under noise (LongMemEval-style long horizon)
+        topics = ["standup", "refactor", "bug bash", "on-call", "retro", "planning",
+                  "lint", "flaky test", "cache", "auth flow", "docs", "telemetry"]
+        for d in range(400):
+            t = topics[d % len(topics)]
+            ev(5 + d % 90, f"Notes from {t} session number {d}: we discussed the {t} "
+                           f"and agreed to follow up about {t} details later.")
+        _ = rng  # reserved for future noise injection
+    return events, cases
+
+
+def check_answer(ans: str, case: dict) -> tuple[bool, float]:
+    expected = str(case["expected"]).lower()
+    a = ans.lower()
+    hit = expected in a
+    stale_leak = False
+    if case.get("stale"):
+        stale_leak = str(case["stale"]).lower() in a and not hit
+    if not hit or stale_leak:
+        return False, 0.0
+    return True, 1.0
+
+
+def build(seed: int = 42, users: int = 8) -> tuple[list[dict], list[dict]]:
+    path = os.environ.get("MEMD_LME_DATA")
+    if path:
+        return load_real(path)
+    return generate_cases(seed=seed, users=users)
+
+
+_LME_DATE = re.compile(r"\s*(\d{4})/(\d{1,2})/(\d{1,2})(?:.*?(\d{1,2}):(\d{2}))?")
+
+
+def _lme_date_ms(s: str) -> int | None:
+    """'2023/05/20 (Sat) 02:21' -> epoch ms (UTC; the dataset has no zone)."""
+    m = _LME_DATE.match(str(s))
+    if not m:
+        return None
+    y, mo, d, hh, mm = (int(x) if x else 0 for x in m.groups())
+    dt = _dt.datetime(y, mo, d, hh, mm, tzinfo=_dt.timezone.utc)
+    return int(dt.timestamp() * 1000)
+
+
+def _load_lme_row(row: dict, events: list[dict], cases: list[dict]) -> None:
+    """One real LongMemEval question: its own user (haystacks are per
+    question), one event per haystack turn, one case. Abstention variants
+    (question_id ending in `_abs`) are skipped: substring grading of
+    "expected" cannot score an abstention."""
+    qid = str(row["question_id"])
+    if qid.endswith("_abs"):
+        return
+    uid = qid
+    sids = row.get("haystack_session_ids") or []
+    dates = row.get("haystack_dates") or []
+    for k, session in enumerate(row.get("haystack_sessions") or []):
+        sid = str(sids[k]) if k < len(sids) else f"s{k}"
+        t_event = _lme_date_ms(dates[k]) if k < len(dates) else None
+        for turn in session:
+            content = turn.get("content") if isinstance(turn, dict) else turn
+            if not content:
+                continue
+            events.append({
+                "content": str(content), "user_id": uid,
+                # uid-scoped: the same distractor session id recurs across
+                # questions, and the harness closes sessions by id alone
+                "session_id": f"{uid}:{sid}",
+                "role": (turn.get("role") if isinstance(turn, dict) else None) or "user",
+                "t_event": t_event,
+            })
+    cases.append({
+        "id": qid, "qclass": row.get("question_type", "real"), "user_id": uid,
+        "query": str(row["question"]), "expected": str(row["answer"]),
+        "question_date": row.get("question_date"),
+        "answer_session_ids": list(row.get("answer_session_ids") or []),
+    })
+
+
+def load_real(path: str) -> tuple[list[dict], list[dict]]:
+    """Load a LongMemEval-derived dataset and normalize it into
+    (events, cases) with the same schema as the synthetic generator."""
+    with open(path) as f:
+        data = json.load(f)
+    if isinstance(data, dict) and "events" in data and "cases" in data:
+        return data["events"], data["cases"]
+    rows = data if isinstance(data, list) else []
+    events: list[dict] = []
+    cases: list[dict] = []
+    for i, row in enumerate(rows):
+        if isinstance(row, dict) and "haystack_sessions" in row and "question_id" in row:
+            _load_lme_row(row, events, cases)
+            continue
+        uid = f"lme{i}"
+        sess_events = row.get("session") or row.get("haystack") or row.get("context") or []
+        for j, s in enumerate(sess_events):
+            content = s.get("content") if isinstance(s, dict) else str(s)
+            if not content:
+                continue
+            events.append({
+                "content": str(content), "user_id": uid,
+                "session_id": s.get("session_id", f"{uid}-s{j//20}") if isinstance(s, dict) else f"{uid}-s{j//20}",
+                "role": (s.get("role", "user") if isinstance(s, dict) else "user"),
+            })
+        q = row.get("question") or row.get("query")
+        a = row.get("answer") or row.get("expected")
+        if q and a:
+            cases.append({
+                "id": f"lme{i}", "qclass": row.get("qclass", row.get("category", "real")),
+                "user_id": uid, "query": str(q), "expected": str(a),
+                "stale": row.get("stale"),
+            })
+    if not cases:
+        raise ValueError(f"{path}: no usable LongMemEval rows found")
+    return events, cases

memd/hosted/__init__.py

@@ -0,0 +1,20 @@
+"""Hosted mode: tenancy (org -> namespaces -> API keys), usage metering,
+entitlements and Stripe billing.
+
+OFF by default. Only `memd serve --http --hosted` (or MEMD_HOSTED=1) loads
+this package; embedded and self-hosted memd never import it, and nothing here
+imports `stripe` at module level (it is the optional `memd[billing]` extra,
+imported lazily by memd.hosted.billing when a Stripe key is configured).
+"""
+from __future__ import annotations
+
+import os
+
+_TRUE = ("1", "true", "yes", "on")
+
+
+def hosted_enabled(flag: bool | None = None) -> bool:
+    """An explicit flag wins; otherwise MEMD_HOSTED decides (default off)."""
+    if flag is not None:
+        return bool(flag)
+    return os.environ.get("MEMD_HOSTED", "").strip().lower() in _TRUE

memd/hosted/app.py

@@ -0,0 +1,258 @@
+"""Hosted-mode wiring for the REST server: the org-aware key store, the
+metering hooks, the /v1/billing routes and the background billing jobs.
+
+memd.server.http imports this module only when hosted mode is on."""
+from __future__ import annotations
+
+import hmac
+import importlib.util
+import logging
+import os
+import threading
+import time
+from collections import OrderedDict
+from typing import Any
+
+from fastapi import Depends, Request
+from fastapi.concurrency import run_in_threadpool
+from pydantic import BaseModel, Field
+
+from memd.hosted.billing import Billing, BillingConfig, BillingError
+from memd.hosted.metering import Forbidden, Metering, QuotaDenied
+from memd.hosted.plans import Plans
+from memd.hosted.store import AdminStore, parse_scopes, period_bounds, period_of
+from memd.metrics import METRICS
+from memd.server.auth import Principal, hash_secret
+
+_log = logging.getLogger("memd.hosted")
+_OPS = {"ns": "_billing"}
+# Stripe's events are a few KiB; nothing legitimate comes near this
+WEBHOOK_MAX_BYTES = 1 << 20
+
+
+class CheckoutIn(BaseModel):
+    plan: str = Field(min_length=1, max_length=64)
+    interval: str = "month"
+
+
+class HostedKeyStore:
+    """KeyStore-compatible adapter over the admin store: keys belong to an
+    org and are bound to one of its namespaces; only SHA-256 hashes of the
+    secrets are stored. The operator key (MEMD_ADMIN_KEY) stays org-less:
+    it is never metered or billed. Legacy keys.toml.json keys are NOT
+    honoured in hosted mode (they carry no org, so they would bypass
+    metering) - adopt them with `memd key migrate --org ...`."""
+
+    CACHE_TTL_S = 5.0  # bounded staleness for a revocation made by another process
+    CACHE_MAX = 10_000
+
+    def __init__(self, store: AdminStore, admin_key: str | None = None):
+        self.store = store
+        self._admin_hash = hash_secret(admin_key) if admin_key else None
+        self._lock = threading.Lock()
+        self._cache: "OrderedDict[str, tuple[Principal, float]]" = OrderedDict()
+
+    def authenticate(self, bearer: str | None) -> Principal | None:
+        if not bearer:
+            return None
+        h = hash_secret(bearer)
+        if self._admin_hash and hmac.compare_digest(h, self._admin_hash):
+            return Principal(key_id="admin", namespace="*", scope_override=True)
+        now = time.monotonic()
+        with self._lock:
+            hit = self._cache.get(h)
+            if hit is not None and now - hit[1] < self.CACHE_TTL_S:
+                return hit[0]
+        rec = self.store.key_for_bearer(bearer)
+        if rec is None:
+            with self._lock:
+                self._cache.pop(h, None)
+            return None
+        scopes = frozenset(rec["scopes"].split())
+        p = Principal(key_id=rec["key_id"], namespace=rec["ns"], pinned_user=rec.get("pinned_user"),
+                      scope_override="override" in scopes, org_id=rec["org_id"], scopes=scopes)
+        with self._lock:
+            self._cache[h] = (p, now)
+            self._cache.move_to_end(h)
+            while len(self._cache) > self.CACHE_MAX:
+                self._cache.popitem(last=False)
+        return p
+
+    def create(self, namespace: str, name: str = "", pinned_user: str | None = None,
+               scope_override: bool = False, *, org_id: str, scopes: str | list[str] | None = None
+               ) -> tuple[str, str]:
+        sc = parse_scopes(scopes)
+        if scope_override and "override" not in sc:
+            sc.append("override")
+        return self.store.create_key(org_id, namespace, name=name, scopes=sc, pinned_user=pinned_user)
+
+    def revoke(self, kid: str) -> bool:
+        ok = self.store.revoke_key(kid)
+        with self._lock:
+            for h, (p, _) in list(self._cache.items()):
+                if p.key_id == kid:
+                    self._cache.pop(h, None)
+        return ok
+
+    def list_keys(self) -> list[dict]:
+        return self.store.list_keys()
+
+
+class BillingJobs:
+    """Background loop: the daily gauge snapshot, the hourly meter push, the
+    settling of usage too old to push on its own, and the daily drift
+    report. Every step is idempotent, so a restart mid-step is safe."""
+
+    def __init__(self, hosted: "Hosted"):
+        self.hosted = hosted
+        self.push_interval_s = float(os.environ.get("MEMD_BILLING_PUSH_INTERVAL_S", "3600"))
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._snapshot_day: str | None = None
+        self._reconciled_day: int | None = None
+
+    def run_once(self, now: float | None = None) -> dict:
+        h = self.hosted
+        now = time.time() if now is None else now
+        out: dict[str, Any] = {}
+        if h.router is not None:
+            # cluster: every node records the size of the namespaces IT holds
+            # (only the leaseholder can read one without taking it over), and
+            # one node - the rendezvous choice for "_billing" among the live
+            # nodes - runs the singleton steps below
+            out["measured"] = h.metering.measure_open()
+            if not h.router.is_leader("_billing"):
+                out["leader"] = False
+                return out
+        day = time.strftime("%Y-%m-%d", time.gmtime(now))
+        if day != self._snapshot_day:
+            out["snapshot"] = h.metering.snapshot_gauges(now=now)
+            self._snapshot_day = day
+        if h.billing.cfg.configured:
+            out["push"] = h.billing.push_usage(now=now)
+            # every tick: usage too old to push on its own is settled against
+            # what Stripe verifiably holds (a no-op when there is none)
+            out["reconcile_pending"] = h.billing.reconcile_pending(now=now)
+            today = int(now // 86400) * 86400
+            if self._reconciled_day != today:
+                # the drift report covers yesterday's quota period to date
+                out["reconcile"] = h.billing.reconcile(period_bounds(period_of(today - 1))[0], today)
+                self._reconciled_day = today
+        return out
+
+    def _loop(self) -> None:
+        while not self._stop.is_set():
+            try:
+                self.run_once()
+            except Exception:
+                METRICS.inc("memd_billing_job_failures_total", **_OPS)
+                _log.exception("memd billing: background job failed")
+            self._stop.wait(self.push_interval_s)
+
+    def start(self) -> None:
+        if self._thread is None and os.environ.get("MEMD_BILLING_JOBS", "1") != "0":
+            self._thread = threading.Thread(target=self._loop, daemon=True, name="memd-billing-jobs")
+            self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=10)
+
+
+class Hosted:
+    def __init__(self, data_dir: str, engine: Any, *, admin_key: str | None = None,
+                 plans: Plans | None = None, billing_config: BillingConfig | None = None,
+                 stripe_client: Any = None, router: Any = None):
+        # from_env runs the sk_live_ guard: a refused key stops startup here
+        self.cfg = billing_config or BillingConfig.from_env()
+        if (self.cfg.configured or self.cfg.webhook_secret) and importlib.util.find_spec("stripe") is None:
+            raise RuntimeError("hosted billing is configured but the Stripe SDK is missing: "
+                               "pip install 'memd-engine[billing]'")
+        self.plans = plans or Plans.from_env()
+        self.store = AdminStore.for_data_root(data_dir)
+        # the cluster router (memd.server.cluster.Router) when this process
+        # is one node of a fleet sharing this admin database, else None
+        self.router = router
+        if router is None:
+            # one process per data root: reservations on disk are a dead process's
+            self.store.clear_reservations()
+        # (a cluster shares the database: the other nodes' reservations are
+        # live, and a dead node's expire by their TTL)
+        self.keystore = HostedKeyStore(self.store, admin_key=admin_key)
+        self.metering = Metering(self.store, self.plans, engine)
+        self.billing = Billing(self.store, self.plans, self.cfg, client=stripe_client)
+        self.jobs = BillingJobs(self)
+        if not self.cfg.configured:
+            _log.warning("memd hosted: no MEMD_STRIPE_SECRET_KEY - tenancy and quotas are on, "
+                         "billing endpoints answer 503 billing_not_configured")
+
+    def close(self) -> None:
+        self.jobs.stop()
+        self.store.close()
+
+    def install(self, app: Any, auth: Any) -> None:
+        """Register the /v1/billing routes and the 402/403 handlers."""
+        from memd.server.http import ApiError, _error
+
+        hosted = self
+
+        @app.exception_handler(QuotaDenied)
+        async def quota_denied(_request: Request, exc: QuotaDenied):
+            return _error(402, exc.detail, exc.code, extra=exc.extra)
+
+        @app.exception_handler(Forbidden)
+        async def forbidden(_request: Request, exc: Forbidden):
+            return _error(403, exc.detail)
+
+        @app.exception_handler(BillingError)
+        async def billing_error(_request: Request, exc: BillingError):
+            return _error(exc.status, exc.detail, exc.code, extra=exc.extra)
+
+        def billing_org(p: Principal = Depends(auth)) -> dict:
+            if p.org_id is None:
+                raise ApiError(400, "billing routes act on the calling key's org; this key has none",
+                               code="no_org")
+            if "billing" not in p.scopes:  # exact: `override` does not imply it
+                raise ApiError(403, "key lacks the 'billing' scope")
+            org = hosted.store.get_org(p.org_id)
+            if org is None:
+                raise ApiError(404, "org not found")
+            return org
+
+        @app.post("/v1/billing/checkout")
+        def checkout(body: CheckoutIn, org: dict = Depends(billing_org)):
+            return hosted.billing.checkout(org, body.plan, body.interval)
+
+        @app.post("/v1/billing/portal")
+        def portal(org: dict = Depends(billing_org)):
+            return hosted.billing.portal(org)
+
+        @app.get("/v1/billing/usage")
+        def usage(org: dict = Depends(billing_org)):
+            return hosted.metering.usage_report(org["id"])
+
+        @app.post("/v1/billing/webhook")
+        async def webhook(request: Request):
+            # unauthenticated by design: the Stripe-Signature HMAC over the
+            # RAW body is the authentication (parsed JSON would not verify).
+            # Capped by bytes actually READ: a chunked body has no
+            # Content-Length for the global 8 MiB guard to see.
+            buf = bytearray()
+            async for chunk in request.stream():
+                buf += chunk
+                if len(buf) > WEBHOOK_MAX_BYTES:
+                    METRICS.inc("memd_billing_webhook_rejected_total", reason="too_large", **_OPS)
+                    return _error(413, "webhook body too large")
+            event = hosted.billing.verify_webhook(bytes(buf), request.headers.get("stripe-signature"))
+            try:
+                result = await run_in_threadpool(hosted.billing.handle_event, event)
+            except BillingError:
+                raise
+            except Exception:
+                _log.exception("memd billing: webhook handler failed for %s", event.get("id"))
+                METRICS.inc("memd_billing_webhook_total", type=str(event.get("type")), result="error", **_OPS)
+                return _error(500, "webhook handler failed")  # Stripe re-delivers
+            if result.get("retry"):
+                return _error(500, "event not applicable yet; retry", "retry", extra={"result": result})
+            return {"received": True, **result}