Skip to main content

chat_core/chat/
embed.rs

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}