shira-hermes/api/services/kb/voyage.py

"""Voyage embeddings client — voyage-multilingual-2, 1024-dim."""

from __future__ import annotations

import asyncio
import logging
import os
import random
from typing import Literal

import httpx

logger = logging.getLogger("shira.kb.voyage")

_VOYAGE_URL = "https://api.voyageai.com/v1/embeddings"
_VOYAGE_RERANK_URL = "https://api.voyageai.com/v1/rerank"
_MAX_BATCH = 128
# Voyage caps each request at 120K tokens total. Keep a safety margin so our
# char-based estimate doesn't push over when the real tokenizer disagrees.
_MAX_BATCH_TOKENS = 90_000
# Rough chars→tokens for Hebrew/mixed text, used only to pre-batch requests.
# Voyage's multilingual tokenizer splits Hebrew more aggressively than ASCII.
_TOKENS_PER_CHAR_EST = 0.55

# Retry parameters for transient failures (429 + 5xx). Dev and prod share a
# single VOYAGE_API_KEY; without retry, simultaneous bulk re-ingests on both
# environments hit the per-key RPM limit and fail one of them. With
# exponential backoff (jittered) the burst is absorbed in seconds.
_MAX_RETRIES = 5
_BASE_BACKOFF_S = 2.0
_MAX_BACKOFF_S = 30.0


class VoyageError(RuntimeError):
    pass


def _retry_after_seconds(resp: httpx.Response, attempt: int) -> float:
    """Honor `Retry-After` header if present, else exponential backoff."""
    hdr = resp.headers.get("retry-after")
    if hdr:
        try:
            return min(_MAX_BACKOFF_S, max(0.1, float(hdr)))
        except ValueError:
            pass
    backoff = min(_MAX_BACKOFF_S, _BASE_BACKOFF_S * (2 ** attempt))
    return backoff + random.uniform(0, backoff * 0.25)


async def _post_with_retry(
    http: httpx.AsyncClient, url: str, *, headers: dict, json: dict, what: str,
) -> httpx.Response:
    """POST with exponential backoff on 429 and 5xx.

    Returns the final Response (caller still checks status_code for
    non-retryable errors, e.g. 4xx other than 429).
    """
    last_resp: httpx.Response | None = None
    for attempt in range(_MAX_RETRIES + 1):
        resp = await http.post(url, headers=headers, json=json)
        last_resp = resp
        if resp.status_code == 200:
            return resp
        if resp.status_code != 429 and resp.status_code < 500:
            return resp
        if attempt == _MAX_RETRIES:
            return resp
        delay = _retry_after_seconds(resp, attempt)
        logger.warning(
            "[kb.voyage] %s got %d, retry %d/%d in %.1fs",
            what, resp.status_code, attempt + 1, _MAX_RETRIES, delay,
        )
        await asyncio.sleep(delay)
    return last_resp  # type: ignore[return-value]


async def embed(
    texts: list[str],
    input_type: Literal["document", "query"] = "document",
    model: str | None = None,
) -> list[list[float]]:
    if not texts:
        return []

    api_key = os.environ.get("VOYAGE_API_KEY")
    if not api_key:
        raise VoyageError("VOYAGE_API_KEY is not set")

    model = model or os.environ.get("VOYAGE_MODEL", "voyage-multilingual-2")

    # Pre-batch by estimated token budget so we never exceed Voyage's 120K cap.
    batches: list[list[str]] = []
    current: list[str] = []
    current_tokens = 0
    for t in texts:
        est = max(1, int(len(t) * _TOKENS_PER_CHAR_EST))
        if current and (current_tokens + est > _MAX_BATCH_TOKENS or len(current) >= _MAX_BATCH):
            batches.append(current)
            current, current_tokens = [], 0
        current.append(t)
        current_tokens += est
    if current:
        batches.append(current)

    out: list[list[float]] = []
    async with httpx.AsyncClient(timeout=120) as http:
        for bi, batch in enumerate(batches):
            resp = await _post_with_retry(
                http,
                _VOYAGE_URL,
                headers={"Authorization": f"Bearer {api_key}"},
                json={"model": model, "input": batch, "input_type": input_type},
                what=f"embed batch {bi+1}/{len(batches)}",
            )
            if resp.status_code != 200:
                raise VoyageError(f"Voyage API {resp.status_code}: {resp.text[:200]}")
            data = resp.json()
            for item in sorted(data["data"], key=lambda d: d["index"]):
                out.append(item["embedding"])
            logger.info(
                "[kb.voyage] batch %d/%d size=%d type=%s",
                bi + 1, len(batches), len(batch), input_type,
            )
    return out


async def rerank(
    query: str,
    documents: list[str],
    top_k: int | None = None,
    model: str | None = None,
) -> list[dict]:
    """Score each document against the query using a Voyage reranker.

    Returns a list of {index, relevance_score} sorted high→low. `index`
    refers back into the original `documents` list so the caller can map
    scores back to its own metadata.

    Empty input → empty output. Upstream failures raise VoyageError; callers
    that want a soft fallback should catch it.
    """
    if not query or not documents:
        return []
    api_key = os.environ.get("VOYAGE_API_KEY")
    if not api_key:
        raise VoyageError("VOYAGE_API_KEY is not set")
    model = model or os.environ.get("VOYAGE_RERANK_MODEL", "rerank-2.5")

    payload: dict = {
        "model": model,
        "query": query,
        "documents": documents,
    }
    if top_k is not None:
        payload["top_k"] = int(top_k)

    async with httpx.AsyncClient(timeout=60) as http:
        resp = await _post_with_retry(
            http,
            _VOYAGE_RERANK_URL,
            headers={"Authorization": f"Bearer {api_key}"},
            json=payload,
            what="rerank",
        )
    if resp.status_code != 200:
        raise VoyageError(f"Voyage rerank {resp.status_code}: {resp.text[:200]}")
    data = resp.json().get("data", [])
    return [
        {"index": item["index"], "relevance_score": item["relevance_score"]}
        for item in data
    ]