69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
import logging
|
|
|
|
import httpx
|
|
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
OPENROUTER_BASE = "https://openrouter.ai/api/v1"
|
|
EMBEDDING_MODEL = settings.openrouter_embedding_model
|
|
IS_NVIDIA = "nvidia" in EMBEDDING_MODEL
|
|
|
|
_client: httpx.AsyncClient | None = None
|
|
|
|
|
|
def get_client() -> httpx.AsyncClient:
|
|
global _client
|
|
if _client is None:
|
|
kwargs: dict = {
|
|
"base_url": OPENROUTER_BASE,
|
|
"headers": {
|
|
"Authorization": f"Bearer {settings.openrouter_api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
"timeout": 30,
|
|
}
|
|
if settings.proxy_url:
|
|
proxy = settings.proxy_url
|
|
if not proxy.startswith(("http://", "https://")):
|
|
proxy = "http://" + proxy
|
|
if settings.proxy_login and settings.proxy_pass:
|
|
proxy = proxy.replace("://", f"://{settings.proxy_login}:{settings.proxy_pass}@")
|
|
kwargs["proxy"] = proxy
|
|
_client = httpx.AsyncClient(**kwargs)
|
|
return _client
|
|
|
|
|
|
def _prefix_text(text: str, prefix: str) -> str:
|
|
if IS_NVIDIA:
|
|
return f"{prefix}: {text}"
|
|
return text
|
|
|
|
|
|
async def embed_texts(texts: list[str]) -> list[list[float]]:
|
|
prefixed = [_prefix_text(t, "passage") for t in texts]
|
|
client = get_client()
|
|
response = await client.post(
|
|
"/embeddings",
|
|
json={"model": EMBEDDING_MODEL, "input": prefixed},
|
|
)
|
|
data = response.json()
|
|
if not data.get("data"):
|
|
logger.error("Empty embedding response: %s", data)
|
|
return [[0.0] * 2048 for _ in texts]
|
|
return [item["embedding"] for item in data["data"]]
|
|
|
|
|
|
async def embed_query(text: str) -> list[float]:
|
|
prefixed = _prefix_text(text, "query")
|
|
client = get_client()
|
|
response = await client.post(
|
|
"/embeddings",
|
|
json={"model": EMBEDDING_MODEL, "input": [prefixed]},
|
|
)
|
|
data = response.json()
|
|
if not data.get("data"):
|
|
logger.error("Empty embedding response for query: %s", data)
|
|
return [0.0] * 2048
|
|
return data["data"][0]["embedding"]
|