6 Commits

Author SHA1 Message Date
Joseph Doherty 013b563f21 merge: T93 cross-chat search service 2026-04-27 02:32:53 -04:00
Joseph Doherty 62d5cdd826 merge: T92 pure-Python cosine vector search service 2026-04-27 02:32:53 -04:00
Joseph Doherty a25c166174 merge: T91 embedding generation service (pseudo-embedding) 2026-04-27 02:32:53 -04:00
Joseph Doherty 8f66e1123a feat: cross-chat search service (T93) 2026-04-27 02:31:31 -04:00
Joseph Doherty caa17b4174 feat: embedding generation service (Phase 4 pseudo-embedding) (T91) 2026-04-27 02:31:07 -04:00
Joseph Doherty c7cb0eb01e feat: pure-Python cosine vector search service (T92) 2026-04-27 02:31:06 -04:00
6 changed files with 750 additions and 0 deletions
+75
View File
@@ -0,0 +1,75 @@
"""Cross-chat search service (T93, Phase 4).
FTS5-based search across ALL owners and ALL chats. Used by the
top-bar search UX (T100) for "where did I last see this character
mention X?" queries. NO witness filter -- this is intentionally a
power-user surface that surfaces memories across POVs.
Mirrors the FTS5 access pattern of ``chat.state.memory.search_memories``
but drops both the ``owner_id = ?`` and the per-witness predicates so a
single query can sweep every chat in the database. The composite
re-rank is also dropped: callers want raw BM25 ordering for the
"highest match strength wins" semantics expected of a global search box.
"""
from __future__ import annotations
from sqlite3 import Connection
def search_all_memories(
conn: Connection,
*,
query: str,
k: int = 20,
) -> list[dict]:
"""Search FTS5 across all owners and chats.
Returns rows with ``{memory_id, owner_id, chat_id, scene_id,
pov_summary, significance, ts, fts_rank}``, sorted by FTS5 BM25
rank ascending (lower rank = stronger match, surfaced first).
The ``memories`` table has no ``ts`` column; we expose ``created_at``
(the projector-side row insertion timestamp) under that key so the
UI does not have to know the storage name.
An empty / whitespace-only ``query`` short-circuits to ``[]`` to
avoid an FTS5 ``MATCH ''`` syntax error and to keep the top-bar
"no input yet" state from triggering a full-table scan.
"""
if not query or not query.strip():
return []
# FTS5 MATCH against the same ``memories_fts`` virtual table that
# backs ``state.memory.search_memories``; the JOIN pulls metadata
# from the content table because the FTS index only stores
# ``pov_summary``. ORDER BY rank ASC because BM25 in FTS5 returns
# negative scores where lower is better.
rows = conn.execute(
"SELECT m.id, m.owner_id, m.chat_id, m.scene_id, "
" m.pov_summary, m.significance, m.created_at, "
" memories_fts.rank "
"FROM memories_fts "
"JOIN memories m ON m.id = memories_fts.rowid "
"WHERE memories_fts MATCH ? "
"ORDER BY memories_fts.rank ASC "
"LIMIT ?",
(query.strip(), k),
).fetchall()
return [
{
"memory_id": r[0],
"owner_id": r[1],
"chat_id": r[2],
"scene_id": r[3],
"pov_summary": r[4],
"significance": r[5],
"ts": r[6],
"fts_rank": r[7],
}
for r in rows
]
__all__ = ["search_all_memories"]
+108
View File
@@ -0,0 +1,108 @@
"""Embedding generation service (T91, Phase 4).
Wraps the embedding API call. For Phase 4's first cut we ship a
deterministic local pseudo-embedding (hash-derived) so the vector
retrieval pipeline can land without an external embedding endpoint
or heavy local dependency. Phase 4.5+ swaps to a real model — the
EmbeddingResult shape stays the same, only the generator changes.
"""
from __future__ import annotations
import hashlib
import math
import struct
from pydantic import BaseModel
from chat.llm.client import LLMClient
DEFAULT_EMBEDDING_DIM = 384
DEFAULT_EMBEDDING_MODEL = "pseudo-sha256-384"
FALLBACK_EMBEDDING_MODEL = "fallback"
class EmbeddingResult(BaseModel):
vector: list[float]
model: str
dim: int
def _pseudo_embed(text: str, dim: int = DEFAULT_EMBEDDING_DIM) -> list[float]:
"""Deterministic pseudo-embedding for Phase 4 first cut.
Hashes the text with SHA-256, then expands by re-hashing each
successive block with the previous block + a counter — this gives
``dim * 4`` bytes of fresh entropy per input rather than naively
repeating the 32-byte digest (which would collapse the vector onto
only 8 unique floats and make distinct inputs cosine-similar).
Bytes are unpacked as little-endian int32s and rescaled to [-1, 1]
so we sidestep the float32 NaN/denormal values that ``struct.unpack
'f'`` would otherwise produce on raw hash bytes. The result is
unit-normalized so cosine similarity reduces to a dot product.
NOT semantically meaningful — just consistent for testing the
pipeline. Phase 4.5 should swap to a real embedding model.
"""
needed = dim * 4 # 4 bytes per int32
seed = text.encode("utf-8")
chunks: list[bytes] = []
counter = 0
while sum(len(c) for c in chunks) < needed:
block = hashlib.sha256(seed + counter.to_bytes(4, "big")).digest()
chunks.append(block)
counter += 1
full = b"".join(chunks)[:needed]
ints = struct.unpack(f"<{dim}i", full)
# Map int32 to roughly [-1, 1] — exact bound doesn't matter since we
# normalize, but keeps values numerically tame.
raw = [x / 2147483648.0 for x in ints]
norm = math.sqrt(sum(x * x for x in raw)) or 1.0
return [x / norm for x in raw]
async def generate_embedding(
client: LLMClient,
*,
text: str,
model: str = DEFAULT_EMBEDDING_MODEL,
dim: int = DEFAULT_EMBEDDING_DIM,
timeout_s: float = 30.0,
) -> EmbeddingResult:
"""Generate an embedding for the given text.
Phase 4 default uses a deterministic local pseudo-embedding. If
the LLMClient grows an ``embed(...)`` method in Phase 4.5, this
wrapper will route to it when ``model != "pseudo-sha256-384"``.
Falls back to a zero vector with ``model="fallback"`` on any
failure (callers detect the sentinel and skip indexing). For the
pseudo path, failure is structurally impossible — it's pure local
computation.
"""
if not text or not text.strip():
# Empty input — return fallback so caller doesn't index empty rows.
return EmbeddingResult(
vector=[0.0] * dim, model=FALLBACK_EMBEDDING_MODEL, dim=dim
)
if model == DEFAULT_EMBEDDING_MODEL:
# Pure-local pseudo path — no LLMClient call.
return EmbeddingResult(vector=_pseudo_embed(text, dim), model=model, dim=dim)
# Future: real embedding via client.embed(...). Phase 4.5 work.
# For Phase 4, any non-default model falls through to fallback.
return EmbeddingResult(
vector=[0.0] * dim, model=FALLBACK_EMBEDDING_MODEL, dim=dim
)
__all__ = [
"DEFAULT_EMBEDDING_DIM",
"DEFAULT_EMBEDDING_MODEL",
"FALLBACK_EMBEDDING_MODEL",
"EmbeddingResult",
"generate_embedding",
]
+79
View File
@@ -0,0 +1,79 @@
"""Vector search service (T92, Phase 4).
Pure-Python cosine similarity over the embeddings table. Phase 4 ships
this without sqlite-vec because the host Python build doesn't support
loadable extensions. For single-user scale (< few thousand memories
per owner), iterating in Python is sub-millisecond.
Phase 4.5+ may swap to sqlite-vec when the host Python supports
enable_load_extension; the public API stays stable.
"""
from __future__ import annotations
import math
from sqlite3 import Connection
from chat.state.embeddings import list_embeddings_for_owner
_VALID_WITNESS_ROLES = {"you", "host", "guest"}
def _cosine_similarity(a: list[float], b: list[float]) -> float:
"""Cosine similarity. Assumes both vectors are non-zero."""
if len(a) != len(b):
return 0.0
dot = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a)) or 1.0
norm_b = math.sqrt(sum(x * x for x in b)) or 1.0
return dot / (norm_a * norm_b)
def vector_search(
conn: Connection,
*,
owner_id: str,
witness_role: str, # "you" | "host" | "guest"
query_vector: list[float],
k: int = 4,
) -> list[dict]:
"""Return top-K memories by cosine similarity to query_vector,
witness-filtered for the viewer's POV. Returns rows with
{memory_id, pov_summary, significance, score} sorted by score
DESC. Empty list if no embeddings indexed for this owner.
"""
if witness_role not in _VALID_WITNESS_ROLES:
raise ValueError(
f"witness_role must be one of {_VALID_WITNESS_ROLES}, got {witness_role!r}"
)
rows = list_embeddings_for_owner(conn, owner_id)
if not rows:
return []
# Witness-filter by the requesting role.
witness_key = f"witness_{witness_role}"
filtered = [r for r in rows if r.get(witness_key) == 1]
if not filtered:
return []
scored: list[tuple[float, dict]] = []
for row in filtered:
score = _cosine_similarity(query_vector, row["vector"])
scored.append(
(
score,
{
"memory_id": row["memory_id"],
"pov_summary": row["pov_summary"],
"significance": row["significance"],
"score": score,
},
)
)
scored.sort(key=lambda t: t[0], reverse=True)
return [item for _, item in scored[:k]]
__all__ = ["vector_search"]
+155
View File
@@ -0,0 +1,155 @@
"""T93 (Phase 4): cross-chat FTS5 search across all owners and chats.
Verifies that ``chat.services.cross_chat_search.search_all_memories``:
* surfaces matches across multiple owner_ids (the per-owner restriction
used by ``state.memory.search_memories`` is intentionally absent),
* applies no witness filter (admin/power-user surface),
* orders results by FTS5 BM25 rank (lower = stronger match, surfaced
first), and
* honours the ``k`` LIMIT and the empty-query fast-path.
"""
from __future__ import annotations
from chat.db.connection import open_db
from chat.db.migrate import apply_migrations
from chat.eventlog.log import append_event
from chat.eventlog.projector import project
from chat.services.cross_chat_search import search_all_memories
import chat.state.memory # noqa: F401 (registers memory_written handler)
def _seed(db, *, memory_specs):
"""Apply migrations + project a list of memory_written events."""
apply_migrations(db)
with open_db(db) as conn:
for spec in memory_specs:
payload = {
"owner_id": spec.get("owner_id", "bot_a"),
"chat_id": spec.get("chat_id", "chat_bot_a"),
"pov_summary": spec["pov_summary"],
"witness_you": spec.get("witness_you", 1),
"witness_host": spec.get("witness_host", 1),
"witness_guest": spec.get("witness_guest", 0),
"source": "direct",
"reliability": 1.0,
"significance": spec.get("significance", 1),
"pinned": 0,
"auto_pinned": 0,
}
append_event(conn, kind="memory_written", payload=payload)
project(conn)
def test_search_all_memories_returns_matches_across_owners(tmp_path):
"""Cross-owner: a single query must surface memories from every owner.
The per-owner ``owner_id = ?`` predicate that ``search_memories`` uses
is intentionally absent here, so a "rabbit" memory under ``bot_a`` and
one under ``bot_b`` should both come back from a single call.
"""
db = tmp_path / "t.db"
_seed(
db,
memory_specs=[
{
"owner_id": "bot_a",
"chat_id": "chat_bot_a",
"pov_summary": "the rabbit darted into the brambles",
},
{
"owner_id": "bot_b",
"chat_id": "chat_bot_b",
"pov_summary": "a white rabbit watched from the hedge",
},
# Distractor: must not appear for "rabbit".
{
"owner_id": "bot_a",
"chat_id": "chat_bot_a",
"pov_summary": "the kettle whistled",
},
],
)
with open_db(db) as conn:
out = search_all_memories(conn, query="rabbit")
owners = {row["owner_id"] for row in out}
assert owners == {"bot_a", "bot_b"}
assert len(out) == 2
# Returned shape contract.
for row in out:
assert set(row.keys()) >= {
"memory_id",
"owner_id",
"chat_id",
"scene_id",
"pov_summary",
"significance",
"ts",
"fts_rank",
}
def test_search_all_memories_orders_by_fts_rank(tmp_path):
"""Stronger BM25 match must come first (rank ASC = lower is better)."""
db = tmp_path / "t.db"
_seed(
db,
memory_specs=[
# Single occurrence -> weaker BM25 score.
{
"owner_id": "bot_a",
"chat_id": "chat_bot_a",
"pov_summary": "a rabbit appeared",
},
# Triple occurrence in a short row -> stronger BM25 score.
{
"owner_id": "bot_b",
"chat_id": "chat_bot_b",
"pov_summary": "rabbit rabbit rabbit",
},
],
)
with open_db(db) as conn:
out = search_all_memories(conn, query="rabbit", k=5)
assert len(out) == 2
# Stronger match first; fts_rank monotonically non-decreasing
# (lower-is-better, so ASC).
assert out[0]["pov_summary"] == "rabbit rabbit rabbit"
assert out[0]["fts_rank"] <= out[1]["fts_rank"]
def test_search_all_memories_respects_k_limit(tmp_path):
"""LIMIT ? must cap result count even when more matches exist."""
db = tmp_path / "t.db"
_seed(
db,
memory_specs=[
{
"owner_id": f"bot_{i}",
"chat_id": f"chat_{i}",
"pov_summary": f"rabbit sighting number {i}",
}
for i in range(10)
],
)
with open_db(db) as conn:
out = search_all_memories(conn, query="rabbit", k=3)
assert len(out) == 3
def test_search_all_memories_empty_query_returns_empty(tmp_path):
"""Empty / whitespace-only query must short-circuit to []."""
db = tmp_path / "t.db"
_seed(
db,
memory_specs=[
{
"owner_id": "bot_a",
"chat_id": "chat_bot_a",
"pov_summary": "the rabbit darted into the brambles",
},
],
)
with open_db(db) as conn:
assert search_all_memories(conn, query="") == []
assert search_all_memories(conn, query=" ") == []
+91
View File
@@ -0,0 +1,91 @@
"""Tests for the embedding generation service (T91, Phase 4).
Phase 4's first cut ships a deterministic local pseudo-embedding so the
vector retrieval pipeline can land without an external embeddings API
or a heavy local model dependency. These tests pin the contract:
* the result has the right shape (vector length, ``dim`` metadata),
* the default ``model`` string is reported back unchanged,
* output is byte-identical for the same input (deterministic),
* distinct inputs produce distinct vectors (so cosine actually
discriminates),
* empty / whitespace-only input collapses to the ``"fallback"`` sentinel
with a zero vector — callers detect this and skip indexing,
* the vector is unit-normalized so cosine similarity behaves.
The pseudo path doesn't touch the LLMClient, so we pass an empty
``MockLLMClient`` — any accidental call into it would raise
``IndexError`` and surface as a regression.
"""
from __future__ import annotations
import math
import pytest
from chat.llm.mock import MockLLMClient
from chat.services.embeddings import (
DEFAULT_EMBEDDING_DIM,
DEFAULT_EMBEDDING_MODEL,
FALLBACK_EMBEDDING_MODEL,
EmbeddingResult,
generate_embedding,
)
def _client() -> MockLLMClient:
# Pseudo path never calls the client — empty canned list ensures any
# accidental call raises and surfaces the regression loudly.
return MockLLMClient(canned=[])
@pytest.mark.asyncio
async def test_generate_embedding_returns_vector_of_correct_dim():
result = await generate_embedding(_client(), text="hello")
assert isinstance(result, EmbeddingResult)
assert isinstance(result.vector, list)
assert len(result.vector) == DEFAULT_EMBEDDING_DIM == 384
assert result.dim == 384
assert all(isinstance(x, float) for x in result.vector)
@pytest.mark.asyncio
async def test_generate_embedding_returns_correct_model_metadata():
result = await generate_embedding(_client(), text="hello")
assert result.model == DEFAULT_EMBEDDING_MODEL == "pseudo-sha256-384"
@pytest.mark.asyncio
async def test_generate_embedding_is_deterministic():
a = await generate_embedding(_client(), text="hello world")
b = await generate_embedding(_client(), text="hello world")
assert a.vector == b.vector
@pytest.mark.asyncio
async def test_generate_embedding_distinct_text_produces_distinct_vectors():
a = await generate_embedding(_client(), text="hello world")
b = await generate_embedding(_client(), text="totally different content")
assert a.vector != b.vector
# Sanity-check cosine similarity — both vectors are unit-normalized,
# so this reduces to a plain dot product.
cosine = sum(x * y for x, y in zip(a.vector, b.vector))
assert cosine < 0.99
@pytest.mark.asyncio
async def test_generate_embedding_empty_text_returns_fallback():
for empty in ("", " ", "\n\t"):
result = await generate_embedding(_client(), text=empty)
assert result.model == FALLBACK_EMBEDDING_MODEL == "fallback"
assert result.dim == DEFAULT_EMBEDDING_DIM
assert len(result.vector) == DEFAULT_EMBEDDING_DIM
assert all(x == 0.0 for x in result.vector)
@pytest.mark.asyncio
async def test_generate_embedding_unit_normalized():
result = await generate_embedding(_client(), text="some non-empty text")
norm_sq = sum(x * x for x in result.vector)
assert math.isclose(norm_sq, 1.0, abs_tol=1e-6)
+242
View File
@@ -0,0 +1,242 @@
from __future__ import annotations
import pytest
from chat.db.connection import open_db
from chat.db.migrate import apply_migrations
from chat.eventlog.log import append_event
from chat.eventlog.projector import project
import chat.state.memory # registers memory_written handler
import chat.state.embeddings # registers embedding handlers
from chat.services.vector_search import vector_search
def _base_memory(**overrides):
payload = {
"owner_id": "bot_a",
"chat_id": "chat_bot_a",
"scene_id": 1,
"pov_summary": "She laughed at his joke about owls.",
"witness_you": 1,
"witness_host": 1,
"witness_guest": 0,
"chat_clock_at": "2026-04-26T10:00:00",
"source": "direct",
"reliability": 1.0,
"significance": 1,
"pinned": 0,
"auto_pinned": 0,
}
payload.update(overrides)
return payload
def _one_hot(dim: int, idx: int) -> list[float]:
"""Return a one-hot vector of length ``dim`` with 1.0 at ``idx``."""
v = [0.0] * dim
v[idx] = 1.0
return v
def _seed_memory_with_embedding(
conn,
*,
owner_id: str,
pov_summary: str,
vector: list[float],
significance: int = 1,
witness_you: int = 1,
witness_host: int = 1,
witness_guest: int = 0,
model: str = "test-model",
) -> int:
append_event(
conn,
kind="memory_written",
payload=_base_memory(
owner_id=owner_id,
pov_summary=pov_summary,
significance=significance,
witness_you=witness_you,
witness_host=witness_host,
witness_guest=witness_guest,
),
)
project(conn)
memory_id = conn.execute(
"SELECT id FROM memories WHERE pov_summary = ?", (pov_summary,)
).fetchone()[0]
append_event(
conn,
kind="embedding_indexed",
payload={
"memory_id": memory_id,
"vector": vector,
"model": model,
"dim": len(vector),
},
)
project(conn)
return memory_id
def test_vector_search_returns_nearest_neighbors(tmp_path):
db = tmp_path / "t.db"
apply_migrations(db)
with open_db(db) as conn:
dim = 8
ids = []
for i in range(5):
mid = _seed_memory_with_embedding(
conn,
owner_id="bot_a",
pov_summary=f"Memory {i}.",
vector=_one_hot(dim, i),
)
ids.append(mid)
# Query close to memory index 3 (one-hot at position 3, plus tiny noise).
query = _one_hot(dim, 3)
query[2] = 0.01
results = vector_search(
conn,
owner_id="bot_a",
witness_role="you",
query_vector=query,
k=3,
)
assert len(results) == 3
# Top-1 must be memory at index 3.
assert results[0]["memory_id"] == ids[3]
assert results[0]["pov_summary"] == "Memory 3."
# Score for the near-perfect match should be very close to 1.0.
assert results[0]["score"] > 0.99
# Results sorted by score DESC.
scores = [r["score"] for r in results]
assert scores == sorted(scores, reverse=True)
# Second place should be memory index 2 (the small noise component).
assert results[1]["memory_id"] == ids[2]
def test_vector_search_respects_witness_filter(tmp_path):
db = tmp_path / "t.db"
apply_migrations(db)
with open_db(db) as conn:
dim = 4
# Memory visible to you=1, host=1, guest=0.
_seed_memory_with_embedding(
conn,
owner_id="bot_a",
pov_summary="Restricted.",
vector=_one_hot(dim, 0),
witness_you=1,
witness_host=1,
witness_guest=0,
)
# Guest sees nothing.
guest_results = vector_search(
conn,
owner_id="bot_a",
witness_role="guest",
query_vector=_one_hot(dim, 0),
k=4,
)
assert guest_results == []
# Host sees the memory.
host_results = vector_search(
conn,
owner_id="bot_a",
witness_role="host",
query_vector=_one_hot(dim, 0),
k=4,
)
assert len(host_results) == 1
assert host_results[0]["pov_summary"] == "Restricted."
# You also see it.
you_results = vector_search(
conn,
owner_id="bot_a",
witness_role="you",
query_vector=_one_hot(dim, 0),
k=4,
)
assert len(you_results) == 1
def test_vector_search_respects_owner_filter(tmp_path):
db = tmp_path / "t.db"
apply_migrations(db)
with open_db(db) as conn:
dim = 4
_seed_memory_with_embedding(
conn,
owner_id="bot_a",
pov_summary="Owner A memory.",
vector=_one_hot(dim, 0),
)
_seed_memory_with_embedding(
conn,
owner_id="bot_b",
pov_summary="Owner B memory.",
vector=_one_hot(dim, 0),
)
a_results = vector_search(
conn,
owner_id="bot_a",
witness_role="you",
query_vector=_one_hot(dim, 0),
k=10,
)
assert len(a_results) == 1
assert a_results[0]["pov_summary"] == "Owner A memory."
b_results = vector_search(
conn,
owner_id="bot_b",
witness_role="you",
query_vector=_one_hot(dim, 0),
k=10,
)
assert len(b_results) == 1
assert b_results[0]["pov_summary"] == "Owner B memory."
def test_vector_search_invalid_witness_role_raises(tmp_path):
db = tmp_path / "t.db"
apply_migrations(db)
with open_db(db) as conn:
with pytest.raises(ValueError, match="witness_role"):
vector_search(
conn,
owner_id="bot_a",
witness_role="invalid",
query_vector=[1.0, 0.0, 0.0],
k=4,
)
def test_vector_search_empty_when_no_embeddings_indexed(tmp_path):
db = tmp_path / "t.db"
apply_migrations(db)
with open_db(db) as conn:
# Seed a memory but don't index an embedding for it.
append_event(
conn,
kind="memory_written",
payload=_base_memory(owner_id="bot_a", pov_summary="No embedding here."),
)
project(conn)
results = vector_search(
conn,
owner_id="bot_a",
witness_role="you",
query_vector=[1.0, 0.0, 0.0, 0.0],
k=4,
)
assert results == []