|
| 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 | +# Conservative character cap (~8 k tokens at ~4 chars/token) to avoid |
| 19 | +# "Failed to encode text" 400s on models with limited context windows. |
| 20 | +_MAX_TEXT_CHARS = 32_000 |
| 21 | + |
| 22 | +# Native output dimensions for known models. The jina-code-embeddings family |
| 23 | +# supports Matryoshka truncation via the `dimensions` API parameter — |
| 24 | +# override with JINA_API_DIMENSIONS to one of the supported sizes: |
| 25 | +# - jina-code-embeddings-0.5b: 64, 128, 256, 512, 896 (native) |
| 26 | +# https://jina.ai/models/jina-code-embeddings-0.5b |
| 27 | +# - jina-code-embeddings-1.5b: 128, 256, 512, 1024, 1536 (native) |
| 28 | +# https://jina.ai/models/jina-code-embeddings-1.5b |
| 29 | +# jina-embeddings-v2-base-code is fixed-size and does not support truncation. |
| 30 | +# https://jina.ai/models/jina-embeddings-v2-base-code |
| 31 | +_NATIVE_DIMENSIONS: dict[str, int] = { |
| 32 | + "jina-embeddings-v2-base-code": 768, |
| 33 | + "jina-code-embeddings-0.5b": 896, |
| 34 | + "jina-code-embeddings-1.5b": 1536, |
| 35 | +} |
| 36 | + |
| 37 | +# Models that accept the `task` parameter (asymmetric retrieval). The v2 model |
| 38 | +# is single-mode and rejects `task`, so we omit it. |
| 39 | +_TASK_AWARE_PREFIXES = ("jina-code-embeddings-",) |
| 40 | + |
| 41 | +# jina-code-embeddings models use a different task vocabulary than the generic |
| 42 | +# "retrieval.*" tasks accepted by other Jina models. |
| 43 | +_JINA_CODE_TASK_MAP = { |
| 44 | + "retrieval.passage": "nl2code.passage", |
| 45 | + "retrieval.query": "nl2code.query", |
| 46 | +} |
| 47 | + |
| 48 | + |
| 49 | +class JinaApiEmbeddingProvider(EmbeddingProvider): |
| 50 | + """Jina AI hosted embeddings — see https://jina.ai/embeddings/.""" |
| 51 | + |
| 52 | + def __init__(self) -> None: |
| 53 | + if not settings.jina_api_key: |
| 54 | + raise RuntimeError( |
| 55 | + "JINA_API_KEY is not set but EMBEDDINGS_PROVIDER=jina-api." |
| 56 | + ) |
| 57 | + self._api_key = settings.jina_api_key |
| 58 | + self._model = settings.jina_api_model |
| 59 | + self._dims_override = settings.jina_api_dimensions |
| 60 | + if self._dims_override is not None: |
| 61 | + self._dims = self._dims_override |
| 62 | + elif self._model in _NATIVE_DIMENSIONS: |
| 63 | + self._dims = _NATIVE_DIMENSIONS[self._model] |
| 64 | + else: |
| 65 | + raise RuntimeError( |
| 66 | + f"Unknown Jina model {self._model!r} — set JINA_API_DIMENSIONS " |
| 67 | + "to declare the output size, or use a known model " |
| 68 | + f"({', '.join(sorted(_NATIVE_DIMENSIONS))})." |
| 69 | + ) |
| 70 | + self._supports_task = self._model.startswith(_TASK_AWARE_PREFIXES) |
| 71 | + self._uses_code_tasks = self._model.startswith("jina-code-embeddings-") |
| 72 | + self._client = httpx.AsyncClient( |
| 73 | + timeout=120.0, |
| 74 | + headers={ |
| 75 | + "Authorization": f"Bearer {self._api_key}", |
| 76 | + "Content-Type": "application/json", |
| 77 | + }, |
| 78 | + ) |
| 79 | + |
| 80 | + @property |
| 81 | + def dimensions(self) -> int: |
| 82 | + return self._dims |
| 83 | + |
| 84 | + async def embed_batch(self, texts: list[str]) -> list[list[float]]: |
| 85 | + return await self._embed(texts, task="retrieval.passage") |
| 86 | + |
| 87 | + async def embed_query(self, text: str) -> list[float]: |
| 88 | + vectors = await self._embed([text], task="retrieval.query") |
| 89 | + return vectors[0] if vectors else [] |
| 90 | + |
| 91 | + def _sanitize(self, text: str) -> str: |
| 92 | + # Encode to UTF-8 replacing lone surrogates and other unencodable |
| 93 | + # code points, then decode back — this removes anything that would |
| 94 | + # cause Jina's tokenizer to return 400 "Failed to encode text". |
| 95 | + cleaned = text.encode("utf-8", errors="replace").decode("utf-8") |
| 96 | + cleaned = "".join(ch for ch in cleaned if ch >= " " or ch in "\t\n\r") |
| 97 | + return cleaned[:_MAX_TEXT_CHARS].strip() or "." |
| 98 | + |
| 99 | + def _make_body(self, inputs: list[str], task: str) -> dict: |
| 100 | + body: dict = {"model": self._model, "input": inputs} |
| 101 | + if self._supports_task: |
| 102 | + body["task"] = ( |
| 103 | + _JINA_CODE_TASK_MAP.get(task, task) if self._uses_code_tasks else task |
| 104 | + ) |
| 105 | + if self._dims_override is not None: |
| 106 | + body["dimensions"] = self._dims_override |
| 107 | + return body |
| 108 | + |
| 109 | + async def _post_with_retry(self, body: dict) -> dict: |
| 110 | + for attempt in range(4): |
| 111 | + resp = await self._client.post(_API_URL, json=body) |
| 112 | + if resp.status_code != 429: |
| 113 | + break |
| 114 | + retry_after = float(resp.headers.get("Retry-After", 0)) |
| 115 | + wait = retry_after if retry_after > 0 else _BACKOFF_DELAYS[attempt] |
| 116 | + logger.warning( |
| 117 | + "Jina rate-limited (429) — retrying in %.0fs (attempt %d/4)", |
| 118 | + wait, |
| 119 | + attempt + 1, |
| 120 | + ) |
| 121 | + await asyncio.sleep(wait) |
| 122 | + if resp.status_code >= 400: |
| 123 | + logger.error("Jina API error %d: %s", resp.status_code, resp.text[:500]) |
| 124 | + resp.raise_for_status() |
| 125 | + return resp.json() |
| 126 | + |
| 127 | + async def _embed_batch_with_fallback( |
| 128 | + self, batch: list[str], task: str |
| 129 | + ) -> list[list[float]]: |
| 130 | + """Embed one item at a time, halving on failure, substituting '.' only as last resort.""" |
| 131 | + vectors: list[list[float]] = [] |
| 132 | + for idx, text in enumerate(batch): |
| 133 | + candidate = text |
| 134 | + embedded = False |
| 135 | + while candidate: |
| 136 | + try: |
| 137 | + data = await self._post_with_retry( |
| 138 | + self._make_body([candidate], task) |
| 139 | + ) |
| 140 | + vectors.append(data["data"][0]["embedding"]) |
| 141 | + if len(candidate) < len(text): |
| 142 | + logger.info( |
| 143 | + "Encoded truncated text at batch index %d (%d → %d chars)", |
| 144 | + idx, |
| 145 | + len(text), |
| 146 | + len(candidate), |
| 147 | + ) |
| 148 | + embedded = True |
| 149 | + break |
| 150 | + except Exception: |
| 151 | + half = len(candidate) // 2 |
| 152 | + if half < 64: |
| 153 | + break |
| 154 | + logger.warning( |
| 155 | + "Text at batch index %d (len=%d) failed — retrying with first %d chars", |
| 156 | + idx, |
| 157 | + len(candidate), |
| 158 | + half, |
| 159 | + ) |
| 160 | + candidate = candidate[:half] |
| 161 | + if not embedded: |
| 162 | + logger.warning( |
| 163 | + "Skipping unencodable text at batch index %d (original len=%d), using placeholder.", |
| 164 | + idx, |
| 165 | + len(text), |
| 166 | + ) |
| 167 | + data = await self._post_with_retry(self._make_body(["."], task)) |
| 168 | + vectors.append(data["data"][0]["embedding"]) |
| 169 | + return vectors |
| 170 | + |
| 171 | + async def _embed(self, texts: list[str], task: str) -> list[list[float]]: |
| 172 | + if not texts: |
| 173 | + return [] |
| 174 | + sanitized = [self._sanitize(t) for t in texts] |
| 175 | + all_vectors: list[list[float]] = [] |
| 176 | + for i in range(0, len(sanitized), _BATCH_SIZE): |
| 177 | + batch = sanitized[i : i + _BATCH_SIZE] |
| 178 | + try: |
| 179 | + data = await self._post_with_retry(self._make_body(batch, task)) |
| 180 | + except Exception as exc: |
| 181 | + if "400" in str(exc): |
| 182 | + logger.warning( |
| 183 | + "Batch of %d failed with 400 — retrying one-by-one", len(batch) |
| 184 | + ) |
| 185 | + all_vectors.extend( |
| 186 | + await self._embed_batch_with_fallback(batch, task) |
| 187 | + ) |
| 188 | + continue |
| 189 | + raise |
| 190 | + batch_vectors = [item["embedding"] for item in data.get("data", [])] |
| 191 | + if len(batch_vectors) != len(batch): |
| 192 | + raise ValueError( |
| 193 | + f"Jina returned {len(batch_vectors)} vectors for " |
| 194 | + f"{len(batch)} inputs — response may be malformed" |
| 195 | + ) |
| 196 | + all_vectors.extend(batch_vectors) |
| 197 | + return all_vectors |
| 198 | + |
| 199 | + async def close(self) -> None: |
| 200 | + await self._client.aclose() |
0 commit comments