risingwave_expr_impl/scalar/
ai_model.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::{Arc, LazyLock};
16use std::time::{Duration, Instant};
17
18use async_openai::Client;
19use async_openai::config::{Config, OpenAIConfig};
20use async_openai::types::embeddings::{CreateEmbeddingRequestArgs, Embedding, EmbeddingInput};
21use prometheus::{Registry, exponential_buckets};
22use risingwave_common::array::{
23    Array, ArrayBuilder, ArrayImpl, ArrayRef, DataChunk, F32Array, ListArrayBuilder, ListValue,
24};
25use risingwave_common::metrics::*;
26use risingwave_common::monitor::GLOBAL_METRICS_REGISTRY;
27use risingwave_common::row::OwnedRow;
28use risingwave_common::types::{DataType, Datum, F32, ScalarImpl};
29use risingwave_expr::expr::{
30    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo,
31};
32use risingwave_expr::{ExprError, Result, build_function};
33use serde::Deserialize;
34use serde_json::Value;
35use thiserror_ext::AsReport;
36
37const OPENAI_EMBEDDING_RETRY_BASE_DELAY: Duration = Duration::from_millis(100);
38const OPENAI_EMBEDDING_RETRY_MAX_DELAY: Duration = Duration::from_secs(10);
39const OPENAI_EMBEDDING_SLOW_REQUEST_THRESHOLD: Duration = Duration::from_secs(60);
40
41/// `OpenAI` embedding context that holds the client and model configuration
42#[derive(Debug)]
43pub struct OpenAiEmbeddingContext {
44    pub client: Client<OpenAIConfig>,
45    pub model: String,
46    pub api_base: String,
47}
48
49#[derive(Deserialize)]
50struct OpenAiEmbeddingConfig {
51    model: String,
52    api_key: Option<String>,
53    org_id: Option<String>,
54    project_id: Option<String>,
55    api_base: Option<String>,
56}
57
58impl OpenAiEmbeddingContext {
59    /// Create a new `OpenAI` embedding context from `api_key` and model
60    pub fn from_config(config: Value) -> Result<Self> {
61        let param: OpenAiEmbeddingConfig = serde_json::from_value(config).map_err(|err| {
62            invalid_param_err(format!("failed to parse config: {}", err.as_report()))
63        })?;
64
65        let mut config = OpenAIConfig::new();
66        if let Some(api_key) = param.api_key {
67            config = config.with_api_key(api_key);
68        }
69        if let Some(org_id) = param.org_id {
70            config = config.with_org_id(org_id);
71        }
72        if let Some(proj_id) = param.project_id {
73            config = config.with_project_id(proj_id);
74        }
75        if let Some(api_base) = param.api_base {
76            config = config.with_api_base(api_base);
77        }
78
79        let client = Client::with_config(config);
80        let api_base = client.config().api_base().to_owned();
81        Ok(Self {
82            client,
83            model: param.model,
84            api_base,
85        })
86    }
87}
88
89#[derive(Debug)]
90struct OpenAiEmbedding {
91    text_expr: BoxedExpression,
92    context: OpenAiEmbeddingContext,
93    metrics: OpenAiEmbeddingMetrics,
94}
95
96impl OpenAiEmbedding {
97    async fn get_embeddings(
98        &self,
99        input: EmbeddingInput,
100        word_counts: &[usize],
101    ) -> Result<Vec<Embedding>> {
102        let expected_embedding_count = word_counts.len();
103        let request_start = Instant::now();
104
105        self.metrics
106            .input_rows
107            .inc_by(expected_embedding_count as u64);
108        for &word_count in word_counts {
109            self.metrics.input_row_word_count.observe(word_count as f64);
110        }
111
112        let request = CreateEmbeddingRequestArgs::default()
113            .model(&self.context.model)
114            .input(input)
115            .build()
116            .map_err(|e| {
117                self.metrics.failure_count.inc();
118                tracing::error!(error = %e.as_report(), "Failed to build OpenAI embedding request");
119                ExprError::Custom("failed to build OpenAI embedding request".into())
120            })?;
121
122        let mut backoff = OPENAI_EMBEDDING_RETRY_BASE_DELAY;
123        loop {
124            let timer = self.metrics.latency.start_timer();
125            let result = {
126                let _inflight_guard = OpenAiEmbeddingInflightGuard::new(
127                    &self.metrics.inflight_requests,
128                    &self.metrics.inflight_rows,
129                    expected_embedding_count,
130                );
131                self.context
132                    .client
133                    .embeddings()
134                    .create(request.clone())
135                    .await
136                    .map_err(|e| {
137                        ExprError::Custom(format!(
138                            "failed to get embedding from OpenAI: {}",
139                            e.as_report()
140                        ))
141                    })
142                    .and_then(|response| {
143                        if response.data.len() != expected_embedding_count {
144                            return Err(ExprError::Custom(
145                                "number of embeddings returned from OpenAI does not match the number of texts"
146                                    .into(),
147                            ));
148                        }
149                        Ok(response.data)
150                    })
151            };
152
153            match result {
154                Ok(embeddings) => {
155                    timer.stop_and_record();
156                    let request_latency = request_start.elapsed();
157                    if request_latency > OPENAI_EMBEDDING_SLOW_REQUEST_THRESHOLD {
158                        let word_count_stats = WordCountStats::new(word_counts);
159                        tracing::warn!(
160                            model = %self.context.model,
161                            row_count = expected_embedding_count,
162                            word_count_min = word_count_stats.min,
163                            word_count_avg = word_count_stats.avg,
164                            word_count_p50 = word_count_stats.p50,
165                            word_count_p99 = word_count_stats.p99,
166                            word_count_max = word_count_stats.max,
167                            latency_ms = request_latency.as_millis(),
168                            latency_secs = request_latency.as_secs_f64(),
169                            "Slow OpenAI embedding request finished"
170                        );
171                    }
172                    self.metrics.success_count.inc();
173                    return Ok(embeddings);
174                }
175                Err(err) => {
176                    timer.stop_and_discard();
177                    self.metrics.failure_count.inc();
178                    tracing::error!(
179                        ?backoff,
180                        error = %err.as_report(),
181                        "Failed to get embedding from OpenAI, retrying"
182                    );
183                    tokio::time::sleep(backoff).await;
184                    backoff = (backoff * 2).min(OPENAI_EMBEDDING_RETRY_MAX_DELAY);
185                }
186            }
187        }
188    }
189}
190
191#[derive(Debug, Clone)]
192struct WordCountStats {
193    min: usize,
194    avg: f64,
195    p50: usize,
196    p99: usize,
197    max: usize,
198}
199
200impl WordCountStats {
201    fn new(word_counts: &[usize]) -> Self {
202        if word_counts.is_empty() {
203            return Self {
204                min: 0,
205                avg: 0.0,
206                p50: 0,
207                p99: 0,
208                max: 0,
209            };
210        }
211
212        let mut sorted = word_counts.to_vec();
213        sorted.sort_unstable();
214        let total: usize = sorted.iter().sum();
215
216        Self {
217            min: sorted[0],
218            avg: total as f64 / sorted.len() as f64,
219            p50: percentile_nearest_rank(&sorted, 0.50),
220            p99: percentile_nearest_rank(&sorted, 0.99),
221            max: sorted[sorted.len() - 1],
222        }
223    }
224}
225
226fn percentile_nearest_rank(sorted: &[usize], percentile: f64) -> usize {
227    debug_assert!(!sorted.is_empty());
228    let rank = (sorted.len() as f64 * percentile).ceil() as usize;
229    sorted[rank.saturating_sub(1).min(sorted.len() - 1)]
230}
231
232fn count_words(text: &str) -> usize {
233    text.split_whitespace().count()
234}
235
236struct OpenAiEmbeddingInflightGuard {
237    inflight_requests: LabelGuardedIntGauge,
238    inflight_rows: LabelGuardedIntGauge,
239    row_count: i64,
240}
241
242impl OpenAiEmbeddingInflightGuard {
243    fn new(
244        inflight_requests: &LabelGuardedIntGauge,
245        inflight_rows: &LabelGuardedIntGauge,
246        row_count: usize,
247    ) -> Self {
248        let row_count = row_count.try_into().unwrap_or(i64::MAX);
249        inflight_requests.inc();
250        inflight_rows.add(row_count);
251        Self {
252            inflight_requests: inflight_requests.clone(),
253            inflight_rows: inflight_rows.clone(),
254            row_count,
255        }
256    }
257}
258
259impl Drop for OpenAiEmbeddingInflightGuard {
260    fn drop(&mut self) {
261        self.inflight_requests.dec();
262        self.inflight_rows.sub(self.row_count);
263    }
264}
265
266impl ExpressionInfo for OpenAiEmbedding {
267    fn return_type(&self) -> DataType {
268        DataType::Float32.list()
269    }
270}
271
272impl AsyncExpression for OpenAiEmbedding {
273    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
274        let text_array = self.text_expr.eval(input).await?;
275        let text_array = text_array.as_utf8();
276
277        // Collect non-null and non-empty texts
278        let mut texts_to_embed = Vec::new();
279        let mut word_counts = Vec::new();
280
281        for i in 0..input.capacity() {
282            if let Some(text) = text_array.value_at(i)
283                && !text.is_empty()
284            {
285                word_counts.push(count_words(text));
286                texts_to_embed.push(text.to_owned());
287            }
288        }
289
290        // Get embeddings in batch
291        let embeddings = if texts_to_embed.is_empty() {
292            Vec::new()
293        } else {
294            self.get_embeddings(EmbeddingInput::StringArray(texts_to_embed), &word_counts)
295                .await?
296        };
297
298        // Map results back to original positions
299        let mut builder = ListArrayBuilder::with_type(input.capacity(), DataType::Float32.list());
300        let mut embedding_idx = 0;
301
302        for i in 0..input.capacity() {
303            if let Some(text) = text_array.value_at(i) {
304                if !text.is_empty() {
305                    // Non-empty text, use the embedding result
306                    if embedding_idx < embeddings.len() {
307                        let embedding = &embeddings[embedding_idx].embedding;
308                        let float_array =
309                            F32Array::from_iter(embedding.iter().map(|&v| Some(F32::from(v))));
310                        let list_value = ListValue::new(float_array.into());
311                        builder.append_owned(Some(list_value));
312                        embedding_idx += 1;
313                    } else {
314                        builder.append(None);
315                    }
316                } else {
317                    // Empty text returns NULL
318                    builder.append(None);
319                }
320            } else {
321                // Null text returns NULL
322                builder.append(None);
323            }
324        }
325
326        Ok(Arc::new(ArrayImpl::List(builder.finish())))
327    }
328
329    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
330        let text_datum = self.text_expr.eval_row(input).await?;
331
332        if let Some(ScalarImpl::Utf8(text)) = text_datum.as_ref() {
333            if text.is_empty() {
334                return Ok(None);
335            }
336
337            let word_count = count_words(text);
338            let word_counts = [word_count];
339            let embeddings = self
340                .get_embeddings(
341                    EmbeddingInput::String(text.to_owned().into_string()),
342                    &word_counts,
343                )
344                .await?;
345            let embedding = &embeddings[0].embedding;
346            let float_array = F32Array::from_iter(embedding.iter().map(|&v| Some(F32::from(v))));
347            Ok(Some(ListValue::new(float_array.into()).into()))
348        } else {
349            Ok(None)
350        }
351    }
352}
353
354fn invalid_param_err(reason: impl Into<String>) -> ExprError {
355    ExprError::InvalidParam {
356        name: "openai_embedding",
357        reason: reason.into().into(),
358    }
359}
360
361#[build_function("openai_embedding(jsonb, varchar) -> float4[]")]
362fn build_openai_embedding_expr(
363    _: DataType,
364    mut children: Vec<BoxedExpression>,
365) -> Result<BoxedExpression> {
366    if children.len() != 2 {
367        return Err(invalid_param_err("expected 2 arguments"));
368    }
369
370    // Check if the first two parameters are constants
371    let config = if let Ok(Some(config_scalar)) = children[0].eval_const() {
372        if let ScalarImpl::Jsonb(config) = config_scalar {
373            config.take()
374        } else {
375            return Err(invalid_param_err(
376                "`embedding_config` must be a jsonb constant",
377            ));
378        }
379    } else {
380        return Err(invalid_param_err("`embedding_config` must be a constant"));
381    };
382
383    let context = OpenAiEmbeddingContext::from_config(config)?;
384
385    Ok(OpenAiEmbedding {
386        text_expr: children.pop().unwrap(), // Take the second expression
387        metrics: GLOBAL_OPENAI_EMBEDDING_METRICS
388            .with_label_values(&context.model, &context.api_base),
389        context,
390    }
391    .boxed())
392}
393
394/// Monitor metrics for `openai_embedding`.
395#[derive(Debug, Clone)]
396struct OpenAiEmbeddingMetricsVec {
397    /// Number of successful `OpenAI` embedding requests.
398    success_count: LabelGuardedIntCounterVec,
399    /// Number of failed `OpenAI` embedding requests.
400    failure_count: LabelGuardedIntCounterVec,
401    /// The latency of `OpenAI` embedding requests in seconds.
402    latency: LabelGuardedHistogramVec,
403    /// Total number of non-null, non-empty input rows submitted to `OpenAI` embedding requests.
404    input_rows: LabelGuardedIntCounterVec,
405    /// Per-row word count of non-null, non-empty input rows submitted to `OpenAI` embedding requests.
406    input_row_word_count: LabelGuardedHistogramVec,
407    /// Number of in-flight `OpenAI` embedding requests.
408    inflight_requests: LabelGuardedIntGaugeVec,
409    /// Number of input rows in in-flight `OpenAI` embedding requests.
410    inflight_rows: LabelGuardedIntGaugeVec,
411}
412
413/// Monitor metrics for `openai_embedding`.
414#[derive(Debug, Clone)]
415struct OpenAiEmbeddingMetrics {
416    /// Number of successful `OpenAI` embedding requests.
417    success_count: LabelGuardedIntCounter,
418    /// Number of failed `OpenAI` embedding requests.
419    failure_count: LabelGuardedIntCounter,
420    /// The latency of `OpenAI` embedding requests in seconds.
421    latency: LabelGuardedHistogram,
422    /// Total number of non-null, non-empty input rows submitted to `OpenAI` embedding requests.
423    input_rows: LabelGuardedIntCounter,
424    /// Per-row word count of non-null, non-empty input rows submitted to `OpenAI` embedding requests.
425    input_row_word_count: LabelGuardedHistogram,
426    /// Number of in-flight `OpenAI` embedding requests.
427    inflight_requests: LabelGuardedIntGauge,
428    /// Number of input rows in in-flight `OpenAI` embedding requests.
429    inflight_rows: LabelGuardedIntGauge,
430}
431
432/// Global `openai_embedding` metrics.
433static GLOBAL_OPENAI_EMBEDDING_METRICS: LazyLock<OpenAiEmbeddingMetricsVec> =
434    LazyLock::new(|| OpenAiEmbeddingMetricsVec::new(&GLOBAL_METRICS_REGISTRY));
435
436impl OpenAiEmbeddingMetricsVec {
437    fn new(registry: &Registry) -> Self {
438        let labels = &["model", "api_base"];
439        let success_count = register_guarded_int_counter_vec_with_registry!(
440            "openai_embedding_success_count",
441            "Total number of successful OpenAI embedding requests",
442            labels,
443            registry
444        )
445        .unwrap();
446        let failure_count = register_guarded_int_counter_vec_with_registry!(
447            "openai_embedding_failure_count",
448            "Total number of failed OpenAI embedding requests",
449            labels,
450            registry
451        )
452        .unwrap();
453        let latency = register_guarded_histogram_vec_with_registry!(
454            "openai_embedding_latency",
455            "The latency(s) of OpenAI embedding requests",
456            labels,
457            exponential_buckets(0.000001, 2.0, 30).unwrap(), // 1us to 1000s
458            registry
459        )
460        .unwrap();
461        let input_rows = register_guarded_int_counter_vec_with_registry!(
462            "openai_embedding_input_rows",
463            "Total number of non-null, non-empty input rows submitted to OpenAI embedding requests",
464            labels,
465            registry
466        )
467        .unwrap();
468        let input_row_word_count = register_guarded_histogram_vec_with_registry!(
469            "openai_embedding_input_row_word_count",
470            "Per-row word count of non-null, non-empty input rows submitted to OpenAI embedding requests",
471            labels,
472            exponential_buckets(1.0, 2.0, 13).unwrap(), // 1 to 4096 words
473            registry
474        )
475        .unwrap();
476        let inflight_requests = register_guarded_int_gauge_vec_with_registry!(
477            "openai_embedding_inflight_requests",
478            "Number of in-flight OpenAI embedding requests",
479            labels,
480            registry
481        )
482        .unwrap();
483        let inflight_rows = register_guarded_int_gauge_vec_with_registry!(
484            "openai_embedding_inflight_rows",
485            "Number of input rows in in-flight OpenAI embedding requests",
486            labels,
487            registry
488        )
489        .unwrap();
490
491        Self {
492            success_count,
493            failure_count,
494            latency,
495            input_rows,
496            input_row_word_count,
497            inflight_requests,
498            inflight_rows,
499        }
500    }
501
502    fn with_label_values(&self, model: &str, api_base: &str) -> OpenAiEmbeddingMetrics {
503        let labels = &[model, api_base];
504
505        OpenAiEmbeddingMetrics {
506            success_count: self.success_count.with_guarded_label_values(labels),
507            failure_count: self.failure_count.with_guarded_label_values(labels),
508            latency: self.latency.with_guarded_label_values(labels),
509            input_rows: self.input_rows.with_guarded_label_values(labels),
510            input_row_word_count: self.input_row_word_count.with_guarded_label_values(labels),
511            inflight_requests: self.inflight_requests.with_guarded_label_values(labels),
512            inflight_rows: self.inflight_rows.with_guarded_label_values(labels),
513        }
514    }
515}
risingwave_expr_impl/scalar/ai_model.rs

risingwave_expr_impl/scalar/
ai_model.rs
