foundry_local_sdk/openai/
embedding_client.rs

1//! OpenAI-compatible embedding client.
2#![allow(deprecated)] // this module implements the deprecated OpenAI facade
3
4use async_openai::types::embeddings::CreateEmbeddingResponse;
5use serde_json::{json, Value};
6
7use crate::detail::native::NativeModel;
8use crate::detail::session::NativeSession;
9use crate::detail::task::spawn_blocking;
10use crate::error::{FoundryLocalError, Result};
11
12/// Client for OpenAI-compatible embedding generation backed by a local model.
13#[deprecated(
14    since = "2.0.0",
15    note = "The OpenAI direct clients are deprecated; use the Session API instead \
16            (`EmbeddingsSession::new(&model)`)."
17)]
18pub struct EmbeddingClient {
19    model_id: String,
20    model: NativeModel,
21}
22
23impl EmbeddingClient {
24    pub(crate) fn new(model_id: &str, model: NativeModel) -> Self {
25        Self {
26            model_id: model_id.to_owned(),
27            model,
28        }
29    }
30
31    /// Generate embeddings for a single input text.
32    pub async fn generate_embedding(&self, input: &str) -> Result<CreateEmbeddingResponse> {
33        Self::validate_input(input)?;
34        let request = self.build_request(json!(input));
35        self.execute_request(request).await
36    }
37
38    /// Generate embeddings for multiple input texts in a single request.
39    pub async fn generate_embeddings(&self, inputs: &[&str]) -> Result<CreateEmbeddingResponse> {
40        if inputs.is_empty() {
41            return Err(FoundryLocalError::Validation {
42                reason: "inputs must be a non-empty array".into(),
43            });
44        }
45        for input in inputs {
46            Self::validate_input(input)?;
47        }
48        let request = self.build_request(json!(inputs));
49        self.execute_request(request).await
50    }
51
52    async fn execute_request(&self, request: Value) -> Result<CreateEmbeddingResponse> {
53        let request_json = serde_json::to_string(&request)?;
54        let model = self.model.clone();
55
56        let raw = spawn_blocking(move || {
57            let session = NativeSession::create(&model)?;
58            session.run_openai_json(&request_json)
59        })
60        .await?;
61
62        // Patch the response to add fields required by async_openai types
63        // that the server doesn't return (object on each item, usage)
64        let mut response_value: Value = serde_json::from_str(&raw)?;
65        if let Some(data) = response_value
66            .get_mut("data")
67            .and_then(|d| d.as_array_mut())
68        {
69            for item in data {
70                if item.get("object").is_none() {
71                    item.as_object_mut()
72                        .map(|m| m.insert("object".into(), json!("embedding")));
73                }
74            }
75        }
76        if response_value.get("usage").is_none() {
77            response_value.as_object_mut().map(|m| {
78                m.insert(
79                    "usage".into(),
80                    json!({"prompt_tokens": 0, "total_tokens": 0}),
81                )
82            });
83        }
84
85        let parsed: CreateEmbeddingResponse = serde_json::from_value(response_value)?;
86        Ok(parsed)
87    }
88
89    fn build_request(&self, input: Value) -> Value {
90        json!({
91            "model": self.model_id,
92            "input": input,
93        })
94    }
95
96    fn validate_input(input: &str) -> Result<()> {
97        if input.trim().is_empty() {
98            return Err(FoundryLocalError::Validation {
99                reason: "input must be a non-empty string".into(),
100            });
101        }
102        Ok(())
103    }
104}
foundry_local_sdk/openai/embedding_client.rs

foundry_local_sdk/openai/
embedding_client.rs
