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 =
193 crate::runtime_config::openrouter_embeddings_url(DEFAULT_OPENROUTER_EMBEDDINGS_URL);
194 Self::new_with_base_url(api_key, model, dim, base_url)
195 }
196
197 pub fn new_with_base_url(
199 api_key: SecretBox<String>,
200 model: String,
201 dim: usize,
202 base_url: String,
203 ) -> Result<Self, AppError> {
204 let client = reqwest::Client::builder()
205 .timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS))
206 .connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS))
207 .user_agent(concat!("sqlite-graphrag/", env!("CARGO_PKG_VERSION")))
208 .build()
209 .map_err(|e| {
210 AppError::Embedding(crate::i18n::validation::embedding_http_client_build_failed(
211 e,
212 ))
213 })?;
214
215 let default_input_type = model_default_input_type(&model);
216
217 Ok(Self {
218 client,
219 api_key,
220 model,
221 dim,
222 default_input_type,
223 base_url,
224 })
225 }
226
227 #[cfg(test)]
231 fn new_with_url(
232 api_key: SecretBox<String>,
233 model: String,
234 dim: usize,
235 base_url: String,
236 ) -> Result<Self, AppError> {
237 Self::new_with_base_url(api_key, model, dim, base_url)
238 }
239
240 pub fn default_input_type(&self) -> Option<&'static str> {
242 self.default_input_type
243 }
244
245 pub async fn embed_single(
247 &self,
248 text: &str,
249 input_type: Option<&str>,
250 ) -> Result<Vec<f32>, EmbedError> {
251 crate::memory_guard::check_embedding_input_size(text)?;
255
256 let request = EmbeddingRequest {
257 model: &self.model,
258 input: EmbeddingInput::Single(text),
259 dimensions: mrl_wire_dimensions(&self.model, self.dim),
260 encoding_format: "float",
261 input_type,
262 };
263
264 let response = self.execute_with_retry(&request).await?;
265
266 let embedding = response
267 .data
268 .into_iter()
269 .next()
270 .ok_or_else(|| {
271 AppError::Embedding(
272 crate::i18n::validation::embedding_empty_response_from_openrouter(),
273 )
274 })?
275 .embedding;
276
277 Ok(self.truncate_embedding(embedding)?)
278 }
279
280 pub async fn embed_batch(
282 &self,
283 texts: &[&str],
284 input_type: Option<&str>,
285 ) -> Result<Vec<Vec<f32>>, EmbedError> {
286 if texts.is_empty() {
287 return Ok(Vec::new());
288 }
289
290 for text in texts {
294 crate::memory_guard::check_embedding_input_size(text)?;
295 }
296
297 let mut all = Vec::with_capacity(texts.len());
298
299 let batch_size = 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(
313 crate::i18n::validation::embedding_expected_count(
314 chunk.len(),
315 response.data.len(),
316 ),
317 )
318 .into());
319 }
320
321 let mut sorted = response.data;
322 sorted.sort_by_key(|d| d.index);
323
324 for d in sorted {
325 all.push(self.truncate_embedding(d.embedding)?);
326 }
327 }
328
329 Ok(all)
330 }
331
332 fn truncate_embedding(&self, embedding: Vec<f32>) -> Result<Vec<f32>, AppError> {
333 if embedding.len() < self.dim {
334 return Err(AppError::Embedding(
335 crate::i18n::validation::embedding_dimension_less_than_requested(
336 embedding.len(),
337 self.dim,
338 ),
339 ));
340 }
341 if embedding.len() == self.dim {
342 Ok(embedding)
343 } else {
344 Ok(embedding[..self.dim].to_vec())
345 }
346 }
347
348 async fn execute_with_retry(
355 &self,
356 request: &EmbeddingRequest<'_>,
357 ) -> Result<EmbeddingResponse, EmbedError> {
358 let mut last_err: Option<EmbedError> = None;
359
360 for attempt in 0..crate::openrouter_http::MAX_RETRIES {
361 let result = self
362 .client
363 .post(&self.base_url)
364 .header(
365 "Authorization",
366 format!("Bearer {}", self.api_key.expose_secret()),
367 )
368 .json(request)
369 .send()
370 .await;
371
372 let resp = match result {
373 Ok(r) => r,
374 Err(e) if e.is_timeout() => {
375 return Err(EmbedError::new(
376 AppError::Embedding(
377 crate::i18n::validation::embedding_openrouter_request_timed_out(),
378 ),
379 AttemptOutcome::Transient,
380 ));
381 }
382 Err(e) => {
383 last_err = Some(EmbedError::new(
384 AppError::Embedding(
385 crate::i18n::validation::embedding_http_request_failed(e),
386 ),
387 AttemptOutcome::Transient,
388 ));
389 crate::openrouter_http::backoff(attempt).await;
390 continue;
391 }
392 };
393
394 let status = resp.status();
395
396 if status.is_success() {
397 let body = resp.text().await.map_err(|e| {
398 EmbedError::new(
399 AppError::Embedding(
400 crate::i18n::validation::embedding_failed_to_read_response_body(e),
401 ),
402 AttemptOutcome::Transient,
403 )
404 })?;
405 match serde_json::from_str::<EmbeddingEnvelope>(&body) {
406 Ok(env) => {
407 if let Some(api_err) = env.error {
412 let retry_class =
413 crate::openrouter_http::provider_error_retry_class(&api_err);
414 return Err(EmbedError::new(
415 AppError::ProviderError {
416 code: api_err.code_string(),
417 message: api_err.message,
418 },
419 retry_class,
420 ));
421 }
422 match env.data {
423 Some(data) => return Ok(EmbeddingResponse { data }),
424 None => {
425 tracing::warn!(
426 attempt,
427 body_len = body.len(),
428 "HTTP 200 with neither data nor error (retrying)"
429 );
430 last_err = Some(EmbedError::new(
431 AppError::Embedding(
432 crate::i18n::validation::embedding_openrouter_200_neither_data_nor_error(),
433 ),
434 AttemptOutcome::Transient,
435 ));
436 crate::openrouter_http::backoff(attempt).await;
437 continue;
438 }
439 }
440 }
441 Err(e) => {
442 tracing::warn!(
443 attempt,
444 body_len = body.len(),
445 "HTTP 200 but JSON unparseable (retrying): {e}"
446 );
447 last_err = Some(EmbedError::new(
448 AppError::Embedding(
449 crate::i18n::validation::embedding_failed_to_parse_response(e),
450 ),
451 AttemptOutcome::Transient,
452 ));
453 crate::openrouter_http::backoff(attempt).await;
454 continue;
455 }
456 }
457 }
458
459 if status.as_u16() == 401 {
460 return Err(EmbedError::new(
461 AppError::Embedding(
462 crate::i18n::validation::embedding_openrouter_invalid_api_key_401(),
463 ),
464 AttemptOutcome::HardFailure,
465 ));
466 }
467
468 if status.as_u16() == 400 || status.as_u16() == 404 {
469 let body = resp.text().await.unwrap_or_default();
470 return Err(EmbedError::new(
471 AppError::Embedding(crate::i18n::validation::embedding_openrouter_returned(
472 status, &body,
473 )),
474 AttemptOutcome::HardFailure,
475 ));
476 }
477
478 if status.as_u16() == 429 {
479 let retry_after = resp
480 .headers()
481 .get("retry-after")
482 .and_then(|v| v.to_str().ok())
483 .and_then(|v| v.parse::<u64>().ok())
484 .unwrap_or(2);
485 tracing::warn!(
486 attempt,
487 retry_after_secs = retry_after,
488 "OpenRouter rate limited, waiting"
489 );
490 last_err = Some(EmbedError::new(
495 AppError::RateLimited {
496 detail: format!("OpenRouter HTTP 429 (retry-after {retry_after}s)"),
497 },
498 AttemptOutcome::Transient,
499 ));
500 tokio::time::sleep(Duration::from_secs(retry_after)).await;
501 continue;
502 }
503
504 if status.is_server_error() {
505 tracing::warn!(attempt, status = %status, "OpenRouter server error, retrying");
506 last_err = Some(EmbedError::new(
507 AppError::Embedding(
508 crate::i18n::validation::embedding_openrouter_server_error(status),
509 ),
510 AttemptOutcome::Transient,
511 ));
512 crate::openrouter_http::backoff(attempt).await;
513 continue;
514 }
515
516 let body = resp.text().await.unwrap_or_default();
517 return Err(EmbedError::new(
518 AppError::Embedding(crate::i18n::validation::embedding_unexpected_http(
519 status, &body,
520 )),
521 crate::openrouter_http::status_retry_class(status),
522 ));
523 }
524
525 Err(last_err.unwrap_or_else(|| {
530 EmbedError::new(
531 AppError::Embedding(crate::i18n::validation::embedding_openrouter_max_retries()),
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}