from openai import AsyncOpenAI

from app.config import settings

client = AsyncOpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=settings.openrouter_api_key,
)

EMBEDDING_MODEL = settings.openrouter_embedding_model
IS_NVIDIA = "nvidia" in EMBEDDING_MODEL


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]
    response = await client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=prefixed,
    )
    return [item.embedding for item in response.data]


async def embed_query(text: str) -> list[float]:
    prefixed = _prefix_text(text, "query")
    response = await client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=[prefixed],
    )
    return response.data[0].embedding