1use crate::{
2 chat::{Chat, state::Embedded},
3 error::{ChatError, ChatFailure},
4 traits::EmbeddingsProvider,
5 types::{
6 callback::CallbackRetryContext, messages::Messages, metadata::Metadata,
7 response::EmbeddingsResponse,
8 },
9};
10
11impl<CP: EmbeddingsProvider> Chat<CP, Embedded> {
12 pub async fn embed(
13 &mut self,
14 messages: &mut Messages,
15 ) -> Result<EmbeddingsResponse, ChatFailure> {
16 if let Some(strategy) = self.before_strategy.as_mut() {
17 strategy(messages, None).await;
18 }
19 let max_retries = self.max_retries.unwrap_or(1);
20 let mut last_metadata: Option<Metadata> = None;
21
22 for idx in 0..max_retries {
23 match self.model.embed(messages).await {
24 Ok(res) => {
25 return Ok(res);
26 }
27 Err(failure) => {
28 if let Some(metadata) = failure.metadata.as_ref() {
29 match &mut last_metadata {
30 Some(existing) => {
31 existing.extend(metadata);
32 }
33 None => {
34 last_metadata = Some(metadata.clone());
35 }
36 }
37 }
38
39 let ctx = CallbackRetryContext { idx, failure };
40
41 if let Some(strategy) = self.retry_strategy.as_mut() {
42 strategy(messages, last_metadata.as_ref(), ctx).await
43 }
44
45 continue;
46 }
47 }
48 }
49
50 Err(ChatFailure {
51 err: ChatError::RateLimited,
52 metadata: last_metadata,
53 })
54 }
55}