swiftide_integrations/openai/
embed.rs

1use async_openai::types::embeddings::{CreateEmbeddingRequest, CreateEmbeddingResponse};
2use async_trait::async_trait;
3
4use swiftide_core::{
5    EmbeddingModel, Embeddings,
6    chat_completion::{Usage, errors::LanguageModelError},
7};
8
9use super::GenericOpenAI;
10use crate::openai::openai_error_to_language_model_error;
11
12#[async_trait]
13impl<
14    C: async_openai::config::Config
15        + std::default::Default
16        + Sync
17        + Send
18        + std::fmt::Debug
19        + Clone
20        + 'static,
21> EmbeddingModel for GenericOpenAI<C>
22{
23    async fn embed(&self, input: Vec<String>) -> Result<Embeddings, LanguageModelError> {
24        let model = self
25            .default_options
26            .embed_model
27            .as_ref()
28            .ok_or(LanguageModelError::PermanentError("Model not set".into()))?;
29
30        let request = self
31            .embed_request_defaults()
32            .model(model)
33            .input(&input)
34            .build()
35            .map_err(LanguageModelError::permanent)?;
36
37        tracing::debug!(
38            num_chunks = input.len(),
39            model = &model,
40            "[Embed] Request to openai"
41        );
42        let response = self
43            .client
44            .embeddings()
45            .create(request)
46            .await
47            .map_err(openai_error_to_language_model_error)?;
48
49        let usage = Usage::from(&response.usage);
50
51        // Only track usage for embedding calls, as requests and responses are extremely verbose
52        self.track_completion(
53            model,
54            Some(&usage),
55            None::<&CreateEmbeddingRequest>,
56            None::<&CreateEmbeddingResponse>,
57        );
58
59        let num_embeddings = response.data.len();
60        tracing::debug!(num_embeddings = num_embeddings, "[Embed] Response openai");
61
62        // WARN: Naively assumes that the order is preserved. Might not always be the case.
63        Ok(response.data.into_iter().map(|d| d.embedding).collect())
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::openai::OpenAI;
71    use serde_json::json;
72    use wiremock::{
73        Mock, MockServer, Request, Respond, ResponseTemplate,
74        matchers::{method, path},
75    };
76
77    #[test_log::test(tokio::test)]
78    async fn test_embed_returns_error_when_model_missing() {
79        let openai = OpenAI::builder().build().unwrap();
80        let err = openai.embed(vec!["text".into()]).await.unwrap_err();
81        assert!(matches!(err, LanguageModelError::PermanentError(_)));
82    }
83
84    #[allow(clippy::items_after_statements)]
85    #[test_log::test(tokio::test)]
86    async fn test_embed_success() {
87        let mock_server = MockServer::start().await;
88
89        let response_body = json!({
90            "data": [{
91                "embedding": [0.1, 0.2],
92                "index": 0,
93                "object": "embedding"
94            }],
95            "model": "text-embedding-3-small",
96            "object": "list",
97            "usage": {"prompt_tokens": 5, "total_tokens": 5}
98        });
99
100        struct ValidateEmbeddingRequest(serde_json::Value);
101
102        impl Respond for ValidateEmbeddingRequest {
103            fn respond(&self, request: &Request) -> ResponseTemplate {
104                let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap();
105                assert_eq!(body["model"], "text-embedding-3-small");
106                assert!(body["input"].is_array());
107                ResponseTemplate::new(200).set_body_json(self.0.clone())
108            }
109        }
110
111        Mock::given(method("POST"))
112            .and(path("/embeddings"))
113            .respond_with(ValidateEmbeddingRequest(response_body))
114            .mount(&mock_server)
115            .await;
116
117        let config = async_openai::config::OpenAIConfig::new().with_api_base(mock_server.uri());
118        let client = async_openai::Client::with_config(config);
119
120        let openai = OpenAI::builder()
121            .client(client)
122            .default_embed_model("text-embedding-3-small")
123            .build()
124            .unwrap();
125
126        let embeddings = openai
127            .embed(vec!["Hello".into(), "World".into()])
128            .await
129            .unwrap();
130
131        assert_eq!(embeddings.len(), 1);
132        assert_eq!(embeddings[0], vec![0.1, 0.2]);
133    }
134}
swiftide_integrations/openai/embed.rs

swiftide_integrations/openai/
embed.rs
