Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions backend/src/docmind/embeddings/hashing.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from __future__ import annotations

import asyncio
import hashlib
import itertools
import math
Expand Down Expand Up @@ -59,6 +60,11 @@ def embed(self, text: str) -> Vector:
return l2_normalize(vector)

async def embed_documents(self, texts: list[str]) -> list[Vector]:
# CPU-bound hashing + numpy; dispatch to a worker thread so the ingest
# loop can serve other requests while a long document is embedded.
return await asyncio.to_thread(self._embed_many, texts)

def _embed_many(self, texts: list[str]) -> list[Vector]:
return [self.embed(text) for text in texts]

async def embed_query(self, text: str) -> Vector:
Expand Down
28 changes: 28 additions & 0 deletions backend/tests/test_retrieval.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from __future__ import annotations

import asyncio
import contextlib
import json
from pathlib import Path

Expand Down Expand Up @@ -57,6 +59,32 @@ def test_rejects_tiny_dimension(self) -> None:
with pytest.raises(ValueError, match="dim"):
HashingEmbedder(8)

async def test_embed_documents_yields_to_the_event_loop(self) -> None:
# The ingest pipeline runs extract in a worker thread and then awaits
# embed_documents on the same event loop; if embedding keeps the loop
# busy, nothing else on the API worker gets to run meanwhile.
emb = HashingEmbedder(384)
texts = [(CORPUS["vac"] + " ") * 400] * 40 # long enough to be slow

ticks = 0

async def ticker() -> None:
nonlocal ticks
while True:
ticks += 1
await asyncio.sleep(0)

ticker_task = asyncio.create_task(ticker())
await asyncio.sleep(0)
ticks_before = ticks
await emb.embed_documents(texts)
ticks_during = ticks - ticks_before
ticker_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await ticker_task

assert ticks_during > 5, f"event loop was blocked; ticker advanced {ticks_during} times"


async def test_openai_compatible_embedder_batches_and_sorts() -> None:
calls: list[list[str]] = []
Expand Down
Loading