Coverage for src / quber / playground / embedding.py: 75%

1"""Embedding helper for the ADE RAG playground. 

2 

3Reuses quber's local `EmbeddingService` (BAAI/bge-large-en-v1.5, 1024 dims), 

4which runs on GPU when available and falls back to CPU. Both ingest (embed 

5each chunk) and query (embed the question) go through the same model, so the 

6vector spaces match. 

7 

8Local embeddings are used instead of a hosted API so the experiment has no 

9external embedding dependency or quota. 

10""" 

11 

12from __future__ import annotations 

13 

14from functools import lru_cache 

15from typing import List 

16 

17from quber.db.embeddings import get_embedding_service 

18 

19EMBED_DIM = 1024 # bge-large-en-v1.5 

20 

21 

22@lru_cache(maxsize=1) 

23def _service(): 

24 # provider/device resolve from settings (EMBEDDING_PROVIDER / EMBEDDING_DEVICE), 

25 # defaulting to local + cuda-with-cpu-fallback. 

26 return get_embedding_service(provider="local") 

27 

28 

29def embed_texts(texts: List[str]) -> List[List[float]]: 

30 if not texts: 

31 return [] 

32 arr = _service().embed(texts, batch_size=32) 

33 return [row.tolist() for row in arr] 

34 

35 

36def embed_query(text: str) -> List[float]: 

37 return embed_texts([text])[0] 

38 

39 

40def embed_document(doc_key: str, texts: List[str]) -> List[List[float]]: 

41 """Embed a document's chunks: on the GPU worker when this host is set up 

42 to send work there, locally otherwise. The question at query time always 

43 embeds locally; only ingestion, thousands of chunks at once, is worth a 

44 worker.""" 

45 from quber.playground import gpu 

46 

47 if gpu.configured(): 

48 return gpu.embed(doc_key, texts) 

49 return embed_texts(texts)