betterdb-agent-memory 0.1.0__py3-none-any.whl

betterdb_agent_memory/__init__.py

@@ -0,0 +1,88 @@
+"""betterdb-agent-memory: long-term vector memory tier for AI agents on Valkey.
+
+Re-exports everything from ``betterdb-agent-cache`` (the short-term cache tiers)
+alongside the memory tier so the facade is a single import.
+"""
+
+from __future__ import annotations
+
+import betterdb_agent_cache as _agent_cache
+from betterdb_agent_cache import *  # noqa: F401,F403
+
+from .agent_memory import AgentMemory, AgentMemoryOptions
+from .build_recall_query import MATCH_ALL_MEMORY_QUERY
+from .composite_score import (
+    composite_score,
+    recency_decay,
+    similarity_from_distance,
+)
+from .discovery import MEMORY_CACHE_TYPE, MEMORY_CAPABILITIES, MemoryDiscovery
+from .memory_store import MemoryStore
+from .telemetry import (
+    DEFAULT_METRICS_PREFIX,
+    DEFAULT_TRACER_NAME,
+    MemoryMetrics,
+    MemoryTelemetry,
+    MemoryTelemetryOptions,
+    create_memory_telemetry,
+)
+from .types import (
+    AgentMemoryConfig,
+    AgentMemoryRecallConfig,
+    ConsolidateResult,
+    EmbedFn,
+    MemoryConfigRefreshConfig,
+    MemoryConfigSnapshot,
+    MemoryDiscoveryConfig,
+    MemoryHit,
+    MemoryItem,
+    MemoryListOptions,
+    MemoryListResult,
+    MemoryScope,
+    MemoryStats,
+    MemoryStoreClient,
+    RecallWeights,
+    SummarizeFn,
+)
+
+__all__ = [
+    # Memory tier
+    "AgentMemory",
+    "AgentMemoryOptions",
+    "AgentMemoryConfig",
+    "AgentMemoryRecallConfig",
+    "MemoryStore",
+    "MemoryDiscovery",
+    "MEMORY_CACHE_TYPE",
+    "MEMORY_CAPABILITIES",
+    # Telemetry
+    "create_memory_telemetry",
+    "DEFAULT_METRICS_PREFIX",
+    "DEFAULT_TRACER_NAME",
+    "MemoryTelemetry",
+    "MemoryTelemetryOptions",
+    "MemoryMetrics",
+    # Scoring
+    "composite_score",
+    "similarity_from_distance",
+    "recency_decay",
+    # Types
+    "EmbedFn",
+    "MemoryStoreClient",
+    "MemoryScope",
+    "MemoryItem",
+    "MemoryHit",
+    "MemoryListOptions",
+    "MemoryListResult",
+    "MemoryStats",
+    "ConsolidateResult",
+    "SummarizeFn",
+    "RecallWeights",
+    "MemoryConfigSnapshot",
+    "MemoryDiscoveryConfig",
+    "MemoryConfigRefreshConfig",
+    "MATCH_ALL_MEMORY_QUERY",
+]
+
+# Surface everything agent-cache exports so consumers need only one import.
+__all__ += list(getattr(_agent_cache, "__all__", []))

betterdb_agent_memory/_num.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import math
+import re
+
+_FLOAT_RE = re.compile(r"^[+-]?(\d+\.?\d*|\.\d+)([eE][+-]?\d+)?")
+_INT_RE = re.compile(r"^[+-]?\d+")
+
+
+def _coerce_str(value: object) -> str:
+    """Coerce a valkey reply value to text the way JS ``String()`` would.
+
+    valkey-py can hand back ``bytes`` for raw replies, so decode those instead of
+    stringifying them (``str(b'1')`` is ``"b'1'"``, which would parse as NaN).
+    """
+    if isinstance(value, bytes):
+        return value.decode()
+    return str(value)
+
+
+def js_number(value: object) -> float:
+    """Mimic JavaScript ``Number(value)`` for the string inputs we see.
+
+    Empty/whitespace-only strings become ``0`` (as in JS); unparseable strings
+    become ``NaN``. ``None`` becomes ``NaN``.
+    """
+    if value is None:
+        return math.nan
+    text = _coerce_str(value).strip()
+    if text == "":
+        return 0.0
+    try:
+        return float(text)
+    except ValueError:
+        return math.nan
+
+
+def parse_float(value: object) -> float:
+    """Mimic JavaScript ``parseFloat``: parse the leading numeric portion, else NaN."""
+    if value is None:
+        return math.nan
+    match = _FLOAT_RE.match(_coerce_str(value).strip())
+    return float(match.group()) if match else math.nan
+
+
+def parse_int(value: object) -> float:
+    """Mimic JavaScript ``parseInt(value, 10)``: parse the leading integer, else NaN."""
+    if value is None:
+        return math.nan
+    match = _INT_RE.match(_coerce_str(value).strip())
+    return int(match.group()) if match else math.nan

betterdb_agent_memory/agent_memory.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+
+from betterdb_agent_cache import AgentCache, AgentCacheOptions
+
+from .memory_store import MemoryStore
+from .telemetry import MemoryTelemetryOptions
+from .types import AgentMemoryConfig, EmbedFn, _empty_memory_config
+
+DEFAULT_NAME = "betterdb_ac"
+
+
+@dataclass(kw_only=True)
+class AgentMemoryOptions(AgentCacheOptions):
+    """Options for the batteries-included :class:`AgentMemory` facade.
+
+    Extends :class:`AgentCacheOptions` (the three short-term cache tiers) with
+    the long-term memory tier: an ``embed_fn`` to vectorize content plus an
+    optional ``memory`` sub-config.
+    """
+
+    embed_fn: EmbedFn
+    memory: AgentMemoryConfig = field(default_factory=_empty_memory_config)
+
+
+class AgentMemory:
+    """Agent cache (llm/tool/session) plus a long-term :class:`MemoryStore` tier."""
+
+    def __init__(self, options: AgentMemoryOptions) -> None:
+        if not callable(getattr(options, "embed_fn", None)):
+            raise ValueError("AgentMemory requires an embed_fn to back the memory tier")
+
+        # The name lives on the shared options object and defaults identically in
+        # both tiers, so the cache and memory key prefixes can never drift apart.
+        name = options.name
+        self._cache = AgentCache(options)
+        self.llm = self._cache.llm
+        self.tool = self._cache.tool
+        self.session = self._cache.session
+
+        memory = options.memory
+        registry = options.telemetry.registry
+        self.memory = MemoryStore(
+            client=options.client,
+            name=name,
+            embed_fn=options.embed_fn,
+            default_threshold=memory.default_threshold,
+            weights=memory.recall.weights if memory.recall is not None else None,
+            half_life_seconds=(
+                memory.recall.half_life_seconds if memory.recall is not None else None
+            ),
+            max_items_per_scope=memory.max_items_per_scope,
+            # The facade is the batteries-included product: discover the memory
+            # tier alongside the cache tiers by default, unless explicitly disabled.
+            discovery=memory.discovery,
+            config_refresh=memory.config_refresh,
+            telemetry=MemoryTelemetryOptions(registry=registry) if registry else None,
+        )
+
+    async def initialize(self) -> None:
+        # Create the memory index before discovery so a freshly constructed facade
+        # is immediately usable for remember/recall without the caller hand-rolling
+        # the FT index. A create failure surfaces — the tier is unusable without it.
+        await self.memory.ensure_index()
+        # Surface a discovery name-collision from either tier instead of swallowing it.
+        await asyncio.gather(
+            self._cache.ensure_discovery_ready(),
+            self.memory.ensure_discovery_ready(),
+            self.memory.ensure_config_refresh_started(),
+        )
+
+    async def close(self) -> None:
+        # Tear down both tiers even if one fails, so timers and heartbeats can't leak.
+        try:
+            await self.memory.close()
+        finally:
+            await self._cache.shutdown()

betterdb_agent_memory/build_memory_index.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from .build_recall_query import VECTOR_FIELD
+
+MEMORY_INDEX_ALGORITHM = "HNSW"
+
+
+def memory_index_name(name: str) -> str:
+    return f"{name}:mem:idx"
+
+
+def memory_key_prefix(name: str) -> str:
+    return f"{name}:mem:"
+
+
+def build_memory_index_args(name: str, dims: int) -> list[str]:
+    if not isinstance(dims, int) or isinstance(dims, bool) or dims <= 0:
+        raise ValueError(f"memory index dimension must be a positive integer, got: {dims}")
+    return [
+        memory_index_name(name),
+        "ON",
+        "HASH",
+        "PREFIX",
+        "1",
+        memory_key_prefix(name),
+        "SCHEMA",
+        VECTOR_FIELD,
+        "VECTOR",
+        MEMORY_INDEX_ALGORITHM,
+        "6",
+        "TYPE",
+        "FLOAT32",
+        "DIM",
+        str(dims),
+        "DISTANCE_METRIC",
+        "COSINE",
+        "threadId",
+        "TAG",
+        "agentId",
+        "TAG",
+        "namespace",
+        "TAG",
+        "tags",
+        "TAG",
+        "SEPARATOR",
+        ",",
+        "source",
+        "TAG",
+        "importance",
+        "NUMERIC",
+        "created_at",
+        "NUMERIC",
+        "SORTABLE",
+        "last_accessed_at",
+        "NUMERIC",
+        "access_count",
+        "NUMERIC",
+        "content",
+        "TEXT",
+    ]

betterdb_agent_memory/build_memory_record.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+from betterdb_valkey_search_kit import encode_float32
+
+DEFAULT_IMPORTANCE = 0.5
+
+
+@dataclass
+class MemoryWrite:
+    key: str
+    fields: list[str | bytes]
+
+
+def build_memory_record(
+    name: str,
+    id: str,
+    content: str,
+    vector: list[float],
+    *,
+    importance: float | None = None,
+    tags: list[str] | None = None,
+    source: str | None = None,
+    thread_id: str | None = None,
+    agent_id: str | None = None,
+    namespace: str | None = None,
+    now: int,
+) -> MemoryWrite:
+    imp = importance if importance is not None else DEFAULT_IMPORTANCE
+    if not isinstance(imp, (int, float)) or not math.isfinite(imp) or imp < 0 or imp > 1:
+        raise ValueError(f"importance must be a finite number in [0, 1], got: {importance}")
+
+    fields: list[str | bytes] = [
+        "content",
+        content,
+        "vector",
+        encode_float32(vector),
+        "importance",
+        str(imp),
+        "created_at",
+        str(now),
+        "last_accessed_at",
+        str(now),
+        "access_count",
+        "0",
+    ]
+
+    tag_list = tags if tags is not None else []
+    for tag in tag_list:
+        if "," in tag:
+            raise ValueError(
+                f"Tag '{tag}' must not contain a comma; tags are stored comma-separated"
+            )
+    if len(tag_list) > 0:
+        fields.extend(["tags", ",".join(tag_list)])
+
+    if thread_id is not None:
+        fields.extend(["threadId", thread_id])
+    if agent_id is not None:
+        fields.extend(["agentId", agent_id])
+    if namespace is not None:
+        fields.extend(["namespace", namespace])
+    if source is not None:
+        fields.extend(["source", source])
+
+    return MemoryWrite(key=f"{name}:mem:{id}", fields=fields)

betterdb_agent_memory/build_recall_query.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from betterdb_valkey_search_kit import escape_tag
+
+from .types import MemoryScope
+
+SCORE_FIELD = "__score"
+VECTOR_FIELD = "vector"
+
+# valkey-search rejects a bare '*' on VECTOR-schema indexes with
+# "Invalid query string syntax". Every memory record has created_at
+# (set at write time), so this numeric range matches all documents
+# and is accepted by the vector index schema.
+MATCH_ALL_MEMORY_QUERY = "@created_at:[-inf +inf]"
+
+
+def _scope_clauses(scope: MemoryScope, tags: list[str]) -> list[str]:
+    clauses: list[str] = []
+    if scope.thread_id is not None:
+        clauses.append(f"@threadId:{{{escape_tag(scope.thread_id)}}}")
+    if scope.agent_id is not None:
+        clauses.append(f"@agentId:{{{escape_tag(scope.agent_id)}}}")
+    if scope.namespace is not None:
+        clauses.append(f"@namespace:{{{escape_tag(scope.namespace)}}}")
+    for tag in tags:
+        clauses.append(f"@tags:{{{escape_tag(tag)}}}")
+    return clauses
+
+
+def _join_clauses(clauses: list[str]) -> str:
+    if len(clauses) == 0:
+        return MATCH_ALL_MEMORY_QUERY
+    return f"({' '.join(clauses)})"
+
+
+def build_scope_filter(scope: MemoryScope, tags: list[str]) -> str:
+    return _join_clauses(_scope_clauses(scope, tags))
+
+
+@dataclass
+class ConsolidateFilterOptions:
+    max_created_at: int | None = None
+    max_importance: float | None = None
+    exclude_source: str | None = None
+
+
+def build_consolidate_filter(
+    scope: MemoryScope,
+    tags: list[str],
+    options: ConsolidateFilterOptions,
+) -> str:
+    clauses = _scope_clauses(scope, tags)
+    if options.max_created_at is not None:
+        clauses.append(f"@created_at:[-inf {options.max_created_at}]")
+    if options.max_importance is not None:
+        clauses.append(f"@importance:[-inf {options.max_importance}]")
+    if options.exclude_source is not None:
+        clauses.append(f"-@source:{{{escape_tag(options.exclude_source)}}}")
+    return _join_clauses(clauses)
+
+
+def build_recall_query(k: int, scope: MemoryScope, tags: list[str]) -> str:
+    return f"{build_scope_filter(scope, tags)}=>[KNN {k} @{VECTOR_FIELD} $vec AS {SCORE_FIELD}]"

betterdb_agent_memory/composite_score.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import math
+
+from .types import RecallWeights
+
+
+def recency_decay(age_seconds: float, half_life_seconds: float) -> float:
+    """True half-life decay: 1 at age 0, 0.5 at one half_life_seconds, approaching 0 beyond."""
+    return math.exp((-math.log(2) * age_seconds) / half_life_seconds)
+
+
+def composite_score(
+    *,
+    similarity: float,
+    age_seconds: float,
+    importance: float,
+    weights: RecallWeights,
+    half_life_seconds: float,
+) -> float:
+    """Weighted blend of semantic similarity, recency, and importance.
+
+    Recency is a true half-life decay: 0.5 at one half_life_seconds.
+    """
+    recency = recency_decay(age_seconds, half_life_seconds)
+    return (
+        weights.similarity * similarity
+        + weights.recency * recency
+        + weights.importance * importance
+    )
+
+
+def similarity_from_distance(distance: float) -> float:
+    """Map cosine distance (0..2, lower = closer) to a 0..1 similarity score."""
+    return 1 - distance / 2

betterdb_agent_memory/discovery.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import socket
+import warnings
+from datetime import datetime, timezone
+from typing import Any, Callable
+
+from betterdb_agent_cache.discovery import (
+    DEFAULT_HEARTBEAT_INTERVAL_S,
+    HEARTBEAT_KEY_PREFIX,
+    HEARTBEAT_TTL_SECONDS,
+    PROTOCOL_KEY,
+    PROTOCOL_VERSION,
+    REGISTRY_KEY,
+)
+
+from .types import MemoryStoreClient
+
+MEMORY_CACHE_TYPE = "agent_memory"
+MEMORY_CAPABILITIES = ["recall", "consolidate", "reinforce"]
+
+
+class MemoryDiscovery:
+    def __init__(
+        self,
+        *,
+        client: MemoryStoreClient,
+        name: str,
+        version: str,
+        stats_key: str,
+        heartbeat_interval_s: float | None = None,
+        on_write_failed: Callable[[], None] | None = None,
+    ) -> None:
+        self._client = client
+        self._name = name
+        self._version = version
+        self._stats_key = stats_key
+        self._heartbeat_interval_s = (
+            heartbeat_interval_s
+            if heartbeat_interval_s is not None
+            else DEFAULT_HEARTBEAT_INTERVAL_S
+        )
+        # Namespace the marker under `{name}:mem` so a memory store and an
+        # agent-cache sharing the same name register distinct registry fields
+        # and heartbeat keys instead of clobbering each other.
+        self._marker_field = f"{name}:mem"
+        self._heartbeat_key = f"{HEARTBEAT_KEY_PREFIX}{self._marker_field}"
+        self._started_at = datetime.now(timezone.utc).isoformat()
+        self._on_write_failed: Callable[[], None] = on_write_failed or (lambda: None)
+        self._heartbeat_task: asyncio.Task[None] | None = None
+
+    def build_marker(self) -> dict[str, Any]:
+        return {
+            "type": MEMORY_CACHE_TYPE,
+            "prefix": self._name,
+            "version": self._version,
+            "protocol_version": PROTOCOL_VERSION,
+            "capabilities": list(MEMORY_CAPABILITIES),
+            "stats_key": self._stats_key,
+            "started_at": self._started_at,
+            "pid": os.getpid(),
+            "hostname": socket.gethostname(),
+        }
+
+    async def register(self) -> None:
+        # HGET-then-HSET is not atomic (TOCTOU); acceptable for best-effort
+        # discovery — a racing writer just means last-writer-wins on the marker.
+        existing = await self._safe_hget()
+        if existing is not None:
+            self._check_collision(existing)
+        await self._write_marker()
+        await self._safe_call(
+            lambda: self._client.execute_command("SET", PROTOCOL_KEY, str(PROTOCOL_VERSION), "NX")
+        )
+        await self._write_heartbeat()
+        self._start_heartbeat()
+
+    async def stop(self, *, delete_heartbeat: bool) -> None:
+        if self._heartbeat_task is not None:
+            self._heartbeat_task.cancel()
+            try:
+                await self._heartbeat_task
+            except (asyncio.CancelledError, Exception):
+                pass
+            self._heartbeat_task = None
+        if not delete_heartbeat:
+            return
+        try:
+            await self._client.execute_command("DEL", self._heartbeat_key)
+        except Exception:
+            self._on_write_failed()
+
+    async def tick_heartbeat(self) -> None:
+        await self._write_heartbeat()
+        await self._write_marker()
+        # PROTOCOL_KEY is set once in register(); the NX SET is a guaranteed
+        # no-op on every subsequent tick, so it's not re-issued here.
+
+    def _start_heartbeat(self) -> None:
+        async def _loop() -> None:
+            try:
+                while True:
+                    await asyncio.sleep(self._heartbeat_interval_s)
+                    await self.tick_heartbeat()
+            except asyncio.CancelledError:
+                pass
+
+        self._heartbeat_task = asyncio.create_task(_loop())
+
+    async def _write_heartbeat(self) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        try:
+            await self._client.execute_command(
+                "SET", self._heartbeat_key, now, "EX", str(HEARTBEAT_TTL_SECONDS)
+            )
+        except Exception:
+            self._on_write_failed()
+
+    async def _write_marker(self) -> None:
+        try:
+            payload = json.dumps(self.build_marker())
+        except Exception:
+            self._on_write_failed()
+            return
+        await self._safe_call(
+            lambda: self._client.execute_command("HSET", REGISTRY_KEY, self._marker_field, payload)
+        )
+
+    async def _safe_hget(self) -> str | None:
+        try:
+            value = await self._client.execute_command("HGET", REGISTRY_KEY, self._marker_field)
+            if value is None:
+                return None
+            return value.decode() if isinstance(value, bytes) else str(value)
+        except Exception:
+            self._on_write_failed()
+            return None
+
+    async def _safe_call(self, fn: Callable[[], Any]) -> None:
+        try:
+            await fn()
+        except Exception:
+            self._on_write_failed()
+
+    def _check_collision(self, existing_json: str) -> None:
+        try:
+            parsed = json.loads(existing_json)
+        except Exception:
+            return
+        existing_type = parsed.get("type") if isinstance(parsed, dict) else None
+        if existing_type and existing_type != MEMORY_CACHE_TYPE:
+            # The memory marker lives under `{name}:mem`, distinct from
+            # agent-cache's `{name}`, so the two tiers never collide here.
+            # Surface it with a visible warning rather than raising into a
+            # swallowed registration; registration then proceeds
+            # last-writer-wins, matching agent-cache.
+            warnings.warn(
+                f"agent-memory discovery: field '{self._marker_field}' already holds a "
+                f"'{existing_type}' marker; overwriting",
+                stacklevel=2,
+            )