Mab runner

evals/_chunking.py

@@ -0,0 +1,98 @@
+"""Sentence-aware token-budgeted chunker, aligned with official MemoryAgentBench.
+
+Port of ``utils/eval_other_utils.chunk_text_into_sentences`` from
+https://github.com/HUST-AI-HYZ/MemoryAgentBench. Used by ruler_eval,
+lru_eval, and longmem_eval so the ingest chunks all four MAB datasets
+(RULER QA1/QA2, LongMemEval-S, infbench_sum, detective_qa) see are
+the same shape the leaderboard runs use:
+
+  - **Atom**: sentence (NLTK ``punkt``).
+  - **Budget**: 4096 tokens of the ``gpt-4o-mini`` tiktoken encoding —
+    the value pinned by every MAB ``data_conf/**/*.yaml`` we've seen.
+  - **Bundle rule**: greedy pile sentences into the current chunk
+    until adding the next one would push the token count over budget,
+    then flush and start a new chunk with that sentence.
+  - **No mid-sentence split**: a sentence larger than the budget still
+    ships in its own chunk (matches the official behaviour — they have
+    no special case for this either).
+
+Why this matters: char-based 4096 (the previous policy) emitted ~4×
+more chunks than official because 4096 chars ≈ 1024 tokens. That made
+MIRIX's retrieval look worse than apples-to-apples because each
+semantic unit was scattered across multiple memories.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import nltk
+import tiktoken
+
+# Match the official MAB chunker exactly.
+DEFAULT_TOKEN_MODEL = "gpt-4o-mini"
+DEFAULT_CHUNK_TOKENS = 4096
+
+
+def _ensure_nltk() -> None:
+    """Make sure the NLTK punkt tokenizer is available.
+
+    Best-effort — NLTK ships a chatty downloader; quiet=True keeps the
+    eval logs clean. If the network is offline and punkt isn't cached
+    yet, the eval will surface the LookupError at first sent_tokenize
+    call, which is the same failure mode the official MAB script has.
+    """
+    try:
+        nltk.data.find("tokenizers/punkt_tab")
+        return
+    except LookupError:
+        pass
+    try:
+        nltk.data.find("tokenizers/punkt")
+        return
+    except LookupError:
+        pass
+    # punkt_tab is the modern split file; older NLTK versions still ship
+    # 'punkt'. Download both so it works on either.
+    for pkg in ("punkt_tab", "punkt"):
+        try:
+            nltk.download(pkg, quiet=True)
+        except Exception:
+            pass
+
+
+_ensure_nltk()
+_ENCODING = tiktoken.encoding_for_model(DEFAULT_TOKEN_MODEL)
+
+
+def chunk_text_into_sentences(
+    text: str,
+    chunk_size: int = DEFAULT_CHUNK_TOKENS,
+) -> List[str]:
+    """Bundle sentences greedily into <= ``chunk_size``-token chunks.
+
+    Returns a list of plain strings (each chunk's sentences joined by
+    a single space, matching the official script). Empty input returns
+    an empty list.
+    """
+    if not text:
+        return []
+
+    sentences = nltk.sent_tokenize(text)
+    chunks: List[str] = []
+    buf: List[str] = []
+    buf_tokens = 0
+
+    for sentence in sentences:
+        st = len(_ENCODING.encode(sentence, allowed_special={"<|endoftext|>"}))
+        if buf and buf_tokens + st > chunk_size:
+            chunks.append(" ".join(buf))
+            buf = [sentence]
+            buf_tokens = st
+        else:
+            buf.append(sentence)
+            buf_tokens += st
+
+    if buf:
+        chunks.append(" ".join(buf))
+    return chunks

evals/_eval_db.py

@@ -0,0 +1,183 @@
+"""Direct-PG helpers for eval scripts.
+
+Two helpers shared by ruler_eval / longmem_eval / main_eval:
+
+- ``measure_memory_size(sample_id)`` — total stored chars+tokens for a
+  user, plus per-table row counts; reads ``MIRIX_PG_DB`` env so it tracks
+  whatever DB the runner set (e.g. ``mirix_isolate``) instead of being
+  pinned to ``mirix``.
+
+- ``dump_memories(sample_id)`` — full per-user dump of episodic +
+  semantic rows, straight from Postgres. Bypasses
+  ``/memory/components``, which caps each memory type at 50 by default
+  and 200 even with an explicit ``limit`` parameter — so it's the only
+  way to get a ``_memories.json`` snapshot whose token count reflects
+  the real store on conversations with thousands of rows.
+"""
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from typing import Dict, List
+
+
+def _pg_conn() -> Dict[str, str]:
+    return {
+        "bin":  shutil.which("psql") or "/usr/local/opt/postgresql@17/bin/psql",
+        "host": os.environ.get("MIRIX_PG_HOST", "localhost"),
+        "user": os.environ.get("MIRIX_PG_USER", "mirix"),
+        "db":   os.environ.get("MIRIX_PG_DB",   "mirix"),
+        "pwd":  os.environ.get("MIRIX_PG_PASSWORD", "mirix"),
+    }
+
+
+def _pg_scalar(sql: str, timeout: int = 30) -> str:
+    c = _pg_conn()
+    try:
+        out = subprocess.run(
+            [c["bin"], "-h", c["host"], "-U", c["user"], "-d", c["db"], "-tAc", sql],
+            capture_output=True, text=True, timeout=timeout,
+            env={**os.environ, "PGPASSWORD": c["pwd"]},
+        )
+        return out.stdout.strip()
+    except Exception:
+        return ""
+
+
+def _pg_rows(sql: str, timeout: int = 300) -> List[List[str]]:
+    """Run SQL with control-byte separators so embedded newlines in
+    summary/details don't get mistaken for row boundaries.
+
+    SOH (``\\x01``) separates columns, STX (``\\x02``) separates records.
+    Neither appears in normal memory text.
+    """
+    c = _pg_conn()
+    out = subprocess.run(
+        [c["bin"], "-h", c["host"], "-U", c["user"], "-d", c["db"],
+         "-A", "-F", "\x01", "-R", "\x02", "-t", "-c", sql],
+        capture_output=True, text=True, timeout=timeout,
+        env={**os.environ, "PGPASSWORD": c["pwd"]},
+    )
+    rows: List[List[str]] = []
+    for chunk in out.stdout.split("\x02"):
+        if not chunk:
+            continue
+        rows.append(chunk.split("\x01"))
+    return rows
+
+
+def _esc(s: str) -> str:
+    return s.replace("'", "''")
+
+
+def measure_memory_size(sample_id: str) -> Dict:
+    """Total stored chars + tokens for ``sample_id``.
+
+    PG (flat) and Neo4j (graph) backends measured on the same yardstick
+    by concatenating every stored text field.
+    """
+    stats: Dict = {"unit": "chars+tokens"}
+    uid = _esc(sample_id)
+
+    flat: Dict = {}
+    flat_chars = 0
+    for table, cols in (
+        ("episodic_memory", "coalesce(summary,'')||coalesce(details,'')"),
+        ("semantic_memory", "coalesce(name,'')||coalesce(summary,'')||coalesce(details,'')"),
+    ):
+        row_n = _pg_scalar(f"SELECT count(*) FROM {table} WHERE user_id='{uid}';")
+        chars = _pg_scalar(
+            f"SELECT coalesce(sum(length({cols})),0) FROM {table} WHERE user_id='{uid}';"
+        )
+        n = int(row_n) if row_n.isdigit() else 0
+        c = int(chars) if chars.lstrip('-').isdigit() else 0
+        flat[table] = {"rows": n, "chars": c}
+        flat_chars += c
+    stats["flat"] = flat
+    stats["flat_total_chars"] = flat_chars
+
+    try:
+        from neo4j import GraphDatabase
+        from mirix.settings import settings
+        driver = GraphDatabase.driver(
+            settings.neo4j_uri, auth=(settings.neo4j_user, settings.neo4j_password)
+        )
+        with driver.session(database=settings.neo4j_database) as ses:
+            rec = ses.run(
+                """
+                MATCH (n) WHERE n.user_id=$uid
+                WITH count(n) AS nodes,
+                     sum(size(coalesce(n.description,'')) + size(coalesce(n.summary,''))) AS node_chars
+                RETURN nodes, node_chars
+                """, uid=sample_id
+            ).single()
+            erec = ses.run(
+                """
+                MATCH (a)-[r]->(b) WHERE a.user_id=$uid
+                RETURN count(r) AS edges,
+                       sum(size(coalesce(r.description,''))) AS edge_chars
+                """, uid=sample_id
+            ).single()
+        driver.close()
+        node_chars = (rec["node_chars"] or 0) if rec else 0
+        edge_chars = (erec["edge_chars"] or 0) if erec else 0
+        stats["graph"] = {
+            "nodes": (rec["nodes"] or 0) if rec else 0,
+            "edges": (erec["edges"] or 0) if erec else 0,
+            "node_chars": node_chars,
+            "edge_chars": edge_chars,
+        }
+        stats["graph_total_chars"] = node_chars + edge_chars
+    except Exception as exc:
+        stats["graph"] = {"error": str(exc)}
+        stats["graph_total_chars"] = 0
+
+    total_chars = max(stats["flat_total_chars"], stats["graph_total_chars"])
+    stats["total_chars"] = total_chars
+    stats["total_tokens"] = total_chars // 4
+    return stats
+
+
+def dump_memories(sample_id: str) -> Dict:
+    """Full per-user dump of episodic + semantic rows for ``sample_id``.
+
+    Shape matches what ``organize_results.count_memory_tokens`` expects:
+    ``{"memories": {"episodic": {"total_count": N, "items": [...]},
+                    "semantic": {...}}}``.
+    """
+    uid = _esc(sample_id)
+    epi_rows = _pg_rows(
+        "SELECT id, "
+        "coalesce(to_char(occurred_at,'YYYY-MM-DD\"T\"HH24:MI:SSOF'),''), "
+        "coalesce(event_type,''), coalesce(actor,''), "
+        "coalesce(summary,''), coalesce(details,'') "
+        f"FROM episodic_memory WHERE user_id='{uid}' "
+        "ORDER BY occurred_at NULLS LAST, created_at"
+    )
+    sem_rows = _pg_rows(
+        "SELECT id, coalesce(name,''), coalesce(summary,''), "
+        "coalesce(details,''), coalesce(source,'') "
+        f"FROM semantic_memory WHERE user_id='{uid}' ORDER BY created_at"
+    )
+    return {
+        "user_id": sample_id,
+        "memories": {
+            "episodic": {
+                "total_count": len(epi_rows),
+                "items": [
+                    {"id": r[0], "occurred_at": r[1], "event_type": r[2],
+                     "actor": r[3], "summary": r[4], "details": r[5]}
+                    for r in epi_rows
+                ],
+            },
+            "semantic": {
+                "total_count": len(sem_rows),
+                "items": [
+                    {"id": r[0], "name": r[1], "summary": r[2],
+                     "details": r[3], "source": r[4]}
+                    for r in sem_rows
+                ],
+            },
+        },
+    }

evals/configs/mab.yaml

@@ -0,0 +1,46 @@
+# Mirix Configuration — MemoryAgentBench (MAB) profile
+#
+# Used by `evals/run_mab_longmem_eval.sh` and any other MAB-benchmark
+# runner (LongMemEval-S etc., HF dataset ai-hyz/MemoryAgentBench).
+#
+# Independent from the 0201c LoCoMo profile so MAB-specific tuning
+# can diverge later without touching LoCoMo runs.
+
+llm_config:
+  model: "gpt-4.1-mini"
+  model_endpoint_type: "openai"
+  api_key: your-api-key
+  model_endpoint: "https://api.openai.com/v1"
+  context_window: 128000
+
+topic_extraction_llm_config:
+  model: "gpt-4.1-nano"
+  model_endpoint_type: "openai"
+  api_key: your-api-key
+  model_endpoint: "https://api.openai.com/v1"
+  context_window: 128000
+
+embedding_config:
+  embedding_model: "text-embedding-3-small"
+  embedding_endpoint: "https://api.openai.com/v1"
+  api_key: your-api-key
+  embedding_endpoint_type: "openai"
+  embedding_dim: 1536
+
+build_embeddings_for_memory: true
+
+meta_agent_config:
+  system_prompts_folder: prompts/0201a/
+  agents:
+    - core_memory_agent
+    - resource_memory_agent
+    - semantic_memory_agent
+    - episodic_memory_agent
+    - procedural_memory_agent
+    - knowledge_vault_memory_agent
+  memory:
+    core:
+      - label: "human"
+        value: ""
+      - label: "persona"
+        value: "I am a helpful assistant."