Ben Kunkle
5b2d39e3ae
Closes #ISSUE Before you mark this PR as ready for review, make sure that you have: - [ ] Added a solid test coverage and/or screenshots from doing manual testing - [ ] Done a self-review taking into account security and performance aspects - [ ] Aligned any UI changes with the [UI checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) Release Notes: - N/A *or* Added/Fixed/Improved ...
1654 lines
58 KiB
Rust
1654 lines
58 KiB
Rust
use anyhow::{Context as _, Result};
|
|
use flate2::read::GzDecoder;
|
|
use gpui::BackgroundExecutor;
|
|
use http_client::{AsyncBody, HttpClient, Method, Request};
|
|
use indoc::indoc;
|
|
use serde::Deserialize;
|
|
use serde_json::{Value as JsonValue, json};
|
|
use std::fmt::Write as _;
|
|
use std::io::Read;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use telemetry_events::EditPredictionRating;
|
|
|
|
use zeta_prompt::{ZetaFormat, ZetaPromptInput, excerpt_range_for_format};
|
|
|
|
use crate::example::Example;
|
|
use crate::progress::{InfoStyle, Progress, Step};
|
|
const EDIT_PREDICTION_DEPLOYMENT_EVENT: &str = "Edit Prediction Deployment";
|
|
use edit_prediction::example_spec::{ExampleSpec, TelemetrySource};
|
|
|
|
pub(crate) const SNOWFLAKE_SUCCESS_CODE: &str = "090001";
|
|
pub(crate) const SNOWFLAKE_ASYNC_IN_PROGRESS_CODE: &str = "333334";
|
|
const PREDICTIVE_EDIT_REQUESTED_EVENT: &str = "Predictive Edit Requested";
|
|
const PREDICTIVE_EDIT_REJECTED_EVENT: &str = "Predictive Edit Rejected";
|
|
const EDIT_PREDICTION_RATED_EVENT: &str = "Edit Prediction Rated";
|
|
const EDIT_PREDICTION_SETTLED_EVENT: &str = "Edit Prediction Settled";
|
|
|
|
/// Minimum Zed version for filtering captured examples.
|
|
/// For example, `MinCaptureVersion { minor: 224, patch: 1 }` means only pull examples
|
|
/// where `zed_version >= 0.224.1`.
|
|
#[derive(Clone, Copy, Debug)]
|
|
pub struct MinCaptureVersion {
|
|
pub minor: u32,
|
|
pub patch: u32,
|
|
}
|
|
|
|
const DEFAULT_STATEMENT_TIMEOUT_SECONDS: u64 = 240;
|
|
const SETTLED_STATEMENT_TIMEOUT_SECONDS: u64 = 240;
|
|
pub(crate) const POLL_INTERVAL: Duration = Duration::from_secs(2);
|
|
pub(crate) const MAX_POLL_ATTEMPTS: usize = 120;
|
|
|
|
/// Parse an input token of the form `captured-after:{timestamp}`.
|
|
pub fn parse_captured_after_input(input: &str) -> Option<&str> {
|
|
input.strip_prefix("captured-after:")
|
|
}
|
|
|
|
/// Parse an input token of the form `rejected-after:{timestamp}`.
|
|
pub fn parse_rejected_after_input(input: &str) -> Option<&str> {
|
|
input.strip_prefix("rejected-after:")
|
|
}
|
|
|
|
/// Parse an input token of the form `requested-after:{timestamp}`.
|
|
pub fn parse_requested_after_input(input: &str) -> Option<&str> {
|
|
input.strip_prefix("requested-after:")
|
|
}
|
|
|
|
/// Parse an input token of the form `settled-after:{timestamp}`.
|
|
pub fn parse_settled_after_input(input: &str) -> Option<&str> {
|
|
input.strip_prefix("settled-after:")
|
|
}
|
|
|
|
/// Parse an input token of the form `rated-after:{timestamp}`, `rated-positive-after:{timestamp}`,
|
|
/// or `rated-negative-after:{timestamp}`.
|
|
/// Returns `(timestamp, Option<EditPredictionRating>)` where `None` means all ratings.
|
|
pub fn parse_rated_after_input(input: &str) -> Option<(&str, Option<EditPredictionRating>)> {
|
|
if let Some(timestamp) = input.strip_prefix("rated-positive-after:") {
|
|
Some((timestamp, Some(EditPredictionRating::Positive)))
|
|
} else if let Some(timestamp) = input.strip_prefix("rated-negative-after:") {
|
|
Some((timestamp, Some(EditPredictionRating::Negative)))
|
|
} else if let Some(timestamp) = input.strip_prefix("rated-after:") {
|
|
Some((timestamp, None))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub(crate) struct SnowflakeStatementResponse {
|
|
#[serde(default)]
|
|
pub(crate) data: Vec<Vec<JsonValue>>,
|
|
#[serde(default)]
|
|
pub(crate) result_set_meta_data: Option<SnowflakeResultSetMetaData>,
|
|
#[serde(default)]
|
|
pub(crate) code: Option<String>,
|
|
#[serde(default)]
|
|
pub(crate) message: Option<String>,
|
|
#[serde(default)]
|
|
pub(crate) statement_handle: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub(crate) struct SnowflakeResultSetMetaData {
|
|
#[serde(default, rename = "rowType")]
|
|
row_type: Vec<SnowflakeColumnMeta>,
|
|
#[serde(default)]
|
|
num_rows: Option<i64>,
|
|
#[serde(default)]
|
|
partition_info: Vec<SnowflakePartitionInfo>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
struct SnowflakePartitionInfo {}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
struct SnowflakeColumnMeta {
|
|
#[serde(default)]
|
|
name: String,
|
|
}
|
|
|
|
async fn run_sql_with_polling(
|
|
http_client: Arc<dyn HttpClient>,
|
|
base_url: &str,
|
|
token: &str,
|
|
request: &serde_json::Value,
|
|
step_progress: &crate::progress::StepProgress,
|
|
background_executor: BackgroundExecutor,
|
|
) -> Result<SnowflakeStatementResponse> {
|
|
let mut response = run_sql(http_client.clone(), base_url, token, request).await?;
|
|
|
|
if response.code.as_deref() == Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) {
|
|
let statement_handle = response
|
|
.statement_handle
|
|
.as_ref()
|
|
.context("async query response missing statementHandle")?
|
|
.clone();
|
|
|
|
for attempt in 1..=MAX_POLL_ATTEMPTS {
|
|
step_progress.set_substatus(format!("polling ({attempt})"));
|
|
|
|
background_executor.timer(POLL_INTERVAL).await;
|
|
|
|
response =
|
|
fetch_partition(http_client.clone(), base_url, token, &statement_handle, 0).await?;
|
|
|
|
if response.code.as_deref() != Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) {
|
|
break;
|
|
}
|
|
}
|
|
|
|
if response.code.as_deref() == Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) {
|
|
anyhow::bail!(
|
|
"query still running after {} poll attempts ({} seconds)",
|
|
MAX_POLL_ATTEMPTS,
|
|
MAX_POLL_ATTEMPTS as u64 * POLL_INTERVAL.as_secs()
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(response)
|
|
}
|
|
|
|
struct SnowflakeConfig {
|
|
token: String,
|
|
base_url: String,
|
|
role: Option<String>,
|
|
}
|
|
|
|
async fn fetch_examples_with_query(
|
|
http_client: Arc<dyn HttpClient>,
|
|
step_progress: &crate::progress::StepProgress,
|
|
background_executor: BackgroundExecutor,
|
|
statement: &str,
|
|
bindings: JsonValue,
|
|
timeout_seconds: u64,
|
|
required_columns: &[&str],
|
|
parse_response: for<'a> fn(
|
|
&'a SnowflakeStatementResponse,
|
|
&'a std::collections::HashMap<String, usize>,
|
|
) -> Result<Box<dyn Iterator<Item = Example> + 'a>>,
|
|
) -> Result<Vec<Example>> {
|
|
let snowflake = SnowflakeConfig {
|
|
token: std::env::var("EP_SNOWFLAKE_API_KEY")
|
|
.context("missing required environment variable EP_SNOWFLAKE_API_KEY")?,
|
|
base_url: std::env::var("EP_SNOWFLAKE_BASE_URL").context(
|
|
"missing required environment variable EP_SNOWFLAKE_BASE_URL (e.g. https://<account>.snowflakecomputing.com)",
|
|
)?,
|
|
role: std::env::var("EP_SNOWFLAKE_ROLE").ok(),
|
|
};
|
|
let request = json!({
|
|
"statement": statement,
|
|
"timeout": timeout_seconds,
|
|
"database": "EVENTS",
|
|
"schema": "PUBLIC",
|
|
"warehouse": "DBT",
|
|
"role": snowflake.role.as_deref(),
|
|
"bindings": bindings
|
|
});
|
|
|
|
let response = run_sql_with_polling(
|
|
http_client.clone(),
|
|
&snowflake.base_url,
|
|
&snowflake.token,
|
|
&request,
|
|
step_progress,
|
|
background_executor,
|
|
)
|
|
.await?;
|
|
|
|
let total_rows = response
|
|
.result_set_meta_data
|
|
.as_ref()
|
|
.and_then(|meta| meta.num_rows)
|
|
.unwrap_or(response.data.len() as i64);
|
|
let partition_count = response
|
|
.result_set_meta_data
|
|
.as_ref()
|
|
.map(|meta| meta.partition_info.len())
|
|
.unwrap_or(1)
|
|
.max(1);
|
|
|
|
step_progress.set_info(format!("{} rows", total_rows), InfoStyle::Normal);
|
|
step_progress.set_substatus("parsing");
|
|
|
|
let column_indices = get_column_indices(&response.result_set_meta_data, required_columns);
|
|
|
|
let mut parsed_examples = Vec::with_capacity(total_rows as usize);
|
|
parsed_examples.extend(parse_response(&response, &column_indices)?);
|
|
|
|
if partition_count > 1 {
|
|
let statement_handle = response
|
|
.statement_handle
|
|
.as_ref()
|
|
.context("response has multiple partitions but no statementHandle")?;
|
|
|
|
for partition in 1..partition_count {
|
|
step_progress.set_substatus(format!(
|
|
"fetching partition {}/{}",
|
|
partition + 1,
|
|
partition_count
|
|
));
|
|
|
|
let partition_response = fetch_partition(
|
|
http_client.clone(),
|
|
&snowflake.base_url,
|
|
&snowflake.token,
|
|
statement_handle,
|
|
partition,
|
|
)
|
|
.await?;
|
|
|
|
parsed_examples.extend(parse_response(&partition_response, &column_indices)?);
|
|
}
|
|
}
|
|
|
|
step_progress.set_substatus("done");
|
|
Ok(parsed_examples)
|
|
}
|
|
|
|
pub(crate) async fn fetch_partition(
|
|
http_client: Arc<dyn HttpClient>,
|
|
base_url: &str,
|
|
token: &str,
|
|
statement_handle: &str,
|
|
partition: usize,
|
|
) -> Result<SnowflakeStatementResponse> {
|
|
let url = format!(
|
|
"{}/api/v2/statements/{}?partition={}",
|
|
base_url.trim_end_matches('/'),
|
|
statement_handle,
|
|
partition
|
|
);
|
|
|
|
let http_request = Request::builder()
|
|
.method(Method::GET)
|
|
.uri(url.as_str())
|
|
.header("Authorization", format!("Bearer {token}"))
|
|
.header(
|
|
"X-Snowflake-Authorization-Token-Type",
|
|
"PROGRAMMATIC_ACCESS_TOKEN",
|
|
)
|
|
.header("Accept", "application/json")
|
|
.header("Accept-Encoding", "gzip")
|
|
.header("User-Agent", "edit_prediction_cli")
|
|
.body(AsyncBody::empty())?;
|
|
|
|
let response = http_client
|
|
.send(http_request)
|
|
.await
|
|
.context("failed to send partition request to Snowflake SQL API")?;
|
|
|
|
let status = response.status();
|
|
let content_encoding = response
|
|
.headers()
|
|
.get("content-encoding")
|
|
.and_then(|v| v.to_str().ok())
|
|
.map(|s| s.to_lowercase());
|
|
|
|
let body_bytes = {
|
|
use futures::AsyncReadExt as _;
|
|
|
|
let mut body = response.into_body();
|
|
let mut bytes = Vec::new();
|
|
body.read_to_end(&mut bytes)
|
|
.await
|
|
.context("failed to read Snowflake SQL API partition response body")?;
|
|
bytes
|
|
};
|
|
|
|
let body_bytes = if content_encoding.as_deref() == Some("gzip") {
|
|
let mut decoder = GzDecoder::new(&body_bytes[..]);
|
|
let mut decompressed = Vec::new();
|
|
decoder
|
|
.read_to_end(&mut decompressed)
|
|
.context("failed to decompress gzip response")?;
|
|
decompressed
|
|
} else {
|
|
body_bytes
|
|
};
|
|
|
|
if !status.is_success() && status.as_u16() != 202 {
|
|
let body_text = String::from_utf8_lossy(&body_bytes);
|
|
anyhow::bail!(
|
|
"snowflake sql api partition request http {}: {}",
|
|
status.as_u16(),
|
|
body_text
|
|
);
|
|
}
|
|
|
|
if body_bytes.is_empty() {
|
|
anyhow::bail!(
|
|
"snowflake sql api partition {} returned empty response body (http {})",
|
|
partition,
|
|
status.as_u16()
|
|
);
|
|
}
|
|
|
|
serde_json::from_slice::<SnowflakeStatementResponse>(&body_bytes).with_context(|| {
|
|
let body_preview = String::from_utf8_lossy(&body_bytes[..body_bytes.len().min(500)]);
|
|
format!(
|
|
"failed to parse Snowflake SQL API partition {} response JSON (http {}): {}",
|
|
partition,
|
|
status.as_u16(),
|
|
body_preview
|
|
)
|
|
})
|
|
}
|
|
|
|
pub(crate) async fn run_sql(
|
|
http_client: Arc<dyn HttpClient>,
|
|
base_url: &str,
|
|
token: &str,
|
|
request: &serde_json::Value,
|
|
) -> Result<SnowflakeStatementResponse> {
|
|
let url = format!("{}/api/v2/statements", base_url.trim_end_matches('/'));
|
|
|
|
let request_body =
|
|
serde_json::to_vec(request).context("failed to serialize Snowflake SQL API request")?;
|
|
|
|
let http_request = Request::builder()
|
|
.method(Method::POST)
|
|
.uri(url.as_str())
|
|
.header("Authorization", format!("Bearer {token}"))
|
|
.header(
|
|
"X-Snowflake-Authorization-Token-Type",
|
|
"PROGRAMMATIC_ACCESS_TOKEN",
|
|
)
|
|
.header("Content-Type", "application/json")
|
|
.header("Accept", "application/json")
|
|
.header("User-Agent", "edit_prediction_cli")
|
|
.body(AsyncBody::from(request_body.clone()))?;
|
|
|
|
let response = http_client
|
|
.send(http_request)
|
|
.await
|
|
.context("failed to send request to Snowflake SQL API")?;
|
|
|
|
let status = response.status();
|
|
let body_bytes = {
|
|
use futures::AsyncReadExt as _;
|
|
|
|
let mut body = response.into_body();
|
|
let mut bytes = Vec::new();
|
|
body.read_to_end(&mut bytes)
|
|
.await
|
|
.context("failed to read Snowflake SQL API response body")?;
|
|
bytes
|
|
};
|
|
|
|
if !status.is_success() && status.as_u16() != 202 {
|
|
let body_text = String::from_utf8_lossy(&body_bytes);
|
|
anyhow::bail!("snowflake sql api http {}: {}", status.as_u16(), body_text);
|
|
}
|
|
|
|
serde_json::from_slice::<SnowflakeStatementResponse>(&body_bytes)
|
|
.context("failed to parse Snowflake SQL API response JSON")
|
|
}
|
|
|
|
pub async fn fetch_rejected_examples_after(
|
|
http_client: Arc<dyn HttpClient>,
|
|
after_timestamps: &[String],
|
|
max_rows_per_timestamp: usize,
|
|
offset: usize,
|
|
background_executor: BackgroundExecutor,
|
|
min_capture_version: Option<MinCaptureVersion>,
|
|
) -> Result<Vec<Example>> {
|
|
if after_timestamps.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let progress = Progress::global();
|
|
|
|
let mut all_examples = Vec::new();
|
|
|
|
for after_date in after_timestamps.iter() {
|
|
let step_progress_name = format!("rejected>{after_date}");
|
|
let step_progress = progress.start(Step::PullExamples, &step_progress_name);
|
|
step_progress.set_substatus("querying");
|
|
|
|
let min_minor_str = min_capture_version.map(|version| version.minor.to_string());
|
|
let min_patch_str = min_capture_version.map(|version| version.patch.to_string());
|
|
let min_minor_str_ref = min_minor_str.as_deref();
|
|
let min_patch_str_ref = min_patch_str.as_deref();
|
|
|
|
let statement = indoc! {r#"
|
|
SELECT
|
|
req.event_properties:request_id::string AS request_id,
|
|
req.device_id::string AS device_id,
|
|
req.time::string AS time,
|
|
req.event_properties:input AS input,
|
|
req.event_properties:prompt::string AS prompt,
|
|
req.event_properties:output::string AS output,
|
|
rej.event_properties:was_shown::boolean AS was_shown,
|
|
rej.event_properties:reason::string AS reason,
|
|
req.event_properties:zed_version::string AS zed_version
|
|
FROM events req
|
|
INNER JOIN events rej
|
|
ON req.event_properties:request_id = rej.event_properties:request_id
|
|
WHERE req.event_type = ?
|
|
AND rej.event_type = ?
|
|
AND req.event_properties:version = 'V3'
|
|
AND rej.event_properties:was_shown = true
|
|
AND req.event_properties:input:can_collect_data = true
|
|
AND req.time > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND (? IS NULL OR (
|
|
TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) > ?
|
|
OR (
|
|
TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) = ?
|
|
AND TRY_CAST(SPLIT_PART(SPLIT_PART(req.event_properties:zed_version::string, '.', 3), '+', 1) AS INTEGER) >= ?
|
|
)
|
|
))
|
|
ORDER BY req.time ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let bindings = json!({
|
|
"1": { "type": "TEXT", "value": PREDICTIVE_EDIT_REQUESTED_EVENT },
|
|
"2": { "type": "TEXT", "value": PREDICTIVE_EDIT_REJECTED_EVENT },
|
|
"3": { "type": "TEXT", "value": after_date },
|
|
"4": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"5": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"6": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"7": { "type": "FIXED", "value": min_patch_str_ref },
|
|
"8": { "type": "FIXED", "value": max_rows_per_timestamp.to_string() },
|
|
"9": { "type": "FIXED", "value": offset.to_string() }
|
|
});
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
bindings,
|
|
DEFAULT_STATEMENT_TIMEOUT_SECONDS,
|
|
&[
|
|
"request_id",
|
|
"device_id",
|
|
"time",
|
|
"input",
|
|
"prompt",
|
|
"output",
|
|
"was_shown",
|
|
"reason",
|
|
"zed_version",
|
|
],
|
|
rejected_examples_from_response,
|
|
)
|
|
.await?;
|
|
|
|
all_examples.extend(examples);
|
|
}
|
|
|
|
Ok(all_examples)
|
|
}
|
|
|
|
pub async fn fetch_requested_examples_after(
|
|
http_client: Arc<dyn HttpClient>,
|
|
after_timestamps: &[String],
|
|
max_rows_per_timestamp: usize,
|
|
offset: usize,
|
|
background_executor: BackgroundExecutor,
|
|
min_capture_version: Option<MinCaptureVersion>,
|
|
) -> Result<Vec<Example>> {
|
|
if after_timestamps.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let progress = Progress::global();
|
|
|
|
let mut all_examples = Vec::new();
|
|
|
|
for after_date in after_timestamps.iter() {
|
|
let step_progress_name = format!("requested>{after_date}");
|
|
let step_progress = progress.start(Step::PullExamples, &step_progress_name);
|
|
step_progress.set_substatus("querying");
|
|
|
|
let min_minor_str = min_capture_version.map(|version| version.minor.to_string());
|
|
let min_patch_str = min_capture_version.map(|version| version.patch.to_string());
|
|
let min_minor_str_ref = min_minor_str.as_deref();
|
|
let min_patch_str_ref = min_patch_str.as_deref();
|
|
|
|
let statement = indoc! {r#"
|
|
SELECT
|
|
req.event_properties:request_id::string AS request_id,
|
|
req.device_id::string AS device_id,
|
|
req.time::string AS time,
|
|
req.event_properties:input AS input,
|
|
req.event_properties:zed_version::string AS zed_version
|
|
FROM events req
|
|
WHERE req.event_type = ?
|
|
AND req.event_properties:version = 'V3'
|
|
AND req.event_properties:input:can_collect_data = true
|
|
AND req.time > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND (? IS NULL OR (
|
|
TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) > ?
|
|
OR (
|
|
TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) = ?
|
|
AND TRY_CAST(SPLIT_PART(SPLIT_PART(req.event_properties:zed_version::string, '.', 3), '+', 1) AS INTEGER) >= ?
|
|
)
|
|
))
|
|
ORDER BY req.time ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let bindings = json!({
|
|
"1": { "type": "TEXT", "value": PREDICTIVE_EDIT_REQUESTED_EVENT },
|
|
"2": { "type": "TEXT", "value": after_date },
|
|
"3": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"4": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"5": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"6": { "type": "FIXED", "value": min_patch_str_ref },
|
|
"7": { "type": "FIXED", "value": max_rows_per_timestamp.to_string() },
|
|
"8": { "type": "FIXED", "value": offset.to_string() }
|
|
});
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
bindings,
|
|
DEFAULT_STATEMENT_TIMEOUT_SECONDS,
|
|
&["request_id", "device_id", "time", "input", "zed_version"],
|
|
requested_examples_from_response,
|
|
)
|
|
.await?;
|
|
|
|
all_examples.extend(examples);
|
|
}
|
|
|
|
Ok(all_examples)
|
|
}
|
|
|
|
pub async fn fetch_captured_examples_after(
|
|
http_client: Arc<dyn HttpClient>,
|
|
after_timestamps: &[String],
|
|
max_rows_per_timestamp: usize,
|
|
offset: usize,
|
|
background_executor: BackgroundExecutor,
|
|
min_capture_version: Option<MinCaptureVersion>,
|
|
) -> Result<Vec<Example>> {
|
|
if after_timestamps.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let progress = Progress::global();
|
|
|
|
let mut all_examples = Vec::new();
|
|
|
|
for after_date in after_timestamps.iter() {
|
|
let step_progress_name = format!("captured>{after_date}");
|
|
let step_progress = progress.start(Step::PullExamples, &step_progress_name);
|
|
step_progress.set_substatus("querying");
|
|
|
|
let min_minor_str = min_capture_version.map(|version| version.minor.to_string());
|
|
let min_patch_str = min_capture_version.map(|version| version.patch.to_string());
|
|
let min_minor_str_ref = min_minor_str.as_deref();
|
|
let min_patch_str_ref = min_patch_str.as_deref();
|
|
|
|
let statement = indoc! {r#"
|
|
SELECT
|
|
settled.event_properties:request_id::string AS request_id,
|
|
settled.device_id::string AS device_id,
|
|
settled.time::string AS time,
|
|
req.event_properties:input AS input,
|
|
settled.event_properties:settled_editable_region::string AS settled_editable_region,
|
|
settled.event_properties:example AS example,
|
|
req.event_properties:zed_version::string AS zed_version
|
|
FROM events settled
|
|
INNER JOIN events req
|
|
ON settled.event_properties:request_id::string = req.event_properties:request_id::string
|
|
WHERE settled.event_type = ?
|
|
AND req.event_type = ?
|
|
AND req.event_properties:version = 'V3'
|
|
AND req.event_properties:input:can_collect_data = true
|
|
AND settled.event_properties:example IS NOT NULL
|
|
AND TYPEOF(settled.event_properties:example) != 'NULL_VALUE'
|
|
AND settled.time > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND (? IS NULL OR (
|
|
TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) > ?
|
|
OR (
|
|
TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) = ?
|
|
AND TRY_CAST(SPLIT_PART(SPLIT_PART(req.event_properties:zed_version::string, '.', 3), '+', 1) AS INTEGER) >= ?
|
|
)
|
|
))
|
|
ORDER BY settled.time ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let bindings = json!({
|
|
"1": { "type": "TEXT", "value": EDIT_PREDICTION_SETTLED_EVENT },
|
|
"2": { "type": "TEXT", "value": PREDICTIVE_EDIT_REQUESTED_EVENT },
|
|
"3": { "type": "TEXT", "value": after_date },
|
|
"4": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"5": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"6": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"7": { "type": "FIXED", "value": min_patch_str_ref },
|
|
"8": { "type": "FIXED", "value": max_rows_per_timestamp.to_string() },
|
|
"9": { "type": "FIXED", "value": offset.to_string() }
|
|
});
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
bindings,
|
|
DEFAULT_STATEMENT_TIMEOUT_SECONDS,
|
|
&[
|
|
"request_id",
|
|
"device_id",
|
|
"time",
|
|
"input",
|
|
"settled_editable_region",
|
|
"example",
|
|
"zed_version",
|
|
],
|
|
captured_examples_from_response,
|
|
)
|
|
.await?;
|
|
|
|
all_examples.extend(examples);
|
|
}
|
|
|
|
Ok(all_examples)
|
|
}
|
|
|
|
pub async fn fetch_settled_examples_after(
|
|
http_client: Arc<dyn HttpClient>,
|
|
after_timestamps: &[String],
|
|
max_rows_per_timestamp: usize,
|
|
offset: usize,
|
|
background_executor: BackgroundExecutor,
|
|
min_capture_version: Option<MinCaptureVersion>,
|
|
) -> Result<Vec<Example>> {
|
|
if after_timestamps.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let progress = Progress::global();
|
|
|
|
let mut all_examples = Vec::new();
|
|
|
|
for after_date in after_timestamps.iter() {
|
|
let step_progress_name = format!("settled>{after_date}");
|
|
let step_progress = progress.start(Step::PullExamples, &step_progress_name);
|
|
step_progress.set_substatus("querying");
|
|
|
|
let _ = min_capture_version;
|
|
|
|
let statement = indoc! {r#"
|
|
WITH requested AS (
|
|
SELECT
|
|
req.event_properties:request_id::string AS request_id,
|
|
req.device_id::string AS device_id,
|
|
req.time AS req_time,
|
|
req.time::string AS time,
|
|
req.event_properties:input AS input,
|
|
req.event_properties:format::string AS requested_format,
|
|
req.event_properties:output::string AS requested_output,
|
|
req.event_properties:zed_version::string AS zed_version
|
|
FROM events req
|
|
WHERE req.event_type = ?
|
|
AND req.event_properties:version = 'V3'
|
|
AND req.event_properties:input:can_collect_data = true
|
|
AND req.time > TRY_TO_TIMESTAMP_NTZ(?)
|
|
)
|
|
SELECT
|
|
req.request_id AS request_id,
|
|
req.device_id AS device_id,
|
|
req.time AS time,
|
|
req.input AS input,
|
|
req.requested_output AS requested_output,
|
|
settled.event_properties:settled_editable_region::string AS settled_editable_region,
|
|
req.requested_format AS requested_format,
|
|
req.zed_version AS zed_version
|
|
FROM requested req
|
|
INNER JOIN events settled
|
|
ON req.request_id = settled.event_properties:request_id::string
|
|
WHERE settled.event_type = ?
|
|
ORDER BY req.req_time ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let bindings = json!({
|
|
"1": { "type": "TEXT", "value": PREDICTIVE_EDIT_REQUESTED_EVENT },
|
|
"2": { "type": "TEXT", "value": after_date },
|
|
"3": { "type": "TEXT", "value": EDIT_PREDICTION_SETTLED_EVENT },
|
|
"4": { "type": "FIXED", "value": max_rows_per_timestamp.to_string() },
|
|
"5": { "type": "FIXED", "value": offset.to_string() }
|
|
});
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
bindings,
|
|
SETTLED_STATEMENT_TIMEOUT_SECONDS,
|
|
&[
|
|
"request_id",
|
|
"device_id",
|
|
"time",
|
|
"input",
|
|
"requested_output",
|
|
"settled_editable_region",
|
|
"requested_format",
|
|
"zed_version",
|
|
],
|
|
settled_examples_from_response,
|
|
)
|
|
.await?;
|
|
|
|
all_examples.extend(examples);
|
|
}
|
|
|
|
Ok(all_examples)
|
|
}
|
|
|
|
pub async fn fetch_rated_examples_after(
|
|
http_client: Arc<dyn HttpClient>,
|
|
inputs: &[(String, Option<EditPredictionRating>)],
|
|
max_rows_per_timestamp: usize,
|
|
offset: usize,
|
|
background_executor: BackgroundExecutor,
|
|
_min_capture_version: Option<MinCaptureVersion>,
|
|
) -> Result<Vec<Example>> {
|
|
if inputs.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let progress = Progress::global();
|
|
|
|
let mut all_examples = Vec::new();
|
|
|
|
for (after_date, rating_filter) in inputs.iter() {
|
|
let filter_label = match rating_filter {
|
|
None => "",
|
|
Some(EditPredictionRating::Positive) => ":positive",
|
|
Some(EditPredictionRating::Negative) => ":negative",
|
|
};
|
|
let step_progress_name = format!("rated{filter_label}>{after_date}");
|
|
let step_progress = progress.start(Step::PullExamples, &step_progress_name);
|
|
step_progress.set_substatus("querying");
|
|
|
|
let rating_value = rating_filter.as_ref().map(|rating| match rating {
|
|
EditPredictionRating::Positive => "Positive",
|
|
EditPredictionRating::Negative => "Negative",
|
|
});
|
|
|
|
let statement = indoc! {r#"
|
|
SELECT
|
|
rated.event_properties:request_id::string AS request_id,
|
|
rated.event_properties:inputs AS inputs,
|
|
rated.event_properties:output::string AS output,
|
|
rated.event_properties:rating::string AS rating,
|
|
rated.event_properties:feedback::string AS feedback,
|
|
rated.device_id::string AS device_id,
|
|
rated.time::string AS time,
|
|
deploy.event_properties:experiment_name::string AS experiment_name,
|
|
deploy.event_properties:environment::string AS environment,
|
|
rated.event_properties:zed_version::string AS zed_version
|
|
FROM events rated
|
|
LEFT JOIN events req
|
|
ON rated.event_properties:request_id::string = req.event_properties:request_id::string
|
|
AND req.event_type = ?
|
|
LEFT JOIN events deploy
|
|
ON req.event_properties:headers:x_baseten_model_id::string = deploy.event_properties:model_id::string
|
|
AND req.event_properties:headers:x_baseten_model_version_id::string = deploy.event_properties:model_version_id::string
|
|
AND deploy.event_type = ?
|
|
WHERE rated.event_type = ?
|
|
AND (? IS NULL OR rated.event_properties:rating::string = ?)
|
|
AND rated.time > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND rated.event_properties:inputs IS NOT NULL
|
|
AND rated.event_properties:inputs:cursor_excerpt IS NOT NULL
|
|
AND rated.event_properties:output IS NOT NULL
|
|
AND rated.event_properties:inputs:can_collect_data = true
|
|
ORDER BY rated.time ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let bindings = json!({
|
|
"1": { "type": "TEXT", "value": PREDICTIVE_EDIT_REQUESTED_EVENT },
|
|
"2": { "type": "TEXT", "value": EDIT_PREDICTION_DEPLOYMENT_EVENT },
|
|
"3": { "type": "TEXT", "value": EDIT_PREDICTION_RATED_EVENT },
|
|
"4": { "type": "TEXT", "value": rating_value },
|
|
"5": { "type": "TEXT", "value": rating_value },
|
|
"6": { "type": "TEXT", "value": after_date },
|
|
"7": { "type": "FIXED", "value": max_rows_per_timestamp.to_string() },
|
|
"8": { "type": "FIXED", "value": offset.to_string() }
|
|
});
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
bindings,
|
|
DEFAULT_STATEMENT_TIMEOUT_SECONDS,
|
|
&[
|
|
"request_id",
|
|
"inputs",
|
|
"output",
|
|
"rating",
|
|
"feedback",
|
|
"device_id",
|
|
"time",
|
|
"experiment_name",
|
|
"environment",
|
|
"zed_version",
|
|
],
|
|
rated_examples_from_response,
|
|
)
|
|
.await?;
|
|
|
|
all_examples.extend(examples);
|
|
}
|
|
|
|
Ok(all_examples)
|
|
}
|
|
|
|
fn rated_examples_from_response<'a>(
|
|
response: &'a SnowflakeStatementResponse,
|
|
column_indices: &'a std::collections::HashMap<String, usize>,
|
|
) -> Result<Box<dyn Iterator<Item = Example> + 'a>> {
|
|
if let Some(code) = &response.code {
|
|
if code != SNOWFLAKE_SUCCESS_CODE {
|
|
anyhow::bail!(
|
|
"snowflake sql api returned error code={code} message={}",
|
|
response.message.as_deref().unwrap_or("<no message>")
|
|
);
|
|
}
|
|
}
|
|
|
|
let iter = response
|
|
.data
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(move |(row_index, data_row)| {
|
|
let get_string = |name: &str| -> Option<String> {
|
|
let index = column_indices.get(name).copied()?;
|
|
match data_row.get(index)? {
|
|
JsonValue::String(s) => Some(s.clone()),
|
|
JsonValue::Null => None,
|
|
other => Some(other.to_string()),
|
|
}
|
|
};
|
|
|
|
let get_json = |name: &str| -> Option<JsonValue> {
|
|
let index = column_indices.get(name).copied()?;
|
|
let value = data_row.get(index)?;
|
|
if value.is_null() {
|
|
return None;
|
|
}
|
|
match value {
|
|
JsonValue::String(s) => serde_json::from_str(s).ok(),
|
|
other => Some(other.clone()),
|
|
}
|
|
};
|
|
|
|
let request_id = get_string("request_id");
|
|
let inputs_json = get_json("inputs");
|
|
let inputs: Option<ZetaPromptInput> = match &inputs_json {
|
|
Some(v) => match serde_json::from_value(v.clone()) {
|
|
Ok(parsed) => Some(parsed),
|
|
Err(e) => {
|
|
log::warn!(
|
|
"skipping row {row_index}: failed to parse inputs - {e}",
|
|
);
|
|
return None;
|
|
}
|
|
},
|
|
None => None,
|
|
};
|
|
let output = get_string("output");
|
|
let rating = get_string("rating");
|
|
let feedback = get_string("feedback").unwrap_or_default();
|
|
let device_id = get_string("device_id");
|
|
let time = get_string("time");
|
|
let experiment_name = get_string("experiment_name");
|
|
let environment = get_string("environment");
|
|
let zed_version = get_string("zed_version");
|
|
|
|
match (inputs, output.clone(), rating.clone(), time.clone()) {
|
|
(Some(inputs), Some(output), Some(rating), Some(time)) => {
|
|
Some(build_rated_example(
|
|
request_id,
|
|
device_id.unwrap_or_default(),
|
|
time,
|
|
inputs,
|
|
output,
|
|
rating,
|
|
feedback,
|
|
experiment_name,
|
|
environment,
|
|
zed_version,
|
|
))
|
|
}
|
|
_ => {
|
|
log::warn!(
|
|
"skipping row {row_index}: missing fields - inputs={:?} output={:?} rating={:?} time={:?}",
|
|
inputs_json.is_some(),
|
|
output.is_some(),
|
|
rating.is_some(),
|
|
time.is_some(),
|
|
);
|
|
None
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Box::new(iter))
|
|
}
|
|
|
|
fn build_rated_example(
|
|
request_id: Option<String>,
|
|
device_id: String,
|
|
time: String,
|
|
input: ZetaPromptInput,
|
|
output: String,
|
|
rating: String,
|
|
feedback: String,
|
|
experiment_name: Option<String>,
|
|
environment: Option<String>,
|
|
zed_version: Option<String>,
|
|
) -> Example {
|
|
let parsed_rating = if rating == "Positive" {
|
|
EditPredictionRating::Positive
|
|
} else {
|
|
EditPredictionRating::Negative
|
|
};
|
|
let is_positive = parsed_rating == EditPredictionRating::Positive;
|
|
let request_id = request_id.unwrap_or_else(|| format!("rated-{}-{}", device_id, time));
|
|
|
|
let mut tags = Vec::with_capacity(3);
|
|
tags.push(if is_positive {
|
|
"rated:positive".to_string()
|
|
} else {
|
|
"rated:negative".to_string()
|
|
});
|
|
if let Some(experiment) = experiment_name {
|
|
tags.push(format!("experiment:{experiment}"));
|
|
}
|
|
if let Some(env) = environment {
|
|
tags.push(format!("environment:{env}"));
|
|
}
|
|
|
|
let mut example =
|
|
build_example_from_snowflake(request_id, device_id, time, input, tags, None, zed_version);
|
|
|
|
example.spec.rating = Some(parsed_rating);
|
|
|
|
if !feedback.is_empty() {
|
|
example
|
|
.spec
|
|
.human_feedback
|
|
.push(edit_prediction::example_spec::HumanFeedback { message: feedback });
|
|
}
|
|
|
|
if is_positive {
|
|
example.spec.expected_patches = vec![output];
|
|
} else {
|
|
example.spec.rejected_patch = Some(output);
|
|
}
|
|
|
|
example
|
|
}
|
|
|
|
fn requested_examples_from_response<'a>(
|
|
response: &'a SnowflakeStatementResponse,
|
|
column_indices: &'a std::collections::HashMap<String, usize>,
|
|
) -> Result<Box<dyn Iterator<Item = Example> + 'a>> {
|
|
if let Some(code) = &response.code {
|
|
if code != SNOWFLAKE_SUCCESS_CODE {
|
|
anyhow::bail!(
|
|
"snowflake sql api returned error code={code} message={}",
|
|
response.message.as_deref().unwrap_or("<no message>")
|
|
);
|
|
}
|
|
}
|
|
|
|
let iter = response
|
|
.data
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(move |(row_index, data_row)| {
|
|
let get_string = |name: &str| -> Option<String> {
|
|
let index = column_indices.get(name).copied()?;
|
|
match data_row.get(index)? {
|
|
JsonValue::String(s) => Some(s.clone()),
|
|
JsonValue::Null => None,
|
|
other => Some(other.to_string()),
|
|
}
|
|
};
|
|
|
|
let get_json = |name: &str| -> Option<JsonValue> {
|
|
let index = column_indices.get(name).copied()?;
|
|
let value = data_row.get(index)?;
|
|
if value.is_null() {
|
|
return None;
|
|
}
|
|
match value {
|
|
JsonValue::String(s) => serde_json::from_str(s).ok(),
|
|
other => Some(other.clone()),
|
|
}
|
|
};
|
|
|
|
let request_id_str = get_string("request_id");
|
|
let device_id = get_string("device_id");
|
|
let time = get_string("time");
|
|
let input_json = get_json("input");
|
|
let input: Option<ZetaPromptInput> =
|
|
input_json.clone().and_then(|v| serde_json::from_value(v).ok());
|
|
let zed_version = get_string("zed_version");
|
|
|
|
match (request_id_str.clone(), device_id.clone(), time.clone(), input) {
|
|
(Some(request_id), Some(device_id), Some(time), Some(input)) => {
|
|
Some(build_example_from_snowflake(
|
|
request_id,
|
|
device_id,
|
|
time,
|
|
input,
|
|
vec!["requested".to_string()],
|
|
None,
|
|
zed_version,
|
|
))
|
|
}
|
|
_ => {
|
|
log::warn!(
|
|
"skipping row {row_index}: missing fields - request_id={:?} device_id={:?} time={:?} input={:?}",
|
|
request_id_str.is_some(),
|
|
device_id.is_some(),
|
|
time.is_some(),
|
|
input_json.is_some(),
|
|
);
|
|
None
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Box::new(iter))
|
|
}
|
|
|
|
fn settled_examples_from_response<'a>(
|
|
response: &'a SnowflakeStatementResponse,
|
|
column_indices: &'a std::collections::HashMap<String, usize>,
|
|
) -> Result<Box<dyn Iterator<Item = Example> + 'a>> {
|
|
if let Some(code) = &response.code {
|
|
if code != SNOWFLAKE_SUCCESS_CODE {
|
|
anyhow::bail!(
|
|
"snowflake sql api returned error code={code} message={}",
|
|
response.message.as_deref().unwrap_or("<no message>")
|
|
);
|
|
}
|
|
}
|
|
|
|
let iter = response
|
|
.data
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(move |(row_index, data_row)| {
|
|
let get_value = |name: &str| -> Option<JsonValue> {
|
|
let index = column_indices.get(name).copied()?;
|
|
let value = data_row.get(index)?;
|
|
if value.is_null() {
|
|
None
|
|
} else {
|
|
Some(value.clone())
|
|
}
|
|
};
|
|
|
|
let get_string = |name: &str| -> Option<String> {
|
|
match get_value(name)? {
|
|
JsonValue::String(s) => Some(s),
|
|
other => Some(other.to_string()),
|
|
}
|
|
};
|
|
|
|
let parse_json_value = |raw: Option<&JsonValue>| -> Option<JsonValue> {
|
|
let value = raw?;
|
|
match value {
|
|
JsonValue::String(s) => serde_json::from_str::<JsonValue>(s).ok(),
|
|
other => Some(other.clone()),
|
|
}
|
|
};
|
|
|
|
let request_id_str = get_string("request_id");
|
|
let device_id = get_string("device_id");
|
|
let time = get_string("time");
|
|
let input_raw = get_value("input");
|
|
let input_json = parse_json_value(input_raw.as_ref());
|
|
let input: Option<ZetaPromptInput> = input_json
|
|
.as_ref()
|
|
.and_then(|parsed| serde_json::from_value(parsed.clone()).ok());
|
|
let requested_output = get_string("requested_output");
|
|
let settled_editable_region = get_string("settled_editable_region");
|
|
let requested_format =
|
|
get_string("requested_format").and_then(|s| ZetaFormat::parse(&s).ok());
|
|
let zed_version = get_string("zed_version");
|
|
|
|
match (
|
|
request_id_str.clone(),
|
|
device_id.clone(),
|
|
time.clone(),
|
|
input.clone(),
|
|
requested_output.clone(),
|
|
settled_editable_region.clone(),
|
|
requested_format,
|
|
) {
|
|
(
|
|
Some(request_id),
|
|
Some(device_id),
|
|
Some(time),
|
|
Some(input),
|
|
Some(requested_output),
|
|
Some(settled_editable_region),
|
|
Some(requested_format),
|
|
) => Some(build_settled_example(
|
|
request_id,
|
|
device_id,
|
|
time,
|
|
input,
|
|
requested_output,
|
|
settled_editable_region,
|
|
requested_format,
|
|
zed_version,
|
|
)),
|
|
_ => {
|
|
let mut missing_fields = Vec::new();
|
|
|
|
if request_id_str.is_none() {
|
|
missing_fields.push("request_id");
|
|
}
|
|
if device_id.is_none() {
|
|
missing_fields.push("device_id");
|
|
}
|
|
if time.is_none() {
|
|
missing_fields.push("time");
|
|
}
|
|
if input_raw.is_none() || input_json.is_none() || input.is_none() {
|
|
missing_fields.push("input");
|
|
}
|
|
if requested_output.is_none() {
|
|
missing_fields.push("requested_output");
|
|
}
|
|
if settled_editable_region.is_none() {
|
|
missing_fields.push("settled_editable_region");
|
|
}
|
|
if requested_format.is_none() {
|
|
missing_fields.push("requested_format");
|
|
}
|
|
|
|
log::warn!(
|
|
"skipping settled row {row_index}: [{}]",
|
|
missing_fields.join(", "),
|
|
);
|
|
None
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Box::new(iter))
|
|
}
|
|
|
|
fn captured_examples_from_response<'a>(
|
|
response: &'a SnowflakeStatementResponse,
|
|
column_indices: &'a std::collections::HashMap<String, usize>,
|
|
) -> Result<Box<dyn Iterator<Item = Example> + 'a>> {
|
|
if let Some(code) = &response.code {
|
|
if code != SNOWFLAKE_SUCCESS_CODE {
|
|
anyhow::bail!(
|
|
"snowflake sql api returned error code={code} message={}",
|
|
response.message.as_deref().unwrap_or("<no message>")
|
|
);
|
|
}
|
|
}
|
|
|
|
let iter = response
|
|
.data
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(move |(row_index, data_row)| {
|
|
let get_value = |name: &str| -> Option<JsonValue> {
|
|
let index = column_indices.get(name).copied()?;
|
|
let value = data_row.get(index)?;
|
|
if value.is_null() {
|
|
None
|
|
} else {
|
|
Some(value.clone())
|
|
}
|
|
};
|
|
|
|
let get_string = |name: &str| -> Option<String> {
|
|
match get_value(name)? {
|
|
JsonValue::String(s) => Some(s),
|
|
other => Some(other.to_string()),
|
|
}
|
|
};
|
|
|
|
let parse_json_value = |raw: Option<&JsonValue>| -> Option<JsonValue> {
|
|
let value = raw?;
|
|
match value {
|
|
JsonValue::String(s) => serde_json::from_str::<JsonValue>(s).ok(),
|
|
other => Some(other.clone()),
|
|
}
|
|
};
|
|
|
|
let request_id = get_string("request_id");
|
|
let device_id = get_string("device_id");
|
|
let time = get_string("time");
|
|
let input_raw = get_value("input");
|
|
let input_json = parse_json_value(input_raw.as_ref());
|
|
let input: Option<ZetaPromptInput> = input_json
|
|
.as_ref()
|
|
.and_then(|parsed| serde_json::from_value(parsed.clone()).ok());
|
|
let example_raw = get_value("example");
|
|
let example_json = parse_json_value(example_raw.as_ref());
|
|
let example_spec: Option<ExampleSpec> = example_json.as_ref().and_then(|parsed| {
|
|
serde_json::from_value(parsed.clone())
|
|
.or_else(|_| {
|
|
parsed
|
|
.as_str()
|
|
.and_then(|markdown| ExampleSpec::from_markdown(markdown).ok())
|
|
.ok_or_else(|| {
|
|
serde_json::Error::io(std::io::Error::other("not markdown"))
|
|
})
|
|
})
|
|
.ok()
|
|
});
|
|
let has_example_spec = example_spec.is_some();
|
|
let settled_editable_region = get_string("settled_editable_region");
|
|
let zed_version = get_string("zed_version");
|
|
|
|
match (
|
|
request_id.clone(),
|
|
device_id.clone(),
|
|
time.clone(),
|
|
input.clone(),
|
|
example_spec,
|
|
settled_editable_region.clone(),
|
|
) {
|
|
(
|
|
Some(request_id),
|
|
Some(device_id),
|
|
Some(time),
|
|
Some(input),
|
|
Some(example_spec),
|
|
Some(settled_editable_region),
|
|
) => Some(build_captured_example(
|
|
request_id,
|
|
device_id,
|
|
time,
|
|
input,
|
|
example_spec,
|
|
settled_editable_region,
|
|
zed_version,
|
|
)),
|
|
_ => {
|
|
let mut missing_fields = Vec::new();
|
|
|
|
if request_id.is_none() {
|
|
missing_fields.push("request_id");
|
|
}
|
|
if device_id.is_none() {
|
|
missing_fields.push("device_id");
|
|
}
|
|
if time.is_none() {
|
|
missing_fields.push("time");
|
|
}
|
|
if input_raw.is_none() || input_json.is_none() || input.is_none() {
|
|
missing_fields.push("input");
|
|
}
|
|
if example_raw.is_none() || !has_example_spec {
|
|
missing_fields.push("example");
|
|
}
|
|
if settled_editable_region.is_none() {
|
|
missing_fields.push("settled_editable_region");
|
|
}
|
|
|
|
log::warn!(
|
|
"skipping captured row {row_index}: [{}]",
|
|
missing_fields.join(", "),
|
|
);
|
|
None
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Box::new(iter))
|
|
}
|
|
|
|
fn build_settled_example(
|
|
request_id: String,
|
|
device_id: String,
|
|
time: String,
|
|
input: ZetaPromptInput,
|
|
requested_output: String,
|
|
settled_editable_region: String,
|
|
requested_format: ZetaFormat,
|
|
zed_version: Option<String>,
|
|
) -> Example {
|
|
let requested_editable_range =
|
|
excerpt_range_for_format(requested_format, &input.excerpt_ranges).0;
|
|
|
|
let base_cursor_excerpt = input.cursor_excerpt.to_string();
|
|
|
|
let requested_range_is_valid = requested_editable_range.start <= requested_editable_range.end
|
|
&& requested_editable_range.end <= base_cursor_excerpt.len();
|
|
let mut example = build_example_from_snowflake(
|
|
request_id.clone(),
|
|
device_id,
|
|
time,
|
|
input,
|
|
vec!["settled".to_string()],
|
|
None,
|
|
zed_version,
|
|
);
|
|
|
|
if !requested_range_is_valid {
|
|
log::warn!(
|
|
"skipping malformed requested range for request {}: requested={:?} (base_len={})",
|
|
request_id,
|
|
requested_editable_range,
|
|
base_cursor_excerpt.len(),
|
|
);
|
|
return example;
|
|
}
|
|
|
|
let settled_replacement = settled_editable_region.as_str();
|
|
let rejected_patch = build_output_patch(
|
|
&example.spec.cursor_path,
|
|
&base_cursor_excerpt,
|
|
&requested_editable_range,
|
|
&requested_output,
|
|
);
|
|
let expected_patch = build_output_patch(
|
|
&example.spec.cursor_path,
|
|
&base_cursor_excerpt,
|
|
&requested_editable_range,
|
|
settled_replacement,
|
|
);
|
|
|
|
example.spec.expected_patches = vec![expected_patch];
|
|
example.spec.rejected_patch = Some(rejected_patch);
|
|
example
|
|
}
|
|
|
|
fn build_captured_example(
|
|
request_id: String,
|
|
device_id: String,
|
|
time: String,
|
|
input: ZetaPromptInput,
|
|
mut example_spec: ExampleSpec,
|
|
settled_editable_region: String,
|
|
zed_version: Option<String>,
|
|
) -> Example {
|
|
let expected_patch = build_output_patch(
|
|
&input.cursor_path,
|
|
input.cursor_excerpt.as_ref(),
|
|
&input.excerpt_ranges.editable_350,
|
|
settled_editable_region.as_str(),
|
|
);
|
|
|
|
example_spec.expected_patches = vec![expected_patch];
|
|
example_spec.telemetry = Some(TelemetrySource {
|
|
request_id,
|
|
device_id,
|
|
time,
|
|
rejection_reason: String::new(),
|
|
was_shown: false,
|
|
});
|
|
|
|
Example {
|
|
spec: example_spec,
|
|
zed_version,
|
|
prompt_inputs: Some(input),
|
|
prompt: None,
|
|
predictions: Vec::new(),
|
|
score: Vec::new(),
|
|
qa: Vec::new(),
|
|
state: None,
|
|
}
|
|
}
|
|
|
|
fn rejected_examples_from_response<'a>(
|
|
response: &'a SnowflakeStatementResponse,
|
|
column_indices: &'a std::collections::HashMap<String, usize>,
|
|
) -> Result<Box<dyn Iterator<Item = Example> + 'a>> {
|
|
if let Some(code) = &response.code {
|
|
if code != SNOWFLAKE_SUCCESS_CODE {
|
|
anyhow::bail!(
|
|
"snowflake sql api returned error code={code} message={}",
|
|
response.message.as_deref().unwrap_or("<no message>")
|
|
);
|
|
}
|
|
}
|
|
|
|
let iter = response
|
|
.data
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(move |(row_index, data_row)| {
|
|
let get_string = |name: &str| -> Option<String> {
|
|
let index = column_indices.get(name).copied()?;
|
|
match data_row.get(index)? {
|
|
JsonValue::String(s) => Some(s.clone()),
|
|
JsonValue::Null => None,
|
|
other => Some(other.to_string()),
|
|
}
|
|
};
|
|
|
|
let get_json = |name: &str| -> Option<JsonValue> {
|
|
let index = column_indices.get(name).copied()?;
|
|
let value = data_row.get(index)?;
|
|
if value.is_null() {
|
|
return None;
|
|
}
|
|
match value {
|
|
JsonValue::String(s) => serde_json::from_str(s).ok(),
|
|
other => Some(other.clone()),
|
|
}
|
|
};
|
|
|
|
let get_bool = |name: &str| -> Option<bool> {
|
|
let index = column_indices.get(name).copied()?;
|
|
match data_row.get(index)? {
|
|
JsonValue::Bool(b) => Some(*b),
|
|
JsonValue::String(s) => s.parse().ok(),
|
|
_ => None,
|
|
}
|
|
};
|
|
|
|
let request_id_str = get_string("request_id");
|
|
let device_id = get_string("device_id");
|
|
let time = get_string("time");
|
|
let input_json = get_json("input");
|
|
let input: Option<ZetaPromptInput> =
|
|
input_json.clone().and_then(|v| serde_json::from_value(v).ok());
|
|
let output = get_string("output");
|
|
let was_shown = get_bool("was_shown");
|
|
let reason = get_string("reason");
|
|
let zed_version = get_string("zed_version");
|
|
|
|
match (request_id_str.clone(), device_id.clone(), time.clone(), input, output.clone(), was_shown, reason.clone()) {
|
|
(Some(request_id), Some(device_id), Some(time), Some(input), Some(output), Some(was_shown), Some(reason)) => {
|
|
Some(build_rejected_example(
|
|
request_id,
|
|
device_id,
|
|
time,
|
|
input,
|
|
output,
|
|
was_shown,
|
|
reason,
|
|
zed_version,
|
|
))
|
|
}
|
|
_ => {
|
|
log::warn!(
|
|
"skipping row {row_index}: missing fields - request_id={:?} device_id={:?} time={:?} input={:?} output={:?} was_shown={:?} reason={:?}",
|
|
request_id_str.is_some(),
|
|
device_id.is_some(),
|
|
time.is_some(),
|
|
input_json.is_some(),
|
|
output.is_some(),
|
|
was_shown.is_some(),
|
|
reason.is_some()
|
|
);
|
|
None
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Box::new(iter))
|
|
}
|
|
|
|
fn build_rejected_example(
|
|
request_id: String,
|
|
device_id: String,
|
|
time: String,
|
|
input: ZetaPromptInput,
|
|
output: String,
|
|
was_shown: bool,
|
|
reason: String,
|
|
zed_version: Option<String>,
|
|
) -> Example {
|
|
let rejected_patch = build_output_patch(
|
|
&input.cursor_path,
|
|
input.cursor_excerpt.as_ref(),
|
|
&input.excerpt_ranges.editable_350,
|
|
&output,
|
|
);
|
|
let mut example = build_example_from_snowflake(
|
|
request_id,
|
|
device_id,
|
|
time,
|
|
input,
|
|
vec![format!("rejection:{}", reason.to_lowercase())],
|
|
Some(RejectionInfo { reason, was_shown }),
|
|
zed_version,
|
|
);
|
|
example.spec.rejected_patch = Some(rejected_patch);
|
|
example
|
|
}
|
|
|
|
struct RejectionInfo {
|
|
reason: String,
|
|
was_shown: bool,
|
|
}
|
|
|
|
fn build_example_from_snowflake(
|
|
request_id: String,
|
|
device_id: String,
|
|
time: String,
|
|
input: ZetaPromptInput,
|
|
tags: Vec<String>,
|
|
rejection: Option<RejectionInfo>,
|
|
zed_version: Option<String>,
|
|
) -> Example {
|
|
let cursor_excerpt = input.cursor_excerpt.as_ref();
|
|
let cursor_offset = input.cursor_offset_in_excerpt;
|
|
|
|
let mut edit_history = String::new();
|
|
for event in &input.events {
|
|
zeta_prompt::write_event(&mut edit_history, event);
|
|
edit_history.push('\n');
|
|
}
|
|
|
|
let (rejection_reason, was_shown) = match &rejection {
|
|
Some(r) => (r.reason.clone(), r.was_shown),
|
|
None => (String::new(), false),
|
|
};
|
|
|
|
let spec = ExampleSpec {
|
|
name: request_id.clone(),
|
|
repository_url: String::new(),
|
|
revision: String::new(),
|
|
tags,
|
|
reasoning: None,
|
|
uncommitted_diff: String::new(),
|
|
cursor_path: input.cursor_path.clone(),
|
|
cursor_position: build_cursor_position(cursor_excerpt, cursor_offset),
|
|
edit_history,
|
|
expected_patches: Vec::new(),
|
|
rejected_patch: None,
|
|
telemetry: Some(TelemetrySource {
|
|
request_id,
|
|
device_id,
|
|
time,
|
|
rejection_reason,
|
|
was_shown,
|
|
}),
|
|
human_feedback: Vec::new(),
|
|
rating: None,
|
|
};
|
|
|
|
Example {
|
|
spec,
|
|
zed_version,
|
|
prompt_inputs: Some(input),
|
|
prompt: None,
|
|
predictions: Vec::new(),
|
|
score: Vec::new(),
|
|
qa: Vec::new(),
|
|
state: None,
|
|
}
|
|
}
|
|
|
|
fn build_cursor_position(excerpt: &str, cursor_offset: usize) -> String {
|
|
let before = &excerpt[..cursor_offset.min(excerpt.len())];
|
|
let after = &excerpt[cursor_offset.min(excerpt.len())..];
|
|
format!("{}[CURSOR_POSITION]{}", before, after)
|
|
}
|
|
|
|
fn build_output_patch(
|
|
cursor_path: &std::path::Path,
|
|
cursor_excerpt: &str,
|
|
editable_range: &std::ops::Range<usize>,
|
|
model_output: &str,
|
|
) -> String {
|
|
let old_text = &cursor_excerpt[editable_range.clone()];
|
|
|
|
let editable_start_row = cursor_excerpt[..editable_range.start]
|
|
.chars()
|
|
.filter(|&c| c == '\n')
|
|
.count() as u32;
|
|
|
|
let diff_body = language::unified_diff_with_offsets(
|
|
old_text,
|
|
model_output,
|
|
editable_start_row,
|
|
editable_start_row,
|
|
);
|
|
|
|
let mut patch = String::new();
|
|
writeln!(&mut patch, "--- a/{}", cursor_path.display()).ok();
|
|
writeln!(&mut patch, "+++ b/{}", cursor_path.display()).ok();
|
|
patch.push_str(&diff_body);
|
|
patch
|
|
}
|
|
|
|
pub(crate) fn get_column_indices(
|
|
meta: &Option<SnowflakeResultSetMetaData>,
|
|
names: &[&str],
|
|
) -> std::collections::HashMap<String, usize> {
|
|
let mut indices = std::collections::HashMap::new();
|
|
if let Some(meta) = meta {
|
|
for (index, col) in meta.row_type.iter().enumerate() {
|
|
for &name in names {
|
|
if col.name.eq_ignore_ascii_case(name) {
|
|
indices.insert(name.to_string(), index);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
indices
|
|
}
|