feat: defer embedding messages (#704)
* feat: defer embedding messages * fix: rm gauges * feat: embed messages immediately on create with reconciler fallback (#766) Adds embed_messages_now background task so newly created messages are searchable within seconds instead of waiting up to the reconciler interval. Three-phase claim/lease → embed → persist never holds a DB session across the embedding call; the reconciler remains the fallback for failures and stragglers. * fix: harden immediate-embed fast path and cover its error branches Wrap embed_messages_now in a top-level try/except so a failure in the claim or persist phase degrades to "reconciler will retry" instead of escaping into the background-task runner; the rows stay pending+leased and the reconciler heals them. Add tests for the previously-uncovered branches: external-store-unavailable persist path, the file-upload endpoint's embed scheduling, and direct unit tests for the shared compute_chunk_positions / build_message_vector_record helpers. Document the semantic-search eventual-consistency window in search.mdx (keyword matches are immediate; vector matches lag creation by seconds). * fix: don't hold DB session across vector-store upserts * fix: align semantic-search function to filter null rows --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Rajat Ahuja
GitHub
parent
bf494257b8
commit
6aa6033a16
|
|
@ -6,6 +6,10 @@ icon: 'magnifying-glass'
|
|||
|
||||
Honcho's search functionality allows you to find relevant messages and conversations across different scopes - from entire workspaces down to specific peers or sessions.
|
||||
|
||||
<Note>
|
||||
Search is hybrid: it combines full-text (keyword) matching with semantic (vector) similarity. Keyword matches are available the instant a message is created. Semantic matches depend on the message's embedding, which is generated in the background, so a freshly created message may take a few seconds to surface in semantic results. If you need to assert on semantic results immediately after writing (for example in tests), wait briefly or poll.
|
||||
</Note>
|
||||
|
||||
## Search Scopes
|
||||
|
||||
### Workspace Search
|
||||
|
|
|
|||
|
|
@ -704,6 +704,9 @@ class EmbeddingSettings(HonchoSettings):
|
|||
VECTOR_DIMENSIONS: Annotated[int, Field(default=1536, gt=0)] = 1536
|
||||
MAX_INPUT_TOKENS: Annotated[int, Field(default=8192, gt=0)] = 8192
|
||||
MAX_TOKENS_PER_REQUEST: Annotated[int, Field(default=300_000, gt=0)] = 300_000
|
||||
# Caps concurrent message-embedding fan-out on the API request path (the
|
||||
# immediate-embed background task). The reconciler is unaffected.
|
||||
MAX_CONCURRENT_EMBEDDINGS: Annotated[int, Field(default=10, gt=0, le=100)] = 10
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
|
|
|
|||
|
|
@ -4,19 +4,18 @@ from logging import getLogger
|
|||
from typing import Any
|
||||
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy import ColumnElement, Select, and_, func, or_, select, text, update
|
||||
from sqlalchemy import ColumnElement, Select, and_, func, or_, select, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models, schemas
|
||||
from src.config import settings
|
||||
from src.dependencies import tracked_db
|
||||
from src.embedding_client import embedding_client
|
||||
from src.exceptions import VectorStoreError
|
||||
from src.telemetry.events import EmbeddingCallPurpose
|
||||
from src.utils.filter import apply_filter
|
||||
from src.utils.formatting import ILIKE_ESCAPE_CHAR, escape_ilike_pattern
|
||||
from src.utils.types import embedding_call_purpose
|
||||
from src.vector_store import VectorRecord, get_external_vector_store
|
||||
from src.vector_store import get_external_vector_store
|
||||
|
||||
from .session import get_or_create_session
|
||||
|
||||
|
|
@ -276,158 +275,37 @@ async def create_messages(
|
|||
|
||||
db.add_all(message_objects)
|
||||
|
||||
# Commit here to release the advisory lock before generating embeddings
|
||||
await db.commit()
|
||||
try:
|
||||
if settings.EMBED_MESSAGES:
|
||||
id_resource_dict = {
|
||||
message.public_id: message.content
|
||||
for message in message_objects
|
||||
if message.content and message.content.strip()
|
||||
}
|
||||
if id_resource_dict:
|
||||
with embedding_call_purpose(
|
||||
EmbeddingCallPurpose.MESSAGE_CREATE.value,
|
||||
workspace_name=workspace_name,
|
||||
parent_category="api",
|
||||
):
|
||||
embedding_dict = await embedding_client.batch_embed(
|
||||
id_resource_dict
|
||||
)
|
||||
else:
|
||||
embedding_dict = {}
|
||||
|
||||
external_vector_store = get_external_vector_store()
|
||||
|
||||
# Determine if we need to persist embeddings to postgres
|
||||
# True when: TYPE=pgvector OR still migrating (dual-write to both stores)
|
||||
store_embeddings_in_postgres = (
|
||||
settings.VECTOR_STORE.TYPE == "pgvector"
|
||||
or not settings.VECTOR_STORE.MIGRATED
|
||||
)
|
||||
|
||||
# Create MessageEmbedding entries
|
||||
embedding_objects: list[models.MessageEmbedding] = []
|
||||
# Maps emb index -> (chunk_position, embedding vector)
|
||||
pending_embedding_data: dict[int, tuple[int, list[float]]] = {}
|
||||
# If embedding is enabled, locally chunk the content and insert
|
||||
# one pending MessageEmbedding row per chunk in chunk order. The actual
|
||||
# embedding work is deferred to the reconciler
|
||||
if settings.EMBED_MESSAGES:
|
||||
id_resource_dict = {
|
||||
message_obj.public_id: message_obj.content
|
||||
for message_obj in message_objects
|
||||
if message_obj.content and message_obj.content.strip()
|
||||
}
|
||||
if id_resource_dict:
|
||||
chunks_by_id = embedding_client.prepare_chunks(id_resource_dict)
|
||||
peer_by_id = {m.public_id: m.peer_name for m in message_objects}
|
||||
pending_rows: list[models.MessageEmbedding] = []
|
||||
for message_obj in message_objects:
|
||||
embeddings = embedding_dict.get(message_obj.public_id, [])
|
||||
for chunk_position, embedding in enumerate(embeddings):
|
||||
embedding_obj = models.MessageEmbedding(
|
||||
content=message_obj.content,
|
||||
message_id=message_obj.public_id,
|
||||
workspace_name=workspace_name,
|
||||
session_name=session_name,
|
||||
peer_name=message_obj.peer_name,
|
||||
sync_state="pending",
|
||||
embedding=embedding if store_embeddings_in_postgres else None,
|
||||
)
|
||||
emb_idx = len(embedding_objects)
|
||||
pending_embedding_data[emb_idx] = (chunk_position, embedding)
|
||||
embedding_objects.append(embedding_obj)
|
||||
|
||||
# Always create MessageEmbedding rows so reconciliation can track sync state
|
||||
# even when embeddings aren't stored in postgres
|
||||
embedding_ids: list[int] = []
|
||||
if embedding_objects:
|
||||
db.add_all(embedding_objects)
|
||||
await db.flush()
|
||||
embedding_ids = [emb.id for emb in embedding_objects]
|
||||
|
||||
await db.commit()
|
||||
|
||||
# If no external vector store (pgvector-only mode), mark as synced immediately
|
||||
if external_vector_store is None:
|
||||
if embedding_ids:
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id.in_(embedding_ids))
|
||||
.values(
|
||||
sync_state="synced",
|
||||
last_sync_at=func.now(),
|
||||
sync_attempts=0,
|
||||
chunks = chunks_by_id.get(message_obj.public_id, [])
|
||||
for chunk_text in chunks:
|
||||
pending_rows.append(
|
||||
models.MessageEmbedding(
|
||||
content=chunk_text,
|
||||
message_id=message_obj.public_id,
|
||||
workspace_name=workspace_name,
|
||||
session_name=session_name,
|
||||
peer_name=peer_by_id[message_obj.public_id],
|
||||
sync_state="pending",
|
||||
embedding=None,
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
else:
|
||||
# External vector store - build and upsert vector records
|
||||
namespace = external_vector_store.get_vector_namespace(
|
||||
"message", workspace_name
|
||||
)
|
||||
if pending_rows:
|
||||
db.add_all(pending_rows)
|
||||
|
||||
# Build vector records with {message_id}_{chunk_position} as vector ID
|
||||
vector_records: list[VectorRecord] = []
|
||||
for emb_idx, emb in enumerate(embedding_objects):
|
||||
chunk_position, embedding = pending_embedding_data[emb_idx]
|
||||
vector_id = f"{emb.message_id}_{chunk_position}"
|
||||
vector_records.append(
|
||||
VectorRecord(
|
||||
id=vector_id,
|
||||
embedding=list(embedding),
|
||||
metadata={
|
||||
"message_id": emb.message_id,
|
||||
"session_name": emb.session_name,
|
||||
"peer_name": emb.peer_name,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
# Upsert to external vector store and update sync state
|
||||
if vector_records:
|
||||
try:
|
||||
await external_vector_store.upsert_many(
|
||||
namespace, vector_records
|
||||
)
|
||||
# Success: mark as synced if we have DB rows
|
||||
if embedding_ids:
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id.in_(embedding_ids))
|
||||
.values(
|
||||
sync_state="synced",
|
||||
last_sync_at=func.now(),
|
||||
sync_attempts=0,
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
except VectorStoreError:
|
||||
logger.warning(
|
||||
"Vector store unavailable; leaving message vectors unsynced"
|
||||
)
|
||||
if embedding_ids:
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id.in_(embedding_ids))
|
||||
.values(
|
||||
sync_attempts=models.MessageEmbedding.sync_attempts
|
||||
+ 1,
|
||||
last_sync_at=func.now(),
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
except Exception:
|
||||
logger.exception("Unexpected error upserting message vectors")
|
||||
if embedding_ids:
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id.in_(embedding_ids))
|
||||
.values(
|
||||
sync_attempts=models.MessageEmbedding.sync_attempts
|
||||
+ 1,
|
||||
last_sync_at=func.now(),
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to generate message embeddings for %s messages in workspace %s and session %s.",
|
||||
len(message_objects),
|
||||
workspace_name,
|
||||
session_name,
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
return message_objects
|
||||
|
||||
|
|
@ -770,6 +648,9 @@ async def _search_messages_pgvector(
|
|||
models.MessageEmbedding,
|
||||
models.Message.public_id == models.MessageEmbedding.message_id,
|
||||
)
|
||||
# Exclude pending rows that haven't been embedded yet: their NULL
|
||||
# distance sorts last and would pad the window with unranked messages.
|
||||
.where(models.MessageEmbedding.embedding.isnot(None))
|
||||
.where(models.MessageEmbedding.workspace_name == workspace_name)
|
||||
.order_by(models.MessageEmbedding.embedding.cosine_distance(query_embedding))
|
||||
.limit(limit * 2)
|
||||
|
|
|
|||
|
|
@ -250,76 +250,61 @@ class _EmbeddingClient:
|
|||
|
||||
async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Simple batch embedding for a list of text strings.
|
||||
Batch-embed a list of text strings. Each input must already fit within
|
||||
`max_embedding_tokens`; this method does not sub-chunk oversized inputs.
|
||||
|
||||
Internally goes through the same token-aware batching pipeline as
|
||||
`batch_embed()` so the per-request token cap is respected.
|
||||
|
||||
Args:
|
||||
texts: List of text strings to embed
|
||||
|
||||
Returns:
|
||||
List of embedding vectors corresponding to input texts
|
||||
List of embedding vectors, one per input text (in order)
|
||||
|
||||
Raises:
|
||||
ValueError: If any text exceeds token limits
|
||||
"""
|
||||
embeddings: list[list[float]] = []
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
for i in range(0, len(texts), self.max_batch_size):
|
||||
batch = texts[i : i + self.max_batch_size]
|
||||
|
||||
async def _embed_batch(batch: list[str] = batch) -> list[list[float]]:
|
||||
"""One provider call for one batch. Lifted into a closure so
|
||||
_emit_embedding_call can time + emit + propagate errors."""
|
||||
batch_embeddings: list[list[float]] = []
|
||||
if isinstance(self.client, genai.Client):
|
||||
# Type cast needed due to genai type signature complexity
|
||||
response = await self.client.aio.models.embed_content(
|
||||
model=self.model,
|
||||
contents=batch, # pyright: ignore[reportArgumentType]
|
||||
config={"output_dimensionality": self.vector_dimensions},
|
||||
)
|
||||
if response.embeddings:
|
||||
for emb in response.embeddings:
|
||||
if emb.values:
|
||||
batch_embeddings.append(
|
||||
self._validate_embedding_dimensions(emb.values)
|
||||
)
|
||||
else: # openai
|
||||
openai_kwargs: dict[str, Any] = {
|
||||
"input": batch,
|
||||
"model": self.model,
|
||||
}
|
||||
if self.send_dimensions:
|
||||
openai_kwargs["dimensions"] = self.vector_dimensions
|
||||
response = await self.client.embeddings.create(**openai_kwargs)
|
||||
batch_embeddings.extend(
|
||||
[
|
||||
self._validate_embedding_dimensions(data.embedding)
|
||||
for data in response.data
|
||||
]
|
||||
)
|
||||
return batch_embeddings
|
||||
|
||||
try:
|
||||
# Pre-compute the tiktoken estimate ONCE for telemetry; the
|
||||
# batch contents don't change between attempts.
|
||||
tokens_estimate = sum(len(self.encoding.encode(t)) for t in batch)
|
||||
batch_embeddings = await _emit_embedding_call(
|
||||
provider=self.transport,
|
||||
model=self.model,
|
||||
texts=batch,
|
||||
input_tokens_estimate=tokens_estimate,
|
||||
fn=_embed_batch,
|
||||
# Validate per-input token limit and collect token counts for batching
|
||||
token_counts: list[int] = []
|
||||
for idx, text in enumerate(texts):
|
||||
tokens = len(self.encoding.encode(text))
|
||||
if tokens > self.max_embedding_tokens:
|
||||
raise ValueError(
|
||||
f"Text at index {idx} exceeds maximum token limit of {self.max_embedding_tokens} tokens (got {tokens} tokens)"
|
||||
)
|
||||
embeddings.extend(batch_embeddings)
|
||||
except Exception as e:
|
||||
# Check if it's a token limit error and re-raise as ValueError for consistency
|
||||
if "token" in str(e).lower():
|
||||
raise ValueError(
|
||||
f"Text content exceeds maximum token limit of {self.max_embedding_tokens}."
|
||||
) from e
|
||||
raise
|
||||
token_counts.append(tokens)
|
||||
|
||||
return embeddings
|
||||
# Use positional indices as text_ids so we can reassemble in input order.
|
||||
text_chunks: dict[str, list[tuple[str, int]]] = {
|
||||
str(i): [(text, token_counts[i])] for i, text in enumerate(texts)
|
||||
}
|
||||
|
||||
batches = self._create_batches(text_chunks)
|
||||
batch_results = await asyncio.gather(
|
||||
*[self._process_batch(batch) for batch in batches],
|
||||
)
|
||||
|
||||
combined: dict[str, list[list[float]]] = self._accumulate_embeddings(
|
||||
batch_results
|
||||
)
|
||||
return [combined[str(i)][0] for i in range(len(texts))]
|
||||
|
||||
def prepare_chunks(self, id_resource_dict: dict[str, str]) -> dict[str, list[str]]:
|
||||
"""
|
||||
Public helper: tokenize and chunk texts using the same rules as
|
||||
`batch_embed()`. Returns ordered chunk texts per input id.
|
||||
|
||||
Intended for callers that want to persist embeddable chunks
|
||||
before later embedding them off the request path.
|
||||
"""
|
||||
return {
|
||||
text_id: [chunk_text for chunk_text, _ in chunks]
|
||||
for text_id, chunks in self._prepare_chunks(id_resource_dict).items()
|
||||
}
|
||||
|
||||
async def batch_embed(
|
||||
self, id_resource_dict: dict[str, str]
|
||||
|
|
@ -623,9 +608,13 @@ class EmbeddingClient:
|
|||
return await self._get_client().embed(query)
|
||||
|
||||
async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Simple batch embedding for a list of text strings."""
|
||||
"""Batch embed a list of text strings (each must fit token limit)."""
|
||||
return await self._get_client().simple_batch_embed(texts)
|
||||
|
||||
def prepare_chunks(self, id_resource_dict: dict[str, str]) -> dict[str, list[str]]:
|
||||
"""Chunk texts using the same rules as `batch_embed` (no network)."""
|
||||
return self._get_client().prepare_chunks(id_resource_dict)
|
||||
|
||||
async def batch_embed(
|
||||
self, id_resource_dict: dict[str, str]
|
||||
) -> dict[str, list[list[float]]]:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,343 @@
|
|||
"""
|
||||
Immediate message-embedding fast path.
|
||||
|
||||
``create_messages`` writes ``MessageEmbedding`` rows as ``sync_state='pending'``
|
||||
with no vector and defers embedding to the reconciler, which runs on a fixed
|
||||
interval. To keep freshly created messages searchable within seconds (not
|
||||
minutes), the message routers schedule ``embed_messages_now`` as a FastAPI
|
||||
background task right after the response is sent. The reconciler remains the
|
||||
fallback for anything this path leaves pending (failures, process restarts, or
|
||||
rows it could not claim).
|
||||
|
||||
The fast path never holds a DB session across a network call (embedding or
|
||||
external vector store): it claims and leases rows in one short transaction,
|
||||
embeds with no session open, then persists in short transactions with any
|
||||
external-store upserts running between them, not inside them. Running concurrently with the
|
||||
reconciler is safe because the claim uses ``FOR UPDATE SKIP LOCKED`` and leases
|
||||
rows by stamping ``last_sync_at``, which the reconciler's backoff filter then
|
||||
skips.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import and_, func, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models
|
||||
from src.config import settings
|
||||
from src.dependencies import tracked_db
|
||||
from src.embedding_client import embedding_client
|
||||
from src.exceptions import VectorStoreError
|
||||
from src.reconciler.sync_vectors import (
|
||||
_backoff_eligible, # pyright: ignore[reportPrivateUsage]
|
||||
build_message_vector_record,
|
||||
compute_chunk_positions,
|
||||
)
|
||||
from src.telemetry.events import EmbeddingCallPurpose
|
||||
from src.utils.types import embedding_call_purpose
|
||||
from src.vector_store import VectorRecord, VectorStore, get_external_vector_store
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_embed_semaphore: asyncio.Semaphore | None = None
|
||||
|
||||
|
||||
def _get_embed_semaphore() -> asyncio.Semaphore:
|
||||
"""Lazily create the embed-concurrency semaphore.
|
||||
|
||||
Built on first use (not at import time) so it binds to the running event
|
||||
loop rather than whatever loop happened to exist at import.
|
||||
"""
|
||||
global _embed_semaphore
|
||||
if _embed_semaphore is None:
|
||||
_embed_semaphore = asyncio.Semaphore(
|
||||
settings.EMBEDDING.MAX_CONCURRENT_EMBEDDINGS
|
||||
)
|
||||
return _embed_semaphore
|
||||
|
||||
|
||||
def reset_embed_semaphore() -> None:
|
||||
"""Test hook: drop the cached semaphore so the next call rebuilds it on the
|
||||
current event loop and current config."""
|
||||
global _embed_semaphore
|
||||
_embed_semaphore = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ClaimedChunk:
|
||||
"""Plain snapshot of a claimed ``MessageEmbedding`` row.
|
||||
|
||||
Captured before the claim transaction commits — after commit the ORM object
|
||||
is detached and attribute access would lazy-load against a closed session.
|
||||
"""
|
||||
|
||||
id: int
|
||||
message_id: str
|
||||
content: str
|
||||
workspace_name: str
|
||||
session_name: str | None
|
||||
peer_name: str | None
|
||||
|
||||
|
||||
async def embed_messages_now(message_ids: list[str]) -> None:
|
||||
"""Embed freshly created messages immediately, leaving the reconciler as the
|
||||
fallback for anything left pending.
|
||||
|
||||
Args:
|
||||
message_ids: ``Message.public_id`` values (what
|
||||
``MessageEmbedding.message_id`` references). Messages without
|
||||
embeddable content simply have no pending rows to claim.
|
||||
"""
|
||||
if not message_ids:
|
||||
return
|
||||
|
||||
# Runs as a fire-and-forget background task, so guard the whole flow: an
|
||||
# unhandled error here would escape into the server's task runner and be
|
||||
# lost. Any failure just leaves rows pending (claimed rows stay leased),
|
||||
# and the reconciler heals them on its next cycle.
|
||||
try:
|
||||
claimed = await _claim_and_lease(message_ids)
|
||||
if not claimed:
|
||||
return
|
||||
|
||||
vectors = await _embed_chunks(claimed)
|
||||
if vectors is None:
|
||||
# Embedding failed; rows stay pending + leased, reconciler will retry.
|
||||
return
|
||||
|
||||
await _persist(message_ids, claimed, vectors)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Immediate embed failed for %s message(s); reconciler will retry",
|
||||
len(message_ids),
|
||||
)
|
||||
|
||||
|
||||
async def _claim_and_lease(message_ids: list[str]) -> list[_ClaimedChunk]:
|
||||
"""Phase 1 (short txn): claim eligible pending rows with FOR UPDATE SKIP
|
||||
LOCKED, lease them by stamping ``last_sync_at``, and snapshot their data.
|
||||
|
||||
``sync_attempts`` is intentionally left untouched: the reconciler owns retry
|
||||
accounting and the eventual ``sync_state='failed'`` backstop, so a transient
|
||||
embedding failure on this best-effort path never burns that budget.
|
||||
"""
|
||||
async with tracked_db("embed_now_claim") as db:
|
||||
rows_stmt = (
|
||||
select(models.MessageEmbedding)
|
||||
.where(
|
||||
and_(
|
||||
models.MessageEmbedding.message_id.in_(message_ids),
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
_backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
)
|
||||
)
|
||||
.order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id)
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
rows = list((await db.execute(rows_stmt)).scalars().all())
|
||||
if not rows:
|
||||
return []
|
||||
|
||||
claimed = [
|
||||
_ClaimedChunk(
|
||||
id=row.id,
|
||||
message_id=row.message_id,
|
||||
content=row.content,
|
||||
workspace_name=row.workspace_name,
|
||||
session_name=row.session_name,
|
||||
peer_name=row.peer_name,
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id.in_([c.id for c in claimed]))
|
||||
.values(last_sync_at=func.now())
|
||||
)
|
||||
await db.commit()
|
||||
return claimed
|
||||
|
||||
|
||||
async def _embed_chunks(claimed: list[_ClaimedChunk]) -> list[list[float]] | None:
|
||||
"""Phase 2 (no DB session): embed the claimed chunk contents under the
|
||||
concurrency semaphore. Returns vectors in input order, or None on failure."""
|
||||
workspaces = {c.workspace_name for c in claimed}
|
||||
try:
|
||||
async with _get_embed_semaphore():
|
||||
with embedding_call_purpose(
|
||||
EmbeddingCallPurpose.MESSAGE_CREATE.value,
|
||||
workspace_name=workspaces.pop() if len(workspaces) == 1 else None,
|
||||
parent_category="api",
|
||||
):
|
||||
return await embedding_client.simple_batch_embed(
|
||||
[c.content for c in claimed]
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Immediate embedding failed for %s chunk(s); reconciler will retry",
|
||||
len(claimed),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def _persist(
|
||||
message_ids: list[str],
|
||||
claimed: list[_ClaimedChunk],
|
||||
vectors: list[list[float]],
|
||||
) -> None:
|
||||
"""Phase 3: persist vectors and mark rows synced. On failure, rows stay
|
||||
pending (already leased) and the reconciler heals them.
|
||||
|
||||
pgvector mode is one short transaction. External-store mode never holds a
|
||||
DB session across the vector-store network call: positions are read in one
|
||||
short transaction, the upserts run with no session open, and the surviving
|
||||
rows are marked synced in a second short transaction."""
|
||||
if len(vectors) != len(claimed):
|
||||
logger.warning(
|
||||
"Embedding count %s != claimed chunk count %s; skipping immediate persist, reconciler will heal",
|
||||
len(vectors),
|
||||
len(claimed),
|
||||
)
|
||||
return
|
||||
|
||||
vector_by_id = {c.id: vec for c, vec in zip(claimed, vectors, strict=True)}
|
||||
# True for pgvector OR during migration (dual-write to both stores).
|
||||
store_in_postgres = (
|
||||
settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED
|
||||
)
|
||||
external = get_external_vector_store()
|
||||
|
||||
if external is None:
|
||||
async with tracked_db("embed_now_persist") as db:
|
||||
await _persist_pgvector(db, claimed, vector_by_id)
|
||||
await db.commit()
|
||||
return
|
||||
|
||||
synced = await _upsert_external(message_ids, claimed, vector_by_id, external)
|
||||
if not synced:
|
||||
return
|
||||
|
||||
async with tracked_db("embed_now_persist") as db:
|
||||
await _mark_synced(db, synced, vector_by_id, store_in_postgres)
|
||||
await db.commit()
|
||||
|
||||
|
||||
async def _persist_pgvector(
|
||||
db: AsyncSession,
|
||||
claimed: list[_ClaimedChunk],
|
||||
vector_by_id: dict[int, list[float]],
|
||||
) -> None:
|
||||
"""pgvector-only mode: write the vector and mark synced per row. The
|
||||
``sync_state='pending'`` guard keeps us idempotent if the reconciler synced
|
||||
a row in the gap."""
|
||||
for c in claimed:
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(
|
||||
and_(
|
||||
models.MessageEmbedding.id == c.id,
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
)
|
||||
)
|
||||
.values(
|
||||
sync_state="synced",
|
||||
last_sync_at=func.now(),
|
||||
sync_attempts=0,
|
||||
embedding=vector_by_id[c.id],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _upsert_external(
|
||||
message_ids: list[str],
|
||||
claimed: list[_ClaimedChunk],
|
||||
vector_by_id: dict[int, list[float]],
|
||||
external: VectorStore,
|
||||
) -> list[_ClaimedChunk]:
|
||||
"""External-store mode: upsert vectors per namespace with no DB session
|
||||
open, returning the chunks whose namespaces upserted successfully.
|
||||
|
||||
Chunk positions come from the shared helper (full sibling ordering) so vector
|
||||
ids match whatever the reconciler writes for any chunk we skipped; reading
|
||||
them is the only DB work here, done in its own short transaction before any
|
||||
network call."""
|
||||
async with tracked_db("embed_now_positions") as db:
|
||||
chunk_position = await compute_chunk_positions(db, message_ids)
|
||||
|
||||
by_namespace: dict[str, list[_ClaimedChunk]] = {}
|
||||
for c in claimed:
|
||||
ns = external.get_vector_namespace("message", c.workspace_name)
|
||||
by_namespace.setdefault(ns, []).append(c)
|
||||
|
||||
synced: list[_ClaimedChunk] = []
|
||||
for namespace, chunks in by_namespace.items():
|
||||
records: list[VectorRecord] = []
|
||||
synced_chunks: list[_ClaimedChunk] = []
|
||||
for c in chunks:
|
||||
pos = chunk_position.get(c.id)
|
||||
if pos is None:
|
||||
continue
|
||||
records.append(
|
||||
build_message_vector_record(
|
||||
message_id=c.message_id,
|
||||
chunk_position=pos,
|
||||
session_name=c.session_name,
|
||||
peer_name=c.peer_name,
|
||||
embedding=vector_by_id[c.id],
|
||||
)
|
||||
)
|
||||
synced_chunks.append(c)
|
||||
|
||||
if not records:
|
||||
continue
|
||||
|
||||
try:
|
||||
await external.upsert_many(namespace, records)
|
||||
except VectorStoreError:
|
||||
logger.warning(
|
||||
"Vector store unavailable during immediate embed of namespace %s; reconciler will retry",
|
||||
namespace,
|
||||
)
|
||||
continue
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Unexpected error during immediate embed of namespace %s; reconciler will retry",
|
||||
namespace,
|
||||
)
|
||||
continue
|
||||
|
||||
synced.extend(synced_chunks)
|
||||
|
||||
return synced
|
||||
|
||||
|
||||
async def _mark_synced(
|
||||
db: AsyncSession,
|
||||
chunks: list[_ClaimedChunk],
|
||||
vector_by_id: dict[int, list[float]],
|
||||
store_in_postgres: bool,
|
||||
) -> None:
|
||||
"""Mark upserted chunks synced (DB-only). The ``sync_state='pending'``
|
||||
guard keeps us idempotent if the reconciler synced a row in the gap."""
|
||||
for c in chunks:
|
||||
values: dict[str, Any] = {
|
||||
"sync_state": "synced",
|
||||
"last_sync_at": func.now(),
|
||||
"sync_attempts": 0,
|
||||
}
|
||||
if store_in_postgres:
|
||||
values["embedding"] = vector_by_id[c.id]
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(
|
||||
and_(
|
||||
models.MessageEmbedding.id == c.id,
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
)
|
||||
)
|
||||
.values(**values)
|
||||
)
|
||||
|
|
@ -9,7 +9,7 @@ import datetime
|
|||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import cast
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy import and_, delete, or_, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
|
@ -109,28 +109,55 @@ async def _get_message_embeddings_needing_sync(
|
|||
"""
|
||||
Get pending message embeddings that need to be synced to the vector store.
|
||||
|
||||
Returns only pending embeddings (with full data including embedding vectors).
|
||||
The batch_size limits the number of embeddings returned.
|
||||
Claims up to `batch_size` distinct message_ids that have at least one
|
||||
eligible pending row, then loads ALL pending rows for those message_ids.
|
||||
This guarantees a single message's chunks are always processed together in
|
||||
one batch, which keeps vector-ID assignment (`{message_id}_{chunk_index}`,
|
||||
derived from row-id ordering) stable across reconciler cycles.
|
||||
|
||||
Uses FOR UPDATE SKIP LOCKED to prevent concurrent processing and
|
||||
orders by last_sync_at (nulls first) to prioritize never-synced records.
|
||||
Uses FOR UPDATE SKIP LOCKED on the per-row claim so concurrent reconcilers
|
||||
don't double-process the same chunks.
|
||||
|
||||
Note: "synced" = done forever, "failed" = permanent failure (manual intervention)
|
||||
"""
|
||||
stmt = (
|
||||
select(models.MessageEmbedding)
|
||||
# Step 1: pick distinct message_ids with at least one eligible pending row,
|
||||
# prioritizing those with the oldest last_sync_at.
|
||||
msg_id_stmt = (
|
||||
select(
|
||||
models.MessageEmbedding.message_id,
|
||||
func.min(models.MessageEmbedding.last_sync_at).label("oldest_attempt"),
|
||||
)
|
||||
.where(
|
||||
and_(
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
_backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
)
|
||||
)
|
||||
.order_by(models.MessageEmbedding.last_sync_at.asc().nullsfirst())
|
||||
.group_by(models.MessageEmbedding.message_id)
|
||||
.order_by(func.min(models.MessageEmbedding.last_sync_at).asc().nullsfirst())
|
||||
.limit(batch_size)
|
||||
)
|
||||
msg_id_rows = (await db.execute(msg_id_stmt)).all()
|
||||
message_ids = [row[0] for row in msg_id_rows]
|
||||
if not message_ids:
|
||||
return []
|
||||
|
||||
# Step 2: claim all pending rows for those messages. Skip rows another
|
||||
# reconciler holds; if we can't claim every chunk of a message right now,
|
||||
# the message will be retried next cycle.
|
||||
rows_stmt = (
|
||||
select(models.MessageEmbedding)
|
||||
.where(
|
||||
and_(
|
||||
models.MessageEmbedding.message_id.in_(message_ids),
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
_backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
)
|
||||
)
|
||||
.order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id)
|
||||
.with_for_update(skip_locked=True)
|
||||
)
|
||||
|
||||
result = await db.execute(stmt)
|
||||
result = await db.execute(rows_stmt)
|
||||
return list(result.scalars().all())
|
||||
|
||||
|
||||
|
|
@ -177,6 +204,63 @@ async def _bump_message_embedding_sync_attempts(
|
|||
)
|
||||
|
||||
|
||||
async def compute_chunk_positions(
|
||||
db: AsyncSession, message_ids: list[str]
|
||||
) -> dict[int, int]:
|
||||
"""Map each MessageEmbedding row id to its 0-indexed chunk position within
|
||||
its message.
|
||||
|
||||
Positions are derived from the full set of sibling rows for each message,
|
||||
ordered by ``(message_id, id)`` — never from a partial subset — so the
|
||||
``{message_id}_{chunk_position}`` vector id stays stable no matter which
|
||||
rows a given caller claimed. Shared by the reconciler and the immediate
|
||||
embed path so the two writers always agree on vector ids.
|
||||
"""
|
||||
if not message_ids:
|
||||
return {}
|
||||
|
||||
sibling_stmt = (
|
||||
select(models.MessageEmbedding.id, models.MessageEmbedding.message_id)
|
||||
.where(models.MessageEmbedding.message_id.in_(message_ids))
|
||||
.order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id)
|
||||
)
|
||||
sibling_rows = (await db.execute(sibling_stmt)).all()
|
||||
|
||||
embs_by_message: dict[str, list[int]] = {}
|
||||
for emb_id, msg_id in sibling_rows:
|
||||
embs_by_message.setdefault(msg_id, []).append(emb_id)
|
||||
|
||||
chunk_position: dict[int, int] = {}
|
||||
for emb_ids in embs_by_message.values():
|
||||
for pos, emb_id in enumerate(emb_ids):
|
||||
chunk_position[emb_id] = pos
|
||||
return chunk_position
|
||||
|
||||
|
||||
def build_message_vector_record(
|
||||
*,
|
||||
message_id: str,
|
||||
chunk_position: int,
|
||||
session_name: str | None,
|
||||
peer_name: str | None,
|
||||
embedding: list[float],
|
||||
) -> VectorRecord:
|
||||
"""Build the external-store record for one message-embedding chunk.
|
||||
|
||||
Single source of the ``{message_id}_{chunk_position}`` vector id and the
|
||||
metadata shape, shared by the reconciler and the immediate embed path.
|
||||
"""
|
||||
return VectorRecord(
|
||||
id=f"{message_id}_{chunk_position}",
|
||||
embedding=[float(x) for x in embedding],
|
||||
metadata={
|
||||
"message_id": message_id,
|
||||
"session_name": session_name,
|
||||
"peer_name": peer_name,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _sync_documents(
|
||||
db: AsyncSession,
|
||||
documents: list[models.Document],
|
||||
|
|
@ -306,16 +390,20 @@ async def _sync_documents(
|
|||
async def _sync_message_embeddings(
|
||||
db: AsyncSession,
|
||||
embeddings: list[models.MessageEmbedding],
|
||||
external_vector_store: VectorStore,
|
||||
external_vector_store: VectorStore | None,
|
||||
) -> tuple[int, int]:
|
||||
"""
|
||||
Sync a batch of pending message embeddings to the external vector store.
|
||||
Sync a batch of pending message embeddings.
|
||||
|
||||
Handles three cases for each embedding:
|
||||
When `external_vector_store` is provided, handles three cases per embedding:
|
||||
1. Embedding exists in postgres → use it for external upsert
|
||||
2. Embedding missing + need postgres storage → re-embed, write to both stores
|
||||
3. Embedding missing + external-only mode → re-embed, write to external only
|
||||
|
||||
When `external_vector_store` is None (pgvector-only mode), re-embeds any
|
||||
pending row missing a vector, writes the vector to postgres, and marks
|
||||
sync_state='synced'. No external upsert is performed.
|
||||
|
||||
Returns (synced_count, failed_count).
|
||||
"""
|
||||
if not embeddings:
|
||||
|
|
@ -338,8 +426,12 @@ async def _sync_message_embeddings(
|
|||
if embs_needing_embed:
|
||||
try:
|
||||
contents = [emb.content for emb in embs_needing_embed]
|
||||
# MESSAGE_CREATE (not VECTOR_SYNC): these rows come from create_messages
|
||||
# as pending chunks; document re-embeds stay on VECTOR_SYNC below.
|
||||
workspaces = {emb.workspace_name for emb in embs_needing_embed}
|
||||
with embedding_call_purpose(
|
||||
EmbeddingCallPurpose.VECTOR_SYNC.value,
|
||||
EmbeddingCallPurpose.MESSAGE_CREATE.value,
|
||||
workspace_name=workspaces.pop() if len(workspaces) == 1 else None,
|
||||
parent_category="reconciliation",
|
||||
):
|
||||
new_embeddings = await embedding_client.simple_batch_embed(contents)
|
||||
|
|
@ -368,6 +460,32 @@ async def _sync_message_embeddings(
|
|||
await _bump_message_embedding_sync_attempts(db, failed_to_embed)
|
||||
failed_count += len(failed_to_embed)
|
||||
|
||||
# pgvector-only mode: no external store to upsert to. Any row that now
|
||||
# has an embedding (either pre-existing or freshly embedded) is fully
|
||||
# synced. Write embeddings via per-row UPDATE so the vector is persisted
|
||||
# alongside sync_state in a single statement (session has autoflush=False,
|
||||
# so the ORM mutation above isn't enough on its own).
|
||||
if external_vector_store is None:
|
||||
embs_done: list[models.MessageEmbedding] = []
|
||||
for emb in embeddings:
|
||||
new_emb = freshly_embedded.get(emb.id)
|
||||
existing = emb.embedding
|
||||
if new_emb is None and existing is None:
|
||||
continue
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id == emb.id)
|
||||
.values(
|
||||
sync_state="synced",
|
||||
last_sync_at=func.now(),
|
||||
sync_attempts=0,
|
||||
**({"embedding": new_emb} if new_emb is not None else {}),
|
||||
)
|
||||
)
|
||||
embs_done.append(emb)
|
||||
synced_count += len(embs_done)
|
||||
return synced_count, failed_count
|
||||
|
||||
# Step 2: Compute chunk positions for vector IDs
|
||||
# Messages can be split into multiple chunks; we need {message_id}_{chunk_position}
|
||||
#
|
||||
|
|
@ -380,21 +498,7 @@ async def _sync_message_embeddings(
|
|||
# 2. Removing MessageEmbedding table entirely if it becomes unnecessary
|
||||
# See: https://github.com/plastic-labs/honcho/issues/XXX
|
||||
message_ids = list({emb.message_id for emb in embeddings})
|
||||
sibling_stmt = (
|
||||
select(models.MessageEmbedding.id, models.MessageEmbedding.message_id)
|
||||
.where(models.MessageEmbedding.message_id.in_(message_ids))
|
||||
.order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id)
|
||||
)
|
||||
sibling_rows = (await db.execute(sibling_stmt)).all()
|
||||
|
||||
embs_by_message: dict[str, list[int]] = {}
|
||||
for emb_id, msg_id in sibling_rows:
|
||||
embs_by_message.setdefault(msg_id, []).append(emb_id)
|
||||
|
||||
chunk_position: dict[int, int] = {}
|
||||
for emb_ids in embs_by_message.values():
|
||||
for pos, emb_id in enumerate(emb_ids):
|
||||
chunk_position[emb_id] = pos
|
||||
chunk_position = await compute_chunk_positions(db, message_ids)
|
||||
|
||||
# Step 3: Build vector records and upsert to external store (all cases)
|
||||
by_namespace: dict[str, list[models.MessageEmbedding]] = {}
|
||||
|
|
@ -416,14 +520,12 @@ async def _sync_message_embeddings(
|
|||
continue
|
||||
|
||||
vector_records.append(
|
||||
VectorRecord(
|
||||
id=f"{emb.message_id}_{chunk_position[emb.id]}",
|
||||
embedding=[float(x) for x in embedding],
|
||||
metadata={
|
||||
"message_id": emb.message_id,
|
||||
"session_name": emb.session_name,
|
||||
"peer_name": emb.peer_name,
|
||||
},
|
||||
build_message_vector_record(
|
||||
message_id=emb.message_id,
|
||||
chunk_position=chunk_position[emb.id],
|
||||
session_name=emb.session_name,
|
||||
peer_name=emb.peer_name,
|
||||
embedding=embedding,
|
||||
)
|
||||
)
|
||||
embs_to_sync.append(emb)
|
||||
|
|
@ -433,11 +535,23 @@ async def _sync_message_embeddings(
|
|||
|
||||
try:
|
||||
await external_vector_store.upsert_many(namespace, vector_records)
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id.in_([e.id for e in embs_to_sync]))
|
||||
.values(sync_state="synced", last_sync_at=func.now(), sync_attempts=0)
|
||||
)
|
||||
# Per-row UPDATEs so freshly-embedded rows persist the vector
|
||||
# alongside sync_state. Session has autoflush=False so the ORM
|
||||
# mutation above isn't sufficient on its own.
|
||||
for emb in embs_to_sync:
|
||||
new_emb = freshly_embedded.get(emb.id)
|
||||
values: dict[str, Any] = {
|
||||
"sync_state": "synced",
|
||||
"last_sync_at": func.now(),
|
||||
"sync_attempts": 0,
|
||||
}
|
||||
if new_emb is not None and store_in_postgres:
|
||||
values["embedding"] = new_emb
|
||||
await db.execute(
|
||||
update(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id == emb.id)
|
||||
.values(**values)
|
||||
)
|
||||
synced_count += len(embs_to_sync)
|
||||
except VectorStoreError:
|
||||
logger.warning(
|
||||
|
|
@ -512,7 +626,7 @@ async def _reconcile_documents_batch(
|
|||
|
||||
|
||||
async def _reconcile_message_embeddings_batch(
|
||||
external_vector_store: VectorStore,
|
||||
external_vector_store: VectorStore | None,
|
||||
metrics: ReconciliationMetrics,
|
||||
) -> bool:
|
||||
"""
|
||||
|
|
@ -592,11 +706,18 @@ async def run_vector_reconciliation_cycle() -> ReconciliationMetrics:
|
|||
external_vector_store = get_external_vector_store()
|
||||
deadline = time.monotonic() + RECONCILIATION_TIME_BUDGET_SECONDS
|
||||
|
||||
# If no external vector store (pgvector mode), only clean up soft-deleted documents
|
||||
# pgvector-only mode: still need to embed pending MessageEmbedding rows
|
||||
# (create_messages defers embedding to the reconciler), then clean up.
|
||||
if external_vector_store is None:
|
||||
while time.monotonic() < deadline:
|
||||
did_work = await _cleanup_pgvector_batch(metrics)
|
||||
if not did_work:
|
||||
embs_work = await _reconcile_message_embeddings_batch(None, metrics)
|
||||
|
||||
if time.monotonic() >= deadline:
|
||||
break
|
||||
|
||||
cleanup_work = await _cleanup_pgvector_batch(metrics)
|
||||
|
||||
if not (embs_work or cleanup_work):
|
||||
break
|
||||
logger.info("Vector reconciliation cycle completed (pgvector mode)")
|
||||
return metrics
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from src.config import settings
|
|||
from src.dependencies import db, read_db
|
||||
from src.deriver import enqueue
|
||||
from src.exceptions import FileTooLargeError, ResourceNotFoundException
|
||||
from src.reconciler.embed_now import embed_messages_now
|
||||
from src.security import require_auth
|
||||
from src.telemetry import prometheus_metrics
|
||||
from src.telemetry.events import FileUploadedEvent, MessageCreatedEvent, emit
|
||||
|
|
@ -140,6 +141,13 @@ async def create_messages_for_session(
|
|||
# Enqueue all messages in one call
|
||||
background_tasks.add_task(enqueue, payloads)
|
||||
|
||||
# Embed immediately so messages are searchable within seconds; the
|
||||
# reconciler is the fallback for anything left pending.
|
||||
if settings.EMBED_MESSAGES and created_messages:
|
||||
background_tasks.add_task(
|
||||
embed_messages_now, [m.public_id for m in created_messages]
|
||||
)
|
||||
|
||||
return created_messages
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to create messages for session {session_id}: {str(e)}")
|
||||
|
|
@ -206,6 +214,14 @@ async def create_messages_with_file(
|
|||
]
|
||||
|
||||
background_tasks.add_task(enqueue, payloads)
|
||||
|
||||
# Embed immediately so messages are searchable within seconds; the
|
||||
# reconciler is the fallback for anything left pending.
|
||||
if settings.EMBED_MESSAGES and created_messages:
|
||||
background_tasks.add_task(
|
||||
embed_messages_now, [m.public_id for m in created_messages]
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"Batch of %s messages created from file uploads and queued for processing",
|
||||
len(created_messages),
|
||||
|
|
|
|||
|
|
@ -164,6 +164,9 @@ class EmbeddingCallPurpose(str, Enum):
|
|||
CREATE_OBSERVATIONS = "create_observations"
|
||||
VECTOR_SYNC = "vector_sync"
|
||||
SUMMARY = "summary"
|
||||
# Pending MessageEmbedding rows from create_messages; embedding runs in the
|
||||
# reconciler (not inline on the API path). Distinct from VECTOR_SYNC, which
|
||||
# covers document re-embeds and other vector-store healing work.
|
||||
MESSAGE_CREATE = "message_create"
|
||||
# Added so previously-unattributed call sites land on a distinct slug
|
||||
# instead of None. Closed taxonomy — coordinate with analytics before
|
||||
|
|
|
|||
|
|
@ -484,6 +484,9 @@ def mock_openai_embeddings(request: pytest.FixtureRequest):
|
|||
patch(
|
||||
"src.embedding_client.embedding_client.simple_batch_embed"
|
||||
) as mock_simple_batch_embed,
|
||||
patch(
|
||||
"src.embedding_client.embedding_client.prepare_chunks"
|
||||
) as mock_prepare_chunks,
|
||||
patch("src.embedding_client.embedding_client.batch_embed") as mock_batch_embed,
|
||||
):
|
||||
# Mock the embed method to return content-dependent embedding
|
||||
|
|
@ -497,6 +500,14 @@ def mock_openai_embeddings(request: pytest.FixtureRequest):
|
|||
|
||||
mock_simple_batch_embed.side_effect = mock_simple_batch_embed_func
|
||||
|
||||
def mock_prepare_chunks_func(
|
||||
id_resource_dict: dict[str, str],
|
||||
) -> dict[str, list[str]]:
|
||||
# No real tokenizer in mocks: treat each input as a single chunk.
|
||||
return {text_id: [text] for text_id, text in id_resource_dict.items()}
|
||||
|
||||
mock_prepare_chunks.side_effect = mock_prepare_chunks_func
|
||||
|
||||
# Mock the batch_embed method to return content-dependent embeddings
|
||||
async def mock_batch_embed_func(
|
||||
id_resource_dict: dict[str, str],
|
||||
|
|
@ -511,6 +522,7 @@ def mock_openai_embeddings(request: pytest.FixtureRequest):
|
|||
yield {
|
||||
"embed": mock_embed,
|
||||
"simple_batch_embed": mock_simple_batch_embed,
|
||||
"prepare_chunks": mock_prepare_chunks,
|
||||
"batch_embed": mock_batch_embed,
|
||||
}
|
||||
|
||||
|
|
@ -790,7 +802,7 @@ def mock_tracked_db(request: pytest.FixtureRequest):
|
|||
yield
|
||||
return
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from contextlib import ExitStack, asynccontextmanager
|
||||
|
||||
db_engine = request.getfixturevalue("db_engine")
|
||||
session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
|
||||
|
|
@ -803,28 +815,35 @@ def mock_tracked_db(request: pytest.FixtureRequest):
|
|||
async with session_factory() as session:
|
||||
yield session
|
||||
|
||||
with (
|
||||
patch("src.dependencies.tracked_db", mock_tracked_db_context),
|
||||
patch("src.deriver.queue_manager.tracked_db", mock_tracked_db_context),
|
||||
patch("src.deriver.consumer.tracked_db", mock_tracked_db_context),
|
||||
patch("src.deriver.enqueue.tracked_db", mock_tracked_db_context),
|
||||
patch("src.routers.peers.tracked_db", mock_tracked_db_context),
|
||||
patch("src.crud.representation.tracked_db", mock_tracked_db_context),
|
||||
patch("src.dreamer.orchestrator.tracked_db", mock_tracked_db_context),
|
||||
patch("src.dreamer.dream_scheduler.tracked_db", mock_tracked_db_context),
|
||||
patch("src.dialectic.chat.tracked_db", mock_tracked_db_context),
|
||||
patch("src.utils.summarizer.tracked_db", mock_tracked_db_context),
|
||||
patch("src.webhooks.events.tracked_db", mock_tracked_db_context),
|
||||
patch("src.webhooks.webhook_delivery.tracked_db", mock_tracked_db_context),
|
||||
patch("src.utils.agent_tools.tracked_db", mock_tracked_db_context),
|
||||
patch("src.utils.search.tracked_db", mock_tracked_db_context),
|
||||
patch("src.crud.document.tracked_db", mock_tracked_db_context),
|
||||
patch("src.crud.message.tracked_db", mock_tracked_db_context),
|
||||
patch("src.reconciler.sync_vectors.tracked_db", mock_tracked_db_context),
|
||||
patch("src.dialectic.core.tracked_db", mock_tracked_db_context),
|
||||
patch("src.dreamer.specialists.tracked_db", mock_tracked_db_context),
|
||||
patch("src.dreamer.surprisal.tracked_db", mock_tracked_db_context),
|
||||
):
|
||||
# Each module imports tracked_db by name, so patch every import site.
|
||||
# Use ExitStack (not a parenthesized `with`) to stay under CPython's
|
||||
# 20-statically-nested-block limit as this list grows.
|
||||
tracked_db_targets = [
|
||||
"src.dependencies.tracked_db",
|
||||
"src.deriver.queue_manager.tracked_db",
|
||||
"src.deriver.consumer.tracked_db",
|
||||
"src.deriver.enqueue.tracked_db",
|
||||
"src.routers.peers.tracked_db",
|
||||
"src.crud.representation.tracked_db",
|
||||
"src.dreamer.orchestrator.tracked_db",
|
||||
"src.dreamer.dream_scheduler.tracked_db",
|
||||
"src.dialectic.chat.tracked_db",
|
||||
"src.utils.summarizer.tracked_db",
|
||||
"src.webhooks.events.tracked_db",
|
||||
"src.webhooks.webhook_delivery.tracked_db",
|
||||
"src.utils.agent_tools.tracked_db",
|
||||
"src.utils.search.tracked_db",
|
||||
"src.crud.document.tracked_db",
|
||||
"src.crud.message.tracked_db",
|
||||
"src.reconciler.sync_vectors.tracked_db",
|
||||
"src.reconciler.embed_now.tracked_db",
|
||||
"src.dialectic.core.tracked_db",
|
||||
"src.dreamer.specialists.tracked_db",
|
||||
"src.dreamer.surprisal.tracked_db",
|
||||
]
|
||||
with ExitStack() as stack:
|
||||
for target in tracked_db_targets:
|
||||
stack.enter_context(patch(target, mock_tracked_db_context))
|
||||
yield
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,268 @@
|
|||
"""
|
||||
Tests for the immediate message-embedding fast path (src/reconciler/embed_now.py).
|
||||
|
||||
These exercise embed_messages_now end-to-end against the test database: it opens
|
||||
its own tracked_db sessions (patched to the test engine in conftest), so each test
|
||||
creates committed fixture rows and asserts on the result via the provided session.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
|
||||
|
||||
from src import models
|
||||
from src.reconciler.embed_now import embed_messages_now, reset_embed_semaphore
|
||||
from src.vector_store import VectorStore
|
||||
|
||||
|
||||
async def _create_message_with_pending_chunks(
|
||||
db_session: AsyncSession,
|
||||
workspace: models.Workspace,
|
||||
peer: models.Peer,
|
||||
chunk_contents: list[str],
|
||||
) -> tuple[str, list[int]]:
|
||||
"""Create a message plus one pending MessageEmbedding row per chunk.
|
||||
|
||||
Returns (message public_id, ordered embedding row ids).
|
||||
"""
|
||||
session = models.Session(name=str(generate_nanoid()), workspace_name=workspace.name)
|
||||
db_session.add(session)
|
||||
await db_session.commit()
|
||||
|
||||
message_id = str(generate_nanoid())
|
||||
message = models.Message(
|
||||
public_id=message_id,
|
||||
session_name=session.name,
|
||||
workspace_name=workspace.name,
|
||||
peer_name=peer.name,
|
||||
content=" ".join(chunk_contents),
|
||||
seq_in_session=1,
|
||||
)
|
||||
db_session.add(message)
|
||||
await db_session.commit()
|
||||
|
||||
rows = [
|
||||
models.MessageEmbedding(
|
||||
content=chunk,
|
||||
message_id=message_id,
|
||||
workspace_name=workspace.name,
|
||||
session_name=session.name,
|
||||
peer_name=peer.name,
|
||||
sync_state="pending",
|
||||
embedding=None,
|
||||
)
|
||||
for chunk in chunk_contents
|
||||
]
|
||||
db_session.add_all(rows)
|
||||
await db_session.commit()
|
||||
for row in rows:
|
||||
await db_session.refresh(row)
|
||||
return message_id, [row.id for row in rows]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_semaphore_fixture():
|
||||
"""Rebuild the module semaphore per test so it binds to the active loop."""
|
||||
reset_embed_semaphore()
|
||||
yield
|
||||
reset_embed_semaphore()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestEmbedMessagesNow:
|
||||
async def test_pgvector_happy_path_marks_synced(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
) -> None:
|
||||
"""pgvector-only mode: rows get a vector and flip to synced immediately."""
|
||||
workspace, peer = sample_data
|
||||
message_id, emb_ids = await _create_message_with_pending_chunks(
|
||||
db_session, workspace, peer, ["hello world"]
|
||||
)
|
||||
|
||||
await embed_messages_now([message_id])
|
||||
|
||||
for emb_id in emb_ids:
|
||||
row = await db_session.get(models.MessageEmbedding, emb_id)
|
||||
assert row is not None
|
||||
await db_session.refresh(row)
|
||||
assert row.sync_state == "synced"
|
||||
assert row.embedding is not None
|
||||
assert row.sync_attempts == 0
|
||||
|
||||
async def test_no_message_ids_is_noop(self) -> None:
|
||||
"""Empty input returns without touching the DB or embedding."""
|
||||
with patch(
|
||||
"src.embedding_client.embedding_client.simple_batch_embed"
|
||||
) as mock_embed:
|
||||
await embed_messages_now([])
|
||||
mock_embed.assert_not_called()
|
||||
|
||||
async def test_already_synced_rows_not_reclaimed(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
) -> None:
|
||||
"""A second run finds no pending rows and does not re-embed."""
|
||||
workspace, peer = sample_data
|
||||
message_id, _ = await _create_message_with_pending_chunks(
|
||||
db_session, workspace, peer, ["first content"]
|
||||
)
|
||||
await embed_messages_now([message_id])
|
||||
|
||||
with patch(
|
||||
"src.embedding_client.embedding_client.simple_batch_embed"
|
||||
) as mock_embed:
|
||||
await embed_messages_now([message_id])
|
||||
mock_embed.assert_not_called()
|
||||
|
||||
async def test_embed_failure_leaves_rows_pending_and_leased(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
) -> None:
|
||||
"""Embedding failure must leave rows pending + leased, attempts untouched,
|
||||
so the reconciler owns retry accounting."""
|
||||
workspace, peer = sample_data
|
||||
message_id, emb_ids = await _create_message_with_pending_chunks(
|
||||
db_session, workspace, peer, ["will fail"]
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.embedding_client.embedding_client.simple_batch_embed",
|
||||
new=AsyncMock(side_effect=RuntimeError("provider down")),
|
||||
):
|
||||
await embed_messages_now([message_id])
|
||||
|
||||
for emb_id in emb_ids:
|
||||
row = await db_session.get(models.MessageEmbedding, emb_id)
|
||||
assert row is not None
|
||||
await db_session.refresh(row)
|
||||
assert row.sync_state == "pending"
|
||||
assert row.embedding is None
|
||||
assert row.sync_attempts == 0 # lease only, no attempt bump
|
||||
assert row.last_sync_at is not None # leased
|
||||
|
||||
async def test_external_store_upserts_with_chunk_positioned_ids(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
mock_vector_store: VectorStore,
|
||||
) -> None:
|
||||
"""External-store mode: upsert each chunk with id {message_id}_{position}
|
||||
and mark rows synced."""
|
||||
workspace, peer = sample_data
|
||||
message_id, emb_ids = await _create_message_with_pending_chunks(
|
||||
db_session, workspace, peer, ["chunk a", "chunk b", "chunk c"]
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.reconciler.embed_now.get_external_vector_store",
|
||||
return_value=mock_vector_store,
|
||||
):
|
||||
await embed_messages_now([message_id])
|
||||
|
||||
upsert_mock: AsyncMock = mock_vector_store.upsert_many # pyright: ignore[reportAssignmentType]
|
||||
upsert_mock.assert_awaited()
|
||||
upserted_ids = {
|
||||
record.id for call in upsert_mock.await_args_list for record in call.args[1]
|
||||
}
|
||||
assert upserted_ids == {
|
||||
f"{message_id}_0",
|
||||
f"{message_id}_1",
|
||||
f"{message_id}_2",
|
||||
}
|
||||
|
||||
for emb_id in emb_ids:
|
||||
row = await db_session.get(models.MessageEmbedding, emb_id)
|
||||
assert row is not None
|
||||
await db_session.refresh(row)
|
||||
assert row.sync_state == "synced"
|
||||
|
||||
async def test_locked_chunk_skipped_keeps_positions_stable(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
db_engine: AsyncEngine,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
mock_vector_store: VectorStore,
|
||||
) -> None:
|
||||
"""If a sibling chunk is locked by another txn, SKIP LOCKED skips it but
|
||||
chunk positions still come from the full sibling ordering — so the claimed
|
||||
chunks keep their {message_id}_0 / _2 ids (not _0 / _1)."""
|
||||
workspace, peer = sample_data
|
||||
message_id, emb_ids = await _create_message_with_pending_chunks(
|
||||
db_session, workspace, peer, ["chunk a", "chunk b", "chunk c"]
|
||||
)
|
||||
locked_id = emb_ids[1] # middle chunk -> position 1
|
||||
|
||||
# Hold a row lock on the middle chunk from an independent transaction.
|
||||
lock_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
|
||||
lock_session = lock_factory()
|
||||
await lock_session.execute(
|
||||
select(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.id == locked_id)
|
||||
.with_for_update()
|
||||
)
|
||||
try:
|
||||
with patch(
|
||||
"src.reconciler.embed_now.get_external_vector_store",
|
||||
return_value=mock_vector_store,
|
||||
):
|
||||
await embed_messages_now([message_id])
|
||||
finally:
|
||||
await lock_session.rollback()
|
||||
await lock_session.close()
|
||||
|
||||
upsert_mock: AsyncMock = mock_vector_store.upsert_many # pyright: ignore[reportAssignmentType]
|
||||
upserted_ids = {
|
||||
record.id for call in upsert_mock.await_args_list for record in call.args[1]
|
||||
}
|
||||
assert upserted_ids == {f"{message_id}_0", f"{message_id}_2"}
|
||||
|
||||
# The locked chunk stays pending; the other two are synced.
|
||||
locked_row = await db_session.get(models.MessageEmbedding, locked_id)
|
||||
assert locked_row is not None
|
||||
await db_session.refresh(locked_row)
|
||||
assert locked_row.sync_state == "pending"
|
||||
for emb_id in (emb_ids[0], emb_ids[2]):
|
||||
row = await db_session.get(models.MessageEmbedding, emb_id)
|
||||
assert row is not None
|
||||
await db_session.refresh(row)
|
||||
assert row.sync_state == "synced"
|
||||
|
||||
async def test_external_store_unavailable_leaves_rows_pending(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
mock_vector_store: VectorStore,
|
||||
) -> None:
|
||||
"""External-store mode: if upsert_many raises VectorStoreError, rows must
|
||||
stay pending with no vector and untouched attempts, so the reconciler
|
||||
heals them. embed_now never bumps sync_attempts."""
|
||||
from src.exceptions import VectorStoreError
|
||||
|
||||
workspace, peer = sample_data
|
||||
message_id, emb_ids = await _create_message_with_pending_chunks(
|
||||
db_session, workspace, peer, ["chunk a", "chunk b"]
|
||||
)
|
||||
|
||||
upsert_mock: AsyncMock = mock_vector_store.upsert_many # pyright: ignore[reportAssignmentType]
|
||||
upsert_mock.side_effect = VectorStoreError("vector store down")
|
||||
|
||||
with patch(
|
||||
"src.reconciler.embed_now.get_external_vector_store",
|
||||
return_value=mock_vector_store,
|
||||
):
|
||||
await embed_messages_now([message_id])
|
||||
|
||||
for emb_id in emb_ids:
|
||||
row = await db_session.get(models.MessageEmbedding, emb_id)
|
||||
assert row is not None
|
||||
await db_session.refresh(row)
|
||||
assert row.sync_state == "pending"
|
||||
assert row.embedding is None
|
||||
assert row.sync_attempts == 0 # embed_now never bumps attempts
|
||||
|
|
@ -23,6 +23,8 @@ from src.reconciler.sync_vectors import (
|
|||
_reconcile_message_embeddings_batch, # pyright: ignore[reportPrivateUsage]
|
||||
_sync_documents, # pyright: ignore[reportPrivateUsage]
|
||||
_sync_message_embeddings, # pyright: ignore[reportPrivateUsage]
|
||||
build_message_vector_record,
|
||||
compute_chunk_positions,
|
||||
run_vector_reconciliation_cycle,
|
||||
)
|
||||
from src.vector_store import (
|
||||
|
|
@ -873,6 +875,82 @@ class TestMessageEmbeddings:
|
|||
assert pending_emb.sync_attempts == 0
|
||||
assert pending_emb.last_sync_at is None
|
||||
|
||||
async def test_pgvector_only_mode_embeds_and_marks_synced(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
) -> None:
|
||||
"""In pgvector-only mode, the reconciler must still embed pending rows."""
|
||||
workspace, peer = sample_data
|
||||
pending_emb = await self._create_pending_message_embedding(
|
||||
db_session, workspace, peer
|
||||
)
|
||||
|
||||
# external_vector_store=None == pgvector-only mode. The reconciler should
|
||||
# re-embed the pending row, write the vector to postgres, and mark synced.
|
||||
synced, failed = await _sync_message_embeddings(db_session, [pending_emb], None)
|
||||
|
||||
await db_session.commit()
|
||||
await db_session.refresh(pending_emb)
|
||||
|
||||
assert synced == 1
|
||||
assert failed == 0
|
||||
assert pending_emb.sync_state == "synced"
|
||||
assert pending_emb.sync_attempts == 0
|
||||
assert pending_emb.embedding is not None
|
||||
|
||||
async def test_all_chunks_of_a_message_claimed_together(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
) -> None:
|
||||
"""A single message's chunks must always be claimed in one batch.
|
||||
|
||||
Selecting by message_id (not row) keeps `{message_id}_{chunk_index}`
|
||||
vector IDs stable across reconciler cycles.
|
||||
"""
|
||||
workspace, peer = sample_data
|
||||
|
||||
# Create one message with 5 chunks.
|
||||
session = models.Session(
|
||||
name=str(generate_nanoid()), workspace_name=workspace.name
|
||||
)
|
||||
db_session.add(session)
|
||||
await db_session.commit()
|
||||
|
||||
message_id = str(generate_nanoid())
|
||||
message = models.Message(
|
||||
public_id=message_id,
|
||||
session_name=session.name,
|
||||
workspace_name=workspace.name,
|
||||
peer_name=peer.name,
|
||||
content="full message content",
|
||||
seq_in_session=1,
|
||||
)
|
||||
db_session.add(message)
|
||||
await db_session.commit()
|
||||
|
||||
chunk_count = 5
|
||||
for i in range(chunk_count):
|
||||
db_session.add(
|
||||
models.MessageEmbedding(
|
||||
content=f"chunk-{i}",
|
||||
message_id=message_id,
|
||||
workspace_name=workspace.name,
|
||||
session_name=session.name,
|
||||
peer_name=peer.name,
|
||||
sync_state="pending",
|
||||
embedding=None,
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
# Even with batch_size=1, all 5 chunks for the message should be claimed
|
||||
# together because the query selects by distinct message_id first.
|
||||
claimed = await _get_message_embeddings_needing_sync(db_session, batch_size=1)
|
||||
assert len(claimed) == chunk_count
|
||||
assert all(emb.message_id == message_id for emb in claimed)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestEndToEndReconciliation:
|
||||
|
|
@ -913,3 +991,90 @@ class TestEndToEndReconciliation:
|
|||
mock_reconcile_docs.assert_awaited_once()
|
||||
mock_reconcile_embs.assert_awaited_once()
|
||||
mock_cleanup_docs.assert_awaited_once()
|
||||
|
||||
|
||||
def test_build_message_vector_record() -> None:
|
||||
"""The shared vector-id/metadata builder: id is {message_id}_{position},
|
||||
embeddings are coerced to float, metadata shape is fixed."""
|
||||
record = build_message_vector_record(
|
||||
message_id="msg_abc",
|
||||
chunk_position=2,
|
||||
session_name="sess",
|
||||
peer_name="peer",
|
||||
embedding=[1, 2, 3], # ints, must be coerced
|
||||
)
|
||||
assert record.id == "msg_abc_2"
|
||||
assert record.embedding == [1.0, 2.0, 3.0]
|
||||
assert all(isinstance(x, float) for x in record.embedding)
|
||||
assert record.metadata == {
|
||||
"message_id": "msg_abc",
|
||||
"session_name": "sess",
|
||||
"peer_name": "peer",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestComputeChunkPositions:
|
||||
"""Direct coverage for compute_chunk_positions, the source of truth for
|
||||
{message_id}_{position} vector ids shared by the reconciler and embed_now."""
|
||||
|
||||
async def test_empty_input_returns_empty(self, db_session: AsyncSession) -> None:
|
||||
assert await compute_chunk_positions(db_session, []) == {}
|
||||
|
||||
async def test_positions_are_per_message_zero_indexed(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
) -> None:
|
||||
"""Each message's rows are numbered from 0 in (message_id, id) order,
|
||||
independent of how rows from other messages interleave."""
|
||||
workspace, peer = sample_data
|
||||
session = models.Session(
|
||||
name=str(generate_nanoid()), workspace_name=workspace.name
|
||||
)
|
||||
db_session.add(session)
|
||||
await db_session.commit()
|
||||
|
||||
# msg_a has 2 chunks, msg_b has 1 chunk.
|
||||
msg_a = str(generate_nanoid())
|
||||
msg_b = str(generate_nanoid())
|
||||
for seq, mid in enumerate((msg_a, msg_b), start=1):
|
||||
db_session.add(
|
||||
models.Message(
|
||||
public_id=mid,
|
||||
session_name=session.name,
|
||||
workspace_name=workspace.name,
|
||||
peer_name=peer.name,
|
||||
content="content",
|
||||
seq_in_session=seq,
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
rows = [
|
||||
models.MessageEmbedding(
|
||||
content=content,
|
||||
message_id=mid,
|
||||
workspace_name=workspace.name,
|
||||
session_name=session.name,
|
||||
peer_name=peer.name,
|
||||
sync_state="pending",
|
||||
embedding=None,
|
||||
)
|
||||
for mid, content in (
|
||||
(msg_a, "a0"),
|
||||
(msg_a, "a1"),
|
||||
(msg_b, "b0"),
|
||||
)
|
||||
]
|
||||
db_session.add_all(rows)
|
||||
await db_session.commit()
|
||||
for row in rows:
|
||||
await db_session.refresh(row)
|
||||
a0, a1, b0 = (row.id for row in rows)
|
||||
|
||||
positions = await compute_chunk_positions(db_session, [msg_a, msg_b])
|
||||
|
||||
assert positions[a0] == 0
|
||||
assert positions[a1] == 1
|
||||
assert positions[b0] == 0
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from src.config import settings
|
|||
from src.crud import create_messages
|
||||
from src.crud import message as message_crud
|
||||
from src.models import Message, Peer, Workspace
|
||||
from src.reconciler.sync_vectors import run_vector_reconciliation_cycle
|
||||
from src.schemas import MessageCreate
|
||||
from src.utils.search import search
|
||||
|
||||
|
|
@ -161,9 +162,12 @@ async def test_blank_messages_are_not_sent_for_embedding(
|
|||
nonblank_content,
|
||||
]
|
||||
|
||||
mock_openai_embeddings["batch_embed"].assert_awaited_once()
|
||||
batch_arg = mock_openai_embeddings["batch_embed"].await_args.args[0]
|
||||
assert batch_arg == {created_messages[1].public_id: nonblank_content}
|
||||
# Inline embedding is gone: create_messages should chunk via prepare_chunks
|
||||
# (no network) and never call batch_embed.
|
||||
mock_openai_embeddings["batch_embed"].assert_not_awaited()
|
||||
mock_openai_embeddings["prepare_chunks"].assert_called_once()
|
||||
prepare_arg = mock_openai_embeddings["prepare_chunks"].call_args.args[0]
|
||||
assert prepare_arg == {created_messages[1].public_id: nonblank_content}
|
||||
|
||||
stmt = select(models.MessageEmbedding).where(
|
||||
models.MessageEmbedding.message_id.in_(
|
||||
|
|
@ -176,6 +180,8 @@ async def test_blank_messages_are_not_sent_for_embedding(
|
|||
assert len(embedding_records) == 1
|
||||
assert embedding_records[0].message_id == created_messages[1].public_id
|
||||
assert embedding_records[0].content == nonblank_content
|
||||
assert embedding_records[0].sync_state == "pending"
|
||||
assert embedding_records[0].embedding is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -327,13 +333,23 @@ async def test_semantic_search_when_embeddings_enabled(
|
|||
assert len(created_messages) == 1
|
||||
created_message = created_messages[0]
|
||||
|
||||
# Verify the embedding was created
|
||||
stmt = select(models.MessageEmbedding).where(
|
||||
models.MessageEmbedding.message_id == created_message.public_id
|
||||
# The pending row exists, but the embedding is generated by the reconciler.
|
||||
# Drive a reconciliation cycle so the row gets an embedding before search.
|
||||
await db_session.commit()
|
||||
await run_vector_reconciliation_cycle()
|
||||
|
||||
# Verify the row was created and reconciled. expire_on_commit=False keeps
|
||||
# stale cached ORM rows, so use populate_existing() to force a reload from
|
||||
# the DB (the reconciler wrote in a different session).
|
||||
stmt = (
|
||||
select(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.message_id == created_message.public_id)
|
||||
.execution_options(populate_existing=True)
|
||||
)
|
||||
result = await db_session.execute(stmt)
|
||||
embedding_record = result.scalar_one_or_none()
|
||||
assert embedding_record is not None
|
||||
assert embedding_record.sync_state == "synced"
|
||||
|
||||
# Now test semantic search without explicitly setting semantic=True
|
||||
# This should use semantic search because EMBED_MESSAGES is True
|
||||
|
|
@ -361,6 +377,77 @@ async def test_semantic_search_when_embeddings_enabled(
|
|||
assert created_message.public_id in found_message_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pgvector_search_excludes_pending_unembedded_rows(
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[Workspace, Peer],
|
||||
):
|
||||
"""Pending MessageEmbedding rows (embedding=None, awaiting the immediate
|
||||
path or reconciler) must not appear in pgvector semantic search results:
|
||||
their NULL distance sorts last and would pad the window with unranked
|
||||
messages."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
embedded_id = str(generate_nanoid())
|
||||
pending_id = str(generate_nanoid())
|
||||
for seq, (mid, content) in enumerate(
|
||||
((embedded_id, "embedded message"), (pending_id, "pending message")), start=1
|
||||
):
|
||||
db_session.add(
|
||||
models.Message(
|
||||
public_id=mid,
|
||||
session_name=test_session.name,
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
content=content,
|
||||
seq_in_session=seq,
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
dims = settings.EMBEDDING.VECTOR_DIMENSIONS
|
||||
db_session.add_all(
|
||||
[
|
||||
models.MessageEmbedding(
|
||||
content="embedded message",
|
||||
message_id=embedded_id,
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer.name,
|
||||
sync_state="synced",
|
||||
embedding=[0.1] * dims,
|
||||
),
|
||||
models.MessageEmbedding(
|
||||
content="pending message",
|
||||
message_id=pending_id,
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer.name,
|
||||
sync_state="pending",
|
||||
embedding=None,
|
||||
),
|
||||
]
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
snippets = await message_crud._search_messages_pgvector( # pyright: ignore[reportPrivateUsage]
|
||||
db_session,
|
||||
test_workspace.name,
|
||||
test_session.name,
|
||||
query_embedding=[0.1] * dims,
|
||||
limit=10,
|
||||
)
|
||||
|
||||
matched_ids = {msg.public_id for matched, _context in snippets for msg in matched}
|
||||
assert embedded_id in matched_ids
|
||||
assert pending_id not in matched_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_merged_snippets_batches_context_query_across_sessions():
|
||||
"""Context expansion should not issue one DB query per matched session."""
|
||||
|
|
@ -638,15 +725,14 @@ async def test_message_chunking_creates_multiple_embeddings(
|
|||
|
||||
test_message_content = "This is a very long message that should be chunked into multiple pieces because it exceeds the token limit that we set for testing purposes. This message contains many words and should definitely be split into multiple chunks."
|
||||
|
||||
def mock_batch_embed_chunked(
|
||||
id_resource_dict: dict[str, str],
|
||||
) -> dict[str, list[list[float]]]:
|
||||
return {
|
||||
text_id: [[0.1] * 1536, [0.2] * 1536, [0.3] * 1536] # 3 chunks per message
|
||||
for text_id in id_resource_dict
|
||||
}
|
||||
chunk_texts = ["chunk-a", "chunk-b", "chunk-c"]
|
||||
|
||||
mock_openai_embeddings["batch_embed"].side_effect = mock_batch_embed_chunked
|
||||
def mock_prepare_chunks_chunked(
|
||||
id_resource_dict: dict[str, str],
|
||||
) -> dict[str, list[str]]:
|
||||
return {text_id: list(chunk_texts) for text_id in id_resource_dict}
|
||||
|
||||
mock_openai_embeddings["prepare_chunks"].side_effect = mock_prepare_chunks_chunked
|
||||
|
||||
messages = [
|
||||
MessageCreate(
|
||||
|
|
@ -666,22 +752,26 @@ async def test_message_chunking_creates_multiple_embeddings(
|
|||
assert len(created_messages) == 1
|
||||
created_message = created_messages[0]
|
||||
|
||||
# Query the MessageEmbedding table to verify multiple embeddings were created
|
||||
stmt = select(models.MessageEmbedding).where(
|
||||
models.MessageEmbedding.message_id == created_message.public_id
|
||||
# batch_embed is no longer called inline; embedding is deferred to reconciler.
|
||||
mock_openai_embeddings["batch_embed"].assert_not_awaited()
|
||||
|
||||
# Query the MessageEmbedding table to verify multiple pending rows were created,
|
||||
# one per chunk, in chunk order (id ascending).
|
||||
stmt = (
|
||||
select(models.MessageEmbedding)
|
||||
.where(models.MessageEmbedding.message_id == created_message.public_id)
|
||||
.order_by(models.MessageEmbedding.id)
|
||||
)
|
||||
result = await db_session.execute(stmt)
|
||||
embedding_records = list(result.scalars().all())
|
||||
|
||||
# Verify multiple embedding records were created (one per chunk)
|
||||
# Embedding vectors are now stored externally in the vector store
|
||||
assert len(embedding_records) == 3 # Should have 3 embeddings for 3 chunks
|
||||
assert len(embedding_records) == 3
|
||||
assert [r.content for r in embedding_records] == chunk_texts
|
||||
|
||||
for _, embedding_record in enumerate(embedding_records):
|
||||
for embedding_record in embedding_records:
|
||||
assert embedding_record.message_id == created_message.public_id
|
||||
assert (
|
||||
embedding_record.content == test_message_content
|
||||
) # Full content stored in each
|
||||
assert embedding_record.workspace_name == test_workspace.name
|
||||
assert embedding_record.session_name == test_session.name
|
||||
assert embedding_record.peer_name == test_peer.name
|
||||
assert embedding_record.sync_state == "pending"
|
||||
assert embedding_record.embedding is None
|
||||
|
|
|
|||
|
|
@ -339,3 +339,108 @@ def test_resolve_send_dimensions_never_returns_false_regardless(
|
|||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_batch_embed_respects_token_budget_per_request(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""simple_batch_embed must split inputs across requests so per-request token cap holds."""
|
||||
fake_embeddings = FakeOpenAIEmbeddingsAPI([0.5] * 4)
|
||||
|
||||
class FakeOpenAIClient:
|
||||
def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
|
||||
self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
|
||||
|
||||
monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient)
|
||||
|
||||
# max_input_tokens=100 per single input; max_tokens_per_request=120 total,
|
||||
# so two ~80-token inputs must end up in *separate* requests.
|
||||
client = _EmbeddingClient(
|
||||
EmbeddingModelConfig(
|
||||
transport="openai",
|
||||
model="text-embedding-3-small",
|
||||
api_key="test-key",
|
||||
base_url=None,
|
||||
),
|
||||
vector_dimensions=4,
|
||||
max_input_tokens=100,
|
||||
max_tokens_per_request=120,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
# "word " * 80 produces ~80 tokens with cl100k_base/the model encoding.
|
||||
long_a = ("alpha " * 80).strip()
|
||||
long_b = ("beta " * 80).strip()
|
||||
|
||||
out = await client.simple_batch_embed([long_a, long_b])
|
||||
assert len(out) == 2
|
||||
# Per-request token cap forces two separate requests.
|
||||
assert len(fake_embeddings.calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_batch_embed_rejects_oversized_input(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Inputs that exceed max_embedding_tokens must raise ValueError immediately."""
|
||||
fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4)
|
||||
|
||||
class FakeOpenAIClient:
|
||||
def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
|
||||
self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
|
||||
|
||||
monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient)
|
||||
|
||||
client = _EmbeddingClient(
|
||||
EmbeddingModelConfig(
|
||||
transport="openai",
|
||||
model="text-embedding-3-small",
|
||||
api_key="test-key",
|
||||
base_url=None,
|
||||
),
|
||||
vector_dimensions=4,
|
||||
max_input_tokens=10,
|
||||
max_tokens_per_request=1000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
too_long = ("word " * 50).strip()
|
||||
with pytest.raises(ValueError, match="maximum token limit"):
|
||||
await client.simple_batch_embed([too_long])
|
||||
|
||||
|
||||
def test_prepare_chunks_returns_ordered_chunks(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""prepare_chunks must split oversized inputs using the same rules as batch_embed."""
|
||||
fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4)
|
||||
|
||||
class FakeOpenAIClient:
|
||||
def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
|
||||
self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
|
||||
|
||||
monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient)
|
||||
|
||||
client = _EmbeddingClient(
|
||||
EmbeddingModelConfig(
|
||||
transport="openai",
|
||||
model="text-embedding-3-small",
|
||||
api_key="test-key",
|
||||
base_url=None,
|
||||
),
|
||||
vector_dimensions=4,
|
||||
max_input_tokens=10,
|
||||
max_tokens_per_request=1000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
short_text = "hello"
|
||||
long_text = ("word " * 50).strip()
|
||||
|
||||
out = client.prepare_chunks({"short": short_text, "long": long_text})
|
||||
|
||||
assert out["short"] == [short_text]
|
||||
assert len(out["long"]) > 1
|
||||
# Order preserved
|
||||
assert isinstance(out["long"][0], str)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import datetime
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
|
@ -46,6 +46,92 @@ async def test_create_message(
|
|||
assert "id" in message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_message_schedules_immediate_embed(
|
||||
client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
"""Creating messages should schedule the immediate-embed background task with
|
||||
the created messages' public ids."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
with (
|
||||
patch("src.config.settings.EMBED_MESSAGES", True),
|
||||
patch(
|
||||
"src.routers.messages.embed_messages_now", new=AsyncMock()
|
||||
) as mock_embed_now,
|
||||
):
|
||||
response = client.post(
|
||||
f"/v3/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={"messages": [{"content": "hello", "peer_id": test_peer.name}]},
|
||||
)
|
||||
assert response.status_code == 201
|
||||
public_id = response.json()[0]["id"]
|
||||
mock_embed_now.assert_awaited_once_with([public_id])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_message_skips_embed_when_disabled(
|
||||
client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
"""When EMBED_MESSAGES is disabled, the immediate-embed task is not scheduled."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
with (
|
||||
patch("src.config.settings.EMBED_MESSAGES", False),
|
||||
patch(
|
||||
"src.routers.messages.embed_messages_now", new=AsyncMock()
|
||||
) as mock_embed_now,
|
||||
):
|
||||
response = client.post(
|
||||
f"/v3/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={"messages": [{"content": "hello", "peer_id": test_peer.name}]},
|
||||
)
|
||||
assert response.status_code == 201
|
||||
mock_embed_now.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_upload_schedules_immediate_embed(
|
||||
client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
"""The file-upload path schedules the immediate-embed task with the created
|
||||
messages' public ids, mirroring the session-message path."""
|
||||
import io
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
with (
|
||||
patch("src.config.settings.EMBED_MESSAGES", True),
|
||||
patch(
|
||||
"src.routers.messages.embed_messages_now", new=AsyncMock()
|
||||
) as mock_embed_now,
|
||||
):
|
||||
files = {"file": ("note.txt", io.BytesIO(b"hello world"), "text/plain")}
|
||||
response = client.post(
|
||||
f"/v3/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/upload",
|
||||
files=files,
|
||||
data={"peer_id": test_peer.name},
|
||||
)
|
||||
assert response.status_code == 201
|
||||
expected_ids = [m["id"] for m in response.json()]
|
||||
mock_embed_now.assert_awaited_once_with(expected_ids)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_messages_with_metadata(
|
||||
client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer]
|
||||
|
|
|
|||
Loading…
Reference in New Issue