bubble-memory 0.1.0__py3-none-any.whl

bubble/__init__.py

@@ -0,0 +1,101 @@
+"""
+bubble — Hierarchical Memory Consolidation System
+
+Typical agent usage
+-------------------
+    import bubble
+
+    # Once, when a user session starts:
+    await bubble.init_graph(user_id)
+
+    # On every user message — retrieve and store in one call (preferred):
+    result = await bubble.observe(user_id, message, prior=agent_reply)
+    context = result["retrieved"]  # SnapshotNode results relevant to this message
+    stored  = result["stored"]     # ingested node descriptors
+
+    # Or separately:
+    await bubble.process(user_id, message, prior=agent_reply)
+    context = await bubble.retrieve(user_id, query)
+
+    # Periodically (runs HDBSCAN + promotion):
+    await bubble.consolidate(user_id)
+
+    # retrieved is a list of dicts:
+    #   {id, summary, members: [{id, summary, confidence_label}],
+    #    context: [{rel, id, summary, confidence_label}]}
+"""
+
+import asyncio
+
+from .db import get_graph, init_graph
+from .embed import embed as _embed
+from .decomposer import decompose as _decompose
+from .ingest import _route_segments, ingest, replay
+from .promote import promote
+from .retrieve import _retrieve_from_vecs, retrieve
+
+
+async def observe(user_id: str, message: str, prior: str | None = None, top_k: int = 3, verbose: bool = False) -> dict:
+    """
+    Decompose once, retrieve relevant memories, then store — all in a single call.
+
+    Shares the decompose+embed step between retrieval and ingestion.
+    Retrieval runs before storage so newly ingested segments don't appear in results.
+
+    Returns:
+      {
+        "retrieved": [...],  # same format as retrieve()
+        "stored":    [...],  # same format as process()
+      }
+    """
+    segments = await _decompose(message, prior)
+    embeddings = list(await asyncio.gather(*[_embed(s["text"]) for s in segments]))
+
+    g = get_graph(user_id)
+    stored = await _route_segments(user_id, segments, embeddings, prior)
+    retrieved = await _retrieve_from_vecs(g, message, embeddings, top_k, verbose)
+    return {"retrieved": retrieved, "stored": stored}
+
+
+async def process(user_id: str, message: str, prior: str | None = None) -> list[dict]:
+    """
+    Ingest a message into the user's memory graph.
+
+    Routes each segment to:
+      - Episodic Episode (intensity >= 0.6): JSONL + Layer 1 node immediately
+      - Layer 0 active pool (everything else): waits for consolidate()
+
+    prior: optional conversational context the user is responding to.
+    Returns the list of created node descriptors.
+    """
+    nodes = await ingest(user_id, message, prior)
+    await promote(user_id)
+    return nodes
+
+
+async def consolidate(user_id: str) -> dict:
+    """
+    Run the full consolidation pipeline on a user's graph:
+      1. HDBSCAN on the Layer 0 active pool
+      2. Promote clusters crossing the t_promo_score threshold to Episodes
+         (includes JSONL archival, SegmentNode deletion, L2 assignment)
+
+    Returns:
+      {"promoted": [...]}  # newly created Episode descriptors
+
+    Call periodically rather than on every message.
+    """
+    promoted = await promote(user_id)
+    return {"promoted": promoted}
+
+
+__all__ = [
+    "init_graph",
+    "observe",
+    "process",
+    "consolidate",
+    "retrieve",
+    "ingest",
+    "promote",
+    "replay",
+]

bubble/_shared.py

@@ -0,0 +1,49 @@
+import os
+from datetime import datetime, timezone
+
+import numpy as np
+from anthropic import AsyncAnthropic
+from dotenv import load_dotenv
+
+load_dotenv()
+
+MODEL = os.getenv("BUBBLE_MODEL", "claude-sonnet-4-6")
+_client = AsyncAnthropic()
+
+_SUMMARIZE_SYSTEM = """\
+You distill one or more user statements into a single memory record.
+
+Rules:
+- Capture the belief, preference, event, or tendency the statements express.
+- When multiple statements are given, identify the common pattern they share.
+- Write exactly one sentence with no grammatical subject.
+- Start with a verb or descriptor that names the belief, event, or pattern.
+- Do not explain, qualify, or ask for clarification.\
+"""
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _normalize(vec: np.ndarray) -> list[float]:
+    """L2-normalize a numpy vector and return as a Python list."""
+    norm = np.linalg.norm(vec)
+    return (vec / norm if norm > 0 else vec).tolist()
+
+
+def _centroid(nodes: list[dict]) -> list[float]:
+    """Mean of source embeddings, L2-normalized."""
+    matrix = np.array([n["embedding"] for n in nodes], dtype=np.float32)
+    return _normalize(matrix.mean(axis=0))
+
+
+async def _summarize(nodes: list[dict]) -> str:
+    texts = "\n".join(f"- {n['raw_text']}" for n in nodes)
+    response = await _client.messages.create(
+        model=MODEL,
+        max_tokens=128,
+        system=_SUMMARIZE_SYSTEM,
+        messages=[{"role": "user", "content": texts}],
+    )
+    return response.content[0].text.strip()

bubble/archive.py

@@ -0,0 +1,48 @@
+import json
+import os
+from pathlib import Path
+
+_ARCHIVE_DIR = os.getenv("BUBBLE_ARCHIVE_DIR", "./data/archive")
+_MKDIR_DONE = False
+
+
+def _path(user_id: str) -> Path:
+    global _MKDIR_DONE
+    p = Path(_ARCHIVE_DIR)
+    if not _MKDIR_DONE:
+        p.mkdir(parents=True, exist_ok=True)
+        _MKDIR_DONE = True
+    return p / f"{user_id}.jsonl"
+
+
+def read_segments(user_id: str):
+    """Yield all archived segment records for a user."""
+    path = _path(user_id)
+    if not path.exists():
+        return
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                yield json.loads(line)
+
+
+def write_segment(
+    user_id: str,
+    *,
+    text: str,
+    prior: str | None,
+    intensity: float,
+    valence: str,
+    timestamp: str,
+) -> None:
+    """Append one segment record to the user's JSONL archive."""
+    entry = {
+        "text":      text,
+        "prior":     prior,
+        "intensity": intensity,
+        "valence":   valence,
+        "timestamp": timestamp,
+    }
+    with _path(user_id).open("a", encoding="utf-8") as f:
+        f.write(json.dumps(entry) + "\n")

bubble/chain.py

@@ -0,0 +1,317 @@
+"""
+Chain assignment pipeline — spec §3.8.
+
+When a new Episode is created, assign it to a topic chain:
+  1 — ANN top-1 against Episode centroids.
+  2 — Similarity threshold check. Below → new isolated SnapshotNode.
+  3 — LLM relatedness check (binary). Not related → new isolated SnapshotNode.
+  4 — Related → traverse FOLLOWED_BY edges to chain tail, wire FOLLOWED_BY edge, join SnapshotNode.
+"""
+
+import os
+import uuid
+
+import numpy as np
+
+from ._shared import MODEL, _client, _normalize, _now
+from .db import get_graph
+
+_CHAIN_MAX_DISTANCE = float(os.getenv("BUBBLE_CHAIN_MAX_DISTANCE", "0.4"))
+_NLI_ENABLED = os.getenv("BUBBLE_ENABLE_NLI", "false").lower() == "true"
+_NLI_ENDPOINT = os.getenv("BUBBLE_NLI_ENDPOINT", "http://localhost:8999/predict")
+
+_SNAPSHOT_SYSTEM = """\
+A sequence of memory records about the user is listed below, from earliest to most recent.
+
+Rules:
+- The most recent memory takes precedence over earlier ones.
+- Earlier memory provides historical context.
+- Synthesize all memory into a single simplified coherent narrative that represents the full arc.
+- Output one concise paragraph.
+- No subject, start with verb.
+- Do not explain or justify.\
+"""
+
+# ---------------------------------------------------------------------------
+# Relatedness check (LLM or NLI)
+# ---------------------------------------------------------------------------
+
+
+async def _related(summary_a: str, summary_b: str) -> bool:
+    """True if the two summaries are about the same topic.
+
+    BUBBLE_CHAIN_BACKEND=nli  — local NLI model, argmax != neutral → related.
+    BUBBLE_CHAIN_BACKEND=llm  — LLM yes/no (default).
+    """
+    if _NLI_ENABLED:
+        import httpx
+
+        async with httpx.AsyncClient() as client:
+            resp = await client.post(
+                _NLI_ENDPOINT,
+                json={"inputs": [summary_a, summary_b]},
+                headers={"Content-Type": "application/json", "Accept": "application/json"},
+            )
+            resp.raise_for_status()
+            scores = resp.json()  # [{"label": ..., "score": ...}, ...]
+        return max(scores, key=lambda x: x["score"])["label"].lower() != "neutral"
+
+    answer = (
+        (
+            await _client.messages.create(
+                model=MODEL,
+                max_tokens=8,
+                system="You are a memory topic classifier. Reply with exactly 'yes' or 'no'.",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": (f"Are these two beliefs about the same topic or subject?\n\nA: {summary_a}\nB: {summary_b}"),
+                    }
+                ],
+            )
+        )
+        .content[0]
+        .text.strip()
+        .lower()
+    )
+    return answer.startswith("y")
+
+
+# ---------------------------------------------------------------------------
+# Chain traversal
+# ---------------------------------------------------------------------------
+
+
+async def _traverse_to_tail(g, start_id: str) -> str:
+    """
+    Follow FOLLOWED_BY edges from start_id to the chain tail
+    (the node with no outgoing FOLLOWED_BY edge).
+    """
+    result = await g.query(
+        "MATCH (start:Episode {id: $id})-[:FOLLOWED_BY*0..]->(tail:Episode) WHERE NOT (tail)-[:FOLLOWED_BY]->() RETURN tail.id LIMIT 1",
+        {"id": start_id},
+    )
+    if result.result_set:
+        return result.result_set[0][0]
+    return start_id
+
+
+# ---------------------------------------------------------------------------
+# Graph writers
+# ---------------------------------------------------------------------------
+
+
+async def _wire_follows(g, from_id: str, to_id: str) -> None:
+    await g.query(
+        "MATCH (a:Episode {id: $a}), (b:Episode {id: $b}) CREATE (a)-[:FOLLOWED_BY]->(b)",
+        {"a": from_id, "b": to_id},
+    )
+
+
+# ---------------------------------------------------------------------------
+# Graph loaders
+# ---------------------------------------------------------------------------
+
+
+def _rows_to_episode_dicts(rows) -> list[dict]:
+    return [
+        {
+            "id": r[0],
+            "summary": r[1],
+            "centroid": r[2],
+            "episodic": bool(r[3]),
+            "timestamp": r[4],
+            "valence": r[5] or "neu",
+        }
+        for r in rows
+        if r[2] is not None
+    ]
+
+
+async def _load_node(g, episode_id: str) -> dict | None:
+    result = await g.query(
+        "MATCH (t:Episode {id: $id}) RETURN t.id, t.summary, t.centroid, t.episodic, t.timestamp, t.valence",
+        {"id": episode_id},
+    )
+    rows = _rows_to_episode_dicts(result.result_set)
+    return rows[0] if rows else None
+
+
+# ---------------------------------------------------------------------------
+# SnapshotNode management
+# ---------------------------------------------------------------------------
+
+
+async def _recompute_snapshot_centroid(g, snap_id: str) -> None:
+    result = await g.query(
+        "MATCH (snap:SnapshotNode {id: $id})-[:SYNTHESIZES]->(t:Episode) RETURN t.centroid",
+        {"id": snap_id},
+    )
+    centroids = [row[0] for row in result.result_set if row[0] is not None]
+    if not centroids:
+        return
+    new_centroid = _normalize(np.array(centroids, dtype=np.float32).mean(axis=0))
+    await g.query(
+        "MATCH (snap:SnapshotNode {id: $id}) SET snap.centroid = vecf32($centroid)",
+        {"id": snap_id, "centroid": new_centroid},
+    )
+
+
+async def _join_snapshot(g, snap_id: str, new_node_id: str) -> None:
+    await g.query(
+        "MATCH (snap:SnapshotNode {id: $snap_id}), (t:Episode {id: $tid}) CREATE (snap)-[:SYNTHESIZES]->(t) SET snap.valid = false",
+        {"snap_id": snap_id, "tid": new_node_id},
+    )
+    await _recompute_snapshot_centroid(g, snap_id)
+
+
+async def _create_snapshot(g, episode_id: str, summary: str, centroid: list[float]) -> None:
+    snap_id = str(uuid.uuid4())
+    await g.query(
+        "CREATE (snap:SnapshotNode {  id: $id, summary: $summary, centroid: vecf32($centroid), valid: true, timestamp: $ts})",
+        {"id": snap_id, "summary": summary, "centroid": centroid, "ts": _now()},
+    )
+    await g.query(
+        "MATCH (snap:SnapshotNode {id: $snap_id}), (t:Episode {id: $tid}) CREATE (snap)-[:SYNTHESIZES]->(t)",
+        {"snap_id": snap_id, "tid": episode_id},
+    )
+
+
+# ---------------------------------------------------------------------------
+# Lazy snapshot summary generation
+# ---------------------------------------------------------------------------
+
+
+async def ensure_snapshot_summary(
+    g,
+    snap_id: str,
+    *,
+    valid: bool | None = None,
+    summary: str | None = None,
+) -> str | None:
+    """
+    Return the SnapshotNode summary, generating it lazily if valid=false or summary is null.
+
+    Pass `valid` and `summary` when already fetched (e.g. from ANN query) to skip the
+    redundant DB round-trip.
+
+    Single-member chains: copy Episode summary directly (no LLM).
+    Multi-member chains: LLM synthesis ordered by timestamp, episodic members last.
+    """
+    if valid is None or summary is None:
+        result = await g.query(
+            "MATCH (snap:SnapshotNode {id: $id}) RETURN snap.summary, snap.valid",
+            {"id": snap_id},
+        )
+        if not result.result_set:
+            return None
+        summary, valid = result.result_set[0]
+    if valid and summary:
+        return summary
+
+    result = await g.query(
+        "MATCH (snap:SnapshotNode {id: $id})-[:SYNTHESIZES]->(t:Episode) RETURN t.id, t.summary, t.episodic, t.timestamp",
+        {"id": snap_id},
+    )
+    members = [
+        {
+            "id": r[0],
+            "summary": r[1],
+            "episodic": bool(r[2]),
+            "timestamp": r[3] or "",
+        }
+        for r in result.result_set
+    ]
+
+    if not members:
+        return None
+
+    if len(members) == 1:
+        new_summary = members[0]["summary"]
+    else:
+        non_ep = sorted([m for m in members if not m["episodic"]], key=lambda m: m["timestamp"])
+        ep = sorted([m for m in members if m["episodic"]], key=lambda m: m["timestamp"])
+        ordered = non_ep + ep
+
+        count = len(ordered)
+        lines = []
+        for i, m in enumerate(ordered):
+            tag = " [episodic]" if m["episodic"] else ""
+            if i == 0:
+                label = f"Belief 1 (earliest){tag}"
+            elif i == count - 1:
+                label = f"Belief {i + 1} (most recent){tag}"
+            else:
+                label = f"Belief {i + 1}{tag}"
+            lines.append(f"{label}: {m['summary']}")
+
+        new_summary = (
+            (
+                await _client.messages.create(
+                    model=MODEL,
+                    max_tokens=256,
+                    system=_SNAPSHOT_SYSTEM,
+                    messages=[{"role": "user", "content": "\n\n".join(lines)}],
+                )
+            )
+            .content[0]
+            .text.strip()
+        )
+
+    await g.query(
+        "MATCH (snap:SnapshotNode {id: $id}) SET snap.summary = $summary, snap.valid = true",
+        {"id": snap_id, "summary": new_summary},
+    )
+    return new_summary
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+async def check_new(user_id: str, new_episode_id: str) -> None:
+    """
+    Assign a newly created Episode to a topic chain.
+    Called by promote() and ingest._store_episodic() after Episode creation.
+    """
+    g = get_graph(user_id)
+    new_node = await _load_node(g, new_episode_id)
+    if new_node is None or new_node["centroid"] is None:
+        return
+
+    # Step 1 — ANN top-1 (k=3: self may occupy a slot if already indexed)
+    result = await g.query(
+        "CALL db.idx.vector.queryNodes('Episode', 'centroid', 2, vecf32($vec)) YIELD node, score WHERE node.id <> $id RETURN node.id, node.summary, score LIMIT 1",
+        {"vec": new_node["centroid"], "id": new_episode_id},
+    )
+
+    if not result.result_set:
+        # No other Episodes exist yet
+        await _create_snapshot(g, new_episode_id, new_node["summary"], new_node["centroid"])
+        return
+
+    closest_id, closest_summary, score = result.result_set[0]
+
+    # Step 2 — Similarity threshold
+    if score > _CHAIN_MAX_DISTANCE:
+        await _create_snapshot(g, new_episode_id, new_node["summary"], new_node["centroid"])
+        return
+
+    # Step 3 — relatedness check
+    if not await _related(new_node["summary"], closest_summary):
+        await _create_snapshot(g, new_episode_id, new_node["summary"], new_node["centroid"])
+        return
+
+    # Step 4 — Append to chain
+    tail_id = await _traverse_to_tail(g, closest_id)
+    await _wire_follows(g, tail_id, new_episode_id)
+
+    snap_result = await g.query(
+        "MATCH (snap:SnapshotNode)-[:SYNTHESIZES]->(t:Episode {id: $id}) RETURN snap.id",
+        {"id": closest_id},
+    )
+    if snap_result.result_set:
+        await _join_snapshot(g, snap_result.result_set[0][0], new_episode_id)
+    else:
+        await _create_snapshot(g, new_episode_id, new_node["summary"], new_node["centroid"])

bubble/cluster.py

@@ -0,0 +1,66 @@
+import os
+
+import numpy as np
+from sklearn.cluster import HDBSCAN
+
+from .db import get_graph
+
+_CLUSTER_MIN_SIZE = int(os.getenv("BUBBLE_CLUSTER_MIN_SIZE", "3"))
+_CLUSTER_DIMS = int(os.getenv("BUBBLE_CLUSTER_DIMS", "128"))
+
+
+async def get_clusters(user_id: str) -> dict[int, list[dict]]:
+    """
+    Run HDBSCAN on all SegmentNodes for a user.
+    All SegmentNodes in the graph are the active pool — promoted nodes are deleted at promotion.
+
+    Returns {cluster_label: [node_dicts]} — noise (label -1) is excluded.
+    Returns empty dict if the pool is smaller than _CLUSTER_MIN_SIZE.
+
+    Each node dict contains: id, raw_text, embedding (768-dim), intensity, valence,
+    prior (str|None), timestamp (str).
+    """
+    g = get_graph(user_id)
+    result = await g.query(
+        "MATCH (n:SegmentNode) "
+        "RETURN n.id, n.raw_text, n.embedding, n.intensity, n.valence, n.prior, n.timestamp"
+    )
+
+    if not result.result_set:
+        return {}
+
+    nodes = [
+        {
+            "id":        row[0],
+            "raw_text":  row[1],
+            "embedding": row[2],
+            "intensity": row[3],
+            "valence":   row[4],
+            "prior":     row[5],
+            "timestamp": row[6],
+        }
+        for row in result.result_set
+    ]
+
+    if len(nodes) < _CLUSTER_MIN_SIZE:
+        return {}
+
+    # Truncate to 128 dims (Matryoshka), then re-normalize.
+    # The full 768-dim vector is unit-norm, but the truncated prefix is not —
+    # re-normalizing restores the cosine distance relationship for euclidean HDBSCAN.
+    raw = np.array([n["embedding"][:_CLUSTER_DIMS] for n in nodes], dtype=np.float32)
+    norms = np.linalg.norm(raw, axis=1, keepdims=True)
+    matrix = raw / np.where(norms > 0, norms, 1.0)
+
+    labels = HDBSCAN(
+        min_cluster_size=_CLUSTER_MIN_SIZE,
+        metric="euclidean",  # equiv to cosine on unit vectors
+    ).fit_predict(matrix)
+
+    clusters: dict[int, list[dict]] = {}
+    for node, label in zip(nodes, labels):
+        if label == -1:
+            continue  # noise — accumulates at Layer 0, never promotes
+        clusters.setdefault(int(label), []).append(node)
+
+    return clusters

bubble/db.py

@@ -0,0 +1,48 @@
+import os
+from dotenv import load_dotenv
+from falkordb.asyncio import FalkorDB
+
+load_dotenv()
+
+_client: FalkorDB | None = None
+
+
+def get_client() -> FalkorDB:
+    global _client
+    if _client is None:
+        _client = FalkorDB(
+            host=os.getenv("FALKORDB_HOST", "localhost"),
+            port=int(os.getenv("FALKORDB_PORT", "6379")),
+        )
+    return _client
+
+
+def get_graph(user_id: str):
+    return get_client().select_graph(f"bubble:{user_id}")
+
+
+async def init_graph(user_id: str) -> None:
+    """Create indexes for a user graph. Safe to call on an already-initialized graph."""
+    g = get_graph(user_id)
+
+    from .embed import EMBED_DIM
+
+    # HNSW vector index on SnapshotNode.centroid — retrieval (L2 entry point).
+    try:
+        await g.query(
+            "CREATE VECTOR INDEX FOR (n:SnapshotNode) ON (n.centroid) "
+            f"OPTIONS {{dimension: {EMBED_DIM}, similarityFunction: 'cosine'}}"
+        )
+    except Exception:
+        pass  # already exists
+
+    # HNSW vector index on Episode.centroid — chain assignment ANN search.
+    # Used by check_new (chain.py) to find the nearest existing Episode
+    # for the three-gate chain assignment cascade.
+    try:
+        await g.query(
+            "CREATE VECTOR INDEX FOR (n:Episode) ON (n.centroid) "
+            f"OPTIONS {{dimension: {EMBED_DIM}, similarityFunction: 'cosine'}}"
+        )
+    except Exception:
+        pass  # already exists