ai_pipeline/src/pipeline/core/llm.py

61 lines
1.7 KiB
Python

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()