|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import asyncio |
| 4 | +import logging |
| 5 | + |
| 6 | +import httpx |
| 7 | + |
| 8 | +from server.config import settings |
| 9 | +from server.embeddings.base import EmbeddingProvider |
| 10 | + |
| 11 | +logger = logging.getLogger(__name__) |
| 12 | + |
| 13 | +_API_URL = "https://api.jina.ai/v1/embeddings" |
| 14 | +# Jina's hosted API accepts up to 2048 inputs per request; 128 keeps us |
| 15 | +# uniform with the OpenAI/Voyage providers. |
| 16 | +_BATCH_SIZE = 128 |
| 17 | +_BACKOFF_DELAYS = [10, 20, 30, 40] |
| 18 | + |
| 19 | +# Native output dimensions for known models. v3 and the jina-code-embeddings |
| 20 | +# family support Matryoshka truncation via the `dimensions` API parameter — |
| 21 | +# override with JINA_API_DIMENSIONS. |
| 22 | +# |
| 23 | +# Other code-tuned models available on the hosted API: |
| 24 | +# - jina-code-embeddings-0.5b |
| 25 | +# - jina-code-embeddings-1.5b |
| 26 | +# Their native dimensions vary; set JINA_API_DIMENSIONS to declare the size. |
| 27 | +_NATIVE_DIMENSIONS: dict[str, int] = { |
| 28 | + "jina-embeddings-v2-base-code": 768, |
| 29 | + "jina-embeddings-v2-base-en": 768, |
| 30 | + "jina-embeddings-v3": 1024, |
| 31 | + "jina-clip-v2": 1024, |
| 32 | +} |
| 33 | + |
| 34 | +# Models that accept the `task` parameter (asymmetric retrieval). v2 models |
| 35 | +# are single-mode and reject `task`, so we omit it for them. |
| 36 | +_TASK_AWARE_PREFIXES = ("jina-embeddings-v3", "jina-code-embeddings-") |
| 37 | + |
| 38 | + |
| 39 | +class JinaApiEmbeddingProvider(EmbeddingProvider): |
| 40 | + """Jina AI hosted embeddings — see https://jina.ai/embeddings/.""" |
| 41 | + |
| 42 | + def __init__(self) -> None: |
| 43 | + if not settings.jina_api_key: |
| 44 | + raise RuntimeError( |
| 45 | + "JINA_API_KEY is not set but EMBEDDINGS_PROVIDER=jina-api." |
| 46 | + ) |
| 47 | + self._api_key = settings.jina_api_key |
| 48 | + self._model = settings.jina_api_model |
| 49 | + self._dims_override = settings.jina_api_dimensions |
| 50 | + if self._dims_override is not None: |
| 51 | + self._dims = self._dims_override |
| 52 | + elif self._model in _NATIVE_DIMENSIONS: |
| 53 | + self._dims = _NATIVE_DIMENSIONS[self._model] |
| 54 | + else: |
| 55 | + raise RuntimeError( |
| 56 | + f"Unknown Jina model {self._model!r} — set JINA_API_DIMENSIONS " |
| 57 | + "to declare the output size, or use a known model " |
| 58 | + f"({', '.join(sorted(_NATIVE_DIMENSIONS))})." |
| 59 | + ) |
| 60 | + self._supports_task = self._model.startswith(_TASK_AWARE_PREFIXES) |
| 61 | + self._client = httpx.AsyncClient( |
| 62 | + timeout=120.0, |
| 63 | + headers={ |
| 64 | + "Authorization": f"Bearer {self._api_key}", |
| 65 | + "Content-Type": "application/json", |
| 66 | + }, |
| 67 | + ) |
| 68 | + |
| 69 | + @property |
| 70 | + def dimensions(self) -> int: |
| 71 | + return self._dims |
| 72 | + |
| 73 | + async def embed_batch(self, texts: list[str]) -> list[list[float]]: |
| 74 | + return await self._embed(texts, task="retrieval.passage") |
| 75 | + |
| 76 | + async def embed_query(self, text: str) -> list[float]: |
| 77 | + vectors = await self._embed([text], task="retrieval.query") |
| 78 | + return vectors[0] if vectors else [] |
| 79 | + |
| 80 | + async def _embed(self, texts: list[str], task: str) -> list[list[float]]: |
| 81 | + if not texts: |
| 82 | + return [] |
| 83 | + all_vectors: list[list[float]] = [] |
| 84 | + for i in range(0, len(texts), _BATCH_SIZE): |
| 85 | + batch = texts[i : i + _BATCH_SIZE] |
| 86 | + body: dict = {"model": self._model, "input": batch} |
| 87 | + if self._supports_task: |
| 88 | + body["task"] = task |
| 89 | + if self._dims_override is not None: |
| 90 | + body["dimensions"] = self._dims_override |
| 91 | + for attempt in range(4): |
| 92 | + resp = await self._client.post(_API_URL, json=body) |
| 93 | + if resp.status_code != 429: |
| 94 | + break |
| 95 | + retry_after = float(resp.headers.get("Retry-After", 0)) |
| 96 | + wait = retry_after if retry_after > 0 else _BACKOFF_DELAYS[attempt] |
| 97 | + logger.warning( |
| 98 | + "Jina rate-limited (429) — retrying in %.0fs (attempt %d/4)", |
| 99 | + wait, |
| 100 | + attempt + 1, |
| 101 | + ) |
| 102 | + await asyncio.sleep(wait) |
| 103 | + resp.raise_for_status() |
| 104 | + data = resp.json() |
| 105 | + batch_vectors = [item["embedding"] for item in data.get("data", [])] |
| 106 | + if len(batch_vectors) != len(batch): |
| 107 | + raise ValueError( |
| 108 | + f"Jina returned {len(batch_vectors)} vectors for " |
| 109 | + f"{len(batch)} inputs — response may be malformed" |
| 110 | + ) |
| 111 | + all_vectors.extend(batch_vectors) |
| 112 | + return all_vectors |
| 113 | + |
| 114 | + async def close(self) -> None: |
| 115 | + await self._client.aclose() |
0 commit comments