1515# uniform with the OpenAI/Voyage providers.
1616_BATCH_SIZE = 128
1717_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
1821
1922# Native output dimensions for known models. The jina-code-embeddings family
2023# supports Matryoshka truncation via the `dimensions` API parameter —
3538# is single-mode and rejects `task`, so we omit it.
3639_TASK_AWARE_PREFIXES = ("jina-code-embeddings-" ,)
3740
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+
3848
3949class JinaApiEmbeddingProvider (EmbeddingProvider ):
4050 """Jina AI hosted embeddings — see https://jina.ai/embeddings/."""
@@ -58,6 +68,7 @@ def __init__(self) -> None:
5868 f"({ ', ' .join (sorted (_NATIVE_DIMENSIONS ))}
5969 )
6070 self ._supports_task = self ._model .startswith (_TASK_AWARE_PREFIXES )
71+ self ._uses_code_tasks = self ._model .startswith ("jina-code-embeddings-" )
6172 self ._client = httpx .AsyncClient (
6273 timeout = 120.0 ,
6374 headers = {
@@ -77,31 +88,105 @@ async def embed_query(self, text: str) -> list[float]:
7788 vectors = await self ._embed ([text ], task = "retrieval.query" )
7889 return vectors [0 ] if vectors else []
7990
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+
80171 async def _embed (self , texts : list [str ], task : str ) -> list [list [float ]]:
81172 if not texts :
82173 return []
174+ sanitized = [self ._sanitize (t ) for t in texts ]
83175 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 ()
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
105190 batch_vectors = [item ["embedding" ] for item in data .get ("data" , [])]
106191 if len (batch_vectors ) != len (batch ):
107192 raise ValueError (
0 commit comments