From b5199bc59a199875187d808158357a33b09e5f60 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 22 Sep 2026 03:03:23 +0000 Subject: [PATCH] feat: add experimental Jev SQL filtering (#9265) * feat: add experimental Jev SQL filtering Signed-off-by: Lei, HUANG * feat: gate Jev filtering behind an opt-in Cargo feature Signed-off-by: Lei, HUANG * ci: verify Jev registration with default features Run the existing registry regression without the jev feature in both PR tests and merge-queue coverage, alongside the existing feature-enabled test runs. Signed-off-by: Lei, HUANG * docs: clarify Jev concurrency scope and stabilization work Document the per-expression/batch concurrency bound and track process-wide limiting, rate-limit backoff, and request budgets and metrics as stabilization prerequisites. Signed-off-by: Lei, HUANG * feat: enable renamed ai-functions feature by default Rename the Jev Cargo feature across the command, query, and function crates and enable it in their defaults. Update CI and documentation, retaining an isolated no-default-features registry check and the runtime API opt-in. Signed-off-by: Lei, HUANG * ci: remove extra AI feature-off checks Use the regular AI-enabled unit and coverage runs for the default feature configuration. Keep feature-off validation available locally and update the usage guide to match. Signed-off-by: Lei, HUANG * refactor: rename AI feature to ai_functions Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- .github/workflows/integration.yml | 2 +- .github/workflows/rust.yml | 4 +- Cargo.lock | 2 + docs/how-to/jev.md | 119 +++++++++ src/cmd/Cargo.toml | 2 + src/common/function/Cargo.toml | 10 +- src/common/function/src/function_registry.rs | 12 + src/common/function/src/scalars.rs | 2 + src/common/function/src/scalars/jev.rs | 209 +++++++++++++++ src/common/function/src/scalars/jev/tests.rs | 250 ++++++++++++++++++ src/query/Cargo.toml | 2 + .../standalone/common/function/jev.result | 66 +++++ .../cases/standalone/common/function/jev.sql | 32 +++ tests/runner/src/env/bare.rs | 2 +- 14 files changed, 709 insertions(+), 5 deletions(-) create mode 100644 docs/how-to/jev.md create mode 100644 src/common/function/src/scalars/jev.rs create mode 100644 src/common/function/src/scalars/jev/tests.rs create mode 100644 tests/cases/standalone/common/function/jev.result create mode 100644 tests/cases/standalone/common/function/jev.sql diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 2cbfa1e1541..c81dba199e9 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -81,7 +81,7 @@ jobs: # `cargo gc` will invoke `cargo build` with specified args run: | started_epoch="$(date +%s)" - cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,vector_index" + cargo gc -- --bin greptime --bin sqlness-runner --features "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions" elapsed_secs=$(( $(date +%s) - started_epoch )) printf -- '- Binary build: **%ss**\n' "${elapsed_secs}" >>"${GITHUB_STEP_SUMMARY}" - name: Pack greptime binaries diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index b274931bc5b..1a093a2d2a5 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -200,7 +200,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index + run: cargo nextest run --workspace -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 @@ -258,7 +258,7 @@ jobs: run: ../../.github/scripts/pull-test-deps-images.sh && docker compose up -d --wait - name: Run nextest cases - run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index + run: cargo llvm-cov nextest --workspace --lcov --output-path lcov.info -F dashboard -F pg_kvbackend -F mysql_kvbackend -F vector_index -F ai_functions env: CARGO_BUILD_RUSTFLAGS: "-C link-arg=-fuse-ld=mold" RUST_BACKTRACE: 1 diff --git a/Cargo.lock b/Cargo.lock index 7dadf276b7a..2c7f3333329 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2479,6 +2479,7 @@ dependencies = [ "arrow-cast 59.2.0", "arrow-schema 59.2.0", "async-trait", + "axum 0.8.4", "bincode", "catalog", "chrono", @@ -2518,6 +2519,7 @@ dependencies = [ "paste", "pretty_assertions", "regex", + "reqwest 0.13.4", "s2", "serde", "serde_json", diff --git a/docs/how-to/jev.md b/docs/how-to/jev.md new file mode 100644 index 00000000000..2a21f085946 --- /dev/null +++ b/docs/how-to/jev.md @@ -0,0 +1,119 @@ +# Natural-language filtering with Jev (experimental MVP) + +`jev(text, statement, threshold)` returns whether Jev judges the statement true +for the text with probability **greater than or equal to** the threshold. +It is an asynchronous SQL scalar function usable in `WHERE` and `SELECT`. + +## Start GreptimeDB + +Jev is compiled and registered by the **default-enabled Cargo feature +`ai_functions`**. Enable API evaluation separately in the server process: + +```sh +export GREPTIMEDB_EXPERIMENTAL_JEV=true +export JEV_API_KEY='' +cargo run -p cmd -- standalone start +``` + +If your key is already exported in `~/.zshrc`, run `source ~/.zshrc` first. +The optional `JEV_MODEL` defaults to `jev-latest`; `JEV_ENDPOINT` defaults to +`https://api.typesafe.ai/v1/systemone` (the full evaluation endpoint URL). +The MVP uses environment variables rather than TOML configuration. +The Cargo feature includes the SQL function in the build; the runtime environment +variables enable and configure its API calls. Default compilation does not turn +on external API calls. To enable the feature explicitly, use +`--features ai_functions`. + +## Query + +For a new example table, quote `service` and `message` in the DDL because the +current parser treats them as keywords: + +```sql +CREATE TABLE events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); +``` + +```sql +SELECT occurred_at, service, message +FROM events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND jev( + (message), + 'The event reports that a payment still failed after retries.', + 0.8 + ) +ORDER BY occurred_at; +``` + +The first two arguments are strings and the third is a number in `[0, 1]`. +`(message)` is an ordinary parenthesized string expression. For several columns, +combine them explicitly, for example `concat(service, ': ', message)`. +Any null argument produces SQL `NULL`, which `WHERE` excludes, without an API call. + +## MVP behavior + +- Each non-null row makes one HTTP request: the text is `state`, and the statement + is a `noul` question's `instructions`. The returned `noul` probability is compared + with the threshold locally; this is not the separate Choice/Score confidence. +- Up to eight requests run concurrently per expression/batch invocation, not per + query or process. Concurrent partitions and queries can exceed eight requests + in total. Each request has a 30-second timeout. Errors (including rate limits + and invalid responses) fail the query. There is no automatic retry or + cross-query cache. +- Start with small, time-bounded queries. Ordinary SQL predicates can reduce the + candidate set, but SQL does not guarantee left-to-right predicate evaluation or + that `LIMIT` bounds the number of API calls. +- Text passed to the function is sent to the configured TypeSafe endpoint. + The API key remains in the server environment, not in SQL. +- This MVP is intended for standalone use. A distributed deployment would need + the environment on every node evaluating the function and separate validation. + +API reference: + +## Before stabilizing + +The experimental MVP still needs the following controls before stabilization: + +- A process-wide concurrency limit shared by Jev invocations, such as a semaphore. +- Bounded retries with exponential backoff for HTTP `429` and `529`, following + TypeSafe's rate-limit guidance. +- Request budgets and metrics for API calls, latency, retries, and rate-limit errors. + +These are follow-up work, not guarantees provided by the current implementation. + +## Validation + +Check both the default build (where `jev` is registered) and an isolated +`common-function` build without default features. Jev's regular tests use a local +HTTP server and need no API key: + +```sh +cargo nextest run -p common-function +cargo nextest run -p common-function --no-default-features +cargo nextest run -p common-function --no-default-features --features ai_functions +``` + +The sqlness runner explicitly includes `ai_functions` when building its test binary: + +```sh +cargo sqlness bare -t jev +``` + +If using `--bins-dir`, provide a binary built with `ai_functions` (included by +default). CI covers the AI-enabled path through the regular unit tests and +sqlness tests; the feature-off checks above can be run locally when needed. + +An opt-in test calls the real service on synthetic payment events: + +```sh +source ~/.zshrc +GREPTIMEDB_EXPERIMENTAL_JEV=true cargo nextest run -p common-function \ + -E 'test(test_jev_live)' --run-ignored ignored-only +``` diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 722e48ae90a..fff268dea45 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -21,11 +21,13 @@ required-features = ["dev-tools"] [features] default = [ + "ai_functions", "servers/pprof", "servers/mem-prof", "meta-srv/pg_kvbackend", "meta-srv/mysql_kvbackend", ] +ai_functions = ["query/ai_functions"] enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterprise"] # Developer-only helper binaries and diagnostic datanode commands. # Kept out of `default` so normal/release builds don't compile them. diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index a1ea31bc32d..ee6cfac9dfb 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -10,8 +10,9 @@ workspace = true [features] testing = [] enterprise = [] -default = ["geo"] +default = ["geo", "ai_functions"] geo = ["geohash", "h3o", "s2", "wkt", "geo-types", "dep:geo"] +ai_functions = ["dep:futures", "dep:reqwest"] [dependencies] ahash.workspace = true @@ -43,6 +44,7 @@ datafusion-pg-catalog.workspace = true datafusion-physical-expr.workspace = true datatypes.workspace = true derive_more.workspace = true +futures = { workspace = true, optional = true } geo = { version = "0.29", optional = true } geo-types = { version = "0.7", optional = true } geohash = { version = "0.13", optional = true } @@ -57,6 +59,7 @@ num = "0.4" num-traits = "0.2" paste.workspace = true regex.workspace = true +reqwest = { workspace = true, optional = true } s2 = { version = "0.0.12", optional = true } serde.workspace = true serde_json.workspace = true @@ -71,6 +74,7 @@ wkt = { version = "0.11", optional = true } [dev-dependencies] approx = "0.5" +axum.workspace = true criterion.workspace = true futures.workspace = true pretty_assertions.workspace = true @@ -80,3 +84,7 @@ tokio.workspace = true [[bench]] name = "uddsketch" harness = false + +[package.metadata.cargo-udeps.ignore] +# Used only by the feature-gated Jev HTTP tests. +development = ["axum"] diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index c8e9874abc9..037ca5b175a 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -34,6 +34,8 @@ use crate::scalars::date::DateFunction; use crate::scalars::expression::ExpressionFunction; use crate::scalars::hll_count::HllCalcFunction; use crate::scalars::ip::IpFunctions; +#[cfg(feature = "ai_functions")] +use crate::scalars::jev::JevFunction; use crate::scalars::json::JsonFunction; use crate::scalars::matches::MatchesFunction; use crate::scalars::matches_term::MatchesTermFunction; @@ -226,6 +228,8 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| // Full text search function MatchesFunction::register(&function_registry); MatchesTermFunction::register(&function_registry); + #[cfg(feature = "ai_functions")] + JevFunction::register(&function_registry); // System and administration functions SystemFunction::register(&function_registry); @@ -379,6 +383,14 @@ mod tests { assert!(FUNCTION_REGISTRY.get_function("uddsketch_rank").is_some()); } + #[test] + fn test_jev_registration_matches_feature() { + assert_eq!( + FUNCTION_REGISTRY.get_function("jev").is_some(), + cfg!(feature = "ai_functions") + ); + } + #[test] fn test_register_if_absent_registers_new_function() { let registry = FunctionRegistry::default(); diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 306d1cd17e9..61fb7ec632c 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -17,6 +17,8 @@ pub(crate) mod date; pub mod expression; #[cfg(feature = "geo")] pub mod geo; +#[cfg(feature = "ai_functions")] +pub(crate) mod jev; pub mod json; pub mod matches; pub mod matches_term; diff --git a/src/common/function/src/scalars/jev.rs b/src/common/function/src/scalars/jev.rs new file mode 100644 index 00000000000..1c6475e6fac --- /dev/null +++ b/src/common/function/src/scalars/jev.rs @@ -0,0 +1,209 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(test)] +mod tests; + +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::BooleanArray; +use arrow::datatypes::DataType; +use async_trait::async_trait; +use datafusion_common::cast::{as_float64_array, as_string_view_array}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, exec_datafusion_err, exec_err, not_impl_err}; +use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use futures::{StreamExt, TryStreamExt, stream}; +use reqwest::Client; +use serde_json::{Value, json}; + +use crate::function_factory::ScalarFunctionFactory; +use crate::function_registry::FunctionRegistry; + +/// Experimental natural-language predicate: `jev(text, statement, threshold)`. +/// Each non-null row is a Noul question; its probability is compared with `>=`. +#[derive(PartialEq, Eq, Hash)] +pub(crate) struct JevFunction { + signature: Signature, + enabled: bool, + api_key: Option, + endpoint: String, + model: String, +} + +impl JevFunction { + pub(crate) fn register(registry: &FunctionRegistry) { + registry.register(ScalarFunctionFactory { + name: "jev".to_string(), + factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()), + }); + } + + async fn evaluate(&self, client: &Client, text: &str, statement: &str) -> Result { + let response: Value = client + .post(&self.endpoint) + .json(&json!({ + "model": self.model, + "state": text, + "questions": { + "matches": { "type": "noul", "instructions": statement } + } + })) + .send() + .await + .and_then(|response| response.error_for_status()) + .map_err(|e| exec_datafusion_err!("jev request failed: {e}"))? + .json() + .await + .map_err(|e| exec_datafusion_err!("jev response is not valid JSON: {e}"))?; + + let answer = &response["answers"]["matches"]; + let probability = answer["noul"].as_f64().filter(|p| (0.0..=1.0).contains(p)); + match (answer["type"].as_str(), probability) { + (Some("noul"), Some(probability)) => Ok(probability), + _ => exec_err!( + "jev response must contain answers.matches with type noul and a probability in [0, 1]" + ), + } + } + + fn client(&self) -> Result { + if !self.enabled { + return exec_err!( + "jev is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it" + ); + } + let key = self + .api_key + .as_deref() + .filter(|key| !key.trim().is_empty()) + .ok_or_else(|| { + exec_datafusion_err!("jev requires the JEV_API_KEY environment variable") + })?; + let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}")) + .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?; + authorization.set_sensitive(true); + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(reqwest::header::AUTHORIZATION, authorization); + Client::builder() + .default_headers(headers) + .timeout(Duration::from_secs(30)) + .build() + .map_err(|e| exec_datafusion_err!("failed to create jev HTTP client: {e}")) + } +} + +impl Default for JevFunction { + fn default() -> Self { + Self { + // External model evaluations must not be constant-folded during planning. + signature: Signature::exact( + vec![DataType::Utf8View, DataType::Utf8View, DataType::Float64], + Volatility::Volatile, + ), + enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"), + api_key: std::env::var("JEV_API_KEY").ok(), + endpoint: std::env::var("JEV_ENDPOINT") + .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()), + model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()), + } + } +} + +impl fmt::Debug for JevFunction { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Credentials must never appear in plans or diagnostic output. + f.debug_struct("JevFunction") + .field("signature", &self.signature) + .field("enabled", &self.enabled) + .field("model", &self.model) + .finish_non_exhaustive() + } +} + +impl ScalarUDFImpl for JevFunction { + fn name(&self) -> &str { + "jev" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + not_impl_err!("jev can only be called from async contexts") + } +} + +#[async_trait] +impl AsyncScalarUDFImpl for JevFunction { + async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = args + .args + .into_iter() + .map(|arg| arg.into_array(args.number_rows)) + .collect::>>()?; + let [texts, statements, thresholds] = take_function_args(self.name(), arrays)?; + let texts = as_string_view_array(&texts)?; + let statements = as_string_view_array(&statements)?; + let thresholds = as_float64_array(&thresholds)?; + let rows: Vec<_> = texts + .iter() + .zip(statements.iter()) + .zip(thresholds.iter()) + .map(|((text, statement), threshold)| Some((text?, statement?, threshold?))) + .collect(); + + // Validate the whole batch before making any billable requests. + for (_, _, threshold) in rows.iter().flatten() { + if !(0.0..=1.0).contains(threshold) { + return exec_err!("jev threshold must be finite and in [0, 1]"); + } + } + if rows.iter().all(Option::is_none) { + return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( + args.number_rows, + )))); + } + + // Reuse connections within the batch; bound in-flight requests and preserve row order. + let client = self.client()?; + let requests: Vec<_> = rows + .into_iter() + .map(|row| { + let client = &client; + async move { + match row { + Some((text, statement, threshold)) => self + .evaluate(client, text, statement) + .await + .map(|p| Some(p >= threshold)), + None => Ok(None), + } + } + }) + .collect(); + // This limit is per expression/batch invocation, not per query or process. + // Concurrent partitions and queries can each have their own in-flight requests. + let matches: Vec> = stream::iter(requests).buffered(8).try_collect().await?; + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(matches)))) + } +} diff --git a/src/common/function/src/scalars/jev/tests.rs b/src/common/function/src/scalars/jev/tests.rs new file mode 100644 index 00000000000..82d7ab03013 --- /dev/null +++ b/src/common/function/src/scalars/jev/tests.rs @@ -0,0 +1,250 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use axum::http::{HeaderMap, StatusCode}; +use axum::routing::post; +use axum::{Json, Router}; +use datafusion::assert_batches_eq; +use datafusion::prelude::{SessionConfig, SessionContext}; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; + +use super::*; +use crate::function::FunctionContext; +use crate::function_registry::FUNCTION_REGISTRY; + +struct MockServer { + endpoint: String, + task: JoinHandle<()>, +} + +impl MockServer { + async fn start(router: Router) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + Self { endpoint, task } + } + + fn function(&self) -> JevFunction { + JevFunction { + enabled: true, + api_key: Some("test-key".to_string()), + endpoint: self.endpoint.clone(), + model: "test-model".to_string(), + ..Default::default() + } + } +} + +impl Drop for MockServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +fn context(function: JevFunction) -> SessionContext { + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf()); + ctx +} + +#[tokio::test] +async fn test_jev_sql_filters_using_noul_probability() { + let server = MockServer::start(Router::new().route( + "/v1/systemone", + post( + |headers: HeaderMap, Json(request): Json| async move { + assert_eq!(headers["authorization"], "Bearer test-key"); + assert_eq!(request["model"], "test-model"); + assert_eq!(request["questions"]["matches"]["type"], "noul"); + // The SQL statement is passed through verbatim, not rewritten as a prompt. + assert_eq!( + request["questions"]["matches"]["instructions"], + "payment failed" + ); + let probability = match request["state"].as_str().unwrap() { + "failed" => { + // Complete this row after later rows to exercise result alignment. + tokio::time::sleep(Duration::from_millis(20)).await; + 0.95 + } + "boundary" => 0.8, + "recovered" => 0.1, + unexpected => panic!("unexpected state: {unexpected}"), + }; + Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}})) + }, + ), + )) + .await; + let ctx = context(server.function()); + let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL)) AS events(id, message)"; + let batches = ctx + .sql(&format!( + "SELECT id FROM {events} WHERE jev((message), 'payment failed', 0.8) ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + ["+----+", "| id |", "+----+", "| 1 |", "| 2 |", "+----+"], + &batches + ); + + let batches = ctx + .sql(&format!( + "SELECT id, jev(message, 'payment failed', 0.8) AS matched FROM {events} ORDER BY id" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+----+---------+", + "| id | matched |", + "+----+---------+", + "| 1 | true |", + "| 2 | true |", + "| 3 | false |", + "| 4 | |", + "+----+---------+" + ], + &batches + ); +} + +#[tokio::test] +async fn test_jev_nulls_and_invalid_thresholds_need_no_api_key() { + let ctx = context(JevFunction { + enabled: true, + api_key: None, + ..Default::default() + }); + let batches = ctx + .sql("SELECT jev(NULL, 'condition', 0.8) AS a, jev('text', NULL, 0.8) AS b, jev('text', 'condition', NULL) AS c") + .await.unwrap().collect().await.unwrap(); + assert_batches_eq!( + [ + "+---+---+---+", + "| a | b | c |", + "+---+---+---+", + "| | | |", + "+---+---+---+" + ], + &batches + ); + + for threshold in ["-0.1", "1.1", "'NaN'::DOUBLE", "'Infinity'::DOUBLE"] { + let error = ctx + .sql(&format!("SELECT jev('text', 'condition', {threshold})")) + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("jev threshold must be finite and in [0, 1]"), + "{error}" + ); + } + let error = ctx + .sql("SELECT jev('text', 'condition', 0.8)") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains("JEV_API_KEY"), "{error}"); + + let ctx = context(JevFunction { + enabled: false, + api_key: None, + ..Default::default() + }); + let error = ctx + .sql("SELECT jev('text', 'condition', 0.8)") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"), + "{error}" + ); +} + +#[tokio::test] +async fn test_jev_bad_api_responses_fail_the_query() { + for (status, body, expected) in [ + (StatusCode::TOO_MANY_REQUESTS, "{}", "jev request failed"), + (StatusCode::OK, "not json", "jev response is not valid JSON"), + ( + StatusCode::OK, + r#"{"answers":{}}"#, + "jev response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#, + "jev response must contain", + ), + ( + StatusCode::OK, + r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#, + "jev response must contain", + ), + ] { + let server = MockServer::start( + Router::new().route("/v1/systemone", post(move || async move { (status, body) })), + ) + .await; + let ctx = context(server.function()); + let error = ctx + .sql("SELECT jev('text', 'condition', 0.8)") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + } +} + +#[tokio::test] +#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"] +async fn test_jev_live() { + let ctx = SessionContext::new(); + let function = FUNCTION_REGISTRY.get_function("jev").unwrap(); + ctx.register_udf(function.provide(FunctionContext::default())); + let batches = ctx.sql( + "SELECT id FROM (VALUES + (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'), + (2, 'Payment succeeded on the second retry; the payment is complete.'), + (3, 'User logged in successfully.') + ) AS events(id, message) + WHERE jev(message, 'The event reports that a payment still failed after retries.', 0.8) + ORDER BY id", + ).await.unwrap().collect().await.unwrap(); + assert_batches_eq!(["+----+", "| id |", "+----+", "| 1 |", "+----+"], &batches); +} diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index b6dd2d1fdf4..480c95a5adb 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -8,6 +8,8 @@ license.workspace = true workspace = true [features] +default = ["ai_functions"] +ai_functions = ["common-function/ai_functions"] enterprise = ["sql/enterprise"] vector_index = [] diff --git a/tests/cases/standalone/common/function/jev.result b/tests/cases/standalone/common/function/jev.result new file mode 100644 index 00000000000..4f6d91d2867 --- /dev/null +++ b/tests/cases/standalone/common/function/jev.result @@ -0,0 +1,66 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; + ++---------+ +| matched | ++---------+ +| | ++---------+ + +SELECT jev('message', NULL, 0.8) AS matched; + ++---------+ +| matched | ++---------+ +| | ++---------+ + +SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; + ++---------+ +| matched | ++---------+ +| | ++---------+ + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE jev_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +Affected Rows: 0 + +INSERT INTO jev_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +Affected Rows: 2 + +SELECT occurred_at, service, message FROM jev_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) +ORDER BY occurred_at; + ++-------------+---------+---------+ +| occurred_at | service | message | ++-------------+---------+---------+ ++-------------+---------+---------+ + +DROP TABLE jev_events; + +Affected Rows: 0 + +-- Invalid thresholds fail locally before any HTTP request. +SELECT jev('message', 'condition', -0.1); + +Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] + +SELECT jev('message', 'condition', 1.1); + +Error: 3001(EngineExecuteQuery), Execution error: jev threshold must be finite and in [0, 1] + diff --git a/tests/cases/standalone/common/function/jev.sql b/tests/cases/standalone/common/function/jev.sql new file mode 100644 index 00000000000..f4c252c9273 --- /dev/null +++ b/tests/cases/standalone/common/function/jev.sql @@ -0,0 +1,32 @@ +-- No credentials or external service needed: NULL propagates without an API call. +SELECT jev(NULL, 'The event reports a failed payment.', 0.8) AS matched; + +SELECT jev('message', NULL, 0.8) AS matched; + +SELECT jev('message', 'The event reports a failed payment.', NULL) AS matched; + +-- Validate SQL registration, coercion, and asynchronous filtering on a real table. +CREATE TABLE jev_events ( + occurred_at TIMESTAMP TIME INDEX, + "service" STRING, + "message" STRING, + PRIMARY KEY ("service") +); + +INSERT INTO jev_events VALUES + ('2026-09-19T01:00:00Z', 'payments', NULL), + ('2026-09-19T02:00:00Z', 'auth', NULL); + +SELECT occurred_at, service, message FROM jev_events +WHERE occurred_at >= '2026-09-19T00:00:00Z' + AND occurred_at < '2026-09-20T00:00:00Z' + AND service = 'payments' + AND jev((message), 'The event reports that a payment still failed after retries.', 0.8) +ORDER BY occurred_at; + +DROP TABLE jev_events; + +-- Invalid thresholds fail locally before any HTTP request. +SELECT jev('message', 'condition', -0.1); + +SELECT jev('message', 'condition', 1.1); diff --git a/tests/runner/src/env/bare.rs b/tests/runner/src/env/bare.rs index 8fc3855edbb..dccd3673b16 100644 --- a/tests/runner/src/env/bare.rs +++ b/tests/runner/src/env/bare.rs @@ -694,7 +694,7 @@ impl Env { "--bin", "greptime", "--features", - "pg_kvbackend,mysql_kvbackend,vector_index", + "pg_kvbackend,mysql_kvbackend,vector_index,ai_functions", ]) .output() .expect("Failed to start GreptimeDB");