embedding-server / app /services /async_embeddings_service.py
Faysal4200's picture
Upload 45 files
ec855e6 verified
Raw
History Blame Contribute Delete
1.74 kB
"""
AsyncEmbeddingsService:
- Wraps the synchronous EmbeddingsService to provide async interface
- Uses asyncio.to_thread to offload blocking model operations to a thread pool
- Manages concurrency with asyncio.Lock
"""
import asyncio
from typing import List, Optional
from .embeddings_service import EmbeddingsService
class AsyncEmbeddingsService:
def __init__(self, model_name: Optional[str] = None):
# We compose the synchronous service
self._sync_service = EmbeddingsService(model_name)
self._lock = asyncio.Lock()
@property
def is_ready(self) -> bool:
return self._sync_service.is_ready
@property
def device(self):
return self._sync_service.device
async def load_model(self) -> None:
"""Async wrapper for loading the model."""
if self.is_ready:
return
async with self._lock:
if self.is_ready:
return
# Run the heavy lifting in a thread
await asyncio.to_thread(self._sync_service.load_model)
async def generate_single_embedding(self, text: str) -> List[float]:
"""
Generate a single embedding asynchronously.
"""
# Ensure model is loaded first (non-blocking check mostly)
await self.load_model()
return await asyncio.to_thread(self._sync_service.generate_single_embedding, text)
async def generate_batch_embeddings(self, texts: List[str]) -> List[List[float]]:
"""
Efficient batch processing asynchronously.
"""
await self.load_model()
return await asyncio.to_thread(self._sync_service.generate_batch_embeddings, texts)