mem01-engine 0.1.0__py3-none-any.whl

mem01/env.py

@@ -0,0 +1,52 @@
+"""Load environment variables from a local .env file.
+
+Where to put secrets:
+  mem01/.env          ← preferred (gitignored)
+  open-source/.env    ← also checked (parent)
+  process environment ← always wins if already set
+
+Copy the template:
+  cp .env.example .env
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+def load_env(*, override: bool = False) -> list[Path]:
+    """Load .env files if python-dotenv is installed. Returns paths loaded."""
+    try:
+        from dotenv import load_dotenv
+    except ImportError:
+        return []
+
+    here = Path(__file__).resolve()
+    candidates = [
+        here.parents[2] / ".env",  # mem01/.env (src/mem01/env.py → mem01/)
+        here.parents[3] / ".env",  # open-source/.env
+        Path.cwd() / ".env",
+    ]
+    loaded: list[Path] = []
+    seen: set[Path] = set()
+    for path in candidates:
+        path = path.resolve()
+        if path in seen or not path.is_file():
+            continue
+        seen.add(path)
+        load_dotenv(path, override=override)
+        loaded.append(path)
+    return loaded
+
+
+def require_openai_key() -> str:
+    load_env()
+    key = os.environ.get("OPENAI_API_KEY", "").strip()
+    if not key:
+        raise RuntimeError(
+            "OPENAI_API_KEY not set. Create mem01/.env from .env.example:\n"
+            "  cp .env.example .env\n"
+            "  # edit .env and paste your key\n"
+        )
+    return key

mem01/ids.py

@@ -0,0 +1,15 @@
+"""Stable identifiers for beliefs.
+
+Why a helper (not raw uuid in every call site):
+- One place to change id format later (ULID, prefixed ids, …)
+- Call sites stay readable: `new_belief_id()` instead of uuid boilerplate
+"""
+
+from __future__ import annotations
+
+import uuid
+
+
+def new_belief_id() -> str:
+    """Return a new unique belief id (UUID4 hex with bel_ prefix)."""
+    return f"bel_{uuid.uuid4().hex}"

mem01/llm/__init__.py

@@ -0,0 +1,10 @@
+"""LLM adapters for the write path only (extract_ops).
+
+Read path must never import this package for default recall.
+Any provider works if it implements LLMClient — OpenAI-shape, Claude, etc.
+"""
+
+from mem01.llm.base import ChatMessage, LLMClient
+from mem01.llm.fake import FakeLLM
+
+__all__ = ["ChatMessage", "FakeLLM", "LLMClient"]

mem01/llm/anthropic.py

@@ -0,0 +1,93 @@
+"""Optional Anthropic Claude client (native Messages API — not OpenAI-shaped).
+
+Requires ANTHROPIC_API_KEY or api_key=. Uses stdlib urllib; `anthropic` package optional.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import urllib.error
+import urllib.request
+
+from mem01.llm.base import ChatMessage
+
+
+class AnthropicLLM:
+    """Claude via https://api.anthropic.com/v1/messages."""
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        model: str = "claude-sonnet-4-20250514",
+        base_url: str = "https://api.anthropic.com",
+        max_tokens: int = 2048,
+        anthropic_version: str = "2023-06-01",
+    ) -> None:
+        self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY") or ""
+        self.model = model
+        self.base_url = base_url.rstrip("/")
+        self.max_tokens = max_tokens
+        self.anthropic_version = anthropic_version
+
+    def complete(
+        self,
+        messages: list[ChatMessage],
+        *,
+        temperature: float = 0.0,
+    ) -> str:
+        if not self.api_key:
+            raise RuntimeError(
+                "AnthropicLLM requires api_key or ANTHROPIC_API_KEY."
+            )
+
+        system_parts: list[str] = []
+        api_messages: list[dict] = []
+        for m in messages:
+            if m.role == "system":
+                system_parts.append(m.content)
+            else:
+                # Anthropic only allows user/assistant in messages list
+                role = m.role if m.role in ("user", "assistant") else "user"
+                api_messages.append({"role": role, "content": m.content})
+
+        if not api_messages:
+            raise ValueError("AnthropicLLM.complete needs at least one non-system message")
+
+        body: dict = {
+            "model": self.model,
+            "max_tokens": self.max_tokens,
+            "temperature": temperature,
+            "messages": api_messages,
+        }
+        if system_parts:
+            body["system"] = "\n\n".join(system_parts)
+
+        url = f"{self.base_url}/v1/messages"
+        data = json.dumps(body).encode("utf-8")
+        req = urllib.request.Request(
+            url,
+            data=data,
+            method="POST",
+            headers={
+                "Content-Type": "application/json",
+                "x-api-key": self.api_key,
+                "anthropic-version": self.anthropic_version,
+            },
+        )
+        try:
+            with urllib.request.urlopen(req, timeout=120) as resp:
+                payload = json.loads(resp.read().decode("utf-8"))
+        except urllib.error.HTTPError as e:
+            detail = e.read().decode("utf-8", errors="replace")
+            raise RuntimeError(f"Anthropic API error {e.code}: {detail}") from e
+
+        try:
+            blocks = payload["content"]
+            texts = [b["text"] for b in blocks if b.get("type") == "text"]
+            if not texts:
+                raise KeyError("no text blocks")
+            return "\n".join(texts)
+        except (KeyError, TypeError) as e:
+            raise RuntimeError(f"unexpected Anthropic response: {payload!r}") from e

mem01/llm/base.py

@@ -0,0 +1,30 @@
+"""LLMClient protocol — provider-agnostic chat completion.
+
+Why not hardcode OpenAI or Claude:
+- Extractor only needs text out for JSON ops
+- OpenAI-compatible, Anthropic Messages, LiteLLM, local models all fit
+  behind the same complete() method
+"""
+
+from __future__ import annotations
+
+from typing import Literal, Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+
+class ChatMessage(BaseModel):
+    role: Literal["system", "user", "assistant"]
+    content: str = Field(min_length=1)
+
+
+@runtime_checkable
+class LLMClient(Protocol):
+    def complete(
+        self,
+        messages: list[ChatMessage],
+        *,
+        temperature: float = 0.0,
+    ) -> str:
+        """Return assistant text (extractor will parse JSON from this)."""
+        ...

mem01/llm/fake.py

@@ -0,0 +1,35 @@
+"""Scripted LLM for tests — no API keys, no network."""
+
+from __future__ import annotations
+
+from mem01.llm.base import ChatMessage
+
+
+class FakeLLM:
+    """Returns a fixed response, or the next item from a queue of responses."""
+
+    def __init__(self, response: str | list[str] = "[]") -> None:
+        if isinstance(response, str):
+            self._responses = [response]
+        else:
+            self._responses = list(response)
+        if not self._responses:
+            self._responses = ["[]"]
+        self._index = 0
+        self.calls: list[list[ChatMessage]] = []
+
+    def complete(
+        self,
+        messages: list[ChatMessage],
+        *,
+        temperature: float = 0.0,
+    ) -> str:
+        self.calls.append(list(messages))
+        # Last scripted response repeats if we run out (stable for retries)
+        idx = min(self._index, len(self._responses) - 1)
+        self._index += 1
+        return self._responses[idx]
+
+    @property
+    def call_count(self) -> int:
+        return len(self.calls)

mem01/llm/openai_compat.py

@@ -0,0 +1,83 @@
+"""Optional OpenAI-compatible chat client (OpenAI, many proxies, some hosts).
+
+Not required for Claude — use anthropic.py or a router (e.g. LiteLLM) instead.
+Only imported when you construct OpenAICompatLLM; core mem01 does not need openai.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import urllib.error
+import urllib.request
+
+from mem01.llm.base import ChatMessage
+
+
+class OpenAICompatLLM:
+    """Minimal HTTP client for POST /v1/chat/completions.
+
+    Uses stdlib urllib so `openai` package is optional.
+    Set base_url to any OpenAI-compatible endpoint.
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        model: str = "gpt-5.6-sol",
+        base_url: str = "https://api.openai.com/v1",
+    ) -> None:
+        if api_key is None:
+            try:
+                from mem01.env import load_env
+
+                load_env()
+            except Exception:
+                pass
+            self.api_key = os.environ.get("OPENAI_API_KEY") or ""
+        else:
+            # Explicit "" means "no key" (tests); do not fall back to env
+            self.api_key = api_key
+        self.model = model
+        self.base_url = base_url.rstrip("/")
+
+    def complete(
+        self,
+        messages: list[ChatMessage],
+        *,
+        temperature: float = 0.0,
+    ) -> str:
+        if not self.api_key:
+            raise RuntimeError(
+                "OpenAICompatLLM requires api_key or OPENAI_API_KEY. "
+                "For Claude use AnthropicLLM (or another LLMClient adapter)."
+            )
+        url = f"{self.base_url}/chat/completions"
+        body = {
+            "model": self.model,
+            "messages": [{"role": m.role, "content": m.content} for m in messages],
+        }
+        if self.model != "gpt-5.6-sol":
+            body["temperature"] = temperature
+        data = json.dumps(body).encode("utf-8")
+        req = urllib.request.Request(
+            url,
+            data=data,
+            method="POST",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.api_key}",
+            },
+        )
+        try:
+            with urllib.request.urlopen(req, timeout=60) as resp:
+                payload = json.loads(resp.read().decode("utf-8"))
+        except urllib.error.HTTPError as e:
+            detail = e.read().decode("utf-8", errors="replace")
+            raise RuntimeError(f"OpenAI-compatible API error {e.code}: {detail}") from e
+
+        try:
+            return payload["choices"][0]["message"]["content"]
+        except (KeyError, IndexError, TypeError) as e:
+            raise RuntimeError(f"unexpected OpenAI-compatible response: {payload!r}") from e

mem01/metrics.py

@@ -0,0 +1,42 @@
+"""Lightweight in-process metrics for cost/latency gates.
+
+Why:
+- Product success bar needs tokens_used, latency, llm_calls — not accuracy alone
+- No external deps (Prometheus later); attach numbers to recall/remember results
+"""
+
+from __future__ import annotations
+
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Iterator
+
+
+@dataclass
+class TimerResult:
+    latency_ms: float
+
+
+@contextmanager
+def timer() -> Iterator[TimerResult]:
+    """Context manager that records wall-clock latency in milliseconds."""
+    result = TimerResult(latency_ms=0.0)
+    start = time.perf_counter()
+    try:
+        yield result
+    finally:
+        result.latency_ms = (time.perf_counter() - start) * 1000.0
+
+
+@dataclass
+class Counter:
+    """Simple named counters (llm_calls, etc.)."""
+
+    values: dict[str, float] = field(default_factory=dict)
+
+    def incr(self, name: str, amount: float = 1.0) -> None:
+        self.values[name] = self.values.get(name, 0.0) + amount
+
+    def get(self, name: str, default: float = 0.0) -> float:
+        return self.values.get(name, default)

mem01/read/__init__.py

@@ -0,0 +1,15 @@
+"""Read pipeline: search → conflict → rank → pack (no LLM on default path)."""
+
+from mem01.read.conflict import filter_conflicts
+from mem01.read.pack import pack_beliefs
+from mem01.read.rank import rank_candidates
+from mem01.read.recall import recall
+from mem01.read.search import search_beliefs
+
+__all__ = [
+    "filter_conflicts",
+    "pack_beliefs",
+    "rank_candidates",
+    "recall",
+    "search_beliefs",
+]

mem01/read/conflict.py

@@ -0,0 +1,85 @@
+"""Conflict-safe filtering on read candidates — no LLM.
+
+Why this exists (product wedge):
+- Even with SUPERSEDE on write, near-duplicates or slipped ADDs can both be active
+- Default recall must not inject two opposite truths when we can detect them
+- Deterministic rules only — keeps hot path cheap (mem0-class latency/$)
+
+v1 rules (mode=current):
+1. Drop non-active statuses
+2. Drop outside validity window (valid_from / valid_to)
+3. If multiple candidates share metadata topic_key, keep best (score, then confidence, then newer)
+
+mode=history (audit / temporal questions):
+- Keep active + superseded (+ optional invalidated)
+- Do NOT collapse topic_key to one winner — both SF (active) and NY (superseded) surface
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Literal
+
+from mem01.types import BeliefStatus, ScoredBelief, utc_now
+
+HistoryMode = Literal["current", "history"]
+
+_HISTORY_STATUSES = frozenset(
+    {
+        BeliefStatus.ACTIVE,
+        BeliefStatus.SUPERSEDED,
+        BeliefStatus.INVALIDATED,
+    }
+)
+
+
+def filter_conflicts(
+    candidates: list[ScoredBelief],
+    *,
+    at: datetime | None = None,
+    mode: HistoryMode = "current",
+) -> list[ScoredBelief]:
+    """Return a filtered subset of *candidates* (order preserved among survivors)."""
+    when = at or utc_now()
+    if mode == "history":
+        return [c for c in candidates if c.belief.status in _HISTORY_STATUSES]
+    current = [c for c in candidates if _is_eligible(c, when)]
+    return _dedupe_by_topic_key(current)
+
+
+def _is_eligible(scored: ScoredBelief, when: datetime) -> bool:
+    b = scored.belief
+    if b.status != BeliefStatus.ACTIVE:
+        return False
+    # Reuse Belief.is_current for validity window
+    return b.is_current(when)
+
+
+def _dedupe_by_topic_key(candidates: list[ScoredBelief]) -> list[ScoredBelief]:
+    """For each topic_key, keep a single winner; beliefs without topic_key all pass."""
+    best_by_topic: dict[str, ScoredBelief] = {}
+    no_topic: list[ScoredBelief] = []
+
+    for c in candidates:
+        key = c.belief.metadata.get("topic_key")
+        if not key or not isinstance(key, str):
+            no_topic.append(c)
+            continue
+        prev = best_by_topic.get(key)
+        if prev is None or _better(c, prev):
+            best_by_topic[key] = c
+
+    # Preserve original relative order: walk candidates, emit if selected
+    winners = set(id(v) for v in best_by_topic.values())
+    winners.update(id(c) for c in no_topic)
+    return [c for c in candidates if id(c) in winners]
+
+
+def _better(a: ScoredBelief, b: ScoredBelief) -> bool:
+    """True if *a* should replace *b* for the same topic_key."""
+    if a.score != b.score:
+        return a.score > b.score
+    if a.belief.confidence != b.belief.confidence:
+        return a.belief.confidence > b.belief.confidence
+    # Newer updated_at wins
+    return a.belief.updated_at >= b.belief.updated_at

mem01/read/pack.py

@@ -0,0 +1,73 @@
+"""Budgeted packing — fit memories under max_memory_tokens.
+
+Why first-class budgets:
+- Product constraint: same or better quality at mem0-class token cost
+- Unlimited top-k is how memory becomes expensive and slow for the agent LLM
+"""
+
+from __future__ import annotations
+
+from mem01.tokens import estimate_tokens
+from mem01.types import Belief, BeliefStatus, PackedMemory, ScoredBelief
+
+
+def format_belief_line(belief: Belief, *, show_status: bool = False) -> str:
+    """One line injected into the agent prompt for a belief.
+
+    History mode tags status so the agent can tell current vs past facts apart
+    (e.g. medical audit: superseded allergy still visible, clearly labeled).
+    """
+    if show_status or belief.status != BeliefStatus.ACTIVE:
+        return f"- [{belief.status.value}] {belief.content}"
+    return f"- {belief.content}"
+
+
+def pack_beliefs(
+    candidates: list[ScoredBelief],
+    *,
+    max_memory_tokens: int = 800,
+    candidate_count: int | None = None,
+    show_status: bool = False,
+) -> PackedMemory:
+    """Greedy pack highest-score candidates until the token budget is hit.
+
+    Assumes *candidates* are already conflict-filtered and preferably ranked
+    highest-first. If not ranked, we sort by score descending here.
+    """
+    if max_memory_tokens <= 0:
+        return PackedMemory(
+            beliefs=[],
+            text="",
+            tokens_used=0,
+            max_memory_tokens=max_memory_tokens,
+            candidate_count=candidate_count if candidate_count is not None else len(candidates),
+        )
+
+    ordered = sorted(candidates, key=lambda c: c.score, reverse=True)
+    chosen: list[Belief] = []
+    lines: list[str] = []
+    tokens_used = 0
+
+    for c in ordered:
+        line = format_belief_line(c.belief, show_status=show_status)
+        # +1 for newline between lines
+        extra = estimate_tokens(line if not lines else "\n" + line)
+        if tokens_used + extra > max_memory_tokens:
+            continue  # try smaller later items? v1 skip; greedy by score only
+        # Strict: if this one doesn't fit, skip it (don't stop entirely —
+        # a shorter lower-score line might still fit)
+        lines.append(line)
+        chosen.append(c.belief)
+        tokens_used += extra
+
+    text = "\n".join(lines)
+    # Recompute exact estimate on final text for honest metrics
+    tokens_used = estimate_tokens(text) if text else 0
+
+    return PackedMemory(
+        beliefs=chosen,
+        text=text,
+        tokens_used=tokens_used,
+        max_memory_tokens=max_memory_tokens,
+        candidate_count=candidate_count if candidate_count is not None else len(candidates),
+    )

mem01/read/rank.py

@@ -0,0 +1,100 @@
+"""Rank memory candidates for packing.
+
+Why not pure cosine alone:
+- Recency and confidence matter for long-lived agents (stale near-matches lose)
+- Single tunable score keeps pack() simple: sort desc, greedy fill
+"""
+
+from __future__ import annotations
+
+import math
+from datetime import datetime
+
+from mem01.types import ScoredBelief, utc_now
+
+
+def score_candidate(
+    scored: ScoredBelief,
+    *,
+    now: datetime | None = None,
+    similarity_weight: float = 0.6,
+    recency_weight: float = 0.25,
+    confidence_weight: float = 0.15,
+) -> float:
+    """Combine similarity, recency, and confidence into one rank score."""
+    when = now or utc_now()
+    sim = max(0.0, float(scored.score))  # cosine can be negative; floor for ranking
+    conf = max(0.0, min(1.0, scored.belief.confidence))
+
+    # Recency: half-life ~30 days (smooth decay, never zero)
+    age_seconds = max(0.0, (when - scored.belief.updated_at).total_seconds())
+    half_life = 30.0 * 24 * 3600
+    recency = math.exp(-math.log(2) * age_seconds / half_life)
+
+    return (
+        similarity_weight * sim
+        + recency_weight * recency
+        + confidence_weight * conf
+    )
+
+
+def rank_candidates(
+    candidates: list[ScoredBelief],
+    *,
+    now: datetime | None = None,
+) -> list[ScoredBelief]:
+    """Return candidates sorted by composite score (highest first).
+
+    Note: returns new ScoredBeliefs with .score set to the composite rank
+    so downstream pack can sort on score alone if needed. Original similarity
+    is replaced — call rank only after conflict filter.
+    """
+    when = now or utc_now()
+    ranked: list[ScoredBelief] = []
+    for c in candidates:
+        s = score_candidate(c, now=when)
+        ranked.append(ScoredBelief(belief=c.belief, score=s))
+    ranked.sort(key=lambda x: x.score, reverse=True)
+    return ranked
+
+
+def _jaccard(a: set[str], b: set[str]) -> float:
+    if not a or not b:
+        return 0.0
+    return len(a & b) / len(a | b)
+
+
+def mmr_select(
+    ranked: list[ScoredBelief],
+    *,
+    lambda_weight: float = 0.75,
+) -> list[ScoredBelief]:
+    """Greedy maximal-marginal-relevance reorder to reduce near-duplicate picks.
+
+    Multi-hop and listing questions need *different* beliefs packed together;
+    pure score ordering lets one topic cluster crowd out the rest of the
+    budget. Uses token Jaccard between belief contents so it works with any
+    store and never touches embeddings on the hot path. O(n^2), n <= recall k.
+    """
+    if len(ranked) <= 2:
+        return list(ranked)
+
+    from mem01.read.search import tokenize
+
+    tokens = {c.belief.id: set(tokenize(c.belief.content)) for c in ranked}
+    remaining = list(ranked)
+    selected: list[ScoredBelief] = [remaining.pop(0)]
+    while remaining:
+        best_idx = 0
+        best_val = float("-inf")
+        for i, cand in enumerate(remaining):
+            redundancy = max(
+                _jaccard(tokens[cand.belief.id], tokens[s.belief.id])
+                for s in selected
+            )
+            val = lambda_weight * cand.score - (1.0 - lambda_weight) * redundancy
+            if val > best_val:
+                best_val = val
+                best_idx = i
+        selected.append(remaining.pop(best_idx))
+    return selected