self-directed-learning/scripts/ump-watcher.py

#!/usr/bin/env python3
"""
ump-watcher: tail `memory.ump.json` on disk and push new/modified records into
the Qdrant `memories_ump` collection via the ump-recall sidecar.

Why file-tail and not SSE:
- The ump-memory SSE endpoint at /ump/subscribe was confirmed to not emit
  events during live writes in our 2026-07-12 testing.
- The on-disk JSON is the canonical store (JsonFileStore flushes per write).
- File-tail is also how we'd backfill 839 historical records — same code path.

Architecture:
- Bootstrap: on startup, scan memory.ump.json, build id→mtime index.
- Steady state: stat() the file every poll_interval (default 5s); if mtime
  changed, re-read fully, diff against index, push new/changed records.
- Embed: POST {id, text} to http://127.0.0.1:4380/embed (sidecar handles
  Ollama → Qdrant). Idempotent: same urn re-embed gets a new Qdrant pid but
  the urn in payload is the canonical key.
- Backfill mode: --backfill flag processes ALL records and exits, useful for
  Phase 1D initial bulk ingest.

Configuration via env (matches ump-recall sidecar):
  UMP_DIR                  /root/.openclaw/agents/main/workspace/state/ump-local
  UMP_FILE                 memory.ump.json (default)
  SIDECAR_URL              http://127.0.0.1:4380
  POLL_INTERVAL            5 (seconds)
  BATCH_SIZE               16 (concurrent embeds)
  LOG_LEVEL                info (debug|info|warn|error)

Failure handling:
- If sidecar is down, log warn and continue (next poll will retry).
- If a single record fails to embed, log and skip (don't halt the batch).
- If JSON file is being written (truncated), retry next poll.

Usage:
  python3 ump-watcher.py                  # daemon mode (default)
  python3 ump-watcher.py --backfill       # one-shot, process all + exit
  python3 ump-watcher.py --once           # one poll cycle + exit (smoke test)
"""

import argparse
import json
import logging
import os
import signal
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import requests

UMP_DIR = Path(os.getenv("UMP_DIR", "/root/.openclaw/agents/main/workspace/state/ump-local"))
UMP_FILE = os.getenv("UMP_FILE", "memory.ump.json")
SIDECAR_URL = os.getenv("SIDECAR_URL", "http://127.0.0.1:4380").rstrip("/")
POLL_INTERVAL = float(os.getenv("POLL_INTERVAL", "5"))
BATCH_SIZE = int(os.getenv("BATCH_SIZE", "4"))
REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "120"))
LOG_LEVEL = os.getenv("LOG_LEVEL", "info").upper()

# State
running = True
ump_path = UMP_DIR / UMP_FILE
state_path = UMP_DIR / ".ump-watcher-state.json"  # tracks last seen per urn


def signal_handler(signum, _frame):
    global running
    log.info("received signal %d, shutting down", signum)
    running = False


def setup_logging():
    logging.basicConfig(
        level=getattr(logging, LOG_LEVEL, logging.INFO),
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%Y-%m-%dT%H:%M:%S",
    )
    return logging.getLogger("ump-watcher")


log = setup_logging()


def load_records(path: Path) -> list[dict]:
    """Read memory.ump.json and return the records list. Robust to mid-write."""
    try:
        with open(path, "r") as f:
            data = json.load(f)
        if isinstance(data, list):
            return data
        log.error("unexpected JSON shape at %s: %s", path, type(data).__name__)
        return []
    except json.JSONDecodeError as e:
        log.warning("JSON decode error at %s (mid-write?): %s", path, e)
        return []
    except FileNotFoundError:
        log.warning("ump store not found at %s", path)
        return []


def load_state() -> dict:
    """Load {urn: last_seen_epoch_seconds} map."""
    try:
        if state_path.exists():
            with open(state_path, "r") as f:
                return json.load(f)
    except (json.JSONDecodeError, OSError) as e:
        log.warning("could not load state at %s: %s", state_path, e)
    return {}


def save_state(state: dict):
    """Atomic write of state file."""
    tmp = state_path.with_suffix(".json.tmp")
    try:
        with open(tmp, "w") as f:
            json.dump(state, f, indent=2, sort_keys=True)
        os.replace(str(tmp), str(state_path))
    except OSError as e:
        log.warning("could not save state at %s: %s", state_path, e)


def text_for_record(rec: dict) -> str:
    """Build the text string we'll embed. Concatenates subject + topic + body.text
    so the embedding captures the semantic content without being just a label.

    Truncates to MAX_EMBED_CHARS (default 4000) because very long texts slow
    Ollama CPU inference dramatically (>60s for SOUL.md-class documents) and
    don't improve embedding quality much past ~2K tokens.
    """
    body = rec.get("body") or {}
    subject = (body.get("subject") or "").strip()
    topic = (body.get("topic") or "").strip()
    text = (body.get("text") or "").strip()
    parts = []
    if subject:
        parts.append(subject)
    if topic and topic not in subject:
        parts.append(f"[{topic}]")
    if text:
        parts.append(text)
    combined = "\n".join(parts) or "(empty record)"
    max_chars = int(os.getenv("MAX_EMBED_CHARS", "4000"))
    if len(combined) > max_chars:
        combined = combined[:max_chars] + "... [truncated]"
    return combined


def embed_one(urn: str, text: str, retries: int = 3) -> tuple[str, bool, str]:
    """POST to sidecar /embed. Returns (urn, ok, error_msg)."""
    for attempt in range(retries + 1):
        try:
            r = requests.post(
                f"{SIDECAR_URL}/embed",
                json={"id": urn, "text": text},
                timeout=REQUEST_TIMEOUT,
            )
            if r.status_code == 200:
                body = r.json()
                if body.get("status") == "ok":
                    return urn, True, ""
                return urn, False, f"sidecar status={body.get('status')}: {body}"
            return urn, False, f"HTTP {r.status_code}: {r.text[:200]}"
        except (requests.RequestException, json.JSONDecodeError) as e:
            if attempt < retries:
                time.sleep(0.5 * (attempt + 1))
                continue
            return urn, False, f"exception: {e!r}"


def push_batch(records: list[dict]) -> dict:
    """Push a batch of records to the sidecar. Returns summary stats."""
    if not records:
        return {"pushed": 0, "failed": 0, "errors": []}

    started = time.time()
    pushed = 0
    failed = 0
    errors = []

    with ThreadPoolExecutor(max_workers=BATCH_SIZE) as pool:
        futures = {
            pool.submit(embed_one, r["id"], text_for_record(r)): r for r in records
        }
        for fut in as_completed(futures):
            urn, ok, err = fut.result()
            if ok:
                pushed += 1
            else:
                failed += 1
                errors.append((urn, err))
                log.warning("embed failed for %s: %s", urn, err)

    elapsed_ms = (time.time() - started) * 1000
    log.info(
        "batch complete: %d ok, %d failed in %.0fms (%.0fms/record)",
        pushed, failed, elapsed_ms, elapsed_ms / max(pushed + failed, 1),
    )
    return {"pushed": pushed, "failed": failed, "errors": errors}


def diff_records(records: list[dict], state: dict) -> list[dict]:
    """Return records whose urn is new or whose `time.valid_to` is unset
    (revised) since last seen. We use record validity as the change signal
    because ump.json doesn't carry a per-record mtime."""
    out = []
    for r in records:
        urn = r.get("id")
        if not urn or not isinstance(urn, str):
            continue
        time_obj = r.get("time") or {}
        valid_to = time_obj.get("valid_to")
        # Use (urn, valid_to) tuple as change marker. valid_to=null means active.
        marker = json.dumps({"urn": urn, "valid_to": valid_to}, sort_keys=True)
        if state.get(urn) != marker:
            out.append(r)
    return out


def update_state(records: list[dict], state: dict):
    for r in records:
        urn = r.get("id")
        if not urn:
            continue
        time_obj = r.get("time") or {}
        valid_to = time_obj.get("valid_to")
        state[urn] = json.dumps({"urn": urn, "valid_to": valid_to}, sort_keys=True)


def sidecar_healthy() -> bool:
    try:
        r = requests.get(f"{SIDECAR_URL}/health", timeout=3)
        return r.status_code == 200 and r.json().get("status") == "ok"
    except requests.RequestException:
        return False


def main():
    parser = argparse.ArgumentParser(description="Tail ump store → push to ump-recall sidecar")
    parser.add_argument("--backfill", action="store_true", help="process all records and exit")
    parser.add_argument("--once", action="store_true", help="run one poll cycle and exit")
    args = parser.parse_args()

    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

    log.info("ump-watcher starting: ump_path=%s sidecar=%s poll=%.1fs",
             ump_path, SIDECAR_URL, POLL_INTERVAL)
    log.info("mode: %s", "backfill" if args.backfill else "once" if args.once else "daemon")

    if not sidecar_healthy():
        log.error("sidecar at %s is not healthy; aborting", SIDECAR_URL)
        sys.exit(1)

    state = load_state()
    log.info("loaded state: %d urns already tracked", len(state))

    if args.backfill:
        # Backfill mode: ignore state, push ALL records, exit.
        records = load_records(ump_path)
        log.info("backfill: %d total records in store", len(records))
        result = push_batch(records)
        log.info("backfill complete: pushed=%d failed=%d", result["pushed"], result["failed"])
        # Don't update state on backfill — let daemon mode track normally from here.
        sys.exit(0 if result["failed"] == 0 else 2)

    # Daemon / once mode
    while running:
        records = load_records(ump_path)
        if records:
            new_or_changed = diff_records(records, state)
            if new_or_changed:
                log.info("detected %d new/changed records (of %d total)",
                         len(new_or_changed), len(records))
                result = push_batch(new_or_changed)
                if result["pushed"] > 0:
                    update_state(records, state)
                    save_state(state)
                    log.info("state updated: %d urns tracked", len(state))
            else:
                log.debug("no changes (records=%d)", len(records))

        if args.once:
            break

        # Interruptible sleep
        slept = 0.0
        while running and slept < POLL_INTERVAL:
            time.sleep(min(0.5, POLL_INTERVAL - slept))
            slept += 0.5

    log.info("ump-watcher exiting cleanly")


if __name__ == "__main__":
    main()