Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 013b563f21 | |||
| 62d5cdd826 | |||
| a25c166174 | |||
| 8f66e1123a | |||
| caa17b4174 | |||
| c7cb0eb01e |
@@ -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"]
|
||||||
@@ -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",
|
||||||
|
]
|
||||||
@@ -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"]
|
||||||
@@ -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=" ") == []
|
||||||
@@ -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)
|
||||||
@@ -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 == []
|
||||||
Reference in New Issue
Block a user