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:
1 parent
6a424c0909
commit
4aa3d5a86c
21 files changed
+1814
-76
No files matched your search
Binary file not shown.
+1
-1
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
|
||||
@@ -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))
|
||||
@@ -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
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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 "")]
|
||||
@@ -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())
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
@@ -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
@@ -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
@@ -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 == []
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user