1169 lines
40 KiB
Python
1169 lines
40 KiB
Python
"""
|
|
Tests for vector store reconciliation.
|
|
|
|
This module tests the vector reconciliation system that syncs documents and
|
|
message embeddings to the vector store, handling failures and retries.
|
|
"""
|
|
|
|
import datetime
|
|
from contextlib import asynccontextmanager
|
|
from typing import cast
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from nanoid import generate as generate_nanoid
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from src import models
|
|
from src.reconciler.sync_vectors import (
|
|
MAX_SYNC_ATTEMPTS,
|
|
ReconciliationMetrics,
|
|
_get_documents_needing_sync, # pyright: ignore[reportPrivateUsage]
|
|
_get_message_embeddings_needing_sync, # pyright: ignore[reportPrivateUsage]
|
|
_reconcile_documents_batch, # pyright: ignore[reportPrivateUsage]
|
|
_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 (
|
|
VectorRecord,
|
|
VectorStore,
|
|
_hash_namespace_components,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestStateTransitions:
|
|
"""Test document sync_state transitions."""
|
|
|
|
async def test_pending_to_synced_on_success(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test documents transition from pending → synced on successful sync."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create documents in pending state with embeddings
|
|
docs = [
|
|
models.Document(
|
|
content=f"doc_{i}",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
sync_attempts=0,
|
|
embedding=[float(i)] * 1536, # Mock embedding
|
|
)
|
|
for i in range(3)
|
|
]
|
|
db_session.add_all(docs)
|
|
await db_session.commit()
|
|
for doc in docs:
|
|
await db_session.refresh(doc)
|
|
|
|
# Mock vector store to succeed
|
|
mock_vector_store = MagicMock(spec=VectorStore)
|
|
mock_vector_store.get_vector_namespace = MagicMock(
|
|
return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}"
|
|
)
|
|
mock_vector_store.upsert_many = AsyncMock(return_value=None)
|
|
|
|
# Run sync
|
|
synced, failed = await _sync_documents(db_session, docs, mock_vector_store)
|
|
|
|
# Verify results
|
|
assert synced == 3
|
|
assert failed == 0
|
|
|
|
# Check state transitions
|
|
for doc in docs:
|
|
await db_session.refresh(doc)
|
|
assert doc.sync_state == "synced"
|
|
assert doc.sync_attempts == 0
|
|
assert doc.last_sync_at is not None
|
|
|
|
async def test_pending_to_pending_with_incremented_attempts(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test documents remain pending with incremented attempts on partial failure."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create document in pending state
|
|
doc = models.Document(
|
|
content="test doc",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
sync_attempts=2, # Already failed twice
|
|
embedding=[1.0] * 1536,
|
|
)
|
|
db_session.add(doc)
|
|
await db_session.commit()
|
|
await db_session.refresh(doc)
|
|
|
|
# Mock vector store to fail with exception
|
|
mock_vector_store = MagicMock(spec=VectorStore)
|
|
mock_vector_store.get_vector_namespace = MagicMock(
|
|
return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}"
|
|
)
|
|
mock_vector_store.upsert_many = AsyncMock(
|
|
side_effect=Exception("Vector store failed")
|
|
)
|
|
|
|
# Run sync
|
|
synced, failed = await _sync_documents(db_session, [doc], mock_vector_store)
|
|
|
|
# Verify failure recorded
|
|
assert synced == 0
|
|
assert failed == 1
|
|
|
|
# Check sync_attempts incremented
|
|
await db_session.refresh(doc)
|
|
assert doc.sync_state == "pending" # Still pending
|
|
assert doc.sync_attempts == 3 # Incremented
|
|
assert doc.last_sync_at is not None # Updated timestamp
|
|
|
|
async def test_pending_to_failed_after_max_attempts(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test documents transition to failed after MAX_SYNC_ATTEMPTS failures."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create document at max attempts - 1
|
|
doc = models.Document(
|
|
content="failing doc",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
sync_attempts=MAX_SYNC_ATTEMPTS - 1, # One more attempt will hit limit
|
|
embedding=[1.0] * 1536,
|
|
)
|
|
db_session.add(doc)
|
|
await db_session.commit()
|
|
await db_session.refresh(doc)
|
|
|
|
# Mock vector store to fail with exception
|
|
mock_vector_store = MagicMock(spec=VectorStore)
|
|
mock_vector_store.get_vector_namespace = MagicMock(
|
|
return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}"
|
|
)
|
|
mock_vector_store.upsert_many = AsyncMock(
|
|
side_effect=Exception("Vector store failed")
|
|
)
|
|
|
|
# Run sync - this should be the final attempt
|
|
synced, failed = await _sync_documents(db_session, [doc], mock_vector_store)
|
|
|
|
# Verify marked as failed
|
|
assert synced == 0
|
|
assert failed == 1
|
|
|
|
await db_session.refresh(doc)
|
|
assert doc.sync_state == "failed" # Permanently failed
|
|
assert doc.sync_attempts == MAX_SYNC_ATTEMPTS
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestBatchProcessing:
|
|
"""Test batch processing and namespace grouping."""
|
|
|
|
async def test_documents_grouped_by_namespace(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test documents from different collections are grouped by namespace."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create another peer for different observer/observed combinations
|
|
peer2 = models.Peer(
|
|
name="peer2",
|
|
workspace_name=workspace.name,
|
|
)
|
|
db_session.add(peer2)
|
|
await db_session.commit()
|
|
|
|
# Create collections for both observer/observed combinations
|
|
collection1 = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
collection2 = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer2.name,
|
|
)
|
|
db_session.add_all([collection1, collection2])
|
|
await db_session.commit()
|
|
|
|
# Create documents for different namespaces
|
|
docs = [
|
|
# Namespace 1: peer1 → peer1
|
|
models.Document(
|
|
content="doc1",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
embedding=[1.0] * 1536,
|
|
),
|
|
# Namespace 2: peer1 → peer2
|
|
models.Document(
|
|
content="doc2",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer2.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
embedding=[2.0] * 1536,
|
|
),
|
|
# Namespace 1 again: peer1 → peer1
|
|
models.Document(
|
|
content="doc3",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
embedding=[3.0] * 1536,
|
|
),
|
|
]
|
|
db_session.add_all(docs)
|
|
await db_session.commit()
|
|
|
|
# Mock vector store to track calls by namespace
|
|
mock_vector_store = MagicMock(spec=VectorStore)
|
|
namespace_calls: dict[str, list[VectorRecord]] = {}
|
|
|
|
def mock_get_namespace(
|
|
_namespace_type: str, workspace: str, observer: str, observed: str
|
|
) -> str:
|
|
return f"honcho.doc.{_hash_namespace_components(workspace, observer, observed)}"
|
|
|
|
async def mock_upsert(namespace: str, vectors: list[VectorRecord]) -> None:
|
|
if namespace not in namespace_calls:
|
|
namespace_calls[namespace] = []
|
|
namespace_calls[namespace].extend(vectors)
|
|
return
|
|
|
|
mock_vector_store.get_vector_namespace = mock_get_namespace
|
|
mock_vector_store.upsert_many = mock_upsert
|
|
|
|
# Run sync
|
|
synced, failed = await _sync_documents(db_session, docs, mock_vector_store)
|
|
|
|
# Verify all synced
|
|
assert synced == 3
|
|
assert failed == 0
|
|
|
|
# Verify namespaces
|
|
expected_ns1 = f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}"
|
|
expected_ns2 = f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer2.name)}"
|
|
|
|
assert expected_ns1 in namespace_calls
|
|
assert expected_ns2 in namespace_calls
|
|
assert len(namespace_calls[expected_ns1]) == 2 # doc1 and doc3
|
|
assert len(namespace_calls[expected_ns2]) == 1 # doc2
|
|
|
|
async def test_batch_size_respected(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test that batch size limits are respected when fetching pending documents."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create more documents than batch size
|
|
docs = [
|
|
models.Document(
|
|
content=f"doc_{i}",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
embedding=[float(i)] * 1536,
|
|
)
|
|
for i in range(150) # More than RECONCILIATION_BATCH_SIZE (100)
|
|
]
|
|
db_session.add_all(docs)
|
|
await db_session.commit()
|
|
|
|
# Fetch documents with batch size limit
|
|
batch = await _get_documents_needing_sync(db_session, batch_size=100)
|
|
|
|
# Verify batch size respected
|
|
assert len(batch) == 100
|
|
|
|
async def test_documents_respect_retry_backoff(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Pending documents should only be fetched once their backoff has elapsed."""
|
|
workspace, peer1 = sample_data
|
|
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
now = datetime.datetime.now(datetime.timezone.utc)
|
|
ineligible_doc = models.Document(
|
|
content="too soon",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
sync_attempts=1,
|
|
last_sync_at=now - datetime.timedelta(minutes=9, seconds=59),
|
|
embedding=[1.0] * 1536,
|
|
)
|
|
eligible_doc = models.Document(
|
|
content="ready",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
sync_attempts=1,
|
|
last_sync_at=now - datetime.timedelta(minutes=10, seconds=1),
|
|
embedding=[2.0] * 1536,
|
|
)
|
|
db_session.add_all([ineligible_doc, eligible_doc])
|
|
await db_session.commit()
|
|
|
|
pending = await _get_documents_needing_sync(db_session)
|
|
pending_ids = {doc.id for doc in pending}
|
|
|
|
assert eligible_doc.id in pending_ids
|
|
assert ineligible_doc.id not in pending_ids
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestReEmbedding:
|
|
"""Test re-embedding logic for documents with NULL embeddings."""
|
|
|
|
async def test_documents_without_embeddings_are_reembedded(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test documents with NULL embeddings are re-embedded during reconciliation."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create documents without embeddings
|
|
docs = [
|
|
models.Document(
|
|
content=f"doc_{i}",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
embedding=None, # NULL embedding
|
|
)
|
|
for i in range(3)
|
|
]
|
|
db_session.add_all(docs)
|
|
await db_session.commit()
|
|
for doc in docs:
|
|
await db_session.refresh(doc)
|
|
|
|
# Mock embedding client
|
|
with patch("src.reconciler.sync_vectors.embedding_client") as mock_embed_client:
|
|
mock_embed_client.simple_batch_embed = AsyncMock(
|
|
return_value=[[float(i)] * 1536 for i in range(3)]
|
|
)
|
|
|
|
# Mock vector store
|
|
mock_vector_store = MagicMock(spec=VectorStore)
|
|
mock_vector_store.get_vector_namespace = MagicMock(
|
|
return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}"
|
|
)
|
|
mock_vector_store.upsert_many = AsyncMock(return_value=None)
|
|
|
|
# Run sync
|
|
synced, failed = await _sync_documents(db_session, docs, mock_vector_store)
|
|
|
|
# Verify embedding was called
|
|
mock_embed_client.simple_batch_embed.assert_called_once()
|
|
|
|
# Verify documents were synced
|
|
assert synced == 3
|
|
assert failed == 0
|
|
|
|
async def test_large_documents_embedded_in_single_batch(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test that documents are embedded in a single batch (no sub-batching)."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create documents without embeddings
|
|
docs = [
|
|
models.Document(
|
|
content=f"doc_{i}",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
embedding=None, # Will be re-embedded
|
|
)
|
|
for i in range(3)
|
|
]
|
|
db_session.add_all(docs)
|
|
await db_session.commit()
|
|
for doc in docs:
|
|
await db_session.refresh(doc)
|
|
|
|
# Mock embedding client to track batch calls
|
|
batch_call_count = 0
|
|
|
|
async def track_batch_embed(contents: list[str]) -> list[list[float]]:
|
|
nonlocal batch_call_count
|
|
batch_call_count += 1
|
|
return [[1.0] * 1536 for _ in contents]
|
|
|
|
with patch("src.reconciler.sync_vectors.embedding_client") as mock_embed_client:
|
|
mock_embed_client.simple_batch_embed = track_batch_embed
|
|
|
|
# Mock vector store
|
|
mock_vector_store = MagicMock(spec=VectorStore)
|
|
mock_vector_store.get_vector_namespace = MagicMock(
|
|
return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}"
|
|
)
|
|
mock_vector_store.upsert_many = AsyncMock(return_value=None)
|
|
|
|
# Run sync
|
|
await _sync_documents(db_session, docs, mock_vector_store)
|
|
|
|
# Verify single batch call (no sub-batching)
|
|
assert batch_call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestSoftDeleteCleanup:
|
|
"""Test soft delete cleanup functionality."""
|
|
|
|
async def test_cleanup_respects_grace_period(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test cleanup only processes documents deleted_at older than threshold."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create recently soft-deleted document (within grace period)
|
|
recent_doc = models.Document(
|
|
content="recent",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
deleted_at=datetime.datetime.now(datetime.timezone.utc)
|
|
- datetime.timedelta(minutes=2), # Only 2 minutes ago
|
|
)
|
|
db_session.add(recent_doc)
|
|
|
|
# Create old soft-deleted document (outside grace period)
|
|
old_doc = models.Document(
|
|
content="old",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
deleted_at=datetime.datetime.now(datetime.timezone.utc)
|
|
- datetime.timedelta(minutes=10), # 10 minutes ago
|
|
)
|
|
db_session.add(old_doc)
|
|
|
|
await db_session.commit()
|
|
|
|
# Query for documents ready for cleanup (older than 5 minutes)
|
|
cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
|
|
minutes=5
|
|
)
|
|
stmt = (
|
|
select(models.Document)
|
|
.where(models.Document.deleted_at.is_not(None))
|
|
.where(models.Document.deleted_at < cutoff)
|
|
)
|
|
result = await db_session.execute(stmt)
|
|
ready_for_cleanup = result.scalars().all()
|
|
|
|
# Only old_doc should be ready
|
|
assert len(ready_for_cleanup) == 1
|
|
assert ready_for_cleanup[0].id == old_doc.id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestMetricsTracking:
|
|
"""Test reconciliation metrics collection."""
|
|
|
|
async def test_metrics_track_sync_results(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Test ReconciliationMetrics tracks synced/failed counts."""
|
|
workspace, peer1 = sample_data
|
|
|
|
# Create collection (required for documents)
|
|
collection = models.Collection(
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
)
|
|
db_session.add(collection)
|
|
await db_session.commit()
|
|
|
|
# Create session
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
# Create mix of documents that will succeed and fail
|
|
success_doc = models.Document(
|
|
content="success",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
embedding=[1.0] * 1536,
|
|
)
|
|
fail_doc = models.Document(
|
|
content="fail",
|
|
workspace_name=workspace.name,
|
|
observer=peer1.name,
|
|
observed=peer1.name,
|
|
session_name=session.name,
|
|
sync_state="pending",
|
|
sync_attempts=MAX_SYNC_ATTEMPTS - 1, # Will fail on next attempt
|
|
embedding=[2.0] * 1536,
|
|
)
|
|
db_session.add_all([success_doc, fail_doc])
|
|
await db_session.commit()
|
|
|
|
# Note: Since both docs have same namespace, they'll be in one batch
|
|
# This test structure needs adjustment for the actual grouping logic
|
|
# For simplicity, let's test metrics at the function level
|
|
|
|
metrics = ReconciliationMetrics()
|
|
metrics.documents_synced = 5
|
|
metrics.documents_failed = 2
|
|
metrics.message_embeddings_synced = 3
|
|
|
|
assert metrics.total_synced == 8
|
|
assert metrics.total_failed == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestMessageEmbeddings:
|
|
"""Test message embedding reconciliation paths."""
|
|
|
|
async def _create_pending_message_embedding(
|
|
self,
|
|
db_session: AsyncSession,
|
|
workspace: models.Workspace,
|
|
peer: models.Peer,
|
|
) -> models.MessageEmbedding:
|
|
"""Helper to create a pending message embedding with no stored vector."""
|
|
session = models.Session(
|
|
name=str(generate_nanoid()), workspace_name=workspace.name
|
|
)
|
|
db_session.add(session)
|
|
await db_session.commit()
|
|
|
|
message = models.Message(
|
|
public_id=str(generate_nanoid()),
|
|
session_name=session.name,
|
|
workspace_name=workspace.name,
|
|
peer_name=peer.name,
|
|
content="hello world",
|
|
seq_in_session=1,
|
|
)
|
|
db_session.add(message)
|
|
await db_session.commit()
|
|
|
|
emb = models.MessageEmbedding(
|
|
content=message.content,
|
|
message_id=message.public_id,
|
|
workspace_name=workspace.name,
|
|
session_name=session.name,
|
|
peer_name=peer.name,
|
|
sync_state="pending",
|
|
embedding=None,
|
|
)
|
|
db_session.add(emb)
|
|
await db_session.commit()
|
|
await db_session.refresh(emb)
|
|
return emb
|
|
|
|
async def test_pending_embeddings_are_selected_without_vectors(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Pending rows with NULL embeddings should still be reconciled."""
|
|
workspace, peer = sample_data
|
|
pending_emb = await self._create_pending_message_embedding(
|
|
db_session, workspace, peer
|
|
)
|
|
|
|
pending = await _get_message_embeddings_needing_sync(db_session)
|
|
assert any(emb.id == pending_emb.id for emb in pending)
|
|
|
|
async def test_message_embeddings_respect_retry_backoff(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
) -> None:
|
|
"""Pending embeddings should only be fetched once their backoff has elapsed."""
|
|
workspace, peer = sample_data
|
|
ineligible_emb = await self._create_pending_message_embedding(
|
|
db_session, workspace, peer
|
|
)
|
|
eligible_emb = await self._create_pending_message_embedding(
|
|
db_session, workspace, peer
|
|
)
|
|
|
|
now = datetime.datetime.now(datetime.timezone.utc)
|
|
ineligible_emb.sync_attempts = 1
|
|
ineligible_emb.last_sync_at = now - datetime.timedelta(minutes=9, seconds=59)
|
|
eligible_emb.sync_attempts = 1
|
|
eligible_emb.last_sync_at = now - datetime.timedelta(minutes=10, seconds=1)
|
|
await db_session.commit()
|
|
|
|
pending = await _get_message_embeddings_needing_sync(db_session)
|
|
pending_ids = {emb.id for emb in pending}
|
|
|
|
assert eligible_emb.id in pending_ids
|
|
assert ineligible_emb.id not in pending_ids
|
|
|
|
async def test_missing_embeddings_reembedded_and_synced(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
mock_vector_store: VectorStore,
|
|
) -> None:
|
|
"""Reconciliation should re-embed missing payloads and mark them synced."""
|
|
workspace, peer = sample_data
|
|
pending_emb = await self._create_pending_message_embedding(
|
|
db_session, workspace, peer
|
|
)
|
|
|
|
synced, failed = await _sync_message_embeddings(
|
|
db_session, [pending_emb], mock_vector_store
|
|
)
|
|
|
|
await db_session.refresh(pending_emb)
|
|
|
|
assert synced == 1
|
|
assert failed == 0
|
|
assert pending_emb.sync_state == "synced"
|
|
assert pending_emb.sync_attempts == 0
|
|
|
|
# Ensure vector upsert was attempted with a populated embedding
|
|
upsert_mock: AsyncMock = cast(AsyncMock, mock_vector_store.upsert_many)
|
|
await_args = upsert_mock.await_args
|
|
assert await_args is not None
|
|
args = await_args.args
|
|
assert len(args) == 2
|
|
vector_records: list[VectorRecord] = args[1]
|
|
assert vector_records
|
|
assert vector_records[0].embedding
|
|
|
|
async def test_upsert_failure_marks_attempt_and_continues(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
mock_vector_store: VectorStore,
|
|
) -> None:
|
|
"""Failures during upsert should bump attempts and keep row pending/failed."""
|
|
workspace, peer = sample_data
|
|
pending_emb = await self._create_pending_message_embedding(
|
|
db_session, workspace, peer
|
|
)
|
|
|
|
upsert_mock: AsyncMock = cast(AsyncMock, mock_vector_store.upsert_many)
|
|
upsert_mock.side_effect = Exception("boom")
|
|
|
|
synced, failed = await _sync_message_embeddings(
|
|
db_session, [pending_emb], mock_vector_store
|
|
)
|
|
|
|
await db_session.refresh(pending_emb)
|
|
|
|
assert synced == 0
|
|
assert failed == 1
|
|
assert pending_emb.sync_state in {"pending", "failed"}
|
|
assert pending_emb.sync_attempts == 1
|
|
|
|
async def test_unexpected_batch_exception_does_not_bump_unattempted_rows(
|
|
self,
|
|
db_session: AsyncSession,
|
|
sample_data: tuple[models.Workspace, models.Peer],
|
|
mock_vector_store: VectorStore,
|
|
) -> None:
|
|
"""Unexpected wrapper-level failures should not penalize the whole batch."""
|
|
workspace, peer = sample_data
|
|
pending_emb = await self._create_pending_message_embedding(
|
|
db_session, workspace, peer
|
|
)
|
|
|
|
metrics = ReconciliationMetrics()
|
|
with (
|
|
patch(
|
|
"src.reconciler.sync_vectors._sync_message_embeddings",
|
|
side_effect=RuntimeError("unexpected"),
|
|
),
|
|
pytest.raises(RuntimeError, match="unexpected"),
|
|
):
|
|
await _reconcile_message_embeddings_batch(mock_vector_store, metrics)
|
|
|
|
await db_session.refresh(pending_emb)
|
|
assert pending_emb.sync_state == "pending"
|
|
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:
|
|
"""Test full reconciliation cycle."""
|
|
|
|
async def test_reconciliation_cycle_completes_when_no_work_found(self) -> None:
|
|
"""Reconciliation should run each stage once and stop when no work is found."""
|
|
mock_vector_store = MagicMock(spec=VectorStore)
|
|
with (
|
|
patch(
|
|
"src.reconciler.sync_vectors.get_external_vector_store",
|
|
return_value=mock_vector_store,
|
|
),
|
|
patch(
|
|
"src.reconciler.sync_vectors._reconcile_documents_batch",
|
|
new_callable=AsyncMock,
|
|
) as mock_reconcile_docs,
|
|
patch(
|
|
"src.reconciler.sync_vectors._reconcile_message_embeddings_batch",
|
|
new_callable=AsyncMock,
|
|
) as mock_reconcile_embs,
|
|
patch(
|
|
"src.reconciler.sync_vectors._cleanup_documents_batch",
|
|
new_callable=AsyncMock,
|
|
) as mock_cleanup_docs,
|
|
):
|
|
mock_reconcile_docs.return_value = False
|
|
mock_reconcile_embs.return_value = False
|
|
mock_cleanup_docs.return_value = False
|
|
|
|
metrics = await run_vector_reconciliation_cycle()
|
|
|
|
assert isinstance(metrics, ReconciliationMetrics)
|
|
assert metrics.total_synced == 0
|
|
assert metrics.total_failed == 0
|
|
assert metrics.total_cleaned == 0
|
|
|
|
mock_reconcile_docs.assert_awaited_once()
|
|
mock_reconcile_embs.assert_awaited_once()
|
|
mock_cleanup_docs.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestReconcilerTracing:
|
|
"""A Sentry transaction is started only when a sync batch finds real work.
|
|
|
|
Reconciler tasks poll on a fixed interval and usually find nothing; an idle
|
|
cycle must create zero transactions so it doesn't drain Sentry quota.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _fake_tracked_db(db: AsyncMock):
|
|
@asynccontextmanager
|
|
async def _cm(*_args: object, **_kwargs: object):
|
|
yield db
|
|
|
|
return _cm
|
|
|
|
async def test_no_transaction_when_no_embeddings_to_sync(self) -> None:
|
|
"""The no-work path returns before starting a transaction."""
|
|
metrics = ReconciliationMetrics()
|
|
with (
|
|
patch(
|
|
"src.reconciler.sync_vectors.tracked_db",
|
|
self._fake_tracked_db(AsyncMock()),
|
|
),
|
|
patch(
|
|
"src.reconciler.sync_vectors._get_message_embeddings_needing_sync",
|
|
new_callable=AsyncMock,
|
|
return_value=[],
|
|
),
|
|
patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn,
|
|
):
|
|
worked = await _reconcile_message_embeddings_batch(None, metrics)
|
|
|
|
assert worked is False
|
|
txn.assert_not_called()
|
|
|
|
async def test_transaction_started_when_embeddings_present(self) -> None:
|
|
"""A batch with real work starts its own named transaction."""
|
|
metrics = ReconciliationMetrics()
|
|
with (
|
|
patch(
|
|
"src.reconciler.sync_vectors.tracked_db",
|
|
self._fake_tracked_db(AsyncMock()),
|
|
),
|
|
patch(
|
|
"src.reconciler.sync_vectors._get_message_embeddings_needing_sync",
|
|
new_callable=AsyncMock,
|
|
return_value=[MagicMock()],
|
|
),
|
|
patch(
|
|
"src.reconciler.sync_vectors._sync_message_embeddings",
|
|
new_callable=AsyncMock,
|
|
return_value=(1, 0),
|
|
),
|
|
patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn,
|
|
):
|
|
worked = await _reconcile_message_embeddings_batch(None, metrics)
|
|
|
|
assert worked is True
|
|
assert metrics.message_embeddings_synced == 1
|
|
txn.assert_called_once()
|
|
assert txn.call_args.kwargs.get("name") == "reconcile_message_embeddings_batch"
|
|
|
|
async def test_no_transaction_when_no_documents_to_sync(self) -> None:
|
|
"""The document batch also skips tracing when there is nothing to sync."""
|
|
metrics = ReconciliationMetrics()
|
|
with (
|
|
patch(
|
|
"src.reconciler.sync_vectors.tracked_db",
|
|
self._fake_tracked_db(AsyncMock()),
|
|
),
|
|
patch(
|
|
"src.reconciler.sync_vectors._get_documents_needing_sync",
|
|
new_callable=AsyncMock,
|
|
return_value=[],
|
|
),
|
|
patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn,
|
|
):
|
|
worked = await _reconcile_documents_batch(
|
|
MagicMock(spec=VectorStore), metrics
|
|
)
|
|
|
|
assert worked is False
|
|
txn.assert_not_called()
|
|
|
|
|
|
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
|