"""Fact CRUD and search — vector + ILIKE fallback."""
import re
import uuid
from datetime import datetime, timezone
import asyncpg
import structlog
from navi.config import settings
from ._embeddings import _vector_to_str
log = structlog.get_logger()
_SEPARATORS = re.compile(r"[-_/\\.]")
_NOISE = re.compile(r"[^\w\s]", re.UNICODE)
_KEY_SPLIT = re.compile(r"[-_/.\s]+")
_AUTO_DUMP_THRESHOLD = 60
_VECTOR_CUTOFF_DISTANCE = 0.3
# Leading tokens that carry no meaning of their own: `preferred_language`,
# `ui_language` and `language` are one key, and the extractor writes all three.
_MERGE_PREFIX_TOKENS = frozenset({"preferred", "preference", "pref", "user", "my", "ui", "response"})
def _normalize_query(query: str) -> list[str]:
q = _SEPARATORS.sub(" ", query)
q = _NOISE.sub(" ", q)
return [t for t in q.lower().split() if len(t) > 1]
def _key_tokens(key: str) -> frozenset[str]:
"""The tokens of a fact key, with *leading* noise words dropped.
Leading only, and that is load-bearing: `ui_language` is the language key with a
noise prefix, but `dark_ui` is a name of its own and must not collapse into
`dark`. A rule that strips `ui` wherever it stands cannot tell those apart.
"""
tokens = [t for t in _KEY_SPLIT.split(key.strip().lower()) if len(t) > 1]
while tokens and tokens[0] in _MERGE_PREFIX_TOKENS:
tokens.pop(0)
return frozenset(tokens)
class FactMixin:
"""Fact storage operations.
Expected on the composite class:
_get_pool() -> asyncpg.Pool
_generate_embedding(text) -> list[float] | None
_has_pgvector() -> bool
"""
async def upsert_fact(
self,
category: str,
key: str,
value: str,
user_id: str | None = None,
source_session_id: str | None = None,
source: str = "conversation",
confidence: int = 70,
expires_at: datetime | None = None,
source_context: str = "",
dedupe: bool = False,
) -> None:
"""Write a fact, keyed on `(user_id, category, key)`.
With `dedupe`, the same fact already stored under a different spelling of
this key is hidden behind the row that survives — see `_merge_key_aliases`.
Off by default, so callers that own their keys (the `memory` tool) write
exactly the row they asked for.
"""
now = datetime.now(timezone.utc)
embedding = await self._generate_embedding(value)
vec_str = _vector_to_str(embedding) if embedding else None
pool = await self._get_pool()
fact_id = str(uuid.uuid4())
async with pool.acquire() as conn, conn.transaction():
# The dedupe path marks losers before the row they point at exists, so a
# failed insert would leave them hidden behind nothing. The plain path
# needs no boundary of its own, and taking one anyway costs nothing.
if dedupe:
await self._merge_key_aliases(conn, category, key, user_id, fact_id)
if vec_str:
await conn.execute(
"""INSERT INTO memory_facts
(id, user_id, category, key, value, created_at, updated_at, source_session_id,
embedding, source, confidence, expires_at, source_context)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::vector, $10, $11, $12, $13)
ON CONFLICT(user_id, category, key) DO UPDATE SET
value = EXCLUDED.value,
updated_at = EXCLUDED.updated_at,
source_session_id = EXCLUDED.source_session_id,
embedding = EXCLUDED.embedding,
source = EXCLUDED.source,
confidence = EXCLUDED.confidence,
expires_at = COALESCE(EXCLUDED.expires_at, memory_facts.expires_at),
source_context = EXCLUDED.source_context""",
fact_id, user_id, category, key, value, now, now,
source_session_id, vec_str, source, confidence, expires_at, source_context,
)
else:
await conn.execute(
"""INSERT INTO memory_facts
(id, user_id, category, key, value, created_at, updated_at, source_session_id,
source, confidence, expires_at, source_context)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
ON CONFLICT(user_id, category, key) DO UPDATE SET
value = EXCLUDED.value,
updated_at = EXCLUDED.updated_at,
source_session_id = EXCLUDED.source_session_id,
source = EXCLUDED.source,
confidence = EXCLUDED.confidence,
expires_at = COALESCE(EXCLUDED.expires_at, memory_facts.expires_at),
source_context = EXCLUDED.source_context""",
fact_id, user_id, category, key, value, now, now,
source_session_id, source, confidence, expires_at, source_context,
)
async def _merge_key_aliases(
self, conn, category: str, key: str, user_id: str | None, new_id: str
) -> None:
"""Hide the same fact already stored under a different spelling of this key.
The exact-key upsert cannot see `preferred_language` when `language` is what
is already there, which is how the base grew 11 spellings of the same handful
of keys. Compared on key tokens and **not** on value: `ui_language` and
`language` are one key while `dark` and `dark_ui` are two, and that is set
equality, not overlap. Values are deliberately left out of the decision: the
survivor is the row that already holds the key the writer used (else the
incoming one), and since the upsert writes the fresh value into it, the
fresher value is what remains readable either way.
Only *leading* noise tokens are dropped (see `_key_tokens`), so a trailing
one survives: `dark_ui` keeps its `ui` and stays apart from `dark`, while
`ui_dark` loses it and does not. That asymmetry is the price of reading `ui`
as noise in `ui_language` and as part of the name in `dark_ui`; no rule reads
it both ways.
Marked, never deleted: a loser keeps its own key and its own value, so the
merge is reversible with `UPDATE ... SET merged_into = NULL`.
"""
tokens = _key_tokens(key)
if not tokens:
return
# Hidden rows are read too: one of them may hold this exact key, and then the
# upsert below will land on it through ON CONFLICT rather than inserting the
# row we are about to point everything at.
rows = await conn.fetch(
"SELECT id, key, merged_into FROM memory_facts "
"WHERE user_id IS NOT DISTINCT FROM $1 AND category = $2 AND id <> $3",
user_id, category, new_id,
)
exact = [r["id"] for r in rows if r["key"] == key]
# Only live rows are worth hiding; the rest are already behind something.
aliases = [
r["id"] for r in rows
if r["key"] != key and r["merged_into"] is None and _key_tokens(r["key"]) == tokens
]
winner = exact[0] if exact else new_id
if exact:
# A row with this exact key is what the ON CONFLICT will update, so it
# becomes live again — and it has to be un-hidden for that, or the write
# lands in a row no read can see and the losers get pointed at a row that
# was never inserted.
await conn.execute(
"UPDATE memory_facts SET merged_into = NULL WHERE id = ANY($1::text[])", exact
)
if not aliases:
return
# The second clause re-points the losers of an earlier merge, so pointers
# never form a chain through a row that is itself hidden.
await conn.execute(
"UPDATE memory_facts SET merged_into = $2 "
"WHERE (id = ANY($1::text[]) OR merged_into = ANY($1::text[])) AND id <> $2",
aliases, winner,
)
log.info(
"memory.key_alias_merged",
category=category, key=key[:60], merged=len(aliases), winner=winner,
)
async def search_facts(self, query: str, user_id: str | None = None, limit: int = 15) -> list[dict]:
# 1. Try vector search if pgvector + embedding backend are available
if self._embedding_backend and await self._has_pgvector():
query_embedding = await self._generate_embedding(query)
vec_str = _vector_to_str(query_embedding) if query_embedding else None
if vec_str:
try:
pool = await self._get_pool()
async with pool.acquire() as conn:
if user_id is None:
rows = await conn.fetch(
"""SELECT id, category, key, value, updated_at,
source, confidence, expires_at, source_context, merged_into,
embedding <=> $1::vector AS distance
FROM memory_facts
WHERE (expires_at IS NULL OR expires_at > now())
AND merged_into IS NULL
AND user_id IS NULL
ORDER BY embedding <=> $1::vector
LIMIT $2""",
vec_str, limit,
)
else:
rows = await conn.fetch(
"""SELECT id, category, key, value, updated_at,
source, confidence, expires_at, source_context, merged_into,
embedding <=> $1::vector AS distance
FROM memory_facts
WHERE (expires_at IS NULL OR expires_at > now())
AND merged_into IS NULL
AND user_id = $3
ORDER BY embedding <=> $1::vector
LIMIT $2""",
vec_str, limit, user_id,
)
results = [
_row_to_dict(r)
for r in rows
if r["distance"] < _VECTOR_CUTOFF_DISTANCE
]
if results:
log.debug("memory.vector_search", hits=len(results), query=query[:40])
return results
except Exception:
log.warning("memory.vector_search_failed", query=query[:40], exc_info=True)
else:
log.warning("memory.vector_fallback_to_ilike", reason="embed_failed", query=query[:40])
else:
reason = "no_embedding_backend" if not self._embedding_backend else "no_pgvector"
log.warning("memory.vector_fallback_to_ilike", reason=reason, query=query[:40])
# 2. Text fallback (original ILIKE logic)
terms = _normalize_query(query)
if not terms:
return await self.get_all_facts(user_id=user_id, limit=limit)
if await self.fact_count(user_id=user_id) <= _AUTO_DUMP_THRESHOLD:
return await self.get_all_facts(user_id=user_id, limit=limit)
base_params: list = []
term_conds: list[str] = []
for term in terms:
like = f"%{term}%"
i = len(base_params) + 1
term_conds.append(
f"(category ILIKE ${i} OR key ILIKE ${i + 1} OR value ILIKE ${i + 2})"
)
base_params.extend([like, like, like])
user_filter = "user_id IS NULL" if user_id is None else f"user_id = ${len(base_params) + 1}"
if user_id is not None:
base_params.append(user_id)
# This branch used to be the only read path that served expired facts, and
# hidden merge losers used to be served by all of them.
user_filter += " AND merged_into IS NULL AND (expires_at IS NULL OR expires_at > now())"
limit_idx = len(base_params) + 1
and_where = " AND ".join(term_conds)
or_where = " OR ".join(term_conds)
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
f"SELECT id, category, key, value, updated_at, source, confidence, "
f"expires_at, source_context, merged_into FROM memory_facts "
f"WHERE {user_filter} AND {and_where} ORDER BY updated_at DESC LIMIT ${limit_idx}",
*base_params, limit,
)
if rows:
return [_row_to_dict(r) for r in rows]
score_expr = " + ".join(
f"CASE WHEN {c} THEN 1 ELSE 0 END" for c in term_conds
)
rows = await conn.fetch(
f"SELECT id, category, key, value, updated_at, source, confidence, "
f"expires_at, source_context, merged_into, ({score_expr}) AS score "
f"FROM memory_facts WHERE {user_filter} AND ({or_where}) "
f"ORDER BY score DESC, updated_at DESC LIMIT ${limit_idx}",
*base_params, limit,
)
return [_row_to_dict(r) for r in rows]
async def delete_fact(self, key: str, category: str | None = None, user_id: str | None = None) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn, conn.transaction():
await self._delete_merged_dependents(conn, key, category, user_id)
if category:
if user_id is None:
result = await conn.execute(
"DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND LOWER(category)=LOWER($2) AND user_id IS NULL",
key, category,
)
else:
result = await conn.execute(
"DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND LOWER(category)=LOWER($2) AND user_id = $3",
key, category, user_id,
)
else:
if user_id is None:
result = await conn.execute(
"DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND user_id IS NULL", key
)
else:
result = await conn.execute(
"DELETE FROM memory_facts WHERE LOWER(key)=LOWER($1) AND user_id = $2", key, user_id
)
return int(result.split()[1])
async def _delete_merged_dependents(
self, conn, key: str, category: str | None, user_id: str | None
) -> None:
"""Forget drops the hidden aliases of the keys it is about to delete.
A merge loser points at its winner; deleting the winner on its own would
leave those rows hidden with nothing left to point at and no way back to
them. They are deleted, not un-hidden — the user asked for the key gone.
"""
conditions = ["LOWER(key) = LOWER($1)"]
args: list = [key]
if category:
args.append(category)
conditions.append(f"LOWER(category) = LOWER(${len(args)})")
if user_id is None:
conditions.append("user_id IS NULL")
else:
args.append(user_id)
conditions.append(f"user_id = ${len(args)}")
await conn.execute(
"DELETE FROM memory_facts WHERE merged_into IN "
f"(SELECT id FROM memory_facts WHERE {' AND '.join(conditions)})",
*args,
)
async def get_all_facts(
self,
user_id: str | None = None,
limit: int | None = None,
offset: int = 0,
search: str | None = None,
sort_by: str = "category",
sort_order: str = "desc",
all_users: bool = False,
include_merged: bool = False,
) -> list[dict]:
q = (
"SELECT id, category, key, value, updated_at, source, confidence, "
"expires_at, source_context, merged_into FROM memory_facts "
)
params: list = []
param_idx = 0
def add_param(value):
nonlocal param_idx
param_idx += 1
params.append(value)
return f"${param_idx}"
conditions = []
if not all_users:
if user_id is None:
conditions.append("user_id IS NULL")
else:
conditions.append(f"user_id = {add_param(user_id)}")
if not include_merged:
conditions.append("merged_into IS NULL")
if search:
like = f"%{search}%"
conditions.append(
f"(key ILIKE {add_param(like)} OR value ILIKE {add_param(like)} OR category ILIKE {add_param(like)})"
)
if conditions:
q += "WHERE " + " AND ".join(conditions) + " "
allowed_cols = {"category", "key", "updated_at", "confidence", "source"}
col = sort_by if sort_by in allowed_cols else "category"
order = "DESC" if sort_order == "desc" else "ASC"
if col == "category":
q += f"ORDER BY {col} {order}, updated_at DESC"
else:
q += f"ORDER BY {col} {order}"
if limit is not None:
q += f" LIMIT {add_param(limit)}"
if offset:
q += f" OFFSET {add_param(offset)}"
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(q, *params)
return [_row_to_dict(r) for r in rows]
async def fact_count(
self,
user_id: str | None = None,
all_users: bool = False,
search: str | None = None,
include_merged: bool = False,
) -> int:
q = "SELECT COUNT(*) FROM memory_facts"
params: list = []
conditions = []
if not all_users:
if user_id is None:
conditions.append("user_id IS NULL")
else:
conditions.append(f"user_id = ${len(params) + 1}")
params.append(user_id)
if not include_merged:
conditions.append("merged_into IS NULL")
if search:
like = f"%{search}%"
conditions.append(
f"(key ILIKE ${len(params) + 1} OR value ILIKE ${len(params) + 2} OR category ILIKE ${len(params) + 3})"
)
params.extend([like, like, like])
if conditions:
q += " WHERE " + " AND ".join(conditions)
pool = await self._get_pool()
async with pool.acquire() as conn:
return await conn.fetchval(q, *params) or 0
async def get_categories(self, user_id: str | None = None) -> list[str]:
pool = await self._get_pool()
async with pool.acquire() as conn:
if user_id is None:
rows = await conn.fetch(
"SELECT DISTINCT category FROM memory_facts "
"WHERE user_id IS NULL AND merged_into IS NULL ORDER BY category"
)
else:
rows = await conn.fetch(
"SELECT DISTINCT category FROM memory_facts "
"WHERE user_id = $1 AND merged_into IS NULL ORDER BY category",
user_id,
)
return [r["category"] for r in rows]
def _row_to_dict(row: asyncpg.Record) -> dict:
val = row["updated_at"]
updated_at = val.isoformat() if hasattr(val, "isoformat") else str(val)
expires_at = row.get("expires_at")
expires_str = expires_at.isoformat() if hasattr(expires_at, "isoformat") else None
return {
"id": row["id"],
"category": row["category"],
"key": row["key"],
"value": row["value"],
"updated_at": updated_at,
"source": row.get("source", "conversation"),
"confidence": row.get("confidence", 70),
"expires_at": expires_str,
"source_context": row.get("source_context", ""),
"merged_into": row.get("merged_into"),
}
