61 lines
1.7 KiB
Python
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()
|