backend/rag/embedder.py

"""NVIDIA NIM embedding client."""
from openai import OpenAI

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from config import NVIDIA_API_KEY, LLM_BASE_URL, EMBED_MODEL


class Embedder:
    def __init__(self):
        self._client = OpenAI(base_url=LLM_BASE_URL, api_key=NVIDIA_API_KEY)

    def embed_query(self, text: str) -> list[float]:
        response = self._client.embeddings.create(
            model=EMBED_MODEL,
            input=[text],
            extra_body={"input_type": "query", "truncate": "END"},
        )
        return response.data[0].embedding

    def embed_passages(self, texts: list[str]) -> list[list[float]]:
        response = self._client.embeddings.create(
            model=EMBED_MODEL,
            input=texts,
            extra_body={"input_type": "passage", "truncate": "END"},
        )
        return [item.embedding for item in response.data]