1use crate::errors::AppError;
10use crate::retry::AttemptOutcome;
11use secrecy::{ExposeSecret, SecretBox};
12use serde::{Deserialize, Serialize};
13use std::time::Duration;
14
15use crate::constants::DEFAULT_OPENROUTER_EMBEDDINGS_URL;
17const DEFAULT_TIMEOUT_SECS: u64 = 30;
18const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 10;
19const DEFAULT_EMBED_HTTP_BATCH_SIZE: usize = crate::constants::FASTEMBED_BATCH_SIZE;
22
23#[derive(Serialize)]
24struct EmbeddingRequest<'a> {
25 model: &'a str,
26 input: EmbeddingInput<'a>,
27 #[serde(skip_serializing_if = "Option::is_none")]
28 dimensions: Option<usize>,
29 encoding_format: &'a str,
30 #[serde(skip_serializing_if = "Option::is_none")]
31 input_type: Option<&'a str>,
32}
33
34#[derive(Serialize)]
35#[serde(untagged)]
36enum EmbeddingInput<'a> {
37 Single(&'a str),
38 Batch(Vec<&'a str>),
39}
40
41#[derive(Deserialize)]
42struct EmbeddingResponse {
43 data: Vec<EmbeddingData>,
44}
45
46#[derive(Deserialize)]
47struct EmbeddingData {
48 embedding: Vec<f32>,
49 index: usize,
50}
51
52#[derive(Deserialize)]
60struct EmbeddingEnvelope {
61 #[serde(default)]
62 data: Option<Vec<EmbeddingData>>,
63 #[serde(default)]
64 error: Option<ApiError>,
65}
66
67use crate::openrouter_http::ApiError;
72
73#[derive(Debug)]
85pub struct EmbedError {
86 pub source: AppError,
88 pub retry_class: AttemptOutcome,
91}
92
93impl EmbedError {
94 fn new(source: AppError, retry_class: AttemptOutcome) -> Self {
95 Self {
96 source,
97 retry_class,
98 }
99 }
100}
101
102impl std::fmt::Display for EmbedError {
103 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104 std::fmt::Display::fmt(&self.source, f)
105 }
106}
107
108impl std::error::Error for EmbedError {
109 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
110 Some(&self.source)
111 }
112}
113
114impl From<AppError> for EmbedError {
123 fn from(source: AppError) -> Self {
124 Self::new(source, AttemptOutcome::HardFailure)
125 }
126}
127
128impl From<EmbedError> for AppError {
135 fn from(err: EmbedError) -> Self {
136 err.source
137 }
138}
139
140pub struct OpenRouterClient {
142 client: reqwest::Client,
143 api_key: SecretBox<String>,
144 model: String,
145 dim: usize,
146 default_input_type: Option<&'static str>,
147 base_url: String,
150}
151
152fn model_supports_mrl(model: &str) -> bool {
153 model.contains("qwen3-embedding")
154 || model.contains("text-embedding-3")
155 || model.contains("gemini-embedding")
156 || model.contains("llama-nemotron-embed")
157 || model.contains("bge-m3")
158}
159
160fn mrl_wire_dimensions(model: &str, dim: usize) -> Option<usize> {
170 if !model_supports_mrl(model) {
171 return None;
172 }
173 if model.contains("qwen3-embedding") {
174 return None;
175 }
176 Some(dim)
177}
178
179fn model_default_input_type(model: &str) -> Option<&'static str> {
180 if model.contains("llama-nemotron-embed") {
181 Some("passage")
182 } else if model.contains("mistral-embed") {
183 None
184 } else {
185 Some("search_document")
186 }
187}
188
189impl OpenRouterClient {
190 pub fn new(api_key: SecretBox<String>, model: String, dim: usize) -> Result<Self, AppError> {
192 let base_url = crate::runtime_config::openrouter_embeddings_url(
193 DEFAULT_OPENROUTER_EMBEDDINGS_URL,
194 );
195 Self::new_with_base_url(api_key, model, dim, base_url)
196 }
197
198 pub fn new_with_base_url(
200 api_key: SecretBox<String>,
201 model: String,
202 dim: usize,
203 base_url: String,
204 ) -> Result<Self, AppError> {
205 let client = reqwest::Client::builder()
206 .timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS))
207 .connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS))
208 .user_agent(concat!("sqlite-graphrag/", env!("CARGO_PKG_VERSION")))
209 .build()
210 .map_err(|e| {
211 AppError::Embedding(crate::i18n::validation::embedding_http_client_build_failed(e))
212 })?;
213
214 let default_input_type = model_default_input_type(&model);
215
216 Ok(Self {
217 client,
218 api_key,
219 model,
220 dim,
221 default_input_type,
222 base_url,
223 })
224 }
225
226 #[cfg(test)]
230 fn new_with_url(
231 api_key: SecretBox<String>,
232 model: String,
233 dim: usize,
234 base_url: String,
235 ) -> Result<Self, AppError> {
236 Self::new_with_base_url(api_key, model, dim, base_url)
237 }
238
239 pub fn default_input_type(&self) -> Option<&'static str> {
241 self.default_input_type
242 }
243
244 pub async fn embed_single(
246 &self,
247 text: &str,
248 input_type: Option<&str>,
249 ) -> Result<Vec<f32>, EmbedError> {
250 crate::memory_guard::check_embedding_input_size(text)?;
254
255 let request = EmbeddingRequest {
256 model: &self.model,
257 input: EmbeddingInput::Single(text),
258 dimensions: mrl_wire_dimensions(&self.model, self.dim),
259 encoding_format: "float",
260 input_type,
261 };
262
263 let response = self.execute_with_retry(&request).await?;
264
265 let embedding = response
266 .data
267 .into_iter()
268 .next()
269 .ok_or_else(|| {
270 AppError::Embedding(
271 crate::i18n::validation::embedding_empty_response_from_openrouter(),
272 )
273 })?
274 .embedding;
275
276 Ok(self.truncate_embedding(embedding)?)
277 }
278
279 pub async fn embed_batch(
281 &self,
282 texts: &[&str],
283 input_type: Option<&str>,
284 ) -> Result<Vec<Vec<f32>>, EmbedError> {
285 if texts.is_empty() {
286 return Ok(Vec::new());
287 }
288
289 for text in texts {
293 crate::memory_guard::check_embedding_input_size(text)?;
294 }
295
296 let mut all = Vec::with_capacity(texts.len());
297
298 let batch_size =
299 crate::runtime_config::embedding_batch_size(DEFAULT_EMBED_HTTP_BATCH_SIZE);
300 for chunk in texts.chunks(batch_size) {
301 let request = EmbeddingRequest {
302 model: &self.model,
303 input: EmbeddingInput::Batch(chunk.to_vec()),
304 dimensions: mrl_wire_dimensions(&self.model, self.dim),
305 encoding_format: "float",
306 input_type,
307 };
308
309 let response = self.execute_with_retry(&request).await?;
310
311 if response.data.len() != chunk.len() {
312 return Err(AppError::Embedding(crate::i18n::validation::embedding_expected_count(
313 chunk.len(),
314 response.data.len(),
315 ))
316 .into());
317 }
318
319 let mut sorted = response.data;
320 sorted.sort_by_key(|d| d.index);
321
322 for d in sorted {
323 all.push(self.truncate_embedding(d.embedding)?);
324 }
325 }
326
327 Ok(all)
328 }
329
330 fn truncate_embedding(&self, embedding: Vec<f32>) -> Result<Vec<f32>, AppError> {
331 if embedding.len() < self.dim {
332 return Err(AppError::Embedding(
333 crate::i18n::validation::embedding_dimension_less_than_requested(
334 embedding.len(),
335 self.dim,
336 ),
337 ));
338 }
339 if embedding.len() == self.dim {
340 Ok(embedding)
341 } else {
342 Ok(embedding[..self.dim].to_vec())
343 }
344 }
345
346 async fn execute_with_retry(
353 &self,
354 request: &EmbeddingRequest<'_>,
355 ) -> Result<EmbeddingResponse, EmbedError> {
356 let mut last_err: Option<EmbedError> = None;
357
358 for attempt in 0..crate::openrouter_http::MAX_RETRIES {
359 let result = self
360 .client
361 .post(&self.base_url)
362 .header(
363 "Authorization",
364 format!("Bearer {}", self.api_key.expose_secret()),
365 )
366 .json(request)
367 .send()
368 .await;
369
370 let resp = match result {
371 Ok(r) => r,
372 Err(e) if e.is_timeout() => {
373 return Err(EmbedError::new(
374 AppError::Embedding(
375 crate::i18n::validation::embedding_openrouter_request_timed_out(),
376 ),
377 AttemptOutcome::Transient,
378 ));
379 }
380 Err(e) => {
381 last_err = Some(EmbedError::new(
382 AppError::Embedding(
383 crate::i18n::validation::embedding_http_request_failed(e),
384 ),
385 AttemptOutcome::Transient,
386 ));
387 crate::openrouter_http::backoff(attempt).await;
388 continue;
389 }
390 };
391
392 let status = resp.status();
393
394 if status.is_success() {
395 let body = resp.text().await.map_err(|e| {
396 EmbedError::new(
397 AppError::Embedding(
398 crate::i18n::validation::embedding_failed_to_read_response_body(e),
399 ),
400 AttemptOutcome::Transient,
401 )
402 })?;
403 match serde_json::from_str::<EmbeddingEnvelope>(&body) {
404 Ok(env) => {
405 if let Some(api_err) = env.error {
410 let retry_class =
411 crate::openrouter_http::provider_error_retry_class(&api_err);
412 return Err(EmbedError::new(
413 AppError::ProviderError {
414 code: api_err.code_string(),
415 message: api_err.message,
416 },
417 retry_class,
418 ));
419 }
420 match env.data {
421 Some(data) => return Ok(EmbeddingResponse { data }),
422 None => {
423 tracing::warn!(
424 attempt,
425 body_len = body.len(),
426 "HTTP 200 with neither data nor error (retrying)"
427 );
428 last_err = Some(EmbedError::new(
429 AppError::Embedding(
430 crate::i18n::validation::embedding_openrouter_200_neither_data_nor_error(),
431 ),
432 AttemptOutcome::Transient,
433 ));
434 crate::openrouter_http::backoff(attempt).await;
435 continue;
436 }
437 }
438 }
439 Err(e) => {
440 tracing::warn!(
441 attempt,
442 body_len = body.len(),
443 "HTTP 200 but JSON unparseable (retrying): {e}"
444 );
445 last_err = Some(EmbedError::new(
446 AppError::Embedding(
447 crate::i18n::validation::embedding_failed_to_parse_response(e),
448 ),
449 AttemptOutcome::Transient,
450 ));
451 crate::openrouter_http::backoff(attempt).await;
452 continue;
453 }
454 }
455 }
456
457 if status.as_u16() == 401 {
458 return Err(EmbedError::new(
459 AppError::Embedding(
460 crate::i18n::validation::embedding_openrouter_invalid_api_key_401(),
461 ),
462 AttemptOutcome::HardFailure,
463 ));
464 }
465
466 if status.as_u16() == 400 || status.as_u16() == 404 {
467 let body = resp.text().await.unwrap_or_default();
468 return Err(EmbedError::new(
469 AppError::Embedding(
470 crate::i18n::validation::embedding_openrouter_returned(status, &body),
471 ),
472 AttemptOutcome::HardFailure,
473 ));
474 }
475
476 if status.as_u16() == 429 {
477 let retry_after = resp
478 .headers()
479 .get("retry-after")
480 .and_then(|v| v.to_str().ok())
481 .and_then(|v| v.parse::<u64>().ok())
482 .unwrap_or(2);
483 tracing::warn!(
484 attempt,
485 retry_after_secs = retry_after,
486 "OpenRouter rate limited, waiting"
487 );
488 last_err = Some(EmbedError::new(
493 AppError::RateLimited {
494 detail: format!("OpenRouter HTTP 429 (retry-after {retry_after}s)"),
495 },
496 AttemptOutcome::Transient,
497 ));
498 tokio::time::sleep(Duration::from_secs(retry_after)).await;
499 continue;
500 }
501
502 if status.is_server_error() {
503 tracing::warn!(attempt, status = %status, "OpenRouter server error, retrying");
504 last_err = Some(EmbedError::new(
505 AppError::Embedding(
506 crate::i18n::validation::embedding_openrouter_server_error(status),
507 ),
508 AttemptOutcome::Transient,
509 ));
510 crate::openrouter_http::backoff(attempt).await;
511 continue;
512 }
513
514 let body = resp.text().await.unwrap_or_default();
515 return Err(EmbedError::new(
516 AppError::Embedding(
517 crate::i18n::validation::embedding_unexpected_http(status, &body),
518 ),
519 crate::openrouter_http::status_retry_class(status),
520 ));
521 }
522
523 Err(last_err.unwrap_or_else(|| {
528 EmbedError::new(
529 AppError::Embedding(
530 crate::i18n::validation::embedding_openrouter_max_retries(),
531 ),
532 AttemptOutcome::Transient,
533 )
534 }))
535 }
536}
537
538#[cfg(test)]
539mod tests {
540 use super::*;
541
542 #[test]
543 fn test_supports_mrl_detection() {
544 assert!(model_supports_mrl("qwen/qwen3-embedding-8b"));
545 assert!(model_supports_mrl("qwen/qwen3-embedding-4b"));
546 assert!(model_supports_mrl("openai/text-embedding-3-small"));
547 assert!(model_supports_mrl("openai/text-embedding-3-large"));
548 assert!(model_supports_mrl("google/gemini-embedding-001"));
549 assert!(model_supports_mrl("google/gemini-embedding-2"));
550 assert!(model_supports_mrl(
551 "nvidia/llama-nemotron-embed-vl-1b-v2:free"
552 ));
553 assert!(model_supports_mrl("baai/bge-m3"));
554
555 assert!(!model_supports_mrl("perplexity/pplx-embed-v1-0.6b"));
556 assert!(!model_supports_mrl("mistralai/mistral-embed-2312"));
557 assert!(!model_supports_mrl("some-random-model"));
558 }
559
560 #[test]
561 fn test_mrl_wire_dimensions_qwen_omits_wire_dim() {
562 assert_eq!(mrl_wire_dimensions("qwen/qwen3-embedding-8b", 384), None);
564 assert_eq!(mrl_wire_dimensions("qwen/qwen3-embedding-4b", 512), None);
565 assert_eq!(
567 mrl_wire_dimensions("openai/text-embedding-3-small", 384),
568 Some(384)
569 );
570 assert_eq!(mrl_wire_dimensions("baai/bge-m3", 256), Some(256));
571 assert_eq!(
573 mrl_wire_dimensions("mistralai/mistral-embed-2312", 1024),
574 None
575 );
576 }
577
578 #[test]
579 fn test_model_default_input_type() {
580 assert_eq!(
581 model_default_input_type("nvidia/llama-nemotron-embed-vl-1b-v2:free"),
582 Some("passage")
583 );
584 assert_eq!(
585 model_default_input_type("mistralai/mistral-embed-2312"),
586 None
587 );
588 assert_eq!(
589 model_default_input_type("qwen/qwen3-embedding-8b"),
590 Some("search_document")
591 );
592 assert_eq!(
593 model_default_input_type("openai/text-embedding-3-small"),
594 Some("search_document")
595 );
596 assert_eq!(
597 model_default_input_type("baai/bge-m3"),
598 Some("search_document")
599 );
600 }
601
602 #[test]
603 fn test_truncate_embedding() {
604 let api_key = SecretBox::new(Box::new("test-key".to_string()));
605 let client = OpenRouterClient::new(api_key, "test-model".into(), 3).unwrap();
606
607 let full = vec![1.0, 2.0, 3.0, 4.0, 5.0];
608 let truncated = client.truncate_embedding(full).unwrap();
609 assert_eq!(truncated, vec![1.0, 2.0, 3.0]);
610
611 let exact = vec![1.0, 2.0, 3.0];
612 let kept = client.truncate_embedding(exact).unwrap();
613 assert_eq!(kept, vec![1.0, 2.0, 3.0]);
614
615 let short = vec![1.0, 2.0];
616 let err = client.truncate_embedding(short);
617 assert!(err.is_err());
618 }
619
620 #[test]
621 fn embedding_envelope_surfaces_provider_error_not_missing_field() {
622 let body = r#"{"error":{"code":400,"message":"context length exceeded"}}"#;
625
626 let legacy_err = match serde_json::from_str::<EmbeddingResponse>(body) {
629 Ok(_) => panic!("legacy parse should have failed on an error body"),
630 Err(e) => e.to_string(),
631 };
632 assert!(
633 legacy_err.contains("missing field"),
634 "precondition: legacy parse masks the cause as a missing field: {legacy_err}"
635 );
636
637 let env: EmbeddingEnvelope =
639 serde_json::from_str(body).expect("envelope parses an error body");
640 assert!(env.data.is_none());
641 let api_err = env.error.expect("error object captured");
642 assert_eq!(api_err.message, "context length exceeded");
643 assert_eq!(api_err.code_string(), "400");
644 }
645
646 #[test]
647 fn embedding_envelope_parses_success_body() {
648 let body = r#"{"data":[{"embedding":[1.0,2.0,3.0],"index":0}]}"#;
649 let env: EmbeddingEnvelope =
650 serde_json::from_str(body).expect("envelope parses a success body");
651 assert!(env.error.is_none());
652 let data = env.data.expect("data present");
653 assert_eq!(data.len(), 1);
654 assert_eq!(data[0].embedding, vec![1.0, 2.0, 3.0]);
655 }
656
657 #[test]
658 fn api_error_code_string_handles_number_string_and_missing() {
659 let num: ApiError = serde_json::from_str(r#"{"code":429,"message":"slow down"}"#).unwrap();
660 assert_eq!(num.code_string(), "429");
661
662 let s: ApiError =
663 serde_json::from_str(r#"{"code":"rate_limited","message":"slow down"}"#).unwrap();
664 assert_eq!(s.code_string(), "rate_limited");
665
666 let missing: ApiError = serde_json::from_str(r#"{"message":"oops"}"#).unwrap();
667 assert_eq!(missing.code_string(), "unknown");
668 }
669
670 #[tokio::test]
671 async fn embed_single_rejects_oversized_input_before_request() {
672 let api_key = SecretBox::new(Box::new("test-key".to_string()));
677 let client = OpenRouterClient::new(api_key, "qwen/qwen3-embedding-8b".into(), 384).unwrap();
678 let big = "word ".repeat(crate::constants::EMBEDDING_REQUEST_MAX_TOKENS + 5_000);
679 match client.embed_single(&big, None).await {
680 Err(EmbedError {
681 source: AppError::TooManyTokens { tokens, limit },
682 retry_class,
683 }) => {
684 assert!(tokens > limit, "tokens={tokens} limit={limit}");
685 assert_eq!(limit, crate::constants::EMBEDDING_REQUEST_MAX_TOKENS as u64);
686 assert_eq!(
687 retry_class,
688 AttemptOutcome::HardFailure,
689 "an oversized input is a permanent client error"
690 );
691 }
692 other => unreachable!("expected TooManyTokens before request, got: {other:?}"),
693 }
694 }
695
696 async fn client_for(server: &wiremock::MockServer, model: &str) -> OpenRouterClient {
697 OpenRouterClient::new_with_url(
698 SecretBox::new(Box::new("test-key".to_string())),
699 model.to_string(),
700 384,
701 format!("{}/embeddings", server.uri()),
702 )
703 .expect("test client builds")
704 }
705
706 #[tokio::test]
707 async fn embed_single_401_is_hard_failure() {
708 use wiremock::{matchers::method, Mock, MockServer, ResponseTemplate};
711 let server = MockServer::start().await;
712 Mock::given(method("POST"))
713 .respond_with(ResponseTemplate::new(401))
714 .mount(&server)
715 .await;
716
717 let client = client_for(&server, "qwen/qwen3-embedding-8b").await;
718 let err = client
719 .embed_single("hello", None)
720 .await
721 .expect_err("401 is an error");
722 assert_eq!(err.retry_class, AttemptOutcome::HardFailure);
723 }
724
725 #[tokio::test]
726 async fn embed_single_exhausted_5xx_is_transient() {
727 use wiremock::{matchers::method, Mock, MockServer, ResponseTemplate};
731 let server = MockServer::start().await;
732 Mock::given(method("POST"))
733 .respond_with(ResponseTemplate::new(503))
734 .mount(&server)
735 .await;
736
737 let client = client_for(&server, "qwen/qwen3-embedding-8b").await;
738 let err = client
739 .embed_single("hello", None)
740 .await
741 .expect_err("persistent 5xx exhausts retries");
742 assert_eq!(err.retry_class, AttemptOutcome::Transient);
743 }
744
745 #[tokio::test]
746 async fn embed_single_provider_error_code_classifies_by_code_not_message() {
747 use wiremock::{matchers::method, Mock, MockServer, ResponseTemplate};
751 let server = MockServer::start().await;
752 Mock::given(method("POST"))
753 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
754 "error": { "code": "context_length_exceeded", "message": "too many tokens" }
755 })))
756 .mount(&server)
757 .await;
758
759 let client = client_for(&server, "qwen/qwen3-embedding-8b").await;
760 let err = client
761 .embed_single("hello", None)
762 .await
763 .expect_err("provider error must surface");
764 assert_eq!(err.retry_class, AttemptOutcome::HardFailure);
765 }
766}