lean-memory 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- lean_memory/__init__.py +14 -0
- lean_memory/embed/__init__.py +0 -0
- lean_memory/embed/base.py +51 -0
- lean_memory/embed/fake.py +41 -0
- lean_memory/embed/sentence_transformer.py +65 -0
- lean_memory/extract/__init__.py +0 -0
- lean_memory/extract/contradiction.py +329 -0
- lean_memory/extract/gliner_extractor.py +483 -0
- lean_memory/extract/llm_typer.py +444 -0
- lean_memory/extract/router.py +418 -0
- lean_memory/extract/rules.py +130 -0
- lean_memory/extract/salience.py +123 -0
- lean_memory/extract/taxonomy.py +244 -0
- lean_memory/mcp_server.py +150 -0
- lean_memory/memory.py +203 -0
- lean_memory/retrieve/__init__.py +0 -0
- lean_memory/retrieve/rerank.py +59 -0
- lean_memory/retrieve/retriever.py +125 -0
- lean_memory/store/__init__.py +0 -0
- lean_memory/store/base.py +91 -0
- lean_memory/store/schema.py +83 -0
- lean_memory/store/sqlite_store.py +301 -0
- lean_memory/types.py +107 -0
- lean_memory-0.1.0.dist-info/METADATA +228 -0
- lean_memory-0.1.0.dist-info/RECORD +27 -0
- lean_memory-0.1.0.dist-info/WHEEL +4 -0
- lean_memory-0.1.0.dist-info/entry_points.txt +2 -0
|
@@ -0,0 +1,125 @@
|
|
|
1
|
+
"""The Retriever — composes the Phase 0 retrieval pipeline (spec section 6):
|
|
2
|
+
|
|
3
|
+
dense (two-stage Matryoshka) + sparse (BM25)
|
|
4
|
+
→ RRF fuse (k=10)
|
|
5
|
+
→ over-retrieve top 30-50
|
|
6
|
+
→ mandatory cross-encoder rerank
|
|
7
|
+
→ salience-decay re-score
|
|
8
|
+
→ temporal filter (is_latest by default; as_of=T → interval predicate)
|
|
9
|
+
|
|
10
|
+
All local, reproducible, no cloud key, no LLM at rank time.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import math
|
|
16
|
+
from typing import Optional
|
|
17
|
+
|
|
18
|
+
from ..embed.base import Embedder, matryoshka_truncate
|
|
19
|
+
from ..store.base import Store
|
|
20
|
+
from ..types import RetrievedFact, now_ms
|
|
21
|
+
from .rerank import Reranker
|
|
22
|
+
|
|
23
|
+
# spec defaults
|
|
24
|
+
RRF_K = 10
|
|
25
|
+
OVER_RETRIEVE = 40 # top 30-50 fused candidates before rerank
|
|
26
|
+
W_REL, W_REC, W_IMP = 0.6, 0.2, 0.2
|
|
27
|
+
DECAY_LAMBDA = 1.0 / (1000 * 60 * 60 * 24 * 30) # ~1/month in ms; recency=exp(-λ·age)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class Retriever:
|
|
31
|
+
def __init__(self, store: Store, embedder: Embedder, reranker: Reranker) -> None:
|
|
32
|
+
self.store = store
|
|
33
|
+
self.embedder = embedder
|
|
34
|
+
self.reranker = reranker
|
|
35
|
+
|
|
36
|
+
def retrieve(
|
|
37
|
+
self,
|
|
38
|
+
query: str,
|
|
39
|
+
k: int = 5,
|
|
40
|
+
*,
|
|
41
|
+
as_of: Optional[int] = None,
|
|
42
|
+
is_latest_only: bool = True,
|
|
43
|
+
now: Optional[int] = None,
|
|
44
|
+
) -> list[RetrievedFact]:
|
|
45
|
+
now = now if now is not None else now_ms()
|
|
46
|
+
|
|
47
|
+
# 1. embed query (full + coarse for the two-stage dense arm)
|
|
48
|
+
q_full = self.embedder.embed_one(query)
|
|
49
|
+
q_coarse = matryoshka_truncate(q_full, self.embedder.coarse_dim)
|
|
50
|
+
|
|
51
|
+
# 2+3. dense and sparse arms, each over-retrieved
|
|
52
|
+
dense = self.store.dense_search(
|
|
53
|
+
q_coarse, q_full, OVER_RETRIEVE,
|
|
54
|
+
is_latest_only=is_latest_only, as_of=as_of,
|
|
55
|
+
)
|
|
56
|
+
sparse = self.store.sparse_search(
|
|
57
|
+
query, OVER_RETRIEVE, is_latest_only=is_latest_only, as_of=as_of,
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
# 4. RRF fuse (k=10). score(d) = Σ 1/(k + rank_i(d)), rank 1-based.
|
|
61
|
+
ranks: dict[str, dict[str, int]] = {}
|
|
62
|
+
for rank, (fid, _) in enumerate(dense, start=1):
|
|
63
|
+
ranks.setdefault(fid, {})["dense"] = rank
|
|
64
|
+
for rank, (fid, _) in enumerate(sparse, start=1):
|
|
65
|
+
ranks.setdefault(fid, {})["sparse"] = rank
|
|
66
|
+
|
|
67
|
+
fused: list[tuple[str, float]] = []
|
|
68
|
+
for fid, r in ranks.items():
|
|
69
|
+
s = 0.0
|
|
70
|
+
if "dense" in r:
|
|
71
|
+
s += 1.0 / (RRF_K + r["dense"])
|
|
72
|
+
if "sparse" in r:
|
|
73
|
+
s += 1.0 / (RRF_K + r["sparse"])
|
|
74
|
+
fused.append((fid, s))
|
|
75
|
+
fused.sort(key=lambda x: x[1], reverse=True)
|
|
76
|
+
fused = fused[:OVER_RETRIEVE]
|
|
77
|
+
if not fused:
|
|
78
|
+
return []
|
|
79
|
+
|
|
80
|
+
# hydrate the fused candidates
|
|
81
|
+
fact_map = self.store.hydrate([fid for fid, _ in fused])
|
|
82
|
+
|
|
83
|
+
# 5. mandatory rerank over the fused candidate texts
|
|
84
|
+
cand_ids = [fid for fid, _ in fused if fid in fact_map]
|
|
85
|
+
cand_texts = [fact_map[fid].fact_text for fid in cand_ids]
|
|
86
|
+
rel_scores = self.reranker.score(query, cand_texts)
|
|
87
|
+
|
|
88
|
+
# 6. salience-decay re-score: final = w_rel·rel + w_rec·recency + w_imp·importance
|
|
89
|
+
rel_norm = _minmax(rel_scores)
|
|
90
|
+
rrf_lookup = dict(fused)
|
|
91
|
+
dense_lookup = {fid: i + 1 for i, (fid, _) in enumerate(dense)}
|
|
92
|
+
sparse_lookup = {fid: i + 1 for i, (fid, _) in enumerate(sparse)}
|
|
93
|
+
|
|
94
|
+
results: list[RetrievedFact] = []
|
|
95
|
+
for fid, rel, rel_n in zip(cand_ids, rel_scores, rel_norm):
|
|
96
|
+
fact = fact_map[fid]
|
|
97
|
+
age = max(0, now - (fact.last_access or fact.valid_at))
|
|
98
|
+
recency = math.exp(-DECAY_LAMBDA * age)
|
|
99
|
+
importance = fact.salience / 10.0
|
|
100
|
+
final = W_REL * rel_n + W_REC * recency + W_IMP * importance
|
|
101
|
+
results.append(
|
|
102
|
+
RetrievedFact(
|
|
103
|
+
fact=fact, final_score=final, relevance=rel,
|
|
104
|
+
recency=recency, importance=importance,
|
|
105
|
+
dense_rank=dense_lookup.get(fid),
|
|
106
|
+
sparse_rank=sparse_lookup.get(fid),
|
|
107
|
+
rrf_score=rrf_lookup.get(fid),
|
|
108
|
+
)
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
results.sort(key=lambda r: r.final_score, reverse=True)
|
|
112
|
+
top = results[:k]
|
|
113
|
+
for r in top:
|
|
114
|
+
self.store.touch(r.fact.id, now)
|
|
115
|
+
return top
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def _minmax(xs: list[float]) -> list[float]:
|
|
119
|
+
"""Scale scores to [0,1] so they combine sanely with recency/importance."""
|
|
120
|
+
if not xs:
|
|
121
|
+
return []
|
|
122
|
+
lo, hi = min(xs), max(xs)
|
|
123
|
+
if hi - lo < 1e-9:
|
|
124
|
+
return [1.0 for _ in xs]
|
|
125
|
+
return [(x - lo) / (hi - lo) for x in xs]
|
|
File without changes
|
|
@@ -0,0 +1,91 @@
|
|
|
1
|
+
"""The `Store` interface — the single abstraction every other component talks to.
|
|
2
|
+
|
|
3
|
+
Two concrete implementations are planned (SqliteStore = default; LanceStore =
|
|
4
|
+
scale tier). Phase 0 ships SqliteStore only. Per BET 4, each namespace is its own
|
|
5
|
+
backing file (per-tenant isolation), so the interface is opened per-namespace.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from abc import ABC, abstractmethod
|
|
11
|
+
from typing import Optional, Sequence
|
|
12
|
+
|
|
13
|
+
import numpy as np
|
|
14
|
+
|
|
15
|
+
from ..types import Entity, Episode, Fact
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class Store(ABC):
|
|
19
|
+
"""Storage + index abstraction. Implementations own one namespace's data."""
|
|
20
|
+
|
|
21
|
+
# ── provenance ──
|
|
22
|
+
@abstractmethod
|
|
23
|
+
def add_episode(self, episode: Episode) -> None: ...
|
|
24
|
+
|
|
25
|
+
# ── entities ──
|
|
26
|
+
@abstractmethod
|
|
27
|
+
def upsert_entity(self, entity: Entity) -> Entity:
|
|
28
|
+
"""Resolve-or-create. If an entity with the same (namespace, name, type)
|
|
29
|
+
exists, return it; otherwise insert `entity` and return it."""
|
|
30
|
+
|
|
31
|
+
@abstractmethod
|
|
32
|
+
def get_entity(self, entity_id: str) -> Optional[Entity]: ...
|
|
33
|
+
|
|
34
|
+
# ── facts ──
|
|
35
|
+
@abstractmethod
|
|
36
|
+
def add_fact(self, fact: Fact, embedding: np.ndarray, embedding_256: np.ndarray) -> None:
|
|
37
|
+
"""Insert a fact row + its vec0 vectors + its FTS row, in one transaction."""
|
|
38
|
+
|
|
39
|
+
@abstractmethod
|
|
40
|
+
def supersede_fact(self, old_fact_id: str, new_fact_id: str, valid_to: int) -> None:
|
|
41
|
+
"""ADD-only supersession: point old→new, set old.valid_to, flip old.is_latest=0.
|
|
42
|
+
Never deletes."""
|
|
43
|
+
|
|
44
|
+
@abstractmethod
|
|
45
|
+
def get_fact(self, fact_id: str) -> Optional[Fact]: ...
|
|
46
|
+
|
|
47
|
+
@abstractmethod
|
|
48
|
+
def find_latest_in_slot(
|
|
49
|
+
self, subject_id: str, predicate: str
|
|
50
|
+
) -> Sequence[Fact]:
|
|
51
|
+
"""All currently-latest facts in a (subject, predicate) slot — for contradiction
|
|
52
|
+
detection / supersession lookup."""
|
|
53
|
+
|
|
54
|
+
# ── retrieval primitives (the Retriever composes these) ──
|
|
55
|
+
@abstractmethod
|
|
56
|
+
def dense_search(
|
|
57
|
+
self,
|
|
58
|
+
query_256: np.ndarray,
|
|
59
|
+
query_768: np.ndarray,
|
|
60
|
+
k: int,
|
|
61
|
+
*,
|
|
62
|
+
is_latest_only: bool = True,
|
|
63
|
+
as_of: Optional[int] = None,
|
|
64
|
+
) -> list[tuple[str, float]]:
|
|
65
|
+
"""Two-stage Matryoshka dense search. Returns [(fact_id, distance)] best-first."""
|
|
66
|
+
|
|
67
|
+
@abstractmethod
|
|
68
|
+
def sparse_search(
|
|
69
|
+
self, query_text: str, k: int, *, is_latest_only: bool = True,
|
|
70
|
+
as_of: Optional[int] = None,
|
|
71
|
+
) -> list[tuple[str, float]]:
|
|
72
|
+
"""BM25 lexical search. Returns [(fact_id, score)] best-first.
|
|
73
|
+
as_of applies the same interval predicate as the dense arm."""
|
|
74
|
+
|
|
75
|
+
@abstractmethod
|
|
76
|
+
def hydrate(self, fact_ids: Sequence[str]) -> dict[str, Fact]:
|
|
77
|
+
"""Bulk-load Fact rows by id (preserves caller's dedup needs)."""
|
|
78
|
+
|
|
79
|
+
@abstractmethod
|
|
80
|
+
def touch(self, fact_id: str, when_ms: int) -> None:
|
|
81
|
+
"""Record an access (recency/decay bookkeeping)."""
|
|
82
|
+
|
|
83
|
+
# ── lifecycle ──
|
|
84
|
+
@abstractmethod
|
|
85
|
+
def close(self) -> None: ...
|
|
86
|
+
|
|
87
|
+
def __enter__(self) -> "Store":
|
|
88
|
+
return self
|
|
89
|
+
|
|
90
|
+
def __exit__(self, *exc) -> None:
|
|
91
|
+
self.close()
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
"""SQLite + vec0 + FTS5 schema for Phase 0.
|
|
2
|
+
|
|
3
|
+
This is the spec's data model translated from its Postgres-flavored DDL to SQLite:
|
|
4
|
+
- Postgres ENUM / halfvec / pgvector → SQLite TEXT + vec0 INT8[N].
|
|
5
|
+
- `{dim}`/`{coarse_dim}` are filled at connect time from the embedder.
|
|
6
|
+
The monotemporal spine is always present; bi-temporal audit columns exist but are
|
|
7
|
+
only diverged from ingest time when the audit extra is enabled (deferred past Phase 0).
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
SCHEMA_SQL = """
|
|
11
|
+
-- ── PROVENANCE LAYER ────────────────────────────────────────────────
|
|
12
|
+
CREATE TABLE IF NOT EXISTS episode (
|
|
13
|
+
id TEXT PRIMARY KEY,
|
|
14
|
+
namespace TEXT NOT NULL,
|
|
15
|
+
raw TEXT NOT NULL,
|
|
16
|
+
source TEXT,
|
|
17
|
+
t_ref INTEGER NOT NULL,
|
|
18
|
+
created_at INTEGER NOT NULL
|
|
19
|
+
);
|
|
20
|
+
|
|
21
|
+
-- ── ENTITY LAYER ────────────────────────────────────────────────────
|
|
22
|
+
CREATE TABLE IF NOT EXISTS entity (
|
|
23
|
+
id TEXT PRIMARY KEY,
|
|
24
|
+
namespace TEXT NOT NULL,
|
|
25
|
+
name TEXT NOT NULL,
|
|
26
|
+
type TEXT,
|
|
27
|
+
summary TEXT,
|
|
28
|
+
resolved_id TEXT,
|
|
29
|
+
created_at INTEGER NOT NULL
|
|
30
|
+
);
|
|
31
|
+
CREATE INDEX IF NOT EXISTS ix_entity_lookup ON entity(namespace, name, type);
|
|
32
|
+
|
|
33
|
+
-- ── FACT LAYER (monotemporal spine always on; audit axis opt-in) ─────
|
|
34
|
+
CREATE TABLE IF NOT EXISTS fact (
|
|
35
|
+
id TEXT PRIMARY KEY,
|
|
36
|
+
namespace TEXT NOT NULL,
|
|
37
|
+
subject_id TEXT NOT NULL REFERENCES entity(id),
|
|
38
|
+
predicate TEXT NOT NULL,
|
|
39
|
+
object_id TEXT REFERENCES entity(id),
|
|
40
|
+
object_literal TEXT,
|
|
41
|
+
fact_text TEXT NOT NULL,
|
|
42
|
+
|
|
43
|
+
valid_at INTEGER NOT NULL,
|
|
44
|
+
valid_to INTEGER,
|
|
45
|
+
superseded_by TEXT REFERENCES fact(id),
|
|
46
|
+
is_latest INTEGER NOT NULL DEFAULT 1,
|
|
47
|
+
|
|
48
|
+
ingested_at INTEGER NOT NULL,
|
|
49
|
+
expired_at INTEGER,
|
|
50
|
+
invalidated_by TEXT REFERENCES fact(id),
|
|
51
|
+
|
|
52
|
+
confidence REAL NOT NULL DEFAULT 1.0,
|
|
53
|
+
salience REAL NOT NULL DEFAULT 0.0,
|
|
54
|
+
last_access INTEGER,
|
|
55
|
+
access_count INTEGER NOT NULL DEFAULT 0,
|
|
56
|
+
is_inference INTEGER NOT NULL DEFAULT 0,
|
|
57
|
+
tier TEXT NOT NULL DEFAULT 'hot',
|
|
58
|
+
episode_id TEXT NOT NULL REFERENCES episode(id),
|
|
59
|
+
created_at INTEGER NOT NULL
|
|
60
|
+
);
|
|
61
|
+
CREATE INDEX IF NOT EXISTS ix_fact_ns_latest ON fact(namespace, is_latest);
|
|
62
|
+
CREATE INDEX IF NOT EXISTS ix_fact_slot ON fact(namespace, subject_id, predicate);
|
|
63
|
+
CREATE INDEX IF NOT EXISTS ix_fact_valid ON fact(namespace, valid_at, valid_to);
|
|
64
|
+
|
|
65
|
+
-- ── VECTOR INDEX (sqlite-vec vec0) ──────────────────────────────────
|
|
66
|
+
-- full 768 + 256-dim coarse Matryoshka vector. Stored FLOAT32 in Phase 0:
|
|
67
|
+
-- the spec targets int8 (size win, ~0.2pt quality cost per BET 1) but
|
|
68
|
+
-- sqlite-vec 0.1.9's int8 INSERT path is broken; flip to INT8[N] once fixed.
|
|
69
|
+
CREATE VIRTUAL TABLE IF NOT EXISTS fact_vec USING vec0(
|
|
70
|
+
fact_id TEXT PRIMARY KEY,
|
|
71
|
+
is_latest INTEGER,
|
|
72
|
+
tier TEXT,
|
|
73
|
+
namespace TEXT,
|
|
74
|
+
embedding FLOAT[{dim}],
|
|
75
|
+
embedding_256 FLOAT[{coarse_dim}]
|
|
76
|
+
);
|
|
77
|
+
|
|
78
|
+
-- ── LEXICAL INDEX (FTS5, external-content style holding its own text) ─
|
|
79
|
+
CREATE VIRTUAL TABLE IF NOT EXISTS fact_fts USING fts5(
|
|
80
|
+
fact_id UNINDEXED,
|
|
81
|
+
fact_text
|
|
82
|
+
);
|
|
83
|
+
"""
|
|
@@ -0,0 +1,301 @@
|
|
|
1
|
+
"""SqliteStore — the default Phase 0 store. One SQLite file per namespace.
|
|
2
|
+
|
|
3
|
+
Design-spec mapping:
|
|
4
|
+
- sqlite-vec `vec0` virtual table holds the int8 768-dim + 256-dim coarse vectors
|
|
5
|
+
(two-stage Matryoshka dense arm).
|
|
6
|
+
- FTS5 holds `fact_text` for the BM25 sparse arm.
|
|
7
|
+
- The relational `fact`/`entity`/`episode` tables hold the monotemporal spine.
|
|
8
|
+
- Per BET 4: this object backs ONE namespace (one file), turning SQLite's
|
|
9
|
+
single-writer limit into free write-isolation.
|
|
10
|
+
|
|
11
|
+
Quantization: vectors arrive L2-normalized float32 in [-1, 1]; we map to int8 by
|
|
12
|
+
scaling by 127 and rounding. vec0 does the distance math in int8 space.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
|
|
17
|
+
import sqlite3
|
|
18
|
+
from pathlib import Path
|
|
19
|
+
from typing import Optional, Sequence
|
|
20
|
+
|
|
21
|
+
import numpy as np
|
|
22
|
+
from sqlite_vec import serialize_float32
|
|
23
|
+
|
|
24
|
+
from ..types import Entity, Episode, Fact
|
|
25
|
+
from .base import Store
|
|
26
|
+
from .schema import SCHEMA_SQL
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _serialize(vec: np.ndarray) -> bytes:
|
|
30
|
+
"""L2-normalized float32 → vec0's float32 wire format.
|
|
31
|
+
|
|
32
|
+
NOTE (Phase 0 decision): we store float32, not int8. The schema/spec target int8
|
|
33
|
+
(size win, ~0.2pt quality cost per BET 1), but sqlite-vec 0.1.9's int8 *insert*
|
|
34
|
+
path is broken ("expected int8, but float32 provided" even for valid int8 blobs),
|
|
35
|
+
while the float32 path is solid. int8 is a documented future optimization to flip
|
|
36
|
+
once the upstream bug is fixed — it does not affect spine correctness.
|
|
37
|
+
"""
|
|
38
|
+
return serialize_float32(vec.astype(np.float32).tolist())
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class SqliteStore(Store):
|
|
42
|
+
def __init__(self, path: str | Path, *, dim: int = 768, coarse_dim: int = 256) -> None:
|
|
43
|
+
self.path = str(path)
|
|
44
|
+
self.dim = dim
|
|
45
|
+
self.coarse_dim = coarse_dim
|
|
46
|
+
self._db = self._connect()
|
|
47
|
+
self._init_schema()
|
|
48
|
+
|
|
49
|
+
# ── connection / schema ──
|
|
50
|
+
def _connect(self) -> sqlite3.Connection:
|
|
51
|
+
import sqlite_vec # lazy: only needed when a real store is opened
|
|
52
|
+
|
|
53
|
+
db = sqlite3.connect(self.path)
|
|
54
|
+
db.row_factory = sqlite3.Row
|
|
55
|
+
db.enable_load_extension(True)
|
|
56
|
+
sqlite_vec.load(db)
|
|
57
|
+
db.enable_load_extension(False)
|
|
58
|
+
db.execute("PRAGMA journal_mode=WAL") # better single-writer concurrency
|
|
59
|
+
db.execute("PRAGMA foreign_keys=ON")
|
|
60
|
+
return db
|
|
61
|
+
|
|
62
|
+
def _init_schema(self) -> None:
|
|
63
|
+
sql = SCHEMA_SQL.format(dim=self.dim, coarse_dim=self.coarse_dim)
|
|
64
|
+
self._db.executescript(sql)
|
|
65
|
+
self._db.commit()
|
|
66
|
+
|
|
67
|
+
# ── provenance ──
|
|
68
|
+
def add_episode(self, episode: Episode) -> None:
|
|
69
|
+
self._db.execute(
|
|
70
|
+
"INSERT INTO episode(id, namespace, raw, source, t_ref, created_at) "
|
|
71
|
+
"VALUES (?,?,?,?,?,?)",
|
|
72
|
+
(episode.id, episode.namespace, episode.raw, episode.source,
|
|
73
|
+
episode.t_ref, episode.created_at),
|
|
74
|
+
)
|
|
75
|
+
self._db.commit()
|
|
76
|
+
|
|
77
|
+
# ── entities ──
|
|
78
|
+
def upsert_entity(self, entity: Entity) -> Entity:
|
|
79
|
+
row = self._db.execute(
|
|
80
|
+
"SELECT * FROM entity WHERE namespace=? AND name=? AND IFNULL(type,'')=IFNULL(?,'')",
|
|
81
|
+
(entity.namespace, entity.name, entity.type),
|
|
82
|
+
).fetchone()
|
|
83
|
+
if row:
|
|
84
|
+
return _row_to_entity(row)
|
|
85
|
+
self._db.execute(
|
|
86
|
+
"INSERT INTO entity(id, namespace, name, type, summary, resolved_id, created_at) "
|
|
87
|
+
"VALUES (?,?,?,?,?,?,?)",
|
|
88
|
+
(entity.id, entity.namespace, entity.name, entity.type,
|
|
89
|
+
entity.summary, entity.resolved_id, entity.created_at),
|
|
90
|
+
)
|
|
91
|
+
self._db.commit()
|
|
92
|
+
return entity
|
|
93
|
+
|
|
94
|
+
def get_entity(self, entity_id: str) -> Optional[Entity]:
|
|
95
|
+
row = self._db.execute("SELECT * FROM entity WHERE id=?", (entity_id,)).fetchone()
|
|
96
|
+
return _row_to_entity(row) if row else None
|
|
97
|
+
|
|
98
|
+
# ── facts ──
|
|
99
|
+
def add_fact(self, fact: Fact, embedding: np.ndarray, embedding_256: np.ndarray) -> None:
|
|
100
|
+
db = self._db
|
|
101
|
+
db.execute(
|
|
102
|
+
"""INSERT INTO fact(
|
|
103
|
+
id, namespace, subject_id, predicate, object_id, object_literal, fact_text,
|
|
104
|
+
valid_at, valid_to, superseded_by, is_latest,
|
|
105
|
+
ingested_at, expired_at, invalidated_by,
|
|
106
|
+
confidence, salience, last_access, access_count, is_inference, tier,
|
|
107
|
+
episode_id, created_at)
|
|
108
|
+
VALUES (?,?,?,?,?,?,?, ?,?,?,?, ?,?,?, ?,?,?,?,?,?, ?,?)""",
|
|
109
|
+
(fact.id, fact.namespace, fact.subject_id, fact.predicate, fact.object_id,
|
|
110
|
+
fact.object_literal, fact.fact_text,
|
|
111
|
+
fact.valid_at, fact.valid_to, fact.superseded_by, fact.is_latest,
|
|
112
|
+
fact.ingested_at, fact.expired_at, fact.invalidated_by,
|
|
113
|
+
fact.confidence, fact.salience, fact.last_access, fact.access_count,
|
|
114
|
+
fact.is_inference, fact.tier,
|
|
115
|
+
fact.episode_id, fact.created_at),
|
|
116
|
+
)
|
|
117
|
+
db.execute(
|
|
118
|
+
"INSERT INTO fact_vec(fact_id, namespace, is_latest, tier, embedding, embedding_256) "
|
|
119
|
+
"VALUES (?,?,?,?,?,?)",
|
|
120
|
+
(fact.id, fact.namespace, fact.is_latest, fact.tier,
|
|
121
|
+
_serialize(embedding), _serialize(embedding_256)),
|
|
122
|
+
)
|
|
123
|
+
db.execute(
|
|
124
|
+
"INSERT INTO fact_fts(fact_id, fact_text) VALUES (?,?)",
|
|
125
|
+
(fact.id, fact.fact_text),
|
|
126
|
+
)
|
|
127
|
+
db.commit()
|
|
128
|
+
|
|
129
|
+
def supersede_fact(self, old_fact_id: str, new_fact_id: str, valid_to: int) -> None:
|
|
130
|
+
db = self._db
|
|
131
|
+
db.execute(
|
|
132
|
+
"UPDATE fact SET superseded_by=?, valid_to=?, is_latest=0 WHERE id=?",
|
|
133
|
+
(new_fact_id, valid_to, old_fact_id),
|
|
134
|
+
)
|
|
135
|
+
# keep the vec0 metadata filter column in sync so superseded facts drop out
|
|
136
|
+
db.execute("UPDATE fact_vec SET is_latest=0 WHERE fact_id=?", (old_fact_id,))
|
|
137
|
+
db.commit()
|
|
138
|
+
|
|
139
|
+
def get_fact(self, fact_id: str) -> Optional[Fact]:
|
|
140
|
+
row = self._db.execute("SELECT * FROM fact WHERE id=?", (fact_id,)).fetchone()
|
|
141
|
+
return _row_to_fact(row) if row else None
|
|
142
|
+
|
|
143
|
+
def find_latest_in_slot(self, subject_id: str, predicate: str) -> Sequence[Fact]:
|
|
144
|
+
rows = self._db.execute(
|
|
145
|
+
"SELECT * FROM fact WHERE subject_id=? AND predicate=? AND is_latest=1",
|
|
146
|
+
(subject_id, predicate),
|
|
147
|
+
).fetchall()
|
|
148
|
+
return [_row_to_fact(r) for r in rows]
|
|
149
|
+
|
|
150
|
+
# ── retrieval primitives ──
|
|
151
|
+
def dense_search(
|
|
152
|
+
self,
|
|
153
|
+
query_256: np.ndarray,
|
|
154
|
+
query_768: np.ndarray,
|
|
155
|
+
k: int,
|
|
156
|
+
*,
|
|
157
|
+
is_latest_only: bool = True,
|
|
158
|
+
as_of: Optional[int] = None,
|
|
159
|
+
) -> list[tuple[str, float]]:
|
|
160
|
+
"""Two-stage Matryoshka: coarse 256-dim KNN over a wider pool, then re-score
|
|
161
|
+
the survivors at full 768-dim. The coarse pool is k*COARSE_FACTOR wide so the
|
|
162
|
+
cheaper first stage doesn't drop gold before the precise re-score."""
|
|
163
|
+
COARSE_FACTOR = 8
|
|
164
|
+
coarse_k = max(k * COARSE_FACTOR, k)
|
|
165
|
+
|
|
166
|
+
latest_clause = "AND is_latest = 1" if is_latest_only else ""
|
|
167
|
+
# Stage 1: coarse KNN. vec0 KNN must use a single MATCH + LIMIT.
|
|
168
|
+
coarse_rows = self._db.execute(
|
|
169
|
+
f"""SELECT fact_id, distance FROM fact_vec
|
|
170
|
+
WHERE embedding_256 MATCH ? {latest_clause}
|
|
171
|
+
ORDER BY distance LIMIT ?""",
|
|
172
|
+
(_serialize(query_256), coarse_k),
|
|
173
|
+
).fetchall()
|
|
174
|
+
if not coarse_rows:
|
|
175
|
+
return []
|
|
176
|
+
|
|
177
|
+
candidate_ids = [r["fact_id"] for r in coarse_rows]
|
|
178
|
+
# Stage 2: re-score candidates at full 768-dim (exact distance, small set).
|
|
179
|
+
# vec0 doesn't take an IN-list on KNN, so we read the full vectors back and
|
|
180
|
+
# compute cosine here — exact, and the candidate set is tiny (coarse_k).
|
|
181
|
+
placeholders = ",".join("?" * len(candidate_ids))
|
|
182
|
+
vec_rows = self._db.execute(
|
|
183
|
+
f"SELECT fact_id, embedding FROM fact_vec WHERE fact_id IN ({placeholders})",
|
|
184
|
+
candidate_ids,
|
|
185
|
+
).fetchall()
|
|
186
|
+
|
|
187
|
+
q = query_768.astype(np.float32)
|
|
188
|
+
q = q / (np.linalg.norm(q) or 1.0)
|
|
189
|
+
scored: list[tuple[str, float]] = []
|
|
190
|
+
for vr in vec_rows:
|
|
191
|
+
stored = np.frombuffer(vr["embedding"], dtype=np.float32)
|
|
192
|
+
sn = np.linalg.norm(stored) or 1.0
|
|
193
|
+
cos = float(np.dot(q, stored) / sn)
|
|
194
|
+
scored.append((vr["fact_id"], 1.0 - cos)) # distance = 1 - cosine
|
|
195
|
+
|
|
196
|
+
if as_of is not None:
|
|
197
|
+
scored = self._apply_as_of(scored, as_of)
|
|
198
|
+
|
|
199
|
+
scored.sort(key=lambda x: x[1])
|
|
200
|
+
return scored[:k]
|
|
201
|
+
|
|
202
|
+
def _apply_as_of(self, scored: list[tuple[str, float]], as_of: int) -> list[tuple[str, float]]:
|
|
203
|
+
ids = [fid for fid, _ in scored]
|
|
204
|
+
if not ids:
|
|
205
|
+
return scored
|
|
206
|
+
placeholders = ",".join("?" * len(ids))
|
|
207
|
+
valid = {
|
|
208
|
+
r["id"]
|
|
209
|
+
for r in self._db.execute(
|
|
210
|
+
f"""SELECT id FROM fact WHERE id IN ({placeholders})
|
|
211
|
+
AND valid_at <= ? AND (valid_to IS NULL OR valid_to > ?)""",
|
|
212
|
+
(*ids, as_of, as_of),
|
|
213
|
+
).fetchall()
|
|
214
|
+
}
|
|
215
|
+
return [(fid, d) for fid, d in scored if fid in valid]
|
|
216
|
+
|
|
217
|
+
def sparse_search(
|
|
218
|
+
self, query_text: str, k: int, *, is_latest_only: bool = True,
|
|
219
|
+
as_of: Optional[int] = None,
|
|
220
|
+
) -> list[tuple[str, float]]:
|
|
221
|
+
# FTS5 BM25: lower bm25() is better, so we negate to "higher is better".
|
|
222
|
+
needs_row_check = is_latest_only or as_of is not None
|
|
223
|
+
rows = self._db.execute(
|
|
224
|
+
"""SELECT f.fact_id AS fact_id, bm25(fact_fts) AS score
|
|
225
|
+
FROM fact_fts f
|
|
226
|
+
WHERE fact_fts MATCH ?
|
|
227
|
+
ORDER BY score LIMIT ?""",
|
|
228
|
+
(_fts_query(query_text), k * (2 if needs_row_check else 1)),
|
|
229
|
+
).fetchall()
|
|
230
|
+
out: list[tuple[str, float]] = []
|
|
231
|
+
for r in rows:
|
|
232
|
+
if needs_row_check:
|
|
233
|
+
row = self._db.execute(
|
|
234
|
+
"SELECT is_latest, valid_at, valid_to FROM fact WHERE id=?",
|
|
235
|
+
(r["fact_id"],),
|
|
236
|
+
).fetchone()
|
|
237
|
+
if not row:
|
|
238
|
+
continue
|
|
239
|
+
if is_latest_only and not row["is_latest"]:
|
|
240
|
+
continue
|
|
241
|
+
if as_of is not None and not (
|
|
242
|
+
row["valid_at"] <= as_of
|
|
243
|
+
and (row["valid_to"] is None or row["valid_to"] > as_of)
|
|
244
|
+
):
|
|
245
|
+
continue
|
|
246
|
+
out.append((r["fact_id"], -float(r["score"])))
|
|
247
|
+
if len(out) >= k:
|
|
248
|
+
break
|
|
249
|
+
return out
|
|
250
|
+
|
|
251
|
+
def hydrate(self, fact_ids: Sequence[str]) -> dict[str, Fact]:
|
|
252
|
+
if not fact_ids:
|
|
253
|
+
return {}
|
|
254
|
+
placeholders = ",".join("?" * len(fact_ids))
|
|
255
|
+
rows = self._db.execute(
|
|
256
|
+
f"SELECT * FROM fact WHERE id IN ({placeholders})", list(fact_ids)
|
|
257
|
+
).fetchall()
|
|
258
|
+
return {r["id"]: _row_to_fact(r) for r in rows}
|
|
259
|
+
|
|
260
|
+
def touch(self, fact_id: str, when_ms: int) -> None:
|
|
261
|
+
self._db.execute(
|
|
262
|
+
"UPDATE fact SET last_access=?, access_count=access_count+1 WHERE id=?",
|
|
263
|
+
(when_ms, fact_id),
|
|
264
|
+
)
|
|
265
|
+
self._db.commit()
|
|
266
|
+
|
|
267
|
+
def close(self) -> None:
|
|
268
|
+
self._db.close()
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
# ── FTS query sanitization ──
|
|
272
|
+
def _fts_query(text: str) -> str:
|
|
273
|
+
"""Turn free text into a safe FTS5 OR-query of bare terms (avoids syntax errors
|
|
274
|
+
from punctuation/operators in user text)."""
|
|
275
|
+
terms = [t for t in "".join(c if c.isalnum() else " " for c in text).split() if t]
|
|
276
|
+
if not terms:
|
|
277
|
+
return '""'
|
|
278
|
+
return " OR ".join(terms)
|
|
279
|
+
|
|
280
|
+
|
|
281
|
+
# ── row → dataclass ──
|
|
282
|
+
def _row_to_entity(row: sqlite3.Row) -> Entity:
|
|
283
|
+
return Entity(
|
|
284
|
+
id=row["id"], namespace=row["namespace"], name=row["name"], type=row["type"],
|
|
285
|
+
summary=row["summary"], resolved_id=row["resolved_id"], created_at=row["created_at"],
|
|
286
|
+
)
|
|
287
|
+
|
|
288
|
+
|
|
289
|
+
def _row_to_fact(row: sqlite3.Row) -> Fact:
|
|
290
|
+
return Fact(
|
|
291
|
+
id=row["id"], namespace=row["namespace"], subject_id=row["subject_id"],
|
|
292
|
+
predicate=row["predicate"], object_id=row["object_id"],
|
|
293
|
+
object_literal=row["object_literal"], fact_text=row["fact_text"],
|
|
294
|
+
valid_at=row["valid_at"], valid_to=row["valid_to"],
|
|
295
|
+
superseded_by=row["superseded_by"], is_latest=row["is_latest"],
|
|
296
|
+
ingested_at=row["ingested_at"], expired_at=row["expired_at"],
|
|
297
|
+
invalidated_by=row["invalidated_by"], confidence=row["confidence"],
|
|
298
|
+
salience=row["salience"], last_access=row["last_access"],
|
|
299
|
+
access_count=row["access_count"], is_inference=row["is_inference"],
|
|
300
|
+
tier=row["tier"], episode_id=row["episode_id"], created_at=row["created_at"],
|
|
301
|
+
)
|