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::io::Read;use std::sync::Arc;use std::time::Duration;use telemetry_events::EditPredictionRating; use zeta_prompt::ZetaPromptInput; 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};use std::fmt::Write as _; 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"; /// 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 = 120;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 `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)} 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 token = std::env::var("EP_SNOWFLAKE_API_KEY")        .context("missing required environment variable EP_SNOWFLAKE_API_KEY")?;    let base_url = std::env::var("EP_SNOWFLAKE_BASE_URL").context(        "missing required environment variable EP_SNOWFLAKE_BASE_URL (e.g. https://<account>.snowflakecomputing.com)",    )?;    let role = std::env::var("EP_SNOWFLAKE_ROLE").ok();     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");         // Join rejected events with their corresponding request events to get the full context.        // We filter for V3 sampling data which contains the structured input we need.        // We also filter for predictions that were actually shown to the user (was_shown = true)        // to focus on explicit user rejections rather than implicit cancellations.        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 min_minor_str = min_capture_version.map(|v| v.minor.to_string());        let min_patch_str = min_capture_version.map(|v| v.patch.to_string());        let min_minor_str_ref = min_minor_str.as_deref();        let min_patch_str_ref = min_patch_str.as_deref();        let request = json!({            "statement": statement,            "timeout": DEFAULT_STATEMENT_TIMEOUT_SECONDS,            "database": "EVENTS",            "schema": "PUBLIC",            "warehouse": "DBT",            "role": role,            "bindings": {                "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 response = run_sql_with_polling(            http_client.clone(),            &base_url,            &token,            &request,            &step_progress,            background_executor.clone(),        )        .await?;         let total_rows = response            .result_set_meta_data            .as_ref()            .and_then(|m| m.num_rows)            .unwrap_or(response.data.len() as i64);         let num_partitions = response            .result_set_meta_data            .as_ref()            .map(|m| m.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,            &[                "request_id",                "device_id",                "time",                "input",                "prompt",                "output",                "was_shown",                "reason",                "zed_version",            ],        );         all_examples.extend(rejected_examples_from_response(&response, &column_indices)?);         if num_partitions > 1 {            let statement_handle = response                .statement_handle                .as_ref()                .context("response has multiple partitions but no statementHandle")?;             for partition in 1..num_partitions {                step_progress.set_substatus(format!(                    "fetching partition {}/{}",                    partition + 1,                    num_partitions                ));                 let partition_response = fetch_partition(                    http_client.clone(),                    &base_url,                    &token,                    statement_handle,                    partition,                )                .await?;                 all_examples.extend(rejected_examples_from_response(                    &partition_response,                    &column_indices,                )?);            }        }         step_progress.set_substatus("done");    }     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 token = std::env::var("EP_SNOWFLAKE_API_KEY")        .context("missing required environment variable EP_SNOWFLAKE_API_KEY")?;    let base_url = std::env::var("EP_SNOWFLAKE_BASE_URL").context(        "missing required environment variable EP_SNOWFLAKE_BASE_URL (e.g. https://<account>.snowflakecomputing.com)",    )?;    let role = std::env::var("EP_SNOWFLAKE_ROLE").ok();     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 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

Lines 1–500 of 1223. Load more