Ben Kunkle
7bcd0b004d
- limit diagnostics sent to teacher - improve parallelism in `ep` commands when using `--max-parallelism` by only grouping by repo when instructed too (repo grouping is only useful for context collection) - update rejected and rated queries to also fetch settled editable region Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Closes #ISSUE Release Notes: - N/A or Added/Fixed/Improved ...
1865 lines
63 KiB
Rust
1865 lines
63 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::collections::HashMap;
|
|
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::PredictionProvider;
|
|
use crate::example::{Example, ExamplePrompt};
|
|
use crate::progress::{InfoStyle, Progress, Step};
|
|
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 SNOWFLAKE_TIMEOUT_CODE: &str = "000630";
|
|
|
|
/// 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,
|
|
}
|
|
|
|
pub(crate) const POLL_INTERVAL: Duration = Duration::from_secs(2);
|
|
const PARTITION_FETCH_MAX_RETRIES: usize = 3;
|
|
const PARTITION_FETCH_RETRY_DELAYS: [Duration; PARTITION_FETCH_MAX_RETRIES] = [
|
|
Duration::from_millis(500),
|
|
Duration::from_secs(1),
|
|
Duration::from_secs(2),
|
|
];
|
|
|
|
/// 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 0.. {
|
|
step_progress.set_substatus(format!("polling ({attempt})"));
|
|
|
|
background_executor.timer(POLL_INTERVAL).await;
|
|
|
|
response = fetch_partition_with_retries(
|
|
http_client.clone(),
|
|
base_url,
|
|
token,
|
|
&statement_handle,
|
|
0,
|
|
background_executor.clone(),
|
|
)
|
|
.await?;
|
|
|
|
if response.code.as_deref() != Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(response)
|
|
}
|
|
|
|
struct SnowflakeConfig {
|
|
token: String,
|
|
base_url: String,
|
|
role: Option<String>,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct QueryRetryState {
|
|
resume_after: String,
|
|
remaining_limit: Option<usize>,
|
|
offset: usize,
|
|
}
|
|
|
|
async fn fetch_examples_with_query<MakeBindings>(
|
|
http_client: Arc<dyn HttpClient>,
|
|
step_progress: &crate::progress::StepProgress,
|
|
background_executor: BackgroundExecutor,
|
|
statement: &str,
|
|
initial_retry_state: QueryRetryState,
|
|
make_bindings: MakeBindings,
|
|
required_columns: &[&str],
|
|
parse_response: for<'a> fn(
|
|
&'a SnowflakeStatementResponse,
|
|
&'a HashMap<String, usize>,
|
|
) -> Result<Box<dyn Iterator<Item = Example> + 'a>>,
|
|
) -> Result<Vec<Example>>
|
|
where
|
|
MakeBindings: Fn(&QueryRetryState) -> JsonValue,
|
|
{
|
|
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 mut requested_columns = required_columns.to_vec();
|
|
if !requested_columns.contains(&"continuation_time") {
|
|
requested_columns.push("continuation_time");
|
|
}
|
|
|
|
let mut parsed_examples = Vec::new();
|
|
let mut retry_state = initial_retry_state;
|
|
let mut retry_count = 0usize;
|
|
|
|
loop {
|
|
let bindings = make_bindings(&retry_state);
|
|
let request = json!({
|
|
"statement": statement,
|
|
"database": "EVENTS",
|
|
"schema": "PUBLIC",
|
|
"warehouse": "DBT",
|
|
"role": snowflake.role.as_deref(),
|
|
"bindings": bindings
|
|
});
|
|
|
|
let response = match run_sql_with_polling(
|
|
http_client.clone(),
|
|
&snowflake.base_url,
|
|
&snowflake.token,
|
|
&request,
|
|
step_progress,
|
|
background_executor.clone(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(response) => response,
|
|
Err(error) => {
|
|
if is_snowflake_timeout_error(&error) && !parsed_examples.is_empty() {
|
|
retry_count += 1;
|
|
step_progress.set_substatus(format!(
|
|
"retrying from {} ({retry_count})",
|
|
retry_state.resume_after
|
|
));
|
|
continue;
|
|
}
|
|
|
|
return Err(error);
|
|
}
|
|
};
|
|
|
|
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, &requested_columns);
|
|
let mut rows_fetched_this_attempt = 0usize;
|
|
let mut timed_out_fetching_partition = false;
|
|
|
|
parsed_examples.extend(parse_response(&response, &column_indices)?);
|
|
rows_fetched_this_attempt += response.data.len();
|
|
let mut last_continuation_time_this_attempt =
|
|
last_continuation_timestamp_from_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 = match fetch_partition_with_retries(
|
|
http_client.clone(),
|
|
&snowflake.base_url,
|
|
&snowflake.token,
|
|
statement_handle,
|
|
partition,
|
|
background_executor.clone(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(response) => response,
|
|
Err(error) => {
|
|
if is_snowflake_timeout_error(&error) && rows_fetched_this_attempt > 0 {
|
|
timed_out_fetching_partition = true;
|
|
break;
|
|
}
|
|
|
|
return Err(error);
|
|
}
|
|
};
|
|
|
|
parsed_examples.extend(parse_response(&partition_response, &column_indices)?);
|
|
rows_fetched_this_attempt += partition_response.data.len();
|
|
|
|
if let Some(partition_continuation_time) =
|
|
last_continuation_timestamp_from_response(&partition_response, &column_indices)
|
|
{
|
|
last_continuation_time_this_attempt = Some(partition_continuation_time);
|
|
}
|
|
}
|
|
}
|
|
|
|
if rows_fetched_this_attempt == 0 {
|
|
step_progress.set_substatus("done");
|
|
return Ok(parsed_examples);
|
|
}
|
|
|
|
if let Some(remaining_limit_value) = &mut retry_state.remaining_limit {
|
|
*remaining_limit_value =
|
|
remaining_limit_value.saturating_sub(rows_fetched_this_attempt);
|
|
if *remaining_limit_value == 0 {
|
|
step_progress.set_substatus("done");
|
|
return Ok(parsed_examples);
|
|
}
|
|
}
|
|
|
|
if !timed_out_fetching_partition {
|
|
step_progress.set_substatus("done");
|
|
return Ok(parsed_examples);
|
|
}
|
|
|
|
let Some(last_continuation_time_this_attempt) = last_continuation_time_this_attempt else {
|
|
step_progress.set_substatus("done");
|
|
return Ok(parsed_examples);
|
|
};
|
|
|
|
retry_state.resume_after = last_continuation_time_this_attempt;
|
|
retry_state.offset = 0;
|
|
retry_count += 1;
|
|
step_progress.set_substatus(format!(
|
|
"retrying from {} ({retry_count})",
|
|
retry_state.resume_after
|
|
));
|
|
}
|
|
}
|
|
|
|
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
|
|
)
|
|
})
|
|
}
|
|
|
|
async fn fetch_partition_with_retries(
|
|
http_client: Arc<dyn HttpClient>,
|
|
base_url: &str,
|
|
token: &str,
|
|
statement_handle: &str,
|
|
partition: usize,
|
|
background_executor: BackgroundExecutor,
|
|
) -> Result<SnowflakeStatementResponse> {
|
|
let mut last_error = None;
|
|
|
|
for retry_attempt in 0..=PARTITION_FETCH_MAX_RETRIES {
|
|
match fetch_partition(
|
|
http_client.clone(),
|
|
base_url,
|
|
token,
|
|
statement_handle,
|
|
partition,
|
|
)
|
|
.await
|
|
{
|
|
Ok(response) => return Ok(response),
|
|
Err(error) => {
|
|
if retry_attempt == PARTITION_FETCH_MAX_RETRIES
|
|
|| !is_transient_partition_fetch_error(&error)
|
|
{
|
|
return Err(error);
|
|
}
|
|
|
|
last_error = Some(error);
|
|
background_executor
|
|
.timer(PARTITION_FETCH_RETRY_DELAYS[retry_attempt])
|
|
.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
match last_error {
|
|
Some(error) => Err(error),
|
|
None => anyhow::bail!("partition fetch retry loop exited without a result"),
|
|
}
|
|
}
|
|
|
|
fn is_transient_partition_fetch_error(error: &anyhow::Error) -> bool {
|
|
error.chain().any(|cause| {
|
|
let message = cause.to_string();
|
|
message.contains("failed to read Snowflake SQL API partition response body")
|
|
|| message.contains("unexpected EOF")
|
|
|| message.contains("peer closed connection without sending TLS close_notify")
|
|
})
|
|
}
|
|
|
|
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
|
|
};
|
|
|
|
let snowflake_response = serde_json::from_slice::<SnowflakeStatementResponse>(&body_bytes)
|
|
.context("failed to parse Snowflake SQL API response JSON")?;
|
|
|
|
if !status.is_success() && status.as_u16() != 202 && !is_timeout_response(&snowflake_response) {
|
|
let body_text = String::from_utf8_lossy(&body_bytes);
|
|
anyhow::bail!("snowflake sql api http {}: {}", status.as_u16(), body_text);
|
|
}
|
|
|
|
if is_timeout_response(&snowflake_response) {
|
|
anyhow::bail!(
|
|
"snowflake sql api timed out code={} message={}",
|
|
snowflake_response.code.as_deref().unwrap_or("<no code>"),
|
|
snowflake_response
|
|
.message
|
|
.as_deref()
|
|
.unwrap_or("<no message>")
|
|
);
|
|
}
|
|
|
|
Ok(snowflake_response)
|
|
}
|
|
|
|
pub async fn fetch_rejected_examples_after(
|
|
http_client: Arc<dyn HttpClient>,
|
|
after_timestamps: &[String],
|
|
max_rows_per_timestamp: Option<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
|
|
ep_request_id AS request_id,
|
|
device_id AS device_id,
|
|
requested_at::string AS continuation_time,
|
|
requested_at::string AS time,
|
|
input_payload AS input,
|
|
prompt AS prompt,
|
|
requested_output AS output,
|
|
settled_editable_region AS settled_editable_region,
|
|
is_ep_shown_before_rejected AS was_shown,
|
|
ep_rejected_reason AS reason,
|
|
zed_version AS zed_version
|
|
FROM ZED_DBT.DBT_PROD.fct_edit_prediction_examples
|
|
WHERE ep_outcome LIKE 'Rejected%'
|
|
AND is_ep_shown_before_rejected = true
|
|
AND requested_at > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND (? IS NULL OR (
|
|
TRY_CAST(SPLIT_PART(zed_version, '.', 2) AS INTEGER) > ?
|
|
OR (
|
|
TRY_CAST(SPLIT_PART(zed_version, '.', 2) AS INTEGER) = ?
|
|
AND TRY_CAST(SPLIT_PART(SPLIT_PART(zed_version, '.', 3), '+', 1) AS INTEGER) >= ?
|
|
)
|
|
))
|
|
ORDER BY requested_at ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
QueryRetryState {
|
|
resume_after: after_date.clone(),
|
|
remaining_limit: max_rows_per_timestamp,
|
|
offset,
|
|
},
|
|
|retry_state| {
|
|
json!({
|
|
"1": { "type": "TEXT", "value": retry_state.resume_after },
|
|
"2": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"3": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"4": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"5": { "type": "FIXED", "value": min_patch_str_ref },
|
|
"6": { "type": "FIXED", "value": format_limit(retry_state.remaining_limit) },
|
|
"7": { "type": "FIXED", "value": retry_state.offset.to_string() }
|
|
})
|
|
},
|
|
&[
|
|
"request_id",
|
|
"device_id",
|
|
"time",
|
|
"input",
|
|
"prompt",
|
|
"output",
|
|
"settled_editable_region",
|
|
"was_shown",
|
|
"reason",
|
|
"zed_version",
|
|
],
|
|
rejected_examples_from_response,
|
|
)
|
|
.await?;
|
|
|
|
all_examples.extend(examples);
|
|
}
|
|
|
|
Ok(all_examples)
|
|
}
|
|
|
|
fn format_limit(limit: Option<usize>) -> String {
|
|
return limit.map(|l| l.to_string()).unwrap_or("NULL".to_string());
|
|
}
|
|
|
|
pub async fn fetch_requested_examples_after(
|
|
http_client: Arc<dyn HttpClient>,
|
|
after_timestamps: &[String],
|
|
max_rows_per_timestamp: Option<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
|
|
ep_request_id AS request_id,
|
|
device_id AS device_id,
|
|
requested_at::string AS continuation_time,
|
|
requested_at::string AS time,
|
|
input_payload AS input,
|
|
zed_version AS zed_version
|
|
FROM ZED_DBT.DBT_PROD.fct_edit_prediction_examples
|
|
WHERE requested_at > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND (? IS NULL OR (
|
|
TRY_CAST(SPLIT_PART(zed_version, '.', 2) AS INTEGER) > ?
|
|
OR (
|
|
TRY_CAST(SPLIT_PART(zed_version, '.', 2) AS INTEGER) = ?
|
|
AND TRY_CAST(SPLIT_PART(SPLIT_PART(zed_version, '.', 3), '+', 1) AS INTEGER) >= ?
|
|
)
|
|
))
|
|
ORDER BY requested_at ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
QueryRetryState {
|
|
resume_after: after_date.clone(),
|
|
remaining_limit: max_rows_per_timestamp,
|
|
offset,
|
|
},
|
|
|retry_state| {
|
|
json!({
|
|
"1": { "type": "TEXT", "value": retry_state.resume_after },
|
|
"2": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"3": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"4": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"5": { "type": "FIXED", "value": min_patch_str_ref },
|
|
"6": { "type": "FIXED", "value": format_limit(retry_state.remaining_limit) },
|
|
"7": { "type": "FIXED", "value": retry_state.offset.to_string() }
|
|
})
|
|
},
|
|
&["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: Option<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
|
|
ep_request_id AS request_id,
|
|
device_id AS device_id,
|
|
requested_at::string AS continuation_time,
|
|
requested_at::string AS time,
|
|
input_payload AS input,
|
|
settled_editable_region AS settled_editable_region,
|
|
example_payload AS example,
|
|
zed_version AS zed_version
|
|
FROM ZED_DBT.DBT_PROD.fct_edit_prediction_examples
|
|
WHERE settled_editable_region IS NOT NULL
|
|
AND example_payload IS NOT NULL
|
|
AND requested_at > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND (? IS NULL OR (
|
|
TRY_CAST(SPLIT_PART(zed_version, '.', 2) AS INTEGER) > ?
|
|
OR (
|
|
TRY_CAST(SPLIT_PART(zed_version, '.', 2) AS INTEGER) = ?
|
|
AND TRY_CAST(SPLIT_PART(SPLIT_PART(zed_version, '.', 3), '+', 1) AS INTEGER) >= ?
|
|
)
|
|
))
|
|
ORDER BY requested_at ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
QueryRetryState {
|
|
resume_after: after_date.clone(),
|
|
remaining_limit: max_rows_per_timestamp,
|
|
offset,
|
|
},
|
|
|retry_state| {
|
|
json!({
|
|
"1": { "type": "TEXT", "value": retry_state.resume_after },
|
|
"2": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"3": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"4": { "type": "FIXED", "value": min_minor_str_ref },
|
|
"5": { "type": "FIXED", "value": min_patch_str_ref },
|
|
"6": { "type": "FIXED", "value": format_limit(retry_state.remaining_limit) },
|
|
"7": { "type": "FIXED", "value": retry_state.offset.to_string() }
|
|
})
|
|
},
|
|
&[
|
|
"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: Option<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#"
|
|
SELECT
|
|
ep_request_id AS request_id,
|
|
device_id AS device_id,
|
|
requested_at::string AS continuation_time,
|
|
requested_at::string AS time,
|
|
input_payload AS input,
|
|
requested_output AS requested_output,
|
|
settled_editable_region AS settled_editable_region,
|
|
requested_format AS requested_format,
|
|
zed_version AS zed_version
|
|
FROM ZED_DBT.DBT_PROD.fct_edit_prediction_examples
|
|
WHERE settled_editable_region IS NOT NULL
|
|
AND requested_at > TRY_TO_TIMESTAMP_NTZ(?)
|
|
ORDER BY requested_at ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
QueryRetryState {
|
|
resume_after: after_date.clone(),
|
|
remaining_limit: max_rows_per_timestamp,
|
|
offset,
|
|
},
|
|
|retry_state| {
|
|
json!({
|
|
"1": { "type": "TEXT", "value": retry_state.resume_after },
|
|
"2": { "type": "FIXED", "value": format_limit(retry_state.remaining_limit) },
|
|
"3": { "type": "FIXED", "value": retry_state.offset.to_string() }
|
|
})
|
|
},
|
|
&[
|
|
"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: Option<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
|
|
ep_request_id AS request_id,
|
|
rated_inputs AS inputs,
|
|
rated_output AS output,
|
|
settled_editable_region AS settled_editable_region,
|
|
rating AS rating,
|
|
feedback AS feedback,
|
|
device_id AS device_id,
|
|
requested_at::string AS continuation_time,
|
|
requested_at::string AS time,
|
|
NULL AS experiment_name,
|
|
NULL AS environment,
|
|
zed_version AS zed_version
|
|
FROM ZED_DBT.DBT_PROD.fct_edit_prediction_examples
|
|
WHERE rating IS NOT NULL
|
|
AND (? IS NULL OR rating = ?)
|
|
AND requested_at > TRY_TO_TIMESTAMP_NTZ(?)
|
|
AND rated_inputs IS NOT NULL
|
|
AND rated_inputs:cursor_excerpt IS NOT NULL
|
|
AND rated_output IS NOT NULL
|
|
ORDER BY requested_at ASC
|
|
LIMIT ?
|
|
OFFSET ?
|
|
"#};
|
|
|
|
let examples = fetch_examples_with_query(
|
|
http_client.clone(),
|
|
&step_progress,
|
|
background_executor.clone(),
|
|
statement,
|
|
QueryRetryState {
|
|
resume_after: after_date.clone(),
|
|
remaining_limit: max_rows_per_timestamp,
|
|
offset,
|
|
},
|
|
|retry_state| {
|
|
json!({
|
|
"1": { "type": "TEXT", "value": rating_value },
|
|
"2": { "type": "TEXT", "value": rating_value },
|
|
"3": { "type": "TEXT", "value": retry_state.resume_after },
|
|
"4": { "type": "FIXED", "value": format_limit(retry_state.remaining_limit) },
|
|
"5": { "type": "FIXED", "value": retry_state.offset.to_string() }
|
|
})
|
|
},
|
|
&[
|
|
"request_id",
|
|
"inputs",
|
|
"output",
|
|
"settled_editable_region",
|
|
"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 settled_editable_region = get_string("settled_editable_region");
|
|
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,
|
|
settled_editable_region,
|
|
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,
|
|
settled_editable_region: Option<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 expected_patch = settled_editable_region
|
|
.as_ref()
|
|
.map(|settled_editable_region| {
|
|
build_output_patch(
|
|
&input.cursor_path,
|
|
input.cursor_excerpt.as_ref(),
|
|
&input.excerpt_ranges.editable_350,
|
|
settled_editable_region,
|
|
)
|
|
});
|
|
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 let Some(expected_patch) = expected_patch {
|
|
example.spec.expected_patches = vec![expected_patch];
|
|
} else if is_positive {
|
|
example.spec.expected_patches = vec![output.clone()];
|
|
}
|
|
|
|
if !is_positive {
|
|
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 prompt = get_string("prompt");
|
|
let output = get_string("output");
|
|
let settled_editable_region = get_string("settled_editable_region");
|
|
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,
|
|
prompt,
|
|
output,
|
|
settled_editable_region,
|
|
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,
|
|
prompt: Option<String>,
|
|
output: String,
|
|
settled_editable_region: Option<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 expected_patch = settled_editable_region
|
|
.as_ref()
|
|
.map(|settled_editable_region| {
|
|
build_output_patch(
|
|
&input.cursor_path,
|
|
input.cursor_excerpt.as_ref(),
|
|
&input.excerpt_ranges.editable_350,
|
|
settled_editable_region,
|
|
)
|
|
});
|
|
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);
|
|
if let Some(expected_patch) = expected_patch {
|
|
example.spec.expected_patches = vec![expected_patch];
|
|
}
|
|
example.prompt = prompt.map(|prompt| ExamplePrompt {
|
|
input: prompt,
|
|
expected_output: None,
|
|
rejected_output: Some(output),
|
|
prefill: None,
|
|
provider: PredictionProvider::default(),
|
|
});
|
|
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
|
|
}
|
|
|
|
fn is_timeout_response(response: &SnowflakeStatementResponse) -> bool {
|
|
response.code.as_deref() == Some(SNOWFLAKE_TIMEOUT_CODE)
|
|
&& response
|
|
.message
|
|
.as_deref()
|
|
.map(|message| message.to_ascii_lowercase().contains("timeout"))
|
|
.unwrap_or(false)
|
|
}
|
|
|
|
fn is_snowflake_timeout_error(error: &anyhow::Error) -> bool {
|
|
error
|
|
.chain()
|
|
.any(|cause| cause.to_string().contains(SNOWFLAKE_TIMEOUT_CODE))
|
|
}
|
|
|
|
fn last_continuation_timestamp_from_response(
|
|
response: &SnowflakeStatementResponse,
|
|
column_indices: &HashMap<String, usize>,
|
|
) -> Option<String> {
|
|
let continuation_time_index = column_indices.get("continuation_time").copied()?;
|
|
response
|
|
.data
|
|
.iter()
|
|
.rev()
|
|
.find_map(|row| match row.get(continuation_time_index)? {
|
|
JsonValue::String(value) => Some(value.clone()),
|
|
JsonValue::Null => None,
|
|
other => Some(other.to_string()),
|
|
})
|
|
}
|
|
|
|
pub(crate) fn get_column_indices(
|
|
meta: &Option<SnowflakeResultSetMetaData>,
|
|
names: &[&str],
|
|
) -> HashMap<String, usize> {
|
|
let mut indices = 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
|
|
}
|