@bcmyguest/openmem 0.0.0-stage → 0.8.0

package/proxy/openmem_context/settings.py

@@ -0,0 +1,134 @@
+"""Runtime settings for the context proxy, read from environment variables."""
+
+import math
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+from .store import ASCII_WHITESPACE
+from .tokens import DEFAULT_CHARS_PER_TOKEN
+
+DEFAULT_TRIM_AT_TOKENS = 120000
+# Trim thresholds by model family (a substring of the lowercased model id), used while CONTEXT_TRIM_AT_TOKENS is
+# unset. Haiku's prompts over 100k tokens bill at a higher tier, so its sessions trim below it.
+MODEL_TRIM_AT_TOKENS = (("haiku", 95000),)
+# A trim keeps older in-turn exchanges while the request stays under this fraction of trim_at_tokens.
+DEFAULT_TRIM_TO_FRACTION = 0.3
+DEFAULT_COMPRESS_OVER_TOKENS = 4000
+DEFAULT_COMPACTION_MAX_CHARS = 12000
+DEFAULT_RECALL_TIMEOUT_S = 5.0
+DEFAULT_HINDSIGHT_URL = "http://127.0.0.1:8888"
+DEFAULT_TTL_DAYS = 30
+# OpenAI-compatible embeddings base URL (POST <url>/embeddings): Ollama on the host. compose.yaml sets the
+# container's view (host.docker.internal).
+DEFAULT_EMBEDDING_URL = "http://127.0.0.1:11434/v1"
+# Seconds per row of an embeddings call, and characters of a row's head sent (docs/context-store.md, "Vectors").
+DEFAULT_EMBEDDING_TIMEOUT_S = 30.0
+DEFAULT_EMBEDDING_MAX_CHARS = 4000
+
+
+def _int(env, name, default):
+    raw = str(env.get(name, "")).strip()
+    if not raw:
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        return default
+
+
+def _float(env, name, default):
+    raw = str(env.get(name, "")).strip()
+    if not raw:
+        return default
+    try:
+        return float(raw)
+    except ValueError:
+        return default
+
+
+def _positive_float(env, name, default):
+    value = _float(env, name, default)
+    return value if math.isfinite(value) and value > 0 else default
+
+
+def _fraction(env, name, default):
+    """A float strictly between 0 and 1, else default."""
+    value = _float(env, name, default)
+    return value if math.isfinite(value) and 0 < value < 1 else default
+
+
+def _positive_int(env, name, default):
+    value = _int(env, name, default)
+    return value if value > 0 else default
+
+
+def _flag(env, name):
+    """True unless the variable is 0, false, no or off."""
+    return str(env.get(name, "1")).strip().lower() not in ("0", "false", "no", "off")
+
+
+@dataclass(frozen=True)
+class Settings:
+    enabled: bool = True
+    # STANDING_INSTRUCTION in every request's system prompt (CONTEXT_STANDING_INSTRUCTION=0 turns it off).
+    standing_instruction: bool = True
+    trim_at_tokens: int = DEFAULT_TRIM_AT_TOKENS
+    # MODEL_TRIM_AT_TOKENS picks the threshold, trim_at_tokens being the fallback (CONTEXT_TRIM_AT_TOKENS unset).
+    trim_at_by_model: bool = False
+    # A trim lands under this fraction of trim_at_tokens, 0 < x < 1 (CONTEXT_TRIM_TO_FRACTION).
+    trim_to_fraction: float = DEFAULT_TRIM_TO_FRACTION
+    compress_over_tokens: int = DEFAULT_COMPRESS_OVER_TOKENS
+    # Characters per token for the trim estimate until the session reports usage (CONTEXT_CHARS_PER_TOKEN).
+    chars_per_token: float = DEFAULT_CHARS_PER_TOKEN
+    compaction_max_chars: int = DEFAULT_COMPACTION_MAX_CHARS
+    # Bounded provider recall in the trim pointer (CONTEXT_COMPACTION_RECALL=0 turns it off).
+    compaction_recall: bool = True
+    recall_timeout_s: float = DEFAULT_RECALL_TIMEOUT_S
+    hindsight_url: str = DEFAULT_HINDSIGHT_URL
+    context_dir: Path = Path.home() / ".local" / "share" / "openmem" / "context"
+    # Session files older than this are deleted at start, at most once a day; 0 keeps them (CONTEXT_STORE_TTL_DAYS).
+    ttl_days: int = DEFAULT_TTL_DAYS
+    # Embedding model for the session databases' vectors; empty turns vectors off (CONTEXT_EMBEDDING_MODEL).
+    embedding_model: str = ""
+    embedding_url: str = DEFAULT_EMBEDDING_URL
+    embedding_api_key: str = ""
+    # An embeddings call of n rows times out after n times this (CONTEXT_EMBEDDING_TIMEOUT_S).
+    embedding_timeout_s: float = DEFAULT_EMBEDDING_TIMEOUT_S
+    # Each row's first this many characters are embedded (CONTEXT_EMBEDDING_MAX_CHARS).
+    embedding_max_chars: int = DEFAULT_EMBEDDING_MAX_CHARS
+
+    @classmethod
+    def from_env(cls, env=None):
+        env = os.environ if env is None else env
+        state_dir = str(env.get("OPENMEM_STATE_DIR", "")).strip(ASCII_WHITESPACE)
+        root = Path(state_dir) if state_dir else Path.home() / ".local" / "share" / "openmem"
+        return cls(
+            enabled=_flag(env, "CONTEXT_PROXY_ENABLED"),
+            standing_instruction=_flag(env, "CONTEXT_STANDING_INSTRUCTION"),
+            trim_at_tokens=_int(env, "CONTEXT_TRIM_AT_TOKENS", DEFAULT_TRIM_AT_TOKENS),
+            trim_at_by_model=_int(env, "CONTEXT_TRIM_AT_TOKENS", None) is None,
+            trim_to_fraction=_fraction(env, "CONTEXT_TRIM_TO_FRACTION", DEFAULT_TRIM_TO_FRACTION),
+            compress_over_tokens=_int(env, "CONTEXT_COMPRESS_OVER_TOKENS", DEFAULT_COMPRESS_OVER_TOKENS),
+            chars_per_token=_positive_float(env, "CONTEXT_CHARS_PER_TOKEN", DEFAULT_CHARS_PER_TOKEN),
+            compaction_max_chars=_int(env, "CONTEXT_COMPACTION_MAX_CHARS", DEFAULT_COMPACTION_MAX_CHARS),
+            compaction_recall=_flag(env, "CONTEXT_COMPACTION_RECALL"),
+            recall_timeout_s=_float(env, "CONTEXT_RECALL_TIMEOUT_S", DEFAULT_RECALL_TIMEOUT_S),
+            hindsight_url=(str(env.get("HINDSIGHT_URL", "")).strip() or DEFAULT_HINDSIGHT_URL).rstrip("/"),
+            context_dir=root / "context",
+            ttl_days=_int(env, "CONTEXT_STORE_TTL_DAYS", DEFAULT_TTL_DAYS),
+            embedding_model=str(env.get("CONTEXT_EMBEDDING_MODEL", "")).strip(),
+            embedding_url=(str(env.get("CONTEXT_EMBEDDING_URL", "")).strip() or DEFAULT_EMBEDDING_URL).rstrip("/"),
+            embedding_api_key=str(env.get("CONTEXT_EMBEDDING_API_KEY", "")).strip(),
+            embedding_timeout_s=_positive_float(env, "CONTEXT_EMBEDDING_TIMEOUT_S", DEFAULT_EMBEDDING_TIMEOUT_S),
+            embedding_max_chars=_positive_int(env, "CONTEXT_EMBEDDING_MAX_CHARS", DEFAULT_EMBEDDING_MAX_CHARS),
+        )
+
+    def trim_at_for(self, model):
+        """The trim threshold for a request to model (its id, or None)."""
+        if self.trim_at_by_model and isinstance(model, str):
+            name = model.lower()
+            for family, tokens in MODEL_TRIM_AT_TOKENS:
+                if family in name:
+                    return tokens
+        return self.trim_at_tokens

package/proxy/openmem_context/store.py

@@ -0,0 +1,129 @@
+"""Context store files, per docs/context-store.md. Directories 700, files 600."""
+
+import json
+import os
+import re
+import threading
+from pathlib import Path
+
+DIR_MODE = 0o700
+FILE_MODE = 0o600
+TRIM_STATE_VERSION = 1
+_SID_RE = re.compile(r"[^A-Za-z0-9._-]")
+# ASCII whitespace only, so the hooks (JS) and the proxy trim the same characters.
+ASCII_WHITESPACE = " \t\n\v\f\r"
+_locks = {}
+_locks_guard = threading.Lock()
+
+
+class InvalidSessionId(ValueError):
+    """The session id is empty after trimming."""
+
+
+class RegistrationNotFound(LookupError):
+    """No readable registration for the session."""
+
+
+def sanitize_sid(sid):
+    """Contract sanitizer (docs/context-store.md): trim, replace characters outside [A-Za-z0-9._-], reject empty."""
+    clean = _SID_RE.sub("_", str(sid).strip(ASCII_WHITESPACE))
+    if not clean:
+        raise InvalidSessionId("empty session id")
+    return clean
+
+
+def session_lock(sid):
+    """In-process lock per session. The proxy runs one LiteLLM worker."""
+    with _locks_guard:
+        lock = _locks.get(sid)
+        if lock is None:
+            lock = _locks[sid] = threading.Lock()
+        return lock
+
+
+def _ensure_dir(path):
+    path.mkdir(mode=DIR_MODE, parents=True, exist_ok=True)
+    try:
+        os.chmod(path, DIR_MODE)
+    except OSError:
+        pass
+
+
+def _write_atomic(path, data):
+    tmp = path.with_name(f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp")
+    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, FILE_MODE)
+    try:
+        with os.fdopen(fd, "wb") as fh:
+            fh.write(data)
+        os.replace(tmp, path)
+    except BaseException:
+        try:
+            os.unlink(tmp)
+        except OSError:
+            pass
+        raise
+
+
+class Store:
+    def __init__(self, root):
+        self.root = Path(root)
+        self.sessions = self.root / "sessions"
+        self.blobs = self.root / "blobs"
+
+    def _session_path(self, sid, suffix):
+        return self.sessions / f"{sid}{suffix}"
+
+    # Registration (written by the SessionStart hook)
+    def registration(self, sid):
+        """Registration dict; raises RegistrationNotFound when it is missing, unreadable or not an object."""
+        try:
+            with open(self._session_path(sid, ".json"), "rb") as fh:
+                value = json.loads(fh.read().decode("utf-8"))
+        except (OSError, ValueError) as exc:
+            raise RegistrationNotFound(sid) from exc
+        if not isinstance(value, dict):
+            raise RegistrationNotFound(sid)
+        return value
+
+    def db_path(self, sid):
+        """Session database (sessiondb.py)."""
+        return self._session_path(sid, ".db")
+
+    def ensure_sessions(self):
+        _ensure_dir(self.sessions)
+
+    # Force flag (written by the prompt hook)
+    def force_requested(self, sid):
+        return self._session_path(sid, ".force").exists()
+
+    def clear_force(self, sid):
+        try:
+            os.unlink(self._session_path(sid, ".force"))
+        except FileNotFoundError:
+            pass
+
+    # Proxy-private trim state
+    def load_trim(self, sid):
+        try:
+            with open(self._session_path(sid, ".trim.json"), "rb") as fh:
+                value = json.loads(fh.read().decode("utf-8"))
+        except (OSError, ValueError):
+            value = None
+        if not isinstance(value, dict) or value.get("version") != TRIM_STATE_VERSION:
+            return {"version": TRIM_STATE_VERSION, "threads": {}}
+        if not isinstance(value.get("threads"), dict):
+            value["threads"] = {}
+        return value
+
+    def save_trim(self, sid, state):
+        _ensure_dir(self.sessions)
+        data = json.dumps(state, ensure_ascii=False, sort_keys=True).encode("utf-8")
+        _write_atomic(self._session_path(sid, ".trim.json"), data)
+
+    # Blobs
+    def write_blob(self, bid, text):
+        path = self.blobs / f"{bid}.txt"
+        if path.exists():
+            return
+        _ensure_dir(self.blobs)
+        _write_atomic(path, text.encode("utf-8"))

package/proxy/openmem_context/tokens.py

@@ -0,0 +1,123 @@
+"""Request size in characters (flat charge for binary attachments), the prompt-token estimate, and fingerprints.
+
+text_tokens (characters / 4) sizes single tool results for compression. The trim threshold uses
+estimate_prompt_tokens, calibrated per session from the provider's reported usage.
+"""
+
+import hashlib
+import json
+import math
+
+# Block types whose payload is binary (base64) rather than text.
+BINARY_BLOCK_TYPES = frozenset({"image", "document", "input_image", "input_file", "image_url", "file"})
+# Flat charge per binary block, in characters (about 1600 tokens).
+BINARY_BLOCK_CHARS = 6400
+# Characters per prompt token before a session has reported usage. Claude Code requests measured about 2.3.
+DEFAULT_CHARS_PER_TOKEN = 2.5
+# Bounds on an observed ratio, so one odd usage report cannot make the estimate absurd.
+MIN_CHARS_PER_TOKEN = 1.0
+MAX_CHARS_PER_TOKEN = 8.0
+# Usage samples kept per session (one per conversation thread).
+MAX_USAGE_THREADS = 32
+
+
+def text_tokens(text):
+    return (len(text) + 3) // 4
+
+
+def _chars(value):
+    if isinstance(value, str):
+        return len(value)
+    if isinstance(value, dict):
+        if isinstance(value.get("type"), str) and value["type"] in BINARY_BLOCK_TYPES:
+            return BINARY_BLOCK_CHARS
+        return sum(len(str(k)) + _chars(v) for k, v in value.items())
+    if isinstance(value, (list, tuple)):
+        return sum(_chars(v) for v in value)
+    if value is None:
+        return 0
+    return len(str(value))
+
+
+def request_chars(*values):
+    """Characters in the values, as the trim measures a request."""
+    return sum(_chars(v) for v in values)
+
+
+def _sample(value):
+    """(chars, tokens) of a stored usage sample, or None when it is missing or malformed."""
+    if not isinstance(value, dict):
+        return None
+    chars, tokens = value.get("chars"), value.get("tokens")
+    if type(chars) is int and type(tokens) is int and chars > 0 and tokens > 0:
+        return chars, tokens
+    return None
+
+
+def _ratio(sample, default):
+    if sample is None:
+        return default
+    return min(MAX_CHARS_PER_TOKEN, max(MIN_CHARS_PER_TOKEN, sample[0] / sample[1]))
+
+
+def _usage_threads(usage):
+    threads = usage.get("threads") if isinstance(usage, dict) else None
+    return threads if isinstance(threads, dict) else {}
+
+
+def estimate_prompt_tokens(chars, usage, root, default_ratio):
+    """Prompt tokens of a request of `chars` characters in conversation thread `root`.
+
+    usage is the session's calibration (record_usage). With a sample for this thread: its reported tokens plus
+    the characters added since, at that sample's chars-per-token ratio. Otherwise chars at the ratio of the
+    session's latest sample, else default_ratio.
+    """
+    own = _sample(_usage_threads(usage).get(root))
+    if own is not None:
+        return max(0, own[1] + math.ceil((chars - own[0]) / _ratio(own, default_ratio)))
+    last = _sample(usage.get("last")) if isinstance(usage, dict) else None
+    return math.ceil(chars / _ratio(last, default_ratio))
+
+
+def record_usage(state, root, chars, tokens):
+    """Store a reported prompt size (tokens) for a request of `chars` characters in thread `root` into state["usage"].
+
+    Returns False, leaving state alone, when either value is not a positive integer.
+    """
+    if _sample({"chars": chars, "tokens": tokens}) is None:
+        return False
+    usage = state.get("usage")
+    if not isinstance(usage, dict):
+        usage = state["usage"] = {}
+    threads = _usage_threads(usage)
+    sample = {"chars": chars, "tokens": tokens}
+    threads.pop(root, None)
+    threads[root] = sample
+    while len(threads) > MAX_USAGE_THREADS:
+        threads.pop(next(iter(threads)))
+    usage["threads"] = threads
+    usage["last"] = dict(sample)
+    return True
+
+
+def _normalize(value):
+    if isinstance(value, dict):
+        out = {k: _normalize(v) for k, v in value.items() if k != "cache_control"}
+        content = out.get("content")
+        if (isinstance(content, list) and len(content) == 1 and isinstance(content[0], dict)
+                and content[0].get("type") == "text" and set(content[0]) == {"type", "text"}):
+            out["content"] = content[0]["text"]
+        return out
+    if isinstance(value, list):
+        return [_normalize(v) for v in value]
+    return value
+
+
+def fingerprint(value):
+    """Stable hash of a message, ignoring key order, cache_control markers and single-text-block wrapping.
+
+    Claude Code marks the newest message with cache_control and sends its content as one text block; on later
+    requests the marker is gone and the same content is a plain string. Both shapes must hash the same.
+    """
+    blob = json.dumps(_normalize(value), sort_keys=True, ensure_ascii=False, separators=(",", ":"), default=str)
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:32]

package/proxy/openmem_context/ttl.py

@@ -0,0 +1,78 @@
+"""TTL cleanup of session files (docs/context-store.md, "Retention"). Same rules as sweepContextStore in
+cli/coding-agent/session-db.mjs."""
+
+import os
+import time
+from pathlib import Path
+
+DAY_S = 86400
+STAMP = ".ttl-stamp"
+# Longest first, so `x.trim.json` groups under `x`, not `x.trim`.
+SESSION_SUFFIXES = (".log.jsonl", ".trim.json", ".db-wal", ".db-shm", ".force", ".json", ".db")
+
+
+def session_of(name):
+    """(sid, suffix) for a session file name, or None for anything else."""
+    if name.startswith("."):
+        return None
+    for suffix in SESSION_SUFFIXES:
+        if name.endswith(suffix) and len(name) > len(suffix):
+            return name[: -len(suffix)], suffix
+    return None
+
+
+def sweep(root, ttl_days, now=None):
+    """Delete every session whose newest file is older than ttl_days; returns the deleted sids, sorted."""
+    if not ttl_days or ttl_days <= 0:
+        return []
+    now = time.time() if now is None else now
+    sessions = Path(root) / "sessions"
+    groups = {}
+    try:
+        entries = list(os.scandir(sessions))
+    except OSError:
+        return []
+    for entry in entries:
+        parsed = session_of(entry.name)
+        if parsed is None or not entry.is_file(follow_symlinks=False):
+            continue
+        try:
+            mtime = entry.stat(follow_symlinks=False).st_mtime
+        except OSError:
+            continue
+        group = groups.setdefault(parsed[0], [0.0, []])
+        group[0] = max(group[0], mtime)
+        group[1].append(entry.path)
+    cutoff = now - ttl_days * DAY_S
+    deleted = []
+    for sid, (newest, paths) in groups.items():
+        if newest >= cutoff:
+            continue
+        for path in paths:
+            try:
+                os.unlink(path)
+            except FileNotFoundError:
+                pass
+        deleted.append(sid)
+    return sorted(deleted)
+
+
+def maybe_sweep(root, ttl_days, now=None):
+    """sweep() at most once a day, recorded by the mtime of <root>/.ttl-stamp; returns the deleted sids."""
+    if not ttl_days or ttl_days <= 0:
+        return []
+    now = time.time() if now is None else now
+    stamp = Path(root) / STAMP
+    try:
+        if now - stamp.stat().st_mtime < DAY_S:
+            return []
+    except FileNotFoundError:
+        pass
+    except OSError:
+        return []
+    if not Path(root).is_dir():
+        return []
+    fd = os.open(stamp, os.O_WRONLY | os.O_CREAT, 0o600)
+    os.close(fd)
+    os.utime(stamp, (now, now))
+    return sweep(root, ttl_days, now)

package/proxy/openmem_context_hook.py

@@ -0,0 +1,34 @@
+"""LiteLLM callback: `litellm_settings.callbacks: openmem_context_hook.handler`.
+
+Thin wrapper; all logic lives in the dependency-free openmem_context package next to this file.
+"""
+
+import asyncio
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from litellm.integrations.custom_logger import CustomLogger  # noqa: E402
+
+from openmem_context import process_request, process_usage, startup  # noqa: E402
+
+
+class OpenMemContextHook(CustomLogger):
+    async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+        # File and Hindsight I/O run off the event loop; process_request never raises. A pass-through call (the
+        # Codex subscription route) is told apart by its route, and its call id lives on the logging object.
+        route = getattr(user_api_key_dict, "request_route", None)
+        call_id = getattr(data.get("litellm_logging_obj"), "litellm_call_id", None) if isinstance(data, dict) else None
+        return await asyncio.to_thread(process_request, data, call_type, route=route, call_id=call_id)
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        # Streams arrive here once, with the assembled response. kwargs carries the litellm_call_id that
+        # async_pre_call_hook saw in data; process_usage never raises.
+        await asyncio.to_thread(process_usage, kwargs, response_obj)
+
+
+# Session-file TTL cleanup, once per proxy start (and at most once a day); never blocks or raises.
+startup()
+
+handler = OpenMemContextHook()