diff --git a/README.md b/README.md index 434d642..eea6321 100644 --- a/README.md +++ b/README.md @@ -154,7 +154,8 @@ each row must be judged in complete isolation. applied before inference, so `WHERE region = 'eu' AND prompt_jev(...) > 0.8` only asks about the `eu` rows. - **Limits.** Each row's text may be up to 64 KiB after JSON encoding. Longer text - fails the query. Requests are capped at 256 KiB. + fails the query. Requests are capped at 256 KiB. Questions may be up to 4,000 + characters, labels 256, and descriptions 1,024. - **Outages.** Each request is tried three times with a 30-second timeout. If all fail, the affected rows return NULL and the query completes. Bad credentials, invalid options, or a malformed reply fail the query instead. diff --git a/src/sql.rs b/src/sql.rs index f8a2913..54d32c4 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -21,12 +21,25 @@ pub struct Question { pub batch_size: usize, } impl Question { + /// Longest instructions accepted, in characters. The provider rejects + /// requests well before the crate's request-size cap would, and a plan-time + /// error names the actual problem instead of a provider HTTP status. + pub const MAX_INSTRUCTION_CHARS: usize = 4_000; + pub const MAX_LABEL_CHARS: usize = 256; + pub const MAX_DESCRIPTION_CHARS: usize = 1_024; + pub fn validate(&self) -> Result<()> { if self.instructions.trim().is_empty() { return Err(plan_datafusion_err!( "prompt_jev instructions must not be empty" )); } + if self.instructions.chars().count() > Self::MAX_INSTRUCTION_CHARS { + return Err(plan_datafusion_err!( + "prompt_jev instructions exceed {} characters", + Self::MAX_INSTRUCTION_CHARS + )); + } if !(1..=64).contains(&self.batch_size) { return Err(plan_datafusion_err!( "prompt_jev batch_size must be between 1 and 64" @@ -42,6 +55,17 @@ impl Question { "prompt_jev criteria require unique non-empty labels and non-empty descriptions" )); } + if c.label.chars().count() > Self::MAX_LABEL_CHARS + || c.description + .as_ref() + .is_some_and(|d| d.chars().count() > Self::MAX_DESCRIPTION_CHARS) + { + return Err(plan_datafusion_err!( + "prompt_jev labels are limited to {} characters and descriptions to {}", + Self::MAX_LABEL_CHARS, + Self::MAX_DESCRIPTION_CHARS + )); + } } let valid = match self.kind.as_str() { "choice" => (2..=255).contains(&self.criteria.len()), @@ -122,7 +146,15 @@ pub fn rewrite_statement(statement: &mut ast::Statement) -> Result<()> { let Expr::Function(f) = expr else { return ControlFlow::Continue(()); }; - if f.name.to_string().to_lowercase() != "prompt_jev" { + // Match the bare identifier, so `"prompt_jev"(...)` and `PROMPT_JEV(...)` + // are the same call; a schema-qualified name is left to DataFusion. + let is_prompt_jev = match f.name.0.as_slice() { + [part] => part + .as_ident() + .is_some_and(|i| i.value.eq_ignore_ascii_case("prompt_jev")), + _ => false, + }; + if !is_prompt_jev { return ControlFlow::Continue(()); } let result = (|| { diff --git a/tests/sql.rs b/tests/sql.rs index b55b2a0..9b84ce8 100644 --- a/tests/sql.rs +++ b/tests/sql.rs @@ -850,3 +850,44 @@ async fn a_folded_limit_survives_the_filter_split() { "rows failing `id > 1` must not be asked about: {asked}" ); } + +#[tokio::test] +async fn quoted_and_uppercase_names_are_the_same_function() { + let mock = Arc::new(Mock::default()); + let result = run( + mock.clone(), + "SELECT \"prompt_jev\"('a', 'q') AS a, PROMPT_JEV('a', 'q') AS b, \"PROMPT_JEV\"('a', 'q') AS c", + ) + .await + .unwrap(); + assert_eq!(result.iter().map(|b| b.num_rows()).sum::(), 1); + assert!(display(&result).contains("0.9")); +} + +#[tokio::test] +async fn oversized_instructions_and_labels_fail_at_plan_time() { + let long_q = "q".repeat(5_000); + let long_label = "L".repeat(300); + let long_desc = "d".repeat(2_000); + for (sql, needle) in [ + ( + format!("SELECT prompt_jev('a', '{long_q}')"), + "4000 characters", + ), + ( + format!("SELECT prompt_jev('a', 'q', choice := ['{long_label}', 'b'])"), + "labels are limited", + ), + ( + format!( + "SELECT prompt_jev('a', 'q', choice := [{{label: 'a', description: '{long_desc}'}}, 'b'])" + ), + "descriptions to 1024", + ), + ] { + let mock = Arc::new(Mock::default()); + let err = run(mock.clone(), &sql).await.err().unwrap().to_string(); + assert!(err.contains(needle), "{needle}: {err}"); + assert!(mock.calls.lock().unwrap().is_empty()); + } +}