import httpx from src.pipeline.config import config from src.pipeline.core.utils import logger class AsyncLLm: def __init__( self, timeout: float = 30.0, max_connections: int = 100, max_keepalive: int = 20, ): self.embedding_api = config["embedding_api_host"].rstrip("/") + "/embeddings" self.embedding_model = config["embedding_model"] self.api_key = config["embedding_api_key"] logger.debug(self.embedding_api) self.embedding_client = httpx.AsyncClient( http2=False, trust_env=False, timeout=httpx.Timeout(timeout), limits=httpx.Limits( max_connections=max_connections, max_keepalive_connections=max_keepalive, ), headers={ "Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}", }, ) async def embedding(self, text: str) -> list[float]: try: resp = await self.embedding_client.post( self.embedding_api, json={"model": self.embedding_model, "input": text}, ) resp.raise_for_status() data = resp.json() return data["data"][0]["embedding"] except httpx.HTTPStatusError as e: logger.error(e) logger.error(f"Embedding HTTP error: {e.response.text}") except Exception as e: logger.exception("Embedding request failed") return [] async def close(self): await self.embedding_client.aclose() client: AsyncLLm | None = None async def init_client(): global client client = AsyncLLm() async def close_client(): await client.close()