Spaces:
Sleeping
Sleeping
| """ | |
| Unit tests for the document embedding path. | |
| The regression these lock down: `_embed` used to call the blocking | |
| genai.embed_content() directly inside a coroutine, stalling the event loop for | |
| the entire duration of a document's ingestion. | |
| """ | |
| import asyncio | |
| import time | |
| from unittest.mock import patch | |
| import pytest | |
| from app.services.document_service import EMBED_BATCH_SIZE, DocumentService | |
| pytestmark = pytest.mark.unit | |
| def _fake_embed(model, content, task_type, output_dimensionality, **kwargs): | |
| """Stand-in for genai.embed_content that blocks like the real one.""" | |
| time.sleep(0.05) | |
| if isinstance(content, list): | |
| return {"embedding": [[0.1] * 768 for _ in content]} | |
| return {"embedding": [0.1] * 768} | |
| # ── The core regression ──────────────────────────────────────────────────────── | |
| async def test_embedding_does_not_block_the_event_loop(): | |
| """ | |
| While embeddings are in flight, an unrelated coroutine must still be | |
| scheduled. Before the fix this assertion failed: the loop was held for the | |
| whole blocking call and the ticker never advanced. | |
| """ | |
| service = DocumentService() | |
| ticks = 0 | |
| async def ticker(): | |
| nonlocal ticks | |
| while True: | |
| await asyncio.sleep(0.005) | |
| ticks += 1 | |
| task = asyncio.create_task(ticker()) | |
| with patch("app.services.document_service.genai.embed_content", _fake_embed): | |
| await service._embed_batch(["chunk one", "chunk two"]) | |
| task.cancel() | |
| assert ticks > 0, "event loop was blocked during embedding" | |
| # ── Batching ─────────────────────────────────────────────────────────────────── | |
| async def test_batches_instead_of_one_call_per_chunk(): | |
| service = DocumentService() | |
| calls = [] | |
| def counting_embed(model, content, **kwargs): | |
| calls.append(content) | |
| return {"embedding": [[0.1] * 768 for _ in content]} | |
| texts = [f"chunk {i}" for i in range(250)] | |
| with patch("app.services.document_service.genai.embed_content", counting_embed): | |
| vectors = await service._embed_batch(texts) | |
| assert len(vectors) == 250 | |
| # 250 chunks at batch size 100 => 3 calls, not 250. | |
| assert len(calls) == 3 | |
| assert [len(c) for c in calls] == [EMBED_BATCH_SIZE, EMBED_BATCH_SIZE, 50] | |
| async def test_results_stay_aligned_with_input_order(): | |
| service = DocumentService() | |
| def indexed_embed(model, content, **kwargs): | |
| return {"embedding": [[float(len(t))] * 768 for t in content]} | |
| texts = ["a", "bb", "ccc"] | |
| with patch("app.services.document_service.genai.embed_content", indexed_embed): | |
| vectors = await service._embed_batch(texts) | |
| assert [v[0] for v in vectors] == [1.0, 2.0, 3.0] | |
| # ── Degradation ──────────────────────────────────────────────────────────────── | |
| async def test_falls_back_to_per_chunk_when_batch_fails(): | |
| """One failing batch must not lose the whole document.""" | |
| service = DocumentService() | |
| def flaky(model, content, **kwargs): | |
| if isinstance(content, list): | |
| raise RuntimeError("batch endpoint unavailable") | |
| return {"embedding": [0.2] * 768} | |
| with patch("app.services.document_service.genai.embed_content", flaky): | |
| vectors = await service._embed_batch(["a", "b", "c"]) | |
| assert len(vectors) == 3 | |
| assert all(v is not None for v in vectors) | |
| async def test_failed_chunks_become_none_not_exceptions(): | |
| service = DocumentService() | |
| def always_fails(model, content, **kwargs): | |
| raise RuntimeError("quota exceeded") | |
| with patch("app.services.document_service.genai.embed_content", always_fails): | |
| vectors = await service._embed_batch(["a", "b"]) | |
| assert vectors == [None, None] | |
| async def test_rejects_shape_mismatch_from_sdk(): | |
| """A flat vector returned for a batch request must not be zipped blindly.""" | |
| service = DocumentService() | |
| calls = [] | |
| def wrong_shape(model, content, **kwargs): | |
| calls.append(content) | |
| if isinstance(content, list): | |
| return {"embedding": [0.1] * 768} # one flat vector, not N | |
| return {"embedding": [0.3] * 768} | |
| with patch("app.services.document_service.genai.embed_content", wrong_shape): | |
| vectors = await service._embed_batch(["a", "b"]) | |
| # Detected the mismatch and re-embedded each chunk individually. | |
| assert len(vectors) == 2 | |
| assert all(v is not None and len(v) == 768 for v in vectors) | |