feat(p1): watcher, core service, hybrid-search skeleton, embed client, mdmem reindex

A file written directly into the vault is searchable in well under 2 s, the
service's own writes are hash-deduped so the watcher never double-indexes, and
reindex runs standalone with --no-embed for a lexical-only rebuild.
This commit is contained in:
Markus committed 2026-09-12 07:53:41 +00:00
1 parent 6a424c0909
commit 4aa3d5a86c
21 files changed
+1814 -76

No files matched your search

Binary file not shown.
+1 -1
View File
@@ -50,7 +50,7 @@ markers = ["live: hits real Bifrost/llama-server; excluded from the default run"
addopts = "-m 'not live' -q"
[tool.ruff]
line-length = 100
line-length = 110
src = ["src", "tests"]
[tool.ruff.lint]
+4 -8
View File
@@ -105,19 +105,15 @@ class Settings:
mcp_host, mcp_port = _addr("MCP_ADDR", "0.0.0.0:8080")
ui_host, ui_port = _addr("UI_ADDR", "0.0.0.0:8081")
origins = tuple(
o.strip()
for o in os.environ.get("UI_ALLOWED_ORIGINS", "").split(",")
if o.strip()
)
origins = tuple(o.strip() for o in os.environ.get("UI_ALLOWED_ORIGINS", "").split(",") if o.strip())
return cls(
vault_dir=vault,
index_db=Path(os.environ.get("INDEX_DB") or data / "index.db"),
cache_db=Path(os.environ.get("CACHE_DB") or data / "emb_cache.db"),
history_dir=Path(os.environ.get("HISTORY_DIR") or data / ".history"),
bifrost_url=(
os.environ.get("BIFROST_URL", "http://bifrost.apps.svc.cluster.local:8080")
).rstrip("/"),
bifrost_url=(os.environ.get("BIFROST_URL", "http://bifrost.apps.svc.cluster.local:8080")).rstrip(
"/"
),
bifrost_vk=os.environ.get("BIFROST_VK", ""),
embed_model=os.environ.get("EMBED_MODEL", "llamaswap/qwen3-embedding-4b"),
embed_dim=_int("EMBED_DIM", 2560),
+3 -9
View File
@@ -124,22 +124,16 @@ class Database:
def _ensure_vec_table(self) -> bool:
assert self._write is not None
cur = self._write
have = cur.execute(
"SELECT name FROM sqlite_master WHERE name='chunks_vec'"
).fetchone()
have = cur.execute("SELECT name FROM sqlite_master WHERE name='chunks_vec'").fetchone()
prev_dim = self._get_meta("embed_dim")
prev_model = self._get_meta("embed_model")
changed = have is not None and (
prev_dim != str(self.dim) or (prev_model or "") != (self.model or "")
)
changed = have is not None and (prev_dim != str(self.dim) or (prev_model or "") != (self.model or ""))
if have is not None and not changed:
return False
if changed:
cur.execute("DROP TABLE IF EXISTS chunks_vec")
cur.execute("UPDATE chunks SET vec_pending=1")
cur.execute(
f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(embedding float[{self.dim}])"
)
cur.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(embedding float[{self.dim}])")
cur.commit()
return bool(changed)
+221
View File
@@ -0,0 +1,221 @@
"""Embeddings via Bifrost, with a content-hash cache (§7).
Documents get no instruction prefix, queries do [R2]. Vectors are truncated to
EMBED_DIM client-side (MRL prefix truncation) and L2-normalised, so the vec0
L2 index is monotone in cosine.
"""
from __future__ import annotations
import asyncio
import hashlib
import logging
import sqlite3
import time
from abc import ABC, abstractmethod
from pathlib import Path
import httpx
import numpy as np
log = logging.getLogger("mdmem.embed")
QUERY_INSTRUCT = (
"Instruct: Given a question or topic from a user or agent, "
"retrieve personal memory notes that are relevant to it\nQuery:"
)
BATCH = 16
RETRIES = 3
def to_blob(vec: np.ndarray) -> bytes:
return np.asarray(vec, dtype=np.float32).tobytes()
def from_blob(blob: bytes) -> np.ndarray:
return np.frombuffer(blob, dtype=np.float32)
def normalize(v: np.ndarray) -> np.ndarray:
norms = np.linalg.norm(v, axis=1, keepdims=True)
norms[norms == 0] = 1.0
return v / norms
class EmbeddingError(RuntimeError):
pass
class EmbeddingCache:
"""content-hash → vector. Survives index rebuilds, so a rebuild is free (§13 P2)."""
def __init__(self, path: Path | str) -> None:
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.conn = sqlite3.connect(str(self.path), check_same_thread=False)
self.conn.execute("PRAGMA journal_mode=WAL")
self.conn.execute(
"CREATE TABLE IF NOT EXISTS emb_cache("
"key TEXT PRIMARY KEY, model TEXT, dim INTEGER, vec BLOB, created TEXT)"
)
self.conn.commit()
self.hits = 0
self.misses = 0
@staticmethod
def key(model: str, dim: int, text: str) -> str:
return hashlib.sha256(f"{model}\n{dim}\n{text}".encode()).hexdigest()
def get(self, model: str, dim: int, text: str) -> np.ndarray | None:
row = self.conn.execute(
"SELECT vec FROM emb_cache WHERE key=?", (self.key(model, dim, text),)
).fetchone()
if row is None:
self.misses += 1
return None
self.hits += 1
return from_blob(row[0])
def put(self, model: str, dim: int, text: str, vec: np.ndarray) -> None:
self.conn.execute(
"INSERT OR REPLACE INTO emb_cache(key,model,dim,vec,created) VALUES(?,?,?,?,datetime('now'))",
(self.key(model, dim, text), model, dim, to_blob(vec)),
)
self.conn.commit()
def close(self) -> None:
self.conn.close()
class Embedder(ABC):
model: str
dim: int
@abstractmethod
async def embed(self, texts: list[str], *, is_query: bool = False) -> np.ndarray: ...
async def selftest(self) -> dict[str, object]:
started = time.monotonic()
v = await self.embed(["ping"], is_query=False)
return {
"model": self.model,
"dim": int(v.shape[1]),
"ms": round((time.monotonic() - started) * 1000, 1),
}
async def aclose(self) -> None: # pragma: no cover - overridden where needed
return None
class BifrostEmbedder(Embedder):
def __init__(
self,
base_url: str,
api_key: str,
model: str,
dim: int,
cache: EmbeddingCache | None = None,
timeout: float = 60.0,
) -> None:
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.dim = dim
self.cache = cache
self._client = httpx.AsyncClient(timeout=timeout)
self.last_latency_ms: float | None = None
self.errors = 0
async def aclose(self) -> None:
await self._client.aclose()
async def embed(self, texts: list[str], *, is_query: bool = False) -> np.ndarray:
if not texts:
return np.zeros((0, self.dim), dtype=np.float32)
prepared = [QUERY_INSTRUCT + t for t in texts] if is_query else list(texts)
out: list[np.ndarray | None] = [None] * len(prepared)
todo: list[int] = []
for i, t in enumerate(prepared):
cached = self.cache.get(self.model, self.dim, t) if self.cache else None
if cached is not None:
out[i] = cached
else:
todo.append(i)
for start in range(0, len(todo), BATCH):
idxs = todo[start : start + BATCH]
vecs = await self._call([prepared[i] for i in idxs])
for i, v in zip(idxs, vecs, strict=True):
out[i] = v
if self.cache:
self.cache.put(self.model, self.dim, prepared[i], v)
return np.vstack([v for v in out if v is not None]).astype(np.float32)
async def _call(self, batch: list[str]) -> np.ndarray:
delay = 0.5
last: Exception | None = None
for attempt in range(1, RETRIES + 1):
started = time.monotonic()
try:
r = await self._client.post(
f"{self.base_url}/v1/embeddings",
headers={
"Authorization": f"Bearer {self.api_key}",
"content-type": "application/json",
},
json={"model": self.model, "input": batch},
)
if r.status_code >= 400:
raise EmbeddingError(
f"embeddings HTTP {r.status_code} for model {self.model!r}: {r.text[:300]}"
)
data = r.json()["data"]
v = np.asarray([d["embedding"] for d in data], dtype=np.float32)
if v.ndim != 2 or v.shape[0] != len(batch):
raise EmbeddingError(f"unexpected embeddings shape {v.shape}")
if v.shape[1] < self.dim:
raise EmbeddingError(
f"model {self.model!r} returned dim {v.shape[1]} < EMBED_DIM {self.dim} "
"— check EMBED_MODEL carries the Bifrost provider prefix (§7, §14.14)"
)
self.last_latency_ms = round((time.monotonic() - started) * 1000, 1)
return normalize(v[:, : self.dim])
except Exception as exc: # noqa: BLE001 - retried, then re-raised
last = exc
self.errors += 1
log.warning(
"embedding call failed",
extra={"fields": {"attempt": attempt, "error": str(exc)[:200]}},
)
if attempt < RETRIES:
await asyncio.sleep(delay)
delay *= 2
raise EmbeddingError(str(last))
class FakeEmbedder(Embedder):
"""Deterministic seeded vectors — tests run with no network and no GPU (§12a.1)."""
def __init__(self, dim: int = 16, model: str = "fake") -> None:
self.dim = dim
self.model = model
self.calls = 0
async def embed(self, texts: list[str], *, is_query: bool = False) -> np.ndarray:
self.calls += 1
rows = []
for t in texts:
base = QUERY_INSTRUCT + t if is_query else t
# Bag-of-words hashing: texts sharing words end up close, which makes
# ranking tests meaningful without a model.
v = np.zeros(self.dim, dtype=np.float32)
words = [w for w in base.lower().split() if w] or ["∅"]
for w in words:
h = hashlib.sha256(w.encode()).digest()
idx = int.from_bytes(h[:4], "big") % self.dim
sign = 1.0 if h[4] % 2 == 0 else -1.0
v[idx] += sign
rows.append(v)
return normalize(np.vstack(rows).astype(np.float32))
+50
View File
@@ -0,0 +1,50 @@
"""Structured errors shared by the MCP tools and the UI."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
@dataclass
class MdmemError(Exception):
code: str
message: str
details: dict[str, Any] | None = None
def __str__(self) -> str: # pragma: no cover - trivial
return f"{self.code}: {self.message}"
def as_dict(self) -> dict[str, Any]:
d: dict[str, Any] = {"error": self.code, "message": self.message}
if self.details:
d["details"] = self.details
return d
class NotFound(MdmemError):
def __init__(self, note_id: str) -> None:
super().__init__("NOT_FOUND", f"no note with id {note_id}", {"id": note_id})
class VersionConflict(MdmemError):
def __init__(self, note_id: str, expected: str, actual: str) -> None:
super().__init__(
"VERSION_CONFLICT",
"the note changed since you read it; re-read it and retry",
{"id": note_id, "expected_version": expected, "current_version": actual},
)
class Invalid(MdmemError):
def __init__(self, message: str, **details: Any) -> None:
super().__init__("INVALID", message, details or None)
class SecretRejected(MdmemError):
def __init__(self, kinds: list[str]) -> None:
super().__init__(
"SECRET_REJECTED",
"the note looks like it contains a credential; memory must not store secrets",
{"patterns": kinds},
)
+13 -2
View File
@@ -13,8 +13,19 @@ from ulid import ULID
_ULID_RE = re.compile(r"^[0-9A-HJKMNP-TV-Z]{26}$")
_TRANSLIT = str.maketrans(
{
"ä": "ae", "ö": "oe", "ü": "ue", "Ä": "ae", "Ö": "oe", "Ü": "ue",
"ß": "ss", "å": "aa", "æ": "ae", "ø": "oe", "đ": "d", "ł": "l", "þ": "th",
"ä": "ae",
"ö": "oe",
"ü": "ue",
"Ä": "ae",
"Ö": "oe",
"Ü": "ue",
"ß": "ss",
"å": "aa",
"æ": "ae",
"ø": "oe",
"đ": "d",
"ł": "l",
"þ": "th",
}
)
MAX_SLUG = 60
+34 -21
View File
@@ -65,7 +65,7 @@ class Indexer:
(note.id,),
)
}
prev_vecs = self._cached_vectors(w, old) # before the rows go away
prev_vecs = self._cached_vectors(w, old) # before the rows go away
self._delete_chunk_rows(w, list(old))
w.execute(
"""INSERT INTO notes(id,path,state,title,type,summary,tags,keywords,
@@ -81,13 +81,23 @@ class Indexer:
created=excluded.created, updated=excluded.updated,
file_hash=excluded.file_hash, mtime=excluded.mtime""",
(
note.id, rel_path, note.state, note.title, note.type, note.summary,
note.id,
rel_path,
note.state,
note.title,
note.type,
note.summary,
json.dumps(note.tags, ensure_ascii=False),
json.dumps(note.keywords, ensure_ascii=False),
note.superseded_by, note.revises,
note.ts("observed_at"), note.ts("valid_from"), note.ts("valid_until"),
note.ts("created"), note.ts("updated"),
note.file_hash or sha256_text(note.to_text()), note.mtime,
note.superseded_by,
note.revises,
note.ts("observed_at"),
note.ts("valid_from"),
note.ts("valid_until"),
note.ts("created"),
note.ts("updated"),
note.file_hash or sha256_text(note.to_text()),
note.mtime,
),
)
@@ -110,9 +120,7 @@ class Indexer:
)
vec = prev_vecs.get(h)
if vec is not None:
w.execute(
"INSERT INTO chunks_vec(rowid,embedding) VALUES(?,?)", (cid, vec)
)
w.execute("INSERT INTO chunks_vec(rowid,embedding) VALUES(?,?)", (cid, vec))
w.execute("UPDATE chunks SET vec_pending=0 WHERE id=?", (cid,))
else:
pending.append(PendingEmbed(chunk_id=cid, text=text, key_hash=h))
@@ -168,9 +176,7 @@ class Indexer:
if key in seen:
continue
seen.add(key)
w.execute(
"INSERT OR REPLACE INTO links(src,dst,rel,origin) VALUES(?,?,?,?)", r
)
w.execute("INSERT OR REPLACE INTO links(src,dst,rel,origin) VALUES(?,?,?,?)", r)
@staticmethod
def _resolve(w, target: str) -> str:
@@ -184,9 +190,7 @@ class Indexer:
if row:
return row[0]
slug = slugify(target)
row = w.execute(
"SELECT id, title FROM notes WHERE path LIKE ? LIMIT 1", (f"%/{slug}--%",)
).fetchone()
row = w.execute("SELECT id, title FROM notes WHERE path LIKE ? LIMIT 1", (f"%/{slug}--%",)).fetchone()
if row:
return row[0]
return f"{UNRESOLVED}{target}"
@@ -215,7 +219,10 @@ class Indexer:
note = self._read_and_fix(path)
except Exception as exc: # a broken file must not abort the scan
stats.errors.append(f"{path}: {exc}")
log.warning("skipping unreadable note", extra={"fields": {"path": str(path), "error": str(exc)}})
log.warning(
"skipping unreadable note",
extra={"fields": {"path": str(path), "error": str(exc)}},
)
continue
seen.add(note.id)
prev = known.get(note.id)
@@ -246,18 +253,24 @@ class Indexer:
if new_path != path:
# the file had no ULID in its name, so write() could not find the original
path.unlink(missing_ok=True)
log.info("assigned id to hand-written note",
extra={"fields": {"id": note.id, "path": str(new_path)}})
log.info(
"assigned id to hand-written note",
extra={"fields": {"id": note.id, "path": str(new_path)}},
)
note = self.vault.read(new_path)
return note
def pending_embeds(self) -> list[PendingEmbed]:
rows = self.db.read().execute(
"""SELECT c.id, c.heading, c.body, c.embed_text_hash,
rows = (
self.db.read()
.execute(
"""SELECT c.id, c.heading, c.body, c.embed_text_hash,
n.title, n.summary, n.keywords
FROM chunks c JOIN notes n ON n.id = c.note_id
WHERE c.vec_pending = 1"""
).fetchall()
)
.fetchall()
)
out: list[PendingEmbed] = []
for r in rows:
kw = ", ".join(json.loads(r["keywords"] or "[]"))
+19 -7
View File
@@ -24,9 +24,23 @@ _WIKILINK_RE = re.compile(r"\[\[([^\[\]|#]+?)(?:#[^\[\]|]*)?(?:\|[^\[\]]*)?\]\]"
_HEADING_RE = re.compile(r"^(#{1,6})\s+(.*?)\s*#*\s*$", re.MULTILINE)
CANONICAL_ORDER = [
"id", "title", "type", "summary", "keywords", "tags", "relations",
"supersedes", "superseded_by", "revises", "possible_duplicates",
"observed_at", "valid_from", "valid_until", "source", "created", "updated",
"id",
"title",
"type",
"summary",
"keywords",
"tags",
"relations",
"supersedes",
"superseded_by",
"revises",
"possible_duplicates",
"observed_at",
"valid_from",
"valid_until",
"source",
"created",
"updated",
]
@@ -75,7 +89,7 @@ def _as_list(v: Any) -> list[str]:
return []
if isinstance(v, str):
return [v] if v.strip() else []
if isinstance(v, (list, tuple)):
if isinstance(v, list | tuple):
return [str(x) for x in v if x is not None and str(x).strip()]
return []
@@ -92,9 +106,7 @@ class Note:
# -- construction ------------------------------------------------------
@classmethod
def from_text(
cls, text: str, state: str = "inbox", path: Path | None = None, **kw: Any
) -> Note:
def from_text(cls, text: str, state: str = "inbox", path: Path | None = None, **kw: Any) -> Note:
fm, body = parse(text)
return cls(fm=fm, body=body, state=state, path=path, **kw)
+330
View File
@@ -0,0 +1,330 @@
"""Hybrid retrieval: FTS5 + sqlite-vec, fused (§8).
No graph expansion: retrieval is lexical + dense only (§2.7). Links are a
filter and a UI affordance, never a scoring channel.
"""
from __future__ import annotations
import json
import math
from dataclasses import dataclass, field
from typing import Any
from mdmem.config import Settings
from mdmem.db import Database
from mdmem.ftsquery import to_match, to_trigram_match
LEX_N = 50
SEM_N = 100
RRF_K = 60
SNIPPET_CHARS = 800
@dataclass
class SearchResult:
id: str
title: str
type: str
state: str
verified: bool
score: float
snippet: str
tags: list[str]
updated: str | None
path: str
debug: dict[str, Any] | None = None
def as_dict(self) -> dict[str, Any]:
d = {
"id": self.id,
"title": self.title,
"type": self.type,
"state": self.state,
"verified": self.verified,
"score": round(self.score, 6),
"snippet": self.snippet,
"tags": self.tags,
"updated": self.updated,
"path": self.path,
}
if self.debug is not None:
d["debug"] = self.debug
return d
@dataclass
class SearchResponse:
results: list[SearchResult] = field(default_factory=list)
truncated: bool = False
chars: int = 0
fusion: str = "rrf"
def as_dict(self) -> dict[str, Any]:
return {
"results": [r.as_dict() for r in self.results],
"truncated": self.truncated,
"chars": self.chars,
"fusion": self.fusion,
}
def _cosine_from_l2(distance: float) -> float:
"""Vectors are L2-normalised, so ||a-b||² = 2 - 2cos."""
return max(-1.0, min(1.0, 1.0 - (distance * distance) / 2.0))
def _snippet(body: str, limit: int = SNIPPET_CHARS) -> str:
body = body.strip()
if len(body) <= limit:
return body
return body[: limit - 1].rstrip() + "…"
def search(
db: Database,
settings: Settings,
query: str,
*,
query_vec: bytes | None = None,
k: int = 8,
types: list[str] | None = None,
tags: list[str] | None = None,
since: str | None = None,
until: str | None = None,
valid_at: str | None = None,
include_inbox: bool | None = None,
include_superseded: bool = False,
include_archive: bool = False,
max_chars: int | None = None,
min_score: float | None = None,
fusion: str | None = None,
cc_alpha: float | None = None,
inbox_weight: float | None = None,
debug: bool = False,
) -> SearchResponse:
conn = db.read()
include_inbox = settings.include_inbox_default if include_inbox is None else include_inbox
max_chars = settings.max_chars_default if max_chars is None else max_chars
fusion = (fusion or settings.fusion).lower()
cc_alpha = settings.cc_alpha if cc_alpha is None else cc_alpha
inbox_weight = settings.inbox_weight if inbox_weight is None else inbox_weight
min_score = settings.min_score if min_score is None else min_score
# 1. lexical -----------------------------------------------------------
lex: dict[int, float] = {}
match = to_match(query)
if match:
w = settings.bm25_weights
rows = conn.execute(
f"SELECT rowid AS cid, bm25(chunks_fts, {w[0]}, {w[1]}, {w[2]}, {w[3]}) AS s "
"FROM chunks_fts WHERE chunks_fts MATCH ? ORDER BY s LIMIT ?",
(match, LEX_N),
).fetchall()
for r in rows:
lex[r["cid"]] = -float(r["s"]) # bm25() is negative; larger is better here
if settings.fts_trigram:
tri = to_trigram_match(query)
if tri and _has_table(conn, "chunks_tri"):
rows = conn.execute(
"SELECT rowid AS cid, bm25(chunks_tri) AS s FROM chunks_tri "
"WHERE chunks_tri MATCH ? ORDER BY s LIMIT ?",
(tri, LEX_N),
).fetchall()
for r in rows:
lex.setdefault(r["cid"], -float(r["s"]) * 0.5)
# 2. dense -------------------------------------------------------------
sem: dict[int, float] = {}
if query_vec is not None:
rows = conn.execute(
"SELECT rowid AS cid, distance FROM chunks_vec "
"WHERE embedding MATCH ? AND k = ? ORDER BY distance",
(query_vec, SEM_N),
).fetchall()
for r in rows:
sem[r["cid"]] = _cosine_from_l2(float(r["distance"]))
cids = set(lex) | set(sem)
if not cids:
return SearchResponse(fusion=fusion)
# 3. filter ------------------------------------------------------------
placeholders = ",".join("?" * len(cids))
rows = conn.execute(
f"""SELECT c.id AS cid, c.body, c.heading, n.id, n.title, n.type, n.state, n.tags,
n.updated, n.path, n.summary, n.superseded_by, n.valid_from, n.valid_until
FROM chunks c JOIN notes n ON n.id = c.note_id
WHERE c.id IN ({placeholders})""",
tuple(cids),
).fetchall()
type_set = {t.lower() for t in types} if types else None
tag_set = {t.lower() for t in tags} if tags else None
kept: dict[str, dict[str, Any]] = {}
for r in rows:
state = r["state"]
if state == "archive" and not include_archive:
continue
if state == "inbox" and not include_inbox:
continue
if r["superseded_by"] and not include_superseded:
continue
if type_set and (r["type"] or "").lower() not in type_set:
continue
note_tags = json.loads(r["tags"] or "[]")
if tag_set and not tag_set.issubset({t.lower() for t in note_tags}):
continue
upd = r["updated"] or ""
if since and upd and upd < since:
continue
if until and upd and upd > until:
continue
if valid_at and not _valid_at(r["valid_from"], r["valid_until"], valid_at):
continue
kept.setdefault(
r["id"],
{
"row": r,
"chunks": [],
"tags": note_tags,
},
)["chunks"].append(r)
if not kept:
return SearchResponse(fusion=fusion)
# 4. fuse at chunk level ------------------------------------------------
alive = {r["cid"] for v in kept.values() for r in v["chunks"]}
scores = _fuse(
{c: s for c, s in lex.items() if c in alive},
{c: s for c, s in sem.items() if c in alive},
fusion=fusion,
alpha=cc_alpha,
)
lex_rank = _ranks(lex)
sem_rank = _ranks(sem)
# 5. aggregate to notes -------------------------------------------------
out: list[SearchResult] = []
for note_id, data in kept.items():
best_cid, best = max(
((r["cid"], scores.get(r["cid"], 0.0)) for r in data["chunks"]),
key=lambda t: t[1],
)
row = next(r for r in data["chunks"] if r["cid"] == best_cid)
state = row["state"]
verified = state == "active"
score = best * (inbox_weight if state == "inbox" else 1.0)
dbg = None
if debug:
dbg = {
"chunk_id": best_cid,
"lex_rank": lex_rank.get(best_cid),
"sem_rank": sem_rank.get(best_cid),
"bm25": round(-lex[best_cid], 4) if best_cid in lex else None,
"cos": round(sem[best_cid], 4) if best_cid in sem else None,
"fused": round(best, 6),
"via": "both"
if best_cid in lex and best_cid in sem
else ("lex" if best_cid in lex else "sem"),
"heading": row["heading"],
}
out.append(
SearchResult(
id=note_id,
title=row["title"] or "",
type=row["type"] or "note",
state=state,
verified=verified,
score=score,
snippet=_snippet(row["body"] or row["summary"] or ""),
tags=data["tags"],
updated=row["updated"],
path=row["path"] or "",
debug=dbg,
)
)
out.sort(key=lambda r: r.score, reverse=True)
# 8. abstention ---------------------------------------------------------
if min_score is not None and (not out or out[0].score < min_score):
return SearchResponse(fusion=fusion)
# 7. budget -------------------------------------------------------------
picked: list[SearchResult] = []
chars = 0
truncated = False
for r in out[:k]:
cost = len(r.snippet) + len(r.title)
if picked and chars + cost > max_chars:
truncated = True
break
picked.append(r)
chars += cost
if len(out) > k:
truncated = True
return SearchResponse(results=picked, truncated=truncated, chars=chars, fusion=fusion)
def _has_table(conn, name: str) -> bool:
return conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (name,)).fetchone() is not None
def _valid_at(valid_from: str | None, valid_until: str | None, t: str) -> bool:
if valid_from and t < valid_from:
return False
if valid_until and t >= valid_until:
return False
return True
def _ranks(scores: dict[int, float]) -> dict[int, int]:
return {
cid: i + 1 for i, (cid, _) in enumerate(sorted(scores.items(), key=lambda kv: kv[1], reverse=True))
}
def _fuse(
lex: dict[int, float],
sem: dict[int, float],
*,
fusion: str,
alpha: float,
) -> dict[int, float]:
if fusion == "cc":
return _cc(lex, sem, alpha)
return _rrf(lex, sem)
def _rrf(lex: dict[int, float], sem: dict[int, float]) -> dict[int, float]:
"""Reciprocal rank fusion, k=60 [R18]."""
out: dict[int, float] = {}
for scores in (lex, sem):
for rank, (cid, _) in enumerate(sorted(scores.items(), key=lambda kv: kv[1], reverse=True), start=1):
out[cid] = out.get(cid, 0.0) + 1.0 / (RRF_K + rank)
return out
def _cc(lex: dict[int, float], sem: dict[int, float], alpha: float) -> dict[int, float]:
"""Convex combination with theoretical min-max normalisation (TM2C2) [R17].
Lexical: s = -bm25 has a theoretical minimum of 0, so s / max(s).
Dense: cosine has a theoretical minimum of -1, so (cos+1) / (max+1).
"""
lex_max = max(lex.values(), default=0.0)
sem_max = max(sem.values(), default=-1.0)
out: dict[int, float] = {}
for cid in set(lex) | set(sem):
ln = (lex.get(cid, 0.0) / lex_max) if lex_max > 0 else 0.0
cos = sem.get(cid)
sn = ((cos + 1.0) / (sem_max + 1.0)) if (cos is not None and sem_max > -1.0) else 0.0
out[cid] = alpha * sn + (1.0 - alpha) * ln
return out
def normalize_scores(values: list[float]) -> list[float]:
total = math.fsum(values)
return [v / total for v in values] if total else values
+34
View File
@@ -0,0 +1,34 @@
"""Secret detection on the write path (§9, §14.12).
Memory is a place agents write to unattended; a leaked key there is a key in a
plain-text file that every harness can read back. Writes with a hit are rejected.
"""
from __future__ import annotations
import re
PATTERNS: list[tuple[str, re.Pattern[str]]] = [
("private_key", re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----")),
("openai_key", re.compile(r"\bsk-[A-Za-z0-9_-]{20,}\b")),
("anthropic_key", re.compile(r"\bsk-ant-[A-Za-z0-9_-]{20,}\b")),
("github_token", re.compile(r"\bgh[pousr]_[A-Za-z0-9]{30,}\b")),
("slack_token", re.compile(r"\bxox[abposr]-[A-Za-z0-9-]{10,}\b")),
("aws_access_key", re.compile(r"\b(?:AKIA|ASIA)[0-9A-Z]{16}\b")),
("google_api_key", re.compile(r"\bAIza[0-9A-Za-z_-]{35}\b")),
("jwt", re.compile(r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b")),
("bearer_header", re.compile(r"(?i)\bauthorization\s*[:=]\s*bearer\s+\S{16,}")),
("hashicorp_vault_token", re.compile(r"\bhv[sb]\.[A-Za-z0-9_-]{20,}\b")),
(
"generic_assignment",
re.compile(
r"(?i)\b(api[_-]?key|secret|passwo?rd|token|client[_-]?secret)\b\s*[:=]\s*"
r"['\"]?[A-Za-z0-9/+=_-]{20,}['\"]?"
),
),
]
def scan(text: str) -> list[str]:
"""Return the names of every pattern that matched. Never returns the match itself."""
return [name for name, pat in PATTERNS if pat.search(text or "")]
+116
View File
@@ -0,0 +1,116 @@
"""Process entry points: reindex, doctor, serve."""
from __future__ import annotations
import asyncio
import json
import logging
from typing import Any
from mdmem.config import Settings
from mdmem.service import MemoryService
log = logging.getLogger("mdmem.server")
def run_reindex(settings: Settings, args: Any) -> int:
async def _main() -> int:
svc = MemoryService.build(settings)
try:
if getattr(args, "no_embed", False):
svc.embedder = None
else:
await svc.start(watch=False, embed=False)
stats = await svc.reindex(full=bool(args.full), embed=not args.no_embed)
print(json.dumps(stats, ensure_ascii=False, indent=2))
return 1 if stats["errors"] else 0
finally:
await svc.stop()
return asyncio.run(_main())
def run_doctor(settings: Settings) -> int:
async def _main() -> int:
svc = MemoryService.build(settings)
try:
report: dict[str, Any] = {
"vault_dir": str(settings.vault_dir),
"index_db": str(settings.index_db),
"embed_model": settings.embed_model,
"embed_dim": settings.embed_dim,
"fusion": settings.fusion,
}
try:
assert svc.embedder is not None
report["embeddings"] = await svc.embedder.selftest()
except Exception as exc:
report["embeddings"] = {"error": str(exc)[:300]}
ok, checks = svc.ready()
report["ready"] = ok
report["checks"] = checks
health = svc.health()
report["counts"] = health.counts
report["vec_pending"] = health.vec_pending
print(json.dumps(report, ensure_ascii=False, indent=2))
return 0 if ok else 1
finally:
await svc.stop()
return asyncio.run(_main())
def run_serve(settings: Settings, args: Any) -> int:
import uvicorn
from mdmem.mcp_server import build_mcp_app
from mdmem.ui.app import build_ui_app
async def _main() -> int:
svc = MemoryService.build(settings)
await svc.start(watch=not getattr(args, "no_watch", False))
servers: list[uvicorn.Server] = []
if not getattr(args, "ui_only", False):
servers.append(
uvicorn.Server(
uvicorn.Config(
build_mcp_app(svc),
host=settings.mcp_host,
port=settings.mcp_port,
log_level=settings.log_level,
access_log=False,
)
)
)
if not getattr(args, "mcp_only", False):
servers.append(
uvicorn.Server(
uvicorn.Config(
build_ui_app(svc),
host=settings.ui_host,
port=settings.ui_port,
log_level=settings.log_level,
access_log=False,
)
)
)
log.info(
"mdmem serving",
extra={
"fields": {
"mcp": f"{settings.mcp_host}:{settings.mcp_port}"
if not getattr(args, "ui_only", False)
else None,
"ui": f"{settings.ui_host}:{settings.ui_port}"
if not getattr(args, "mcp_only", False)
else None,
}
},
)
try:
await asyncio.gather(*(s.serve() for s in servers))
finally:
await svc.stop()
return 0
return asyncio.run(_main())
+766
View File
@@ -0,0 +1,766 @@
"""Core service. Every write — MCP or UI — goes through exactly these functions (§10)."""
from __future__ import annotations
import asyncio
import json
import logging
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
from mdmem import search as search_mod
from mdmem.chunker import chunk_note, embed_text
from mdmem.config import NOTE_TYPES, STATES, Settings
from mdmem.db import Database
from mdmem.embed import BifrostEmbedder, Embedder, EmbeddingCache, to_blob
from mdmem.errors import Invalid, NotFound, SecretRejected, VersionConflict
from mdmem.ids import is_ulid
from mdmem.index import UNRESOLVED, Indexer
from mdmem.logging import AuditLog
from mdmem.notes import Note, utcnow
from mdmem.secrets import scan as scan_secrets
from mdmem.vaultio import Vault
from mdmem.watcher import VaultWatcher
log = logging.getLogger("mdmem.service")
IDEMPOTENCY_TTL_S = 24 * 3600
@dataclass
class Health:
counts: dict[str, int]
chunks: int
vec_pending: int
embed_model: str
embed_dim: int
embed_ok: bool
embed_last_ms: float | None
embed_errors: int
unresolved_links: int
orphans: int
secrets_in_vault: list[str]
class MemoryService:
def __init__(
self,
settings: Settings,
db: Database,
vault: Vault,
indexer: Indexer,
embedder: Embedder | None,
audit: AuditLog | None = None,
) -> None:
self.settings = settings
self.db = db
self.vault = vault
self.indexer = indexer
self.embedder = embedder
self.audit = audit or AuditLog()
self.watcher: VaultWatcher | None = None
self.embed_ok = embedder is not None
self._write_lock = asyncio.Lock()
self._secret_warnings: list[str] = []
self.metrics = {"search": 0, "writes": 0, "embed_calls": 0}
# -- construction ------------------------------------------------------
@classmethod
def build(cls, settings: Settings, embedder: Embedder | None = None) -> MemoryService:
vault = Vault(settings.vault_dir, settings.history_dir)
vault.ensure_dirs()
db = Database(
settings.index_db,
dim=settings.embed_dim,
trigram=settings.fts_trigram,
model=settings.embed_model,
)
rebuilt = db.open()
if rebuilt:
log.warning(
"embedding model or dim changed — vectors dropped, re-embedding",
extra={"fields": {"model": settings.embed_model, "dim": settings.embed_dim}},
)
if embedder is None:
embedder = BifrostEmbedder(
settings.bifrost_url,
settings.bifrost_vk,
settings.embed_model,
settings.embed_dim,
cache=EmbeddingCache(settings.cache_db),
)
return cls(settings, db, vault, Indexer(db, vault), embedder)
async def start(self, *, watch: bool = True, embed: bool = True) -> None:
if self.embedder is not None and embed:
try:
info = await self.embedder.selftest()
log.info("embedding backend ready", extra={"fields": info})
self.embed_ok = True
except Exception as exc:
self.embed_ok = False
log.error(
"embedding self-test failed — serving lexical-only until it recovers "
"(check EMBED_MODEL carries the provider prefix, §7/§14.14)",
extra={"fields": {"model": self.settings.embed_model, "error": str(exc)[:300]}},
)
stats = await asyncio.to_thread(self.indexer.scan)
log.info(
"vault scanned",
extra={
"fields": {
"scanned": stats.scanned,
"indexed": stats.indexed,
"removed": stats.removed,
"errors": len(stats.errors),
"seconds": round(stats.seconds, 3),
}
},
)
self._secret_warnings = await asyncio.to_thread(self._scan_vault_for_secrets)
if self._secret_warnings:
log.warning(
"vault contains possible secrets",
extra={"fields": {"notes": len(self._secret_warnings)}},
)
if embed:
await self.embed_pending()
if watch:
self.watcher = VaultWatcher(self.vault.root, self._on_changes)
self.watcher.start()
async def stop(self) -> None:
if self.watcher is not None:
await self.watcher.stop()
if self.embedder is not None:
await self.embedder.aclose()
self.db.close()
# -- indexing ----------------------------------------------------------
async def _on_changes(self, paths: set[Path]) -> None:
async with self._write_lock:
pending = []
for path in paths:
try:
pending.extend(await asyncio.to_thread(self._sync_path, path))
except Exception as exc:
log.warning(
"reindex of changed file failed",
extra={"fields": {"path": str(path), "error": str(exc)[:200]}},
)
await self._embed_items(pending)
def _sync_path(self, path: Path) -> list[Any]:
from mdmem.ids import id_from_filename
if not path.exists():
note_id = id_from_filename(path.name)
if note_id:
# A move between lifecycle folders shows up as a delete plus a create;
# only drop the note if no copy of it is left anywhere in the vault.
if self.vault.find(note_id) is None:
self.indexer.remove(note_id)
return []
note = self.vault.read(path)
row = self.db.read().execute("SELECT file_hash, path FROM notes WHERE id=?", (note.id,)).fetchone()
if row and row["file_hash"] == note.file_hash and row["path"] == str(path):
return [] # our own write, already indexed (§14.6)
return self.indexer.index_note(note)
async def reindex(self, *, full: bool = False, embed: bool = True) -> dict[str, Any]:
async with self._write_lock:
stats = await asyncio.to_thread(self.indexer.scan, full=full)
embedded = await self.embed_pending() if embed else 0
return {
"scanned": stats.scanned,
"indexed": stats.indexed,
"removed": stats.removed,
"unchanged": stats.unchanged,
"errors": stats.errors,
"embedded": embedded,
"seconds": round(stats.seconds, 3),
}
async def embed_pending(self) -> int:
items = await asyncio.to_thread(self.indexer.pending_embeds)
return await self._embed_items(items)
async def _embed_items(self, items: list[Any]) -> int:
if not items or self.embedder is None:
return 0
try:
vecs = await self.embedder.embed([i.text for i in items], is_query=False)
self.metrics["embed_calls"] += 1
self.embed_ok = True
except Exception as exc:
self.embed_ok = False
log.warning(
"embedding failed; chunks stay lexical-only and are retried later (§7)",
extra={"fields": {"chunks": len(items), "error": str(exc)[:200]}},
)
return 0
payload = [(i.chunk_id, to_blob(v)) for i, v in zip(items, vecs, strict=True)]
await asyncio.to_thread(self.indexer.store_vectors, payload)
return len(payload)
# -- read --------------------------------------------------------------
async def search(self, query: str, **kw: Any) -> search_mod.SearchResponse:
started = time.monotonic()
vec: bytes | None = None
if self.embedder is not None:
try:
v = await self.embedder.embed([query], is_query=True)
vec = to_blob(v[0])
except Exception as exc:
self.embed_ok = False
log.warning(
"query embedding failed; lexical-only for this search",
extra={"fields": {"error": str(exc)[:200]}},
)
resp = await asyncio.to_thread(search_mod.search, self.db, self.settings, query, query_vec=vec, **kw)
self.metrics["search"] += 1
log.debug(
"search",
extra={
"fields": {
"hits": len(resp.results),
"ms": round((time.monotonic() - started) * 1000, 1),
}
},
)
return resp
def get_note(self, note_id: str) -> Note:
path = self.vault.find(note_id)
if path is None:
raise NotFound(note_id)
return self.vault.read(path)
def get(self, note_id: str) -> dict[str, Any]:
note = self.get_note(note_id)
conn = self.db.read()
out_links = [
self._link_info(conn, r["dst"], r["rel"])
for r in conn.execute("SELECT dst, rel FROM links WHERE src=? ORDER BY rel", (note_id,))
]
backlinks = [
self._link_info(conn, r["src"], r["rel"])
for r in conn.execute("SELECT src, rel FROM links WHERE dst=? ORDER BY rel", (note_id,))
]
return {
"id": note.id,
"title": note.title,
"type": note.type,
"state": note.state,
"verified": note.state == "active",
"version": note.file_hash,
"markdown": note.to_text(),
"body": note.body,
"summary": note.summary,
"keywords": note.keywords,
"tags": note.tags,
"source": note.source,
"updated": note.ts("updated"),
"observed_at": note.ts("observed_at"),
"links": out_links,
"backlinks": backlinks,
"supersede_chain": self.supersede_chain(note),
"possible_duplicates": note.possible_duplicates,
"flags": self.flags_for(note_id),
}
def _link_info(self, conn, other_id: str, rel: str) -> dict[str, Any]:
if other_id.startswith(UNRESOLVED):
return {"id": None, "title": other_id[len(UNRESOLVED) :], "rel": rel, "resolved": False}
row = conn.execute("SELECT title FROM notes WHERE id=?", (other_id,)).fetchone()
return {
"id": other_id,
"title": row["title"] if row else other_id,
"rel": rel,
"resolved": row is not None,
}
def supersede_chain(self, note: Note) -> list[dict[str, Any]]:
conn = self.db.read()
chain: list[dict[str, Any]] = []
for old_id in note.supersedes:
row = conn.execute("SELECT title FROM notes WHERE id=?", (old_id,)).fetchone()
chain.append({"id": old_id, "title": row["title"] if row else old_id, "rel": "supersedes"})
if note.superseded_by:
row = conn.execute("SELECT title FROM notes WHERE id=?", (note.superseded_by,)).fetchone()
chain.append(
{
"id": note.superseded_by,
"title": row["title"] if row else note.superseded_by,
"rel": "superseded_by",
}
)
return chain
def flags_for(self, note_id: str) -> list[dict[str, Any]]:
return [
dict(r)
for r in self.db.read().execute(
"SELECT id, reason, source, created, resolved FROM flags WHERE note_id=? ORDER BY id DESC",
(note_id,),
)
]
def list_notes(
self, state: str | None = None, limit: int = 500, offset: int = 0, order: str = "updated"
) -> list[dict[str, Any]]:
allowed = {"updated", "created", "title", "type", "state"}
order = order if order in allowed else "updated"
sql = (
"SELECT id,title,type,state,tags,keywords,updated,created,summary,superseded_by,revises "
"FROM notes"
)
args: list[Any] = []
if state:
sql += " WHERE state=?"
args.append(state)
direction = "ASC" if order == "title" else "DESC"
sql += f" ORDER BY {order} {direction} LIMIT ? OFFSET ?"
args += [limit, offset]
rows = self.db.read().execute(sql, args).fetchall()
out = []
for r in rows:
d = dict(r)
d["tags"] = json.loads(d.pop("tags") or "[]")
d["keywords"] = json.loads(d.pop("keywords") or "[]")
out.append(d)
return out
# -- write -------------------------------------------------------------
def _validate(self, *, title: str, body: str, type: str) -> None:
if not title or not title.strip():
raise Invalid("title must not be empty")
if type not in NOTE_TYPES:
raise Invalid(f"type must be one of {list(NOTE_TYPES)}", type=type)
if len(body) > self.settings.max_note_chars:
raise Invalid(
f"body exceeds MAX_NOTE_CHARS ({self.settings.max_note_chars}); split it into atomic notes",
length=len(body),
)
hits = scan_secrets(f"{title}\n{body}")
if hits:
raise SecretRejected(hits)
async def write_note(
self,
*,
title: str,
body: str,
type: str = "note",
summary: str = "",
keywords: list[str] | None = None,
tags: list[str] | None = None,
relations: dict[str, list[str]] | None = None,
observed_at: str | None = None,
valid_from: str | None = None,
valid_until: str | None = None,
source: dict[str, Any] | None = None,
supersedes: list[str] | None = None,
revises: str | None = None,
state: str = "inbox",
) -> dict[str, Any]:
self._validate(title=title, body=body, type=type)
note = Note.new(
title=title,
body=body,
type=type,
summary=summary,
keywords=keywords or [],
tags=tags or [],
relations=relations or {},
observed_at=observed_at,
valid_from=valid_from,
valid_until=valid_until,
source=source or {},
supersedes=supersedes or [],
revises=revises,
state=state,
)
dups = await self._duplicates_for(note)
if dups:
note.set("possible_duplicates", [d["id"] for d in dups])
await self._persist(note)
self.metrics["writes"] += 1
return {
"id": note.id,
"state": note.state,
"version": note.file_hash,
"possible_duplicates": dups,
}
async def update_note(
self,
note_id: str,
expected_version: str | None,
*,
body: str | None = None,
title: str | None = None,
summary: str | None = None,
keywords: list[str] | None = None,
tags: list[str] | None = None,
relations: dict[str, list[str]] | None = None,
source: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Active notes are never edited in place by an agent: the edit becomes a
proposed revision in the inbox (§2.4, §9)."""
note = self.get_note(note_id)
if expected_version and expected_version != note.file_hash:
raise VersionConflict(note_id, expected_version, note.file_hash)
new_body = note.body if body is None else body
new_title = note.title if title is None else title
self._validate(title=new_title, body=new_body, type=note.type)
if note.state == "active":
existing = (
self.db.read()
.execute("SELECT id FROM notes WHERE revises=? AND state='inbox'", (note_id,))
.fetchone()
)
proposal = Note.new(
title=new_title,
body=new_body,
type=note.type,
summary=note.summary if summary is None else summary,
keywords=note.keywords if keywords is None else keywords,
tags=note.tags if tags is None else tags,
relations=note.relations if relations is None else relations,
observed_at=note.ts("observed_at"),
valid_from=note.ts("valid_from"),
valid_until=note.ts("valid_until"),
source=source or note.source,
revises=note_id,
state="inbox",
)
await self._persist(proposal)
return {
"id": proposal.id,
"state": "inbox",
"revises": note_id,
"version": proposal.file_hash,
"replaced_proposal": existing["id"] if existing else None,
}
for key, value in (
("title", title),
("summary", summary),
("keywords", keywords),
("tags", tags),
("relations", relations),
("source", source),
):
if value is not None:
note.set(key, value)
if body is not None:
note.body = body
note.touch()
await self._persist(note)
return {"id": note.id, "state": note.state, "version": note.file_hash}
async def supersede(
self, old_id: str, *, title: str, body: str, type: str = "fact", **kw: Any
) -> dict[str, Any]:
old = self.get_note(old_id)
result = await self.write_note(title=title, body=body, type=type, supersedes=[old.id], **kw)
result["supersedes"] = [old.id]
return result
async def flag(self, note_id: str, reason: str, source: str = "agent") -> dict[str, Any]:
self.get_note(note_id) # raises NotFound
async with self._write_lock:
def _insert() -> int:
with self.db.write() as w:
cur = w.execute(
"INSERT INTO flags(note_id,reason,source,created) VALUES(?,?,?,?)",
(note_id, reason[:500], source, utcnow()),
)
return int(cur.lastrowid or 0)
flag_id = await asyncio.to_thread(_insert)
return {"id": note_id, "flag_id": flag_id, "reason": reason[:500]}
# -- review actions (UI only) -----------------------------------------
async def approve(self, note_id: str, *, edited_body: str | None = None) -> dict[str, Any]:
note = self.get_note(note_id)
if note.state != "inbox":
raise Invalid("only inbox items can be approved", state=note.state)
if edited_body is not None:
note.body = edited_body
target_id = note.revises
if target_id:
return await self._merge_revision(note, target_id)
for old_id in note.supersedes:
old_path = self.vault.find(old_id)
if old_path is not None:
old = self.vault.read(old_path)
old.set("superseded_by", note.id)
old.touch()
await self._persist(old)
note.set("possible_duplicates", [])
note.touch()
await self._persist(note, state="active")
return {"id": note.id, "state": "active", "version": note.file_hash}
async def _merge_revision(self, proposal: Note, target_id: str) -> dict[str, Any]:
target = self.get_note(target_id)
target.body = proposal.body
for key in ("title", "summary", "keywords", "tags", "relations"):
target.set(key, proposal.fm.get(key))
target.touch()
await self._persist(target)
proposal.set("merged_into", target_id)
proposal.touch()
await self._persist(proposal, state="archive")
return {
"id": target.id,
"state": target.state,
"merged_from": proposal.id,
"version": target.file_hash,
}
async def reject(self, note_id: str, reason: str = "") -> dict[str, Any]:
note = self.get_note(note_id)
note.set("rejected", True)
if reason:
note.set("rejected_reason", reason[:500])
note.touch()
await self._persist(note, state="archive")
return {"id": note.id, "state": "archive"}
async def merge_into(self, src_id: str, dst_id: str) -> dict[str, Any]:
src = self.get_note(src_id)
dst = self.get_note(dst_id)
dst.body = f"{dst.body.rstrip()}\n\n{src.body.strip()}\n"
merged_kw = list(dict.fromkeys(dst.keywords + src.keywords))
merged_tags = list(dict.fromkeys(dst.tags + src.tags))
dst.set("keywords", merged_kw)
dst.set("tags", merged_tags)
dst.touch()
self._validate(title=dst.title, body=dst.body, type=dst.type)
await self._persist(dst)
src.set("merged_into", dst.id)
src.touch()
await self._persist(src, state="archive")
return {"id": dst.id, "merged_from": src.id, "state": dst.state}
async def set_state(self, note_id: str, state: str) -> dict[str, Any]:
if state not in STATES:
raise Invalid(f"state must be one of {list(STATES)}", state=state)
note = self.get_note(note_id)
note.touch()
await self._persist(note, state=state)
return {"id": note.id, "state": state}
async def edit(
self,
note_id: str,
*,
body: str,
expected_version: str,
frontmatter: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Direct edit from the UI: the owner may change an active note in place."""
note = self.get_note(note_id)
if expected_version and expected_version != note.file_hash:
raise VersionConflict(note_id, expected_version, note.file_hash)
if frontmatter:
for k, v in frontmatter.items():
note.set(k, v)
note.body = body
self._validate(title=note.title, body=note.body, type=note.type)
note.touch()
await self._persist(note)
return {"id": note.id, "version": note.file_hash, "state": note.state}
async def hard_delete(self, note_id: str) -> dict[str, Any]:
note = self.get_note(note_id)
async with self._write_lock:
await asyncio.to_thread(self.vault.delete, note)
await asyncio.to_thread(self.indexer.remove, note_id)
def _cleanup() -> None:
with self.db.write() as w:
w.execute("DELETE FROM flags WHERE note_id=?", (note_id,))
w.execute("DELETE FROM links WHERE dst=?", (note_id,))
await asyncio.to_thread(_cleanup)
return {"id": note_id, "deleted": True}
# -- helpers -----------------------------------------------------------
async def _persist(self, note: Note, state: str | None = None) -> None:
"""Atomic file write → history snapshot → synchronous reindex (§10)."""
async with self._write_lock:
await asyncio.to_thread(self.vault.write, note, state)
pending = await asyncio.to_thread(self.indexer.index_note, note)
await self._embed_items(pending)
async def _duplicates_for(self, note: Note) -> list[dict[str, Any]]:
if self.embedder is None:
return []
chunks = chunk_note(note)
if not chunks:
return []
try:
vec = await self.embedder.embed([embed_text(note, chunks[0])], is_query=False)
except Exception:
return []
blob = to_blob(vec[0])
def _query() -> list[dict[str, Any]]:
rows = (
self.db.read()
.execute(
"SELECT rowid AS cid, distance FROM chunks_vec "
"WHERE embedding MATCH ? AND k = 10 ORDER BY distance",
(blob,),
)
.fetchall()
)
out: list[dict[str, Any]] = []
seen: set[str] = set()
for r in rows:
cos = search_mod._cosine_from_l2(float(r["distance"]))
if cos < self.settings.dup_threshold:
continue
row = (
self.db.read()
.execute(
"SELECT n.id, n.title, n.state FROM chunks c JOIN notes n ON n.id=c.note_id "
"WHERE c.id=?",
(r["cid"],),
)
.fetchone()
)
if row is None or row["state"] != "active" or row["id"] in seen:
continue
seen.add(row["id"])
out.append({"id": row["id"], "title": row["title"], "cosine": round(cos, 4)})
return out
return await asyncio.to_thread(_query)
def _scan_vault_for_secrets(self) -> list[str]:
hits: list[str] = []
for path, _state in self.vault.iter_files():
try:
text = path.read_text(encoding="utf-8")
except OSError:
continue
if scan_secrets(text):
hits.append(str(path.name))
return hits
# -- idempotency -------------------------------------------------------
def idempotent_get(self, request_id: str | None, tool: str) -> dict[str, Any] | None:
if not request_id:
return None
row = (
self.db.read()
.execute(
"SELECT result, created FROM idempotency WHERE request_id=? AND tool=?",
(request_id, tool),
)
.fetchone()
)
if row is None:
return None
try:
age = time.time() - time.mktime(time.strptime(row["created"], "%Y-%m-%dT%H:%M:%SZ"))
except ValueError:
age = 0.0
if age > IDEMPOTENCY_TTL_S:
return None
return json.loads(row["result"])
def idempotent_put(self, request_id: str | None, tool: str, result: dict[str, Any]) -> None:
if not request_id:
return
with self.db.write() as w:
w.execute(
"INSERT OR REPLACE INTO idempotency(request_id,tool,result,created) VALUES(?,?,?,?)",
(request_id, tool, json.dumps(result, ensure_ascii=False), utcnow()),
)
# -- health ------------------------------------------------------------
def health(self) -> Health:
conn = self.db.read()
counts = {
r["state"]: r["n"] for r in conn.execute("SELECT state, count(*) AS n FROM notes GROUP BY state")
}
chunks = conn.execute("SELECT count(*) FROM chunks").fetchone()[0]
pending = conn.execute("SELECT count(*) FROM chunks WHERE vec_pending=1").fetchone()[0]
unresolved = conn.execute("SELECT count(*) FROM links WHERE dst LIKE 'unresolved:%'").fetchone()[0]
orphans = conn.execute(
"SELECT count(*) FROM notes n WHERE NOT EXISTS(SELECT 1 FROM links WHERE src=n.id) "
"AND NOT EXISTS(SELECT 1 FROM links WHERE dst=n.id)"
).fetchone()[0]
emb = self.embedder
return Health(
counts={s: counts.get(s, 0) for s in STATES},
chunks=chunks,
vec_pending=pending,
embed_model=self.settings.embed_model,
embed_dim=self.settings.embed_dim,
embed_ok=self.embed_ok,
embed_last_ms=getattr(emb, "last_latency_ms", None),
embed_errors=getattr(emb, "errors", 0),
unresolved_links=unresolved,
orphans=orphans,
secrets_in_vault=self._secret_warnings,
)
def unresolved_links(self) -> list[dict[str, Any]]:
return [
{
"src": r["src"],
"target": r["dst"][len(UNRESOLVED) :],
"rel": r["rel"],
"title": r["title"],
}
for r in self.db.read().execute(
"SELECT l.src, l.dst, l.rel, n.title FROM links l JOIN notes n ON n.id=l.src "
"WHERE l.dst LIKE 'unresolved:%' ORDER BY n.title"
)
]
def ready(self) -> tuple[bool, dict[str, Any]]:
"""Readiness must not depend on Bifrost (§13a.3, §14.17)."""
checks: dict[str, Any] = {}
ok = True
try:
self.db.read().execute("SELECT count(*) FROM notes").fetchone()
checks["index_db"] = "ok"
except Exception as exc:
checks["index_db"] = str(exc)[:120]
ok = False
try:
self.db.read().execute("SELECT vec_version()").fetchone()
checks["sqlite_vec"] = "ok"
except Exception as exc:
checks["sqlite_vec"] = str(exc)[:120]
ok = False
checks["vault_writable"] = "ok" if self.vault.writable() else "not writable"
ok = ok and checks["vault_writable"] == "ok"
checks["embeddings"] = "ok" if self.embed_ok else "degraded (lexical only)"
return ok, checks
def validate_id(note_id: str) -> str:
if not is_ulid(note_id):
raise Invalid("id must be a ULID", id=note_id)
return note_id
def as_float32(v: Any) -> np.ndarray: # pragma: no cover - convenience
return np.asarray(v, dtype=np.float32)
+69
View File
@@ -0,0 +1,69 @@
"""Vault watcher (§6.2): debounced, hash-deduped, ignores *.tmp and .history."""
from __future__ import annotations
import asyncio
import logging
from collections.abc import Awaitable, Callable
from pathlib import Path
from watchfiles import Change, awatch
from mdmem.vaultio import TMP_SUFFIX
log = logging.getLogger("mdmem.watcher")
DEBOUNCE_MS = 500
def _relevant(_change: Change, path: str) -> bool:
p = Path(path)
if p.name.endswith(TMP_SUFFIX) or p.suffix != ".md":
return False
return ".history" not in p.parts
class VaultWatcher:
def __init__(
self,
root: Path,
on_changes: Callable[[set[Path]], Awaitable[None]],
debounce_ms: int = DEBOUNCE_MS,
) -> None:
self.root = Path(root)
self.on_changes = on_changes
self.debounce_ms = debounce_ms
self._task: asyncio.Task[None] | None = None
self._stop = asyncio.Event()
async def _run(self) -> None:
log.info("watching vault", extra={"fields": {"path": str(self.root)}})
try:
async for batch in awatch(
self.root,
watch_filter=_relevant,
step=self.debounce_ms,
debounce=self.debounce_ms * 3,
stop_event=self._stop,
):
paths = {Path(p) for _c, p in batch}
try:
await self.on_changes(paths)
except Exception: # a bad file must not kill the watcher
log.exception("change handling failed")
except asyncio.CancelledError: # pragma: no cover
raise
def start(self) -> None:
self._stop = asyncio.Event()
self._task = asyncio.create_task(self._run(), name="mdmem-watcher")
async def stop(self) -> None:
self._stop.set()
if self._task is not None:
self._task.cancel()
try:
await self._task
except (asyncio.CancelledError, Exception):
pass
self._task = None
+5 -5
View File
@@ -10,10 +10,7 @@ def test_short_note_is_one_chunk():
def test_long_note_splits_on_h2_and_h3():
body = "\n\n".join(
["intro " * 50]
+ [f"## Section {i}\n\n" + ("text " * 400) for i in range(3)]
)
body = "\n\n".join(["intro " * 50] + [f"## Section {i}\n\n" + ("text " * 400) for i in range(3)])
n = Note.new(title="T", body=body, type="note")
chunks = chunk_note(n)
assert len(chunks) >= 4
@@ -40,7 +37,10 @@ def test_one_huge_paragraph_is_split_without_overlap():
def test_embed_text_merges_keys_at_index_time():
n = Note.new(
title="Titel", body="Körper", type="fact", summary="Zusammenfassung",
title="Titel",
body="Körper",
type="fact",
summary="Zusammenfassung",
keywords=["a", "b"],
)
c = chunk_note(n)[0]
+1 -3
View File
@@ -3,9 +3,7 @@ import sys
def _run(*args):
return subprocess.run(
[sys.executable, "-m", "mdmem", *args], capture_output=True, text=True
)
return subprocess.run([sys.executable, "-m", "mdmem", *args], capture_output=True, text=True)
def test_help_exits_zero():
+16 -3
View File
@@ -10,8 +10,19 @@ def test_schema_created_and_versioned(tmp_path):
db.open()
assert db.meta_get("schema_version") == str(SCHEMA_VERSION)
assert db.meta_get("embed_dim") == "8"
tables = {r[0] for r in db.read().execute("SELECT name FROM sqlite_master WHERE type IN ('table','view')")}
for t in ("notes", "chunks", "chunks_fts", "chunks_tri", "chunks_vec", "links", "flags", "idempotency"):
tables = {
r[0] for r in db.read().execute("SELECT name FROM sqlite_master WHERE type IN ('table','view')")
}
for t in (
"notes",
"chunks",
"chunks_fts",
"chunks_tri",
"chunks_vec",
"links",
"flags",
"idempotency",
):
assert t in tables
db.close()
@@ -39,7 +50,9 @@ def test_dim_change_rebuilds_vector_table(tmp_path):
db.open()
with db.write() as w:
w.execute("INSERT INTO notes(id,path,state,file_hash) VALUES('a','p','active','h')")
w.execute("INSERT INTO chunks(id,note_id,ord,heading,body,embed_text_hash) VALUES(1,'a',0,'','b','h')")
w.execute(
"INSERT INTO chunks(id,note_id,ord,heading,body,embed_text_hash) VALUES(1,'a',0,'','b','h')"
)
w.execute("INSERT INTO chunks_vec(rowid,embedding) VALUES(1,?)", (b"\x00" * 16,))
db.close()
+33 -15
View File
@@ -26,8 +26,14 @@ def _note(vault, **kw):
def test_index_note_writes_all_rows(idx):
indexer, vault, db = idx
n = Note.new(title="llama-swap", body="routet Modelle", type="fact",
summary="s", keywords=["router", "modell"], tags=["homelab"])
n = Note.new(
title="llama-swap",
body="routet Modelle",
type="fact",
summary="s",
keywords=["router", "modell"],
tags=["homelab"],
)
vault.write(n)
n = vault.read(vault.find(n.id))
pending = indexer.index_note(n)
@@ -43,17 +49,18 @@ def test_index_note_writes_all_rows(idx):
def test_fts_finds_keyword_and_body(idx):
indexer, vault, db = idx
n = Note.new(title="Titel", body="Bifrost fronts llama-swap", type="fact",
keywords=["Modellrouter"])
n = Note.new(title="Titel", body="Bifrost fronts llama-swap", type="fact", keywords=["Modellrouter"])
vault.write(n)
indexer.index_note(vault.read(vault.find(n.id)))
hits = db.read().execute(
"SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ?", ('"modellrouter"',)
).fetchall()
hits = (
db.read()
.execute("SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ?", ('"modellrouter"',))
.fetchall()
)
assert hits
assert db.read().execute(
"SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ?", ('"fronts"',)
).fetchall()
assert (
db.read().execute("SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ?", ('"fronts"',)).fetchall()
)
def test_reindex_replaces_chunks_not_duplicates(idx):
@@ -76,12 +83,18 @@ def test_links_are_extracted_and_resolved(idx):
vault.write(target)
indexer.index_note(vault.read(vault.find(target.id)))
src = Note.new(title="llama-swap", body="runs on [[evo-x2]] and [[nowhere]]",
type="fact", relations={"part_of": ["[[evo-x2]]"]})
src = Note.new(
title="llama-swap",
body="runs on [[evo-x2]] and [[nowhere]]",
type="fact",
relations={"part_of": ["[[evo-x2]]"]},
)
vault.write(src)
indexer.index_note(vault.read(vault.find(src.id)))
links = {(r["dst"], r["rel"]) for r in db.read().execute("SELECT dst, rel FROM links WHERE src=?", (src.id,))}
links = {
(r["dst"], r["rel"]) for r in db.read().execute("SELECT dst, rel FROM links WHERE src=?", (src.id,))
}
assert (target.id, "links") in links
assert (target.id, "part_of") in links
assert ("unresolved:nowhere", "links") in links
@@ -92,7 +105,12 @@ def test_forward_reference_resolves_when_target_arrives(idx):
src = Note.new(title="A", body="see [[Later Note]]", type="note")
vault.write(src)
indexer.index_note(vault.read(vault.find(src.id)))
assert db.read().execute("SELECT dst FROM links WHERE src=?", (src.id,)).fetchone()[0].startswith("unresolved:")
assert (
db.read()
.execute("SELECT dst FROM links WHERE src=?", (src.id,))
.fetchone()[0]
.startswith("unresolved:")
)
later = Note.new(title="Later Note", body="here", type="note")
vault.write(later)
@@ -150,7 +168,7 @@ def test_unchanged_embed_text_keeps_its_vector(idx):
pending = indexer.index_note(n)
indexer.store_vectors([(pending[0].chunk_id, b"\x00" * 32)])
n.set("tags", ["neu"]) # tags are not part of the embed text
n.set("tags", ["neu"]) # tags are not part of the embed text
vault.write(n)
again = indexer.index_note(vault.read(vault.find(n.id)))
assert again == []
+27
View File
@@ -0,0 +1,27 @@
import pytest
from mdmem.secrets import scan
@pytest.mark.parametrize(
"text,expected",
[
("nothing to see here", []),
("-----BEGIN OPENSSH PRIVATE KEY-----\nabc", ["private_key"]),
("key is sk-abcdefghijklmnopqrstuvwxyz012345", ["openai_key"]),
("AKIAIOSFODNN7EXAMPLE", ["aws_access_key"]),
("Authorization: Bearer abcdefghijklmnopqrstuvwxyz", ["bearer_header"]),
("password: hunter2", []), # short values are not secrets, just noise
],
)
def test_scan(text, expected):
assert scan(text) == expected
def test_generic_assignment_catches_long_values():
assert "generic_assignment" in scan('client_secret = "Xy3k9QpL2mNvB8sT4uW1zC7dF0gH5jK6"')
def test_scan_never_returns_the_secret_itself():
hits = scan("sk-abcdefghijklmnopqrstuvwxyz012345")
assert all("sk-" not in h for h in hits)
+1 -2
View File
@@ -94,8 +94,7 @@ def test_read_normalises_nfd_from_disk(vault):
p = vault.root / "active" / "x--01JZZZZZZZZZZZZZZZZZZZZZZZ.md"
p.write_text(
"---\nid: 01JZZZZZZZZZZZZZZZZZZZZZZZ\ntitle: t\n---\n"
+ unicodedata.normalize("NFD", "Grüße"),
"---\nid: 01JZZZZZZZZZZZZZZZZZZZZZZZ\ntitle: t\n---\n" + unicodedata.normalize("NFD", "Grüße"),
encoding="utf-8",
)
n = vault.read(p)
+71
View File
@@ -0,0 +1,71 @@
import asyncio
import pytest
from mdmem.embed import FakeEmbedder
from mdmem.notes import Note
from mdmem.service import MemoryService
@pytest.fixture()
async def svc(settings):
s = MemoryService.build(settings, embedder=FakeEmbedder(settings.embed_dim))
await s.start(watch=True)
yield s
await s.stop()
async def _wait_for(fn, timeout=2.0, interval=0.05):
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
while loop.time() < deadline:
got = fn()
if got:
return got
await asyncio.sleep(interval)
return fn()
async def test_file_written_on_disk_becomes_searchable_within_two_seconds(svc):
n = Note.new(
title="Zeitreise",
body="Der Wasserstand am Traunsee",
type="fact",
keywords=["traunsee"],
state="active",
)
path = svc.vault.root / "active" / n.filename()
path.write_text(n.to_text(), encoding="utf-8")
async def _hit():
r = await svc.search("Wasserstand", k=5)
return r.results
results = None
loop = asyncio.get_running_loop()
deadline = loop.time() + 2.0
while loop.time() < deadline:
results = await _hit()
if results:
break
await asyncio.sleep(0.05)
assert results, "note written on disk was not searchable within 2 s"
assert results[0].id == n.id
async def test_deleting_a_file_removes_it_from_the_index(svc):
n = Note.new(title="Weg damit", body="temporär", type="note", state="active")
path = svc.vault.root / "active" / n.filename()
path.write_text(n.to_text(), encoding="utf-8")
await _wait_for(lambda: svc.db.read().execute("SELECT count(*) FROM notes").fetchone()[0])
path.unlink()
gone = await _wait_for(lambda: svc.db.read().execute("SELECT count(*) FROM notes").fetchone()[0] == 0)
assert gone
async def test_service_own_write_does_not_double_index(svc):
res = await svc.write_note(title="Eigener Schreibvorgang", body="Inhalt", type="fact")
await asyncio.sleep(0.9) # let the watcher see our own write
rows = svc.db.read().execute("SELECT count(*) FROM chunks WHERE note_id=?", (res["id"],)).fetchone()[0]
assert rows == 1