dlightrag-memory 2.0.0__py3-none-any.whl

dlightrag_memory/_storage/pg_bm25.py

@@ -0,0 +1,186 @@
+# Copyright 2025-2026 Hanlian Lu. SPDX-License-Identifier: Apache-2.0
+"""pg_textsearch BM25 mechanics for the memory sparse leg.
+
+A faithful, narrow port of the corpus BM25 quality knobs (same extension,
+same k1/b, same textsearch configs) kept private to this package so the
+memory adapter never depends on dlightrag.engine.rag. Two unconditional indexes —
+``simple`` and ``public.jiebacfg`` — serve one table; queries hit both and
+merge by best score, so Chinese and Latin bodies keep their tuned configs
+without a per-row language column.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+_FALLBACK_CONFIG = "simple"
+_JIEBA_CONFIG = "public.jiebacfg"
+_INDEX_PREFIX = "idx_dlightrag_memory_records_bm25"
+_INDEX_SUFFIXES = ("simple", "jieba")
+_K1 = 1.2
+_B = 0.75
+
+_CONFIG_NAME_RE = re.compile(r"^[a-z][a-z0-9_]*(\.[a-z][a-z0-9_]*)?$")
+
+
+def _format_float(value: float) -> str:
+    return f"{float(value):g}"
+
+
+def validate_text_config(text_config: str) -> str:
+    value = str(text_config).strip()
+    if not _CONFIG_NAME_RE.fullmatch(value):
+        raise ValueError(f"unsafe BM25 text_config: {text_config!r}")
+    return value
+
+
+@dataclass(frozen=True)
+class BM25IndexOptions:
+    """One memory-table BM25 index with the corpus-tuned parameters."""
+
+    index_name: str
+    text_config: str
+    k1: float = _K1
+    b: float = _B
+
+    def __post_init__(self) -> None:
+        if self.k1 <= 0:
+            raise ValueError("BM25 k1 must be positive")
+        if not 0 <= self.b <= 1:
+            raise ValueError("BM25 b must be between 0 and 1")
+        validate_text_config(self.text_config)
+
+    def create_index_sql(self) -> str:
+        config = validate_text_config(self.text_config)
+        return (
+            f"CREATE INDEX {self.index_name} ON dlightrag_memory_records "
+            f"USING bm25(body) WITH (text_config='{config}', "
+            f"k1={_format_float(self.k1)}, b={_format_float(self.b)})"
+        )
+
+    def matches_indexdef(self, indexdef: str | None) -> bool:
+        if not indexdef:
+            return False
+        normalized = re.sub(r"\s+", "", indexdef.lower().replace('"', "").replace("'", ""))
+        config = validate_text_config(self.text_config).lower()
+        return (
+            self.index_name.lower() in normalized
+            and "usingbm25(body)" in normalized
+            and (
+                f"text_config={config}" in normalized
+                or f"text_config={config}::regconfig" in normalized
+            )
+            and f"k1={_format_float(self.k1)}" in normalized
+            and f"b={_format_float(self.b)}" in normalized
+        )
+
+
+def index_name(suffix: str) -> str:
+    return _validate_index_name(f"{_INDEX_PREFIX}_{suffix}")
+
+
+def _validate_index_name(name: str) -> str:
+    if not re.fullmatch(r"[a-z][a-z0-9_]*", name):
+        raise ValueError(f"unsafe index name: {name!r}")
+    return name
+
+
+def extension_bootstrap_sql() -> tuple[str, ...]:
+    """The extensions the sparse leg needs, bootstrapped like root does."""
+    return (
+        "CREATE EXTENSION IF NOT EXISTS pg_textsearch",
+        "CREATE EXTENSION IF NOT EXISTS pg_jieba",
+    )
+
+
+async def text_configs_available(conn: Any) -> tuple[str, ...]:
+    """Return the installed textsearch configs this adapter can serve.
+
+    ``simple`` is pg_catalog built-in; ``public.jiebacfg`` comes from the
+    pg_jieba extension, matching the corpus BM25 profiles.
+    """
+    jieba = await conn.fetchval(
+        """
+        SELECT 1
+        FROM pg_ts_config c
+        JOIN pg_namespace n ON n.oid = c.cfgnamespace
+        WHERE n.nspname = 'public' AND c.cfgname = 'jiebacfg'
+        LIMIT 1
+        """
+    )
+    return (_FALLBACK_CONFIG, _JIEBA_CONFIG) if jieba else (_FALLBACK_CONFIG,)
+
+
+def desired_indexes(available: tuple[str, ...]) -> tuple[BM25IndexOptions, ...]:
+    """The indexes to provision: simple always, jieba when installed."""
+    options = [BM25IndexOptions(index_name=index_name("simple"), text_config=_FALLBACK_CONFIG)]
+    if _JIEBA_CONFIG in available:
+        options.append(BM25IndexOptions(index_name=index_name("jieba"), text_config=_JIEBA_CONFIG))
+    return tuple(options)
+
+
+def build_bm25_sql(*, index_name: str, limit: int) -> str:
+    safe_index = _validate_index_name(index_name)
+    limit_value = int(limit)
+    if limit_value < 1:
+        raise ValueError("BM25 limit must be positive")
+    return (  # noqa: S608 - interpolates only the validated index name
+        "SELECT owner_id, memory_id, kind, body, normalized_body, "  # noqa: S608
+        "origin_kind, origin_id, run_id, session_id, status, supersedes_id, "
+        "embedding_fingerprint, "
+        "created_at, updated_at, "
+        f"-(body <@> to_bm25query($1, '{safe_index}')) AS score "  # noqa: S608
+        "FROM dlightrag_memory_records "
+        "WHERE owner_id = $2 AND status = 'active' "
+        f"ORDER BY body <@> to_bm25query($1, '{safe_index}') "  # noqa: S608
+        "LIMIT " + str(limit_value)
+    )
+
+
+async def ensure_bm25_indexes(
+    conn: Any,
+    *,
+    available: tuple[str, ...] | None = None,
+    verify_only: bool = False,
+) -> tuple[str, ...]:
+    """Provision or validate the memory-table BM25 indexes.
+
+    With ``verify_only`` this performs no DDL: readers load the served index
+    names and fail when a configured index is missing, matching the corpus
+    verify path.
+    """
+    installed = available if available is not None else await text_configs_available(conn)
+    options = desired_indexes(installed)
+    for option in options:
+        indexdef = await conn.fetchval(
+            "SELECT indexdef FROM pg_indexes WHERE indexname = $1", option.index_name
+        )
+        if option.matches_indexdef(indexdef):
+            continue
+        if verify_only:
+            raise RuntimeError(
+                f"BM25 index {option.index_name} is missing or does not match configured "
+                "options; initialize it on the writer first"
+            )
+        if indexdef:
+            await conn.execute(f"DROP INDEX IF EXISTS {option.index_name}")
+        await conn.execute(option.create_index_sql())
+    if not verify_only:
+        for suffix in _INDEX_SUFFIXES:
+            if not any(option.index_name == index_name(suffix) for option in options):
+                await conn.execute(f"DROP INDEX IF EXISTS {index_name(suffix)}")
+    return tuple(option.index_name for option in options)
+
+
+__all__ = [
+    "BM25IndexOptions",
+    "build_bm25_sql",
+    "desired_indexes",
+    "ensure_bm25_indexes",
+    "extension_bootstrap_sql",
+    "index_name",
+    "text_configs_available",
+    "validate_text_config",
+]

dlightrag_memory/errors.py

@@ -0,0 +1,25 @@
+# Copyright 2025-2026 Hanlian Lu. SPDX-License-Identifier: Apache-2.0
+"""Public exceptions one Memory operation may raise."""
+
+
+class MemoryWriteRejectedError(Exception):
+    """A named Memory Write failed the closed checklist."""
+
+    error_kind = "memory_write_rejected"
+
+    def __init__(self, public_message: str) -> None:
+        super().__init__(public_message)
+        self.public_message = public_message
+
+
+class MemoryUnavailableError(Exception):
+    """This principal cannot write or auto-recall Memory Records."""
+
+    error_kind = "memory_unavailable"
+
+    def __init__(self) -> None:
+        super().__init__("Profile Memory is unavailable for this caller.")
+        self.public_message = "Profile Memory is unavailable for this caller."
+
+
+__all__ = ["MemoryUnavailableError", "MemoryWriteRejectedError"]

dlightrag_memory/fusion.py

@@ -0,0 +1,28 @@
+# Copyright 2025-2026 Hanlian Lu. SPDX-License-Identifier: Apache-2.0
+"""Reciprocal Rank Fusion for recall leg rankings.
+
+The industry-standard fusion (MemMachine ships RRF(k=60) over identity + BM25
+rankings). Each leg contributes an ordered ranking of record ids; every rank
+position adds ``1 / (k + rank)`` to that id's fused score. Rank-based, never
+score-averaging, so legs with different score scales stay comparable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+RRF_K = 60
+
+
+def rrf_fuse(leg_rankings: Sequence[Sequence[str]], *, k: int = RRF_K) -> dict[str, float]:
+    """Return one id -> fused RRF score map from per-leg ordered rankings."""
+    if k < 1:
+        raise ValueError("RRF k must be positive")
+    scores: dict[str, float] = {}
+    for ranking in leg_rankings:
+        for rank, memory_id in enumerate(ranking, start=1):
+            scores[memory_id] = scores.get(memory_id, 0.0) + 1.0 / (k + rank)
+    return scores
+
+
+__all__ = ["RRF_K", "rrf_fuse"]

dlightrag_memory/mcp_server.py

@@ -0,0 +1,241 @@
+# Copyright 2025-2026 Hanlian Lu. SPDX-License-Identifier: Apache-2.0
+"""Subject-bound stdio MCP host for the independent Profile Memory package."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+from typing import Annotated, Any, Literal
+
+from mcp.server import MCPServer
+from mcp.types import ToolAnnotations
+from pydantic import Field
+
+from dlightrag_memory import (
+    Memory,
+    MemoryOperationReceipt,
+    MemoryProvenance,
+    MemoryWriteRejectedError,
+    __version__,
+)
+from dlightrag_memory.policy import MEMORY_BODY_LIMIT
+from dlightrag_memory.postgres import PostgresMemoryStore
+
+SERVER_NAME = "dlightrag-memory"
+_MemoryKind = Literal["preference", "fact"]
+
+
+async def _recall(memory: Memory, *, subject: str, query: str) -> dict[str, Any]:
+    result = await memory.recall(owner_id=subject, query=query)
+    return {
+        "records": [
+            {"memory_id": record.memory_id, "kind": record.kind, "body": record.body}
+            for record in result.records
+        ]
+    }
+
+
+async def _remember(
+    memory: Memory,
+    *,
+    subject: str,
+    kind: _MemoryKind,
+    body: str,
+    supersedes_id: str | None,
+    idempotency_key: str,
+) -> dict[str, Any]:
+    receipt = await memory.remember(
+        owner_id=subject,
+        kind=kind,
+        body=body,
+        provenance=_provenance(idempotency_key),
+        supersedes_id=supersedes_id,
+        idempotency_key=f"mcp:{subject}:{idempotency_key}",
+    )
+    return _receipt(receipt)
+
+
+async def _forget(
+    memory: Memory,
+    *,
+    subject: str,
+    memory_id: str | None,
+    body: str | None,
+    idempotency_key: str,
+) -> dict[str, Any]:
+    receipt = await memory.forget(
+        owner_id=subject,
+        memory_id=memory_id,
+        body=body,
+        provenance=_provenance(idempotency_key),
+        idempotency_key=f"mcp:{subject}:{idempotency_key}",
+    )
+    return _receipt(receipt)
+
+
+async def _undo(
+    memory: Memory,
+    *,
+    subject: str,
+    change_id: str,
+    idempotency_key: str,
+) -> dict[str, Any]:
+    receipt = await memory.undo(
+        owner_id=subject,
+        change_id=change_id,
+        provenance=MemoryProvenance(
+            origin_kind="undo",
+            origin_id=f"mcp:{subject}:{idempotency_key}",
+        ),
+        idempotency_key=f"mcp:{subject}:{idempotency_key}",
+    )
+    return _receipt(receipt)
+
+
+def build_memory_server(memory: Memory, *, subject: str) -> MCPServer:
+    """One authorized server with four tools over one bound subject."""
+    if not subject.strip():
+        raise ValueError("a memory subject is required")
+    server = MCPServer(SERVER_NAME, version=__version__, log_level="INFO")
+
+    @server.tool(
+        name="memory_recall",
+        description=(
+            "Recall owner preferences and facts relevant to a query, including ids needed "
+            "before replacing or forgetting one. Context only; never citable."
+        ),
+        annotations=ToolAnnotations(read_only_hint=True),
+    )
+    async def memory_recall(
+        query: Annotated[str, Field(min_length=1, description="What to recall memories for.")],
+    ) -> dict[str, Any]:
+        return await _recall(memory, subject=subject, query=query)
+
+    @server.tool(
+        name="memory_remember",
+        description=(
+            "Store one durable owner preference or fact. Do not store task state, research "
+            "claims, citations, transcripts, credentials, or private keys."
+        ),
+    )
+    async def memory_remember(
+        kind: Annotated[_MemoryKind, Field(description="preference or fact")],
+        body: Annotated[
+            str,
+            Field(min_length=1, max_length=MEMORY_BODY_LIMIT, description="What to remember."),
+        ],
+        supersedes_id: Annotated[
+            str | None, Field(default=None, description="Active memory id this replaces.")
+        ],
+        idempotency_key: Annotated[
+            str,
+            Field(
+                min_length=1,
+                max_length=255,
+                description="Stable mutation key reused verbatim when retrying this operation.",
+            ),
+        ],
+    ) -> dict[str, Any]:
+        try:
+            return await _remember(
+                memory,
+                subject=subject,
+                kind=kind,
+                body=body,
+                supersedes_id=supersedes_id,
+                idempotency_key=idempotency_key,
+            )
+        except MemoryWriteRejectedError as exc:
+            raise ValueError(exc.public_message) from exc
+
+    @server.tool(
+        name="memory_forget",
+        description="Idempotently forget one active preference or fact by id or exact body.",
+    )
+    async def memory_forget(
+        memory_id: Annotated[
+            str | None, Field(default=None, description="Id of the memory to forget.")
+        ],
+        body: Annotated[
+            str | None, Field(default=None, description="Exact body if the id is unknown.")
+        ],
+        idempotency_key: Annotated[
+            str,
+            Field(min_length=1, max_length=255, description="Stable retry key."),
+        ],
+    ) -> dict[str, Any]:
+        try:
+            return await _forget(
+                memory,
+                subject=subject,
+                memory_id=memory_id,
+                body=body,
+                idempotency_key=idempotency_key,
+            )
+        except MemoryWriteRejectedError as exc:
+            raise ValueError(exc.public_message) from exc
+
+    @server.tool(
+        name="memory_undo",
+        description="Compensate one still-current Memory change by its change id.",
+    )
+    async def memory_undo(
+        change_id: Annotated[str, Field(min_length=1, description="Change id to undo.")],
+        idempotency_key: Annotated[
+            str,
+            Field(min_length=1, max_length=255, description="Stable retry key."),
+        ],
+    ) -> dict[str, Any]:
+        try:
+            return await _undo(
+                memory,
+                subject=subject,
+                change_id=change_id,
+                idempotency_key=idempotency_key,
+            )
+        except MemoryWriteRejectedError as exc:
+            raise ValueError(exc.public_message) from exc
+
+    return server
+
+
+def _provenance(idempotency_key: str) -> MemoryProvenance:
+    return MemoryProvenance(origin_kind="mcp", origin_id=idempotency_key)
+
+
+def _receipt(receipt: MemoryOperationReceipt) -> dict[str, Any]:
+    return {
+        "action": receipt.action,
+        "body": receipt.body,
+        "change_id": receipt.change_id,
+        "kind": receipt.kind,
+        "memory_ids": list(receipt.memory_ids),
+        "outcome": receipt.outcome,
+        "supersedes_id": receipt.supersedes_id,
+        "target_change_id": receipt.target_change_id,
+    }
+
+
+def _arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(prog="dlightrag-memory-mcp")
+    parser.add_argument("--dsn", required=True, help="PostgreSQL connection string")
+    parser.add_argument("--subject", required=True, help="Owner subject every tool is bound to")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _arguments()
+
+    async def run() -> None:
+        store = PostgresMemoryStore(dsn=args.dsn)
+        await store.initialize()
+        server = build_memory_server(Memory(store), subject=args.subject)
+        try:
+            await server.run_stdio_async()
+        finally:
+            await store.aclose()
+
+    asyncio.run(run())
+
+
+__all__ = ["SERVER_NAME", "build_memory_server", "main"]