Teknium
243 lines
14 KiB
Python
243 lines
14 KiB
Python
"""Hybrid keyword/BM25 retrieval for the memory store: FTS5 candidates reranked with
|
||
Jaccard similarity and HRR vector similarity, trust-weighted (ported from KIK memory_agent.py)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from collections.abc import Callable
|
||
from datetime import datetime, timezone
|
||
from typing import TYPE_CHECKING
|
||
|
||
if TYPE_CHECKING:
|
||
from .store import MemoryStore
|
||
|
||
try:
|
||
from . import holographic as hrr
|
||
except ImportError:
|
||
import holographic as hrr # type: ignore[no-redef]
|
||
|
||
_FACT_COLUMNS = (
|
||
"fact_id, content, category, tags, trust_score, "
|
||
"retrieval_count, helpful_count, created_at, updated_at"
|
||
)
|
||
_ROLE_ENTITY, _ROLE_CONTENT = hrr.ROLE_ENTITY, hrr.ROLE_CONTENT
|
||
_PUNCT = ".,;:!?\"'()[]{}#@<>"
|
||
_FTS_OPERATORS = str.maketrans("", "", '"()*^:-+')
|
||
# Stopwords dropped before FTS5 OR-expansion: short English function words that
|
||
# carry no retrieval signal and force false-negative AND matches.
|
||
_FTS_STOPWORDS = frozenset("""
|
||
a about above after again all am an and any are as at be because been before being between both but by
|
||
can could did do does doing don down during each few for from further had has have having he her here hers
|
||
herself him himself his how i if in into is it its itself just me more most my myself no nor not now of off
|
||
on once only or other our ours ourselves out over own same she should so some such than that the their theirs
|
||
them themselves then there these they this those through to too under until up very was we were what when
|
||
where which while who whom why will with would you your yours yourself yourselves
|
||
""".split())
|
||
|
||
|
||
def _shift(sim: float) -> float:
|
||
"""Cosine similarity [-1, 1] -> [0, 1]."""
|
||
return (sim + 1.0) / 2.0
|
||
|
||
|
||
class FactRetriever:
|
||
"""Multi-strategy fact retrieval with trust-weighted scoring."""
|
||
|
||
def __init__(self, store: MemoryStore, temporal_decay_half_life: int = 0, # days, 0 = disabled
|
||
fts_weight: float = 0.4, jaccard_weight: float = 0.3, hrr_weight: float = 0.3, hrr_dim: int = 1024):
|
||
self.store, self.half_life, self.hrr_dim = store, temporal_decay_half_life, hrr_dim
|
||
if hrr_weight > 0 and not hrr._HAS_NUMPY: # redistribute weights without numpy
|
||
fts_weight, jaccard_weight, hrr_weight = 0.6, 0.4, 0.0
|
||
self.fts_weight, self.jaccard_weight, self.hrr_weight = fts_weight, jaccard_weight, hrr_weight
|
||
|
||
def _atom(self, word: str):
|
||
return hrr.encode_atom(word, self.hrr_dim)
|
||
|
||
def _phases(self, blob: bytes):
|
||
return hrr.bytes_to_phases(blob, dim=self.hrr_dim)
|
||
|
||
def search(self, query: str, category: str | None = None, min_trust: float = 0.3, limit: int = 10) -> list[dict]:
|
||
"""Hybrid search: FTS5 candidates (limit*3) → Jaccard + HRR rerank → trust weighting →
|
||
optional temporal decay 0.5^(age_days / half_life). Returns fact dicts with 'score', sorted desc."""
|
||
candidates = self._fts_candidates(query, category, min_trust, limit * 3)
|
||
if not candidates:
|
||
return []
|
||
query_tokens = self._tokenize(query)
|
||
# Query vector is loop-invariant; encode lazily on the first candidate that carries
|
||
# an HRR vector so stores whose hrr_vector was never backfilled don't pay for it.
|
||
query_vec = None
|
||
for fact in candidates:
|
||
jaccard = self._jaccard_similarity(query_tokens, self._tokenize(fact["content"]) | self._tokenize(fact.get("tags", "")))
|
||
hrr_sim = 0.5 # neutral
|
||
if self.hrr_weight > 0 and fact.get("hrr_vector"):
|
||
fact_vec = self._phases(fact["hrr_vector"])
|
||
if query_vec is None:
|
||
query_vec = hrr.encode_text(query, self.hrr_dim)
|
||
hrr_sim = _shift(hrr.similarity(query_vec, fact_vec))
|
||
relevance = self.fts_weight * fact.get("fts_rank", 0.0) + self.jaccard_weight * jaccard + self.hrr_weight * hrr_sim
|
||
fact["score"] = relevance * fact["trust_score"]
|
||
if self.half_life > 0:
|
||
fact["score"] *= self._temporal_decay(fact.get("updated_at") or fact.get("created_at"))
|
||
results = sorted(candidates, key=lambda x: x["score"], reverse=True)[:limit]
|
||
for fact in results:
|
||
fact.pop("hrr_vector", None) # callers expect JSON-serializable dicts
|
||
return results
|
||
|
||
def _vector_query(self, fallback: str, category: str | None, limit: int, make_sim: Callable[[], Callable]) -> list[dict]:
|
||
"""Rank every fact vector (optionally per category) with the sim fn built by make_sim(); FTS5 fallback
|
||
when no vectors exist. make_sim runs after the rows check so role atoms are encoded once, never for fallback."""
|
||
rows = self._vector_rows(category)
|
||
if not rows:
|
||
return self.search(fallback, category=category, limit=limit)
|
||
return self._rank_by_vector(rows, make_sim(), limit)
|
||
|
||
def probe(self, entity: str, category: str | None = None, limit: int = 10) -> list[dict]:
|
||
"""Compositional entity query: unbind bind(entity, ROLE_ENTITY) from the category bank
|
||
(or each fact vector) to find facts where the entity plays a structural role.
|
||
Not keyword search. Falls back to FTS5 without numpy."""
|
||
if not hrr._HAS_NUMPY:
|
||
return self.search(entity, category=category, limit=limit)
|
||
probe_key = hrr.bind(self._atom(entity.lower()), self._atom(_ROLE_ENTITY))
|
||
if category: # category bank first, then individual fact vectors
|
||
bank_row = self.store._conn.execute("SELECT vector FROM memory_banks WHERE bank_name = ?", (f"cat:{category}",)).fetchone()
|
||
if bank_row:
|
||
extracted = hrr.unbind(self._phases(bank_row["vector"]), probe_key)
|
||
return self._rank_by_vector(self._vector_rows(category), lambda _f, fact_vec: hrr.similarity(extracted, fact_vec), limit)
|
||
|
||
def make_sim():
|
||
role_content = self._atom(_ROLE_CONTENT)
|
||
# Does unbinding the probe key leave the fact's content signal?
|
||
return lambda fact, fact_vec: hrr.similarity(
|
||
hrr.unbind(fact_vec, probe_key), hrr.bind(hrr.encode_text(fact["content"], self.hrr_dim), role_content))
|
||
|
||
return self._vector_query(entity, category, limit, make_sim)
|
||
|
||
def related(self, entity: str, category: str | None = None, limit: int = 10) -> list[dict]:
|
||
"""Facts structurally connected to an entity (shared context), not just facts
|
||
*about* it as in probe. Falls back to FTS5 without numpy."""
|
||
if not hrr._HAS_NUMPY:
|
||
return self.search(entity, category=category, limit=limit)
|
||
entity_vec = self._atom(entity.lower()) # bare atom, not role-bound: ANY structural match
|
||
|
||
def make_sim():
|
||
roles = (self._atom(_ROLE_ENTITY), self._atom(_ROLE_CONTENT))
|
||
# A residual similar to ANY role vector means the entity plays a structural role.
|
||
return lambda _f, fact_vec: max(hrr.similarity(hrr.unbind(fact_vec, entity_vec), role) for role in roles)
|
||
|
||
return self._vector_query(entity, category, limit, make_sim)
|
||
|
||
def reason(self, entities: list[str], category: str | None = None, limit: int = 10) -> list[dict]:
|
||
"""Multi-entity compositional query (vector-space JOIN): facts where ALL entities
|
||
play structural roles. Falls back to FTS5 without numpy."""
|
||
if not hrr._HAS_NUMPY or not entities:
|
||
return self.search(" ".join(entities), category=category, limit=limit)
|
||
role_entity = self._atom(_ROLE_ENTITY)
|
||
probe_keys = [hrr.bind(self._atom(entity.lower()), role_entity) for entity in entities]
|
||
|
||
def make_sim():
|
||
role_content = self._atom(_ROLE_CONTENT)
|
||
# AND semantics via min: high only if EVERY entity is structurally present.
|
||
return lambda _f, fact_vec: min(hrr.similarity(hrr.unbind(fact_vec, key), role_content) for key in probe_keys)
|
||
|
||
return self._vector_query(" ".join(entities), category, limit, make_sim)
|
||
|
||
def contradict(self, category: str | None = None, threshold: float = 0.3, limit: int = 10) -> list[dict]:
|
||
"""Memory hygiene: pairs of facts that share entities (same subject) but have low
|
||
content-vector similarity (different claims). Empty without numpy."""
|
||
if not hrr._HAS_NUMPY:
|
||
return []
|
||
rows = self._vector_rows(category, columns="fact_id, content, category, tags, trust_score, created_at, updated_at, hrr_vector")
|
||
if len(rows) < 2:
|
||
return []
|
||
if len(rows) > 500: # O(n²) guard: only compare the most recently updated facts
|
||
rows = sorted(rows, key=lambda r: r["updated_at"] or r["created_at"], reverse=True)[:500]
|
||
facts = [] # (public dict, lower-cased entity names, phase vector)
|
||
for row in rows:
|
||
fact = dict(row)
|
||
entity_rows = self.store._conn.execute(
|
||
"SELECT e.name FROM entities e JOIN fact_entities fe ON fe.entity_id = e.entity_id WHERE fe.fact_id = ?",
|
||
(fact["fact_id"],),
|
||
).fetchall()
|
||
facts.append((fact, {r["name"].lower() for r in entity_rows}, self._phases(fact.pop("hrr_vector"))))
|
||
contradictions = []
|
||
for i, (f1, ents1, vec1) in enumerate(facts):
|
||
for f2, ents2, vec2 in facts[i + 1:]:
|
||
if not ents1 or not ents2:
|
||
continue
|
||
entity_overlap = len(ents1 & ents2) / len(ents1 | ents2)
|
||
if entity_overlap < 0.3:
|
||
continue # not enough shared subject to be contradictory
|
||
content_sim = hrr.similarity(vec1, vec2)
|
||
contradiction_score = entity_overlap * (1.0 - _shift(content_sim)) # high overlap + low similarity
|
||
if contradiction_score >= threshold:
|
||
contradictions.append({
|
||
"fact_a": f1, "fact_b": f2,
|
||
"entity_overlap": round(entity_overlap, 3),
|
||
"content_similarity": round(content_sim, 3),
|
||
"contradiction_score": round(contradiction_score, 3),
|
||
"shared_entities": sorted(ents1 & ents2),
|
||
})
|
||
return sorted(contradictions, key=lambda x: x["contradiction_score"], reverse=True)[:limit]
|
||
|
||
def _vector_rows(self, category: str | None, columns: str = _FACT_COLUMNS + ", hrr_vector") -> list:
|
||
"""All facts that carry an HRR vector, optionally filtered by category."""
|
||
where = "WHERE hrr_vector IS NOT NULL" + (" AND category = ?" if category else "")
|
||
return self.store._conn.execute(f"SELECT {columns} FROM facts {where}", [category] if category else []).fetchall()
|
||
|
||
def _rank_by_vector(self, rows: list, sim_fn: Callable[[dict, object], float], limit: int) -> list[dict]:
|
||
"""Score each row as (sim + 1) / 2 * trust_score (sim shifted to [0, 1]), sorted desc."""
|
||
scored = [dict(row) for row in rows]
|
||
for fact in scored:
|
||
fact["score"] = _shift(sim_fn(fact, self._phases(fact.pop("hrr_vector")))) * fact["trust_score"]
|
||
return sorted(scored, key=lambda x: x["score"], reverse=True)[:limit]
|
||
|
||
def _fts_candidates(self, query: str, category: str | None, min_trust: float, limit: int) -> list[dict]:
|
||
"""Raw FTS5 MATCH candidates with rank normalized to [0, 1] as 'fts_rank'."""
|
||
category_clause = "AND f.category = ? " if category else ""
|
||
params = [self._sanitize_fts_query(query)] + ([category] if category else []) + [min_trust, limit]
|
||
sql = ("SELECT f.*, facts_fts.rank as fts_rank_raw FROM facts_fts JOIN facts f ON f.fact_id = facts_fts.rowid "
|
||
f"WHERE facts_fts MATCH ? {category_clause}AND f.trust_score >= ? ORDER BY facts_fts.rank LIMIT ?")
|
||
try:
|
||
results = [dict(row) for row in self.store._conn.execute(sql, params).fetchall()]
|
||
except Exception:
|
||
return [] # FTS5 MATCH can fail on malformed queries
|
||
# FTS5 rank is negative (lower = better); normalize |rank| / max to [0, 1] (1e-6 floor avoids div by zero)
|
||
max_rank = max([abs(f["fts_rank_raw"]) for f in results] + [1e-6])
|
||
for fact in results:
|
||
fact["fts_rank"] = abs(fact.pop("fts_rank_raw")) / max_rank
|
||
return results
|
||
|
||
@staticmethod
|
||
def _tokenize(text: str) -> set[str]:
|
||
"""Lowercase whitespace tokens with surrounding punctuation stripped (no stemming)."""
|
||
return {c for c in (w.strip(_PUNCT) for w in text.lower().split()) if c} if text else set()
|
||
|
||
@staticmethod
|
||
def _sanitize_fts_query(query: str) -> str:
|
||
"""Natural-language query -> FTS5-safe OR expression of quoted tokens. FTS5 AND-joins a multi-word
|
||
MATCH by default, which tanks recall on prose: drop stopwords and <2-char tokens, strip FTS5 operator
|
||
chars, phrase-quote each survivor. If nothing survives, return the raw query (zero results, not a SQL error)."""
|
||
if not query:
|
||
return ""
|
||
tokens = [f'"{c}"' for c in (raw.strip(_PUNCT).translate(_FTS_OPERATORS) for raw in query.lower().split())
|
||
if len(c) >= 2 and c not in _FTS_STOPWORDS]
|
||
return " OR ".join(tokens) if tokens else query
|
||
|
||
@staticmethod
|
||
def _jaccard_similarity(set_a: set, set_b: set) -> float:
|
||
"""Jaccard similarity coefficient: |A ∩ B| / |A ∪ B|."""
|
||
return len(set_a & set_b) / len(set_a | set_b) if set_a and set_b else 0.0
|
||
|
||
def _temporal_decay(self, timestamp_str: str | None) -> float:
|
||
"""0.5^(age_days / half_life); 1.0 if disabled, missing, unparseable, or in the future."""
|
||
if not self.half_life or not timestamp_str:
|
||
return 1.0
|
||
try:
|
||
ts = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00")) if isinstance(timestamp_str, str) else timestamp_str
|
||
if ts.tzinfo is None:
|
||
ts = ts.replace(tzinfo=timezone.utc)
|
||
age_days = (datetime.now(timezone.utc) - ts).total_seconds() / 86400
|
||
return 1.0 if age_days < 0 else math.pow(0.5, age_days / self.half_life)
|
||
except (ValueError, TypeError):
|
||
return 1.0
|