perf: batch page embeddings via /api/embed and upsert before pruning stale points
- embed_batch now issues one batched /api/embed request (the old loop made one /api/embeddings round-trip per chunk) with a per-text fallback that preserves None-for-failed semantics - update_from_page embeds all chunks in that single call and stores them in one Qdrant batch upsert (upsert_points) - reindex order reversed: upsert new points first, then prune stale ids (deterministic uuid5 ids make overwrite safe) so a mid-way failure no longer leaves the page with zero vectors - VectorUpdateSummary gains status (success/partial/failed) and chunks_skipped; all-embeddings-failed keeps old vectors and reports failure instead of success=True Measured on a 7-chunk page ingest (local server, llm_tester): ~375ms -> ~181ms median over 3 runs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QbFZyDvYksazX6nYQYZ67L
This commit is contained in:
@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
### Changed (performance)
|
||||
|
||||
- **Batched embeddings + delete-last reindex** - `OllamaClient.embed_batch` now sends ONE batched `/api/embed` request (verified against the live Ollama; the old "batch" looped one `/api/embeddings` call per chunk) with a per-text fallback preserving partial-success semantics. `VectorService.update_from_page` embeds all chunks in that single call and upserts them in one Qdrant batch, and the reindex order is reversed: new points are upserted BEFORE stale points are pruned (deterministic uuid5 chunk ids make the overwrite safe), so a mid-way failure can no longer leave a page with zero vectors — the old order deleted everything first. The summary now reports `status` (`success`/`partial`/`failed`) and `chunks_skipped` instead of unconditional `success=True`; a fully failed embedding pass keeps the old vectors and reports failure. Measured on a real 7-chunk page ingest as `llm_tester` against the local server: ~375ms → ~181ms median (3 runs each).
|
||||
- **Document sync indexing fixed** - `DocumentSyncService._index_vectors` now awaits `ensure_collection` (the coroutine was created but never ran, so fresh tenants had no collection at upsert time), filters out `None` entries from `embed_batch` so one failed chunk embedding no longer aborts the whole document upsert (all-failed still reports failure), and routes the raw `client.delete`/`client.upsert` calls through the async wrapper (`delete_by_filter` / new batch `upsert_points`). Offline unit tests added.
|
||||
- **Async Qdrant client** - `QdrantClientWrapper` now uses `AsyncQdrantClient` with an explicit timeout (`QDRANT_TIMEOUT`, default 30s). Every vector call previously ran on the synchronous client inside async wrapper methods, blocking the FastAPI event loop for the duration of each Qdrant round-trip. The wrapper API is unchanged (all methods were already `async`), so call sites only gained real awaits. The HybridRAG document leg was moved off the deprecated raw `client.search` onto the wrapper's `search_vectors` (fixing a latent `AttributeError`: it called the nonexistent `ollama.embed_text`, so the leg always reported `failed`), and the health check awaits `get_collections`.
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ class OllamaClient:
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.model = model
|
||||
self.embeddings_url = f"{self.base_url}/api/embeddings"
|
||||
self.embed_url = f"{self.base_url}/api/embed"
|
||||
self.generate_url = f"{self.base_url}/api/generate"
|
||||
self.tags_url = f"{self.base_url}/api/tags"
|
||||
self.client = httpx.AsyncClient(timeout=120.0) # Embeddings can be slow
|
||||
@@ -106,8 +107,37 @@ class OllamaClient:
|
||||
>>> len(embeddings)
|
||||
3
|
||||
"""
|
||||
embeddings = []
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
# Single batched request via Ollama's /api/embed (the old
|
||||
# implementation looped one /api/embeddings call per text).
|
||||
try:
|
||||
response = await self.client.post(
|
||||
self.embed_url,
|
||||
json={"model": self.model, "input": texts}
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
embeddings = data.get("embeddings")
|
||||
|
||||
if embeddings is not None and len(embeddings) == len(texts):
|
||||
if show_progress:
|
||||
logger.info(f"Batched embedding complete: {len(embeddings)}/{len(texts)}")
|
||||
return embeddings
|
||||
|
||||
logger.warning(
|
||||
f"Batched embed returned {len(embeddings or [])} vectors for "
|
||||
f"{len(texts)} inputs, falling back to per-text embedding"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Batched embed failed ({e}), falling back to per-text embedding"
|
||||
)
|
||||
|
||||
# Fallback: per-text embedding preserves partial-success semantics
|
||||
# (None entries for texts that failed to embed).
|
||||
embeddings = []
|
||||
for i, text in enumerate(texts):
|
||||
if show_progress and i % 10 == 0:
|
||||
logger.info(f"Embedding progress: {i}/{len(texts)}")
|
||||
|
||||
@@ -67,6 +67,8 @@ class VectorUpdateSummary(BaseModel):
|
||||
chunks_created: int = Field(default=0, description="New chunks created")
|
||||
chunks_updated: int = Field(default=0, description="Existing chunks updated")
|
||||
chunks_deleted: int = Field(default=0, description="Old chunks deleted")
|
||||
chunks_skipped: int = Field(default=0, description="Chunks skipped (embedding failed)")
|
||||
status: str = Field(default="success", description="'success', 'partial' (some chunks skipped), or 'failed'")
|
||||
total_chunks: int = Field(default=0, description="Total chunks for this page")
|
||||
embedding_dim: int = Field(default=768, description="Embedding dimensionality")
|
||||
processing_time_ms: float = Field(..., description="Processing time in milliseconds")
|
||||
|
||||
@@ -154,50 +154,79 @@ class VectorService:
|
||||
chunks = self._chunk_text(content)
|
||||
logger.info(f"Split page {page_id} into {len(chunks)} chunks")
|
||||
|
||||
# Delete existing chunks for this page
|
||||
deleted_count = await self.qdrant.delete_by_filter(
|
||||
collection_name=collection_name,
|
||||
filter_conditions={"page_id": page_id}
|
||||
)
|
||||
# Generate ALL embeddings in one batched /api/embed call
|
||||
# (previously one sequential Ollama round-trip per chunk)
|
||||
embeddings = await self.ollama.embed_batch(chunks)
|
||||
|
||||
# Generate embeddings and upsert chunks
|
||||
chunks_created = 0
|
||||
for idx, chunk_text in enumerate(chunks):
|
||||
# Generate deterministic UUID from page_id and chunk_index
|
||||
chunk_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, f"page_{page_id}_chunk_{idx}"))
|
||||
|
||||
# Generate embedding
|
||||
embedding = await self.ollama.embed(chunk_text)
|
||||
# Build points; deterministic uuid5 IDs mean re-upserting the
|
||||
# same page overwrites its previous chunks in place.
|
||||
points = []
|
||||
chunks_skipped = 0
|
||||
embedding_dim = 768
|
||||
for idx, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
|
||||
if not embedding:
|
||||
logger.error(f"Failed to generate embedding for chunk {chunk_id}")
|
||||
chunks_skipped += 1
|
||||
logger.error(f"Failed to generate embedding for page {page_id} chunk {idx}")
|
||||
continue
|
||||
|
||||
# Prepare metadata
|
||||
metadata = {
|
||||
"page_id": page_id,
|
||||
"page_title": title,
|
||||
"page_path": path,
|
||||
"chunk_index": idx,
|
||||
"chunk_text": chunk_text,
|
||||
"user": user
|
||||
}
|
||||
embedding_dim = len(embedding)
|
||||
chunk_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, f"page_{page_id}_chunk_{idx}"))
|
||||
points.append({
|
||||
"id": chunk_id,
|
||||
"vector": embedding,
|
||||
"payload": {
|
||||
"page_id": page_id,
|
||||
"page_title": title,
|
||||
"page_path": path,
|
||||
"chunk_index": idx,
|
||||
"chunk_text": chunk_text,
|
||||
"user": user
|
||||
}
|
||||
})
|
||||
|
||||
# Upsert to Qdrant
|
||||
success = await self.qdrant.upsert_vector(
|
||||
# Upsert BEFORE deleting stale points. The old order (delete all,
|
||||
# then embed+upsert one by one) left the page with ZERO vectors if
|
||||
# anything failed mid-way; now old vectors survive until the new
|
||||
# ones are safely stored.
|
||||
chunks_created = 0
|
||||
if points:
|
||||
chunks_created = await self.qdrant.upsert_points(
|
||||
collection_name=collection_name,
|
||||
vector_id=chunk_id,
|
||||
vector=embedding,
|
||||
payload=metadata
|
||||
points=points
|
||||
)
|
||||
|
||||
if success:
|
||||
chunks_created += 1
|
||||
# Prune stale points from a previous version of the page (chunk
|
||||
# indexes beyond the new count, or indexes whose new embedding
|
||||
# failed). Only prune if the new upsert actually stored points —
|
||||
# a fully failed embedding pass must not wipe the old vectors.
|
||||
deleted_count = 0
|
||||
if points:
|
||||
new_ids = {p["id"] for p in points}
|
||||
existing = await self.qdrant.scroll_all_points(
|
||||
collection_name=collection_name,
|
||||
filter_conditions={"page_id": page_id},
|
||||
with_payload=False
|
||||
)
|
||||
stale_ids = [pt["id"] for pt in existing if pt["id"] not in new_ids]
|
||||
if stale_ids:
|
||||
deleted_count = await self.qdrant.delete_by_ids(
|
||||
collection_name=collection_name,
|
||||
point_ids=stale_ids
|
||||
)
|
||||
|
||||
processing_time_ms = (time.time() - start_time) * 1000
|
||||
|
||||
if chunks_created == 0:
|
||||
status = "failed"
|
||||
elif chunks_skipped > 0:
|
||||
status = "partial"
|
||||
else:
|
||||
status = "success"
|
||||
|
||||
logger.info(
|
||||
f"Updated vectors for page {page_id}: "
|
||||
f"{chunks_created} chunks created, {deleted_count} old chunks deleted"
|
||||
f"{chunks_created} chunks created, {chunks_skipped} skipped, "
|
||||
f"{deleted_count} stale chunks deleted ({status})"
|
||||
)
|
||||
|
||||
return VectorUpdateSummary(
|
||||
@@ -205,10 +234,16 @@ class VectorService:
|
||||
page_title=title,
|
||||
chunks_created=chunks_created,
|
||||
chunks_deleted=deleted_count,
|
||||
chunks_skipped=chunks_skipped,
|
||||
total_chunks=chunks_created,
|
||||
embedding_dim=len(embedding) if embedding else 768,
|
||||
embedding_dim=embedding_dim,
|
||||
processing_time_ms=processing_time_ms,
|
||||
success=True
|
||||
success=chunks_created > 0,
|
||||
status=status,
|
||||
error_message=(
|
||||
f"{chunks_skipped}/{len(chunks)} chunk embeddings failed"
|
||||
if chunks_skipped else None
|
||||
)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
@@ -41,6 +41,9 @@ def mock_qdrant():
|
||||
qdrant.search_vectors = AsyncMock(return_value=[])
|
||||
qdrant.search_with_expiry_filter = AsyncMock(return_value=[])
|
||||
qdrant.upsert_vector = AsyncMock(return_value=True)
|
||||
qdrant.upsert_points = AsyncMock(side_effect=lambda collection_name, points: len(points))
|
||||
qdrant.scroll_all_points = AsyncMock(return_value=[])
|
||||
qdrant.delete_by_ids = AsyncMock(return_value=0)
|
||||
qdrant.delete_by_filter = AsyncMock(return_value=0)
|
||||
return qdrant
|
||||
|
||||
@@ -49,6 +52,9 @@ def mock_qdrant():
|
||||
def mock_ollama():
|
||||
ollama = MagicMock()
|
||||
ollama.embed = AsyncMock(return_value=[0.1] * 768)
|
||||
ollama.embed_batch = AsyncMock(
|
||||
side_effect=lambda texts, **kw: [[0.1] * 768 for _ in texts]
|
||||
)
|
||||
ollama.embed_text = AsyncMock(return_value=[0.1] * 768)
|
||||
ollama.generate_text = AsyncMock(return_value="{}")
|
||||
return ollama
|
||||
@@ -148,8 +154,8 @@ class TestVectorLegScoping:
|
||||
|
||||
assert summary.success is False
|
||||
assert "outside user" in (summary.error_message or "")
|
||||
mock_qdrant.upsert_vector.assert_not_awaited()
|
||||
mock_qdrant.delete_by_filter.assert_not_awaited()
|
||||
mock_qdrant.upsert_points.assert_not_awaited()
|
||||
mock_qdrant.delete_by_ids.assert_not_awaited()
|
||||
|
||||
async def test_update_from_page_rejects_sibling_prefix_namespace(
|
||||
self, vector_service, mock_qdrant, mock_wiki
|
||||
@@ -163,7 +169,7 @@ class TestVectorLegScoping:
|
||||
summary = await vector_service.update_from_page(page_id=43, user=TENANT)
|
||||
|
||||
assert summary.success is False
|
||||
mock_qdrant.upsert_vector.assert_not_awaited()
|
||||
mock_qdrant.upsert_points.assert_not_awaited()
|
||||
|
||||
async def test_update_from_page_accepts_own_namespace(
|
||||
self, vector_service, mock_qdrant, mock_wiki
|
||||
@@ -177,7 +183,7 @@ class TestVectorLegScoping:
|
||||
|
||||
assert summary.success is True
|
||||
assert (
|
||||
mock_qdrant.upsert_vector.await_args.kwargs["collection_name"]
|
||||
mock_qdrant.upsert_points.await_args.kwargs["collection_name"]
|
||||
== TENANT_COLLECTION
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
Offline unit tests for the batched, delete-last vector reindex path.
|
||||
|
||||
Pins the Phase D fixes in VectorService.update_from_page:
|
||||
- embeddings come from ONE batched embed_batch call, not per-chunk embed()
|
||||
- new points are upserted BEFORE stale points are deleted (deterministic
|
||||
uuid5 ids make the overwrite safe), so a failure can no longer leave
|
||||
the page with zero vectors
|
||||
- the summary reports partial/failed status instead of unconditional
|
||||
success=True when chunks are skipped
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.vector_service import VectorService
|
||||
|
||||
TENANT = "llm_tester"
|
||||
TENANT_COLLECTION = "library_desk_llm_tester"
|
||||
PAGE = {
|
||||
"id": 44,
|
||||
"title": "T",
|
||||
"path": "users/llm_tester/page",
|
||||
"content": " ".join(f"word{i}" for i in range(1200)), # 3 chunks @ 500/50
|
||||
"tags": [],
|
||||
}
|
||||
|
||||
|
||||
def _chunk_id(page_id: int, idx: int) -> str:
|
||||
return str(uuid.uuid5(uuid.NAMESPACE_DNS, f"page_{page_id}_chunk_{idx}"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_qdrant():
|
||||
qdrant = MagicMock()
|
||||
qdrant.collection_exists = AsyncMock(return_value=True)
|
||||
qdrant.ensure_collection = AsyncMock()
|
||||
qdrant.upsert_points = AsyncMock(side_effect=lambda collection_name, points: len(points))
|
||||
qdrant.scroll_all_points = AsyncMock(return_value=[])
|
||||
qdrant.delete_by_ids = AsyncMock(side_effect=lambda collection_name, point_ids: len(point_ids))
|
||||
qdrant.delete_by_filter = AsyncMock(return_value=0)
|
||||
return qdrant
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ollama():
|
||||
ollama = MagicMock()
|
||||
ollama.embed_batch = AsyncMock(
|
||||
side_effect=lambda texts, **kw: [[0.1] * 768 for _ in texts]
|
||||
)
|
||||
return ollama
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_wiki():
|
||||
wiki = MagicMock()
|
||||
wiki.get_page = AsyncMock(return_value=dict(PAGE))
|
||||
return wiki
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def service(mock_qdrant, mock_wiki, mock_ollama):
|
||||
return VectorService(mock_qdrant, mock_wiki, mock_ollama)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestBatchedReindex:
|
||||
async def test_single_batched_embed_and_single_upsert(
|
||||
self, service, mock_qdrant, mock_ollama
|
||||
):
|
||||
summary = await service.update_from_page(page_id=44, user=TENANT)
|
||||
|
||||
assert summary.success is True
|
||||
assert summary.status == "success"
|
||||
assert summary.chunks_created == 3
|
||||
mock_ollama.embed_batch.assert_awaited_once()
|
||||
mock_qdrant.upsert_points.assert_awaited_once()
|
||||
|
||||
async def test_upsert_happens_before_stale_delete(self, service, mock_qdrant):
|
||||
order = []
|
||||
mock_qdrant.upsert_points = AsyncMock(
|
||||
side_effect=lambda collection_name, points: order.append("upsert") or len(points)
|
||||
)
|
||||
mock_qdrant.delete_by_ids = AsyncMock(
|
||||
side_effect=lambda collection_name, point_ids: order.append("delete") or len(point_ids)
|
||||
)
|
||||
# One stale point from a previous, longer version of the page
|
||||
mock_qdrant.scroll_all_points = AsyncMock(return_value=[
|
||||
{"id": _chunk_id(44, i), "payload": {}} for i in range(4)
|
||||
])
|
||||
|
||||
summary = await service.update_from_page(page_id=44, user=TENANT)
|
||||
|
||||
assert order == ["upsert", "delete"]
|
||||
assert summary.chunks_deleted == 1 # only the stale 4th chunk
|
||||
deleted = mock_qdrant.delete_by_ids.await_args.kwargs["point_ids"]
|
||||
assert deleted == [_chunk_id(44, 3)]
|
||||
|
||||
async def test_partial_embedding_failure_marks_partial(
|
||||
self, service, mock_qdrant, mock_ollama
|
||||
):
|
||||
mock_ollama.embed_batch = AsyncMock(
|
||||
side_effect=lambda texts, **kw: [
|
||||
[0.1] * 768 if i != 1 else None for i in range(len(texts))
|
||||
]
|
||||
)
|
||||
|
||||
summary = await service.update_from_page(page_id=44, user=TENANT)
|
||||
|
||||
assert summary.status == "partial"
|
||||
assert summary.success is True
|
||||
assert summary.chunks_created == 2
|
||||
assert summary.chunks_skipped == 1
|
||||
assert "failed" in (summary.error_message or "")
|
||||
|
||||
async def test_total_embedding_failure_keeps_old_vectors(
|
||||
self, service, mock_qdrant, mock_ollama
|
||||
):
|
||||
mock_ollama.embed_batch = AsyncMock(
|
||||
side_effect=lambda texts, **kw: [None for _ in texts]
|
||||
)
|
||||
mock_qdrant.scroll_all_points = AsyncMock(return_value=[
|
||||
{"id": _chunk_id(44, 0), "payload": {}}
|
||||
])
|
||||
|
||||
summary = await service.update_from_page(page_id=44, user=TENANT)
|
||||
|
||||
assert summary.status == "failed"
|
||||
assert summary.success is False
|
||||
# Old vectors are NOT wiped when nothing new was stored
|
||||
mock_qdrant.upsert_points.assert_not_awaited()
|
||||
mock_qdrant.delete_by_ids.assert_not_awaited()
|
||||
Reference in New Issue
Block a user