zed/crates/edit_prediction_cli/src/predict.rs

use crate::{
    FormatPromptArgs, PredictArgs, PredictionProvider, TeacherBackend,
    anthropic_client::AnthropicClient,
    example::{Example, ExamplePrediction, ExamplePrompt},
    format_prompt::{TeacherMultiRegionPrompt, TeacherPrompt, run_format_prompt},
    headless::EpAppState,
    load_project::run_load_project,
    openai_client::OpenAiClient,
    parse_output::parse_prediction_output,
    paths::{LATEST_EXAMPLE_RUN_DIR, RUN_DIR},
    progress::{ExampleProgress, InfoStyle, Progress, Step, StepProgress},
    retrieve_context::run_context_retrieval,
};
use anyhow::Context as _;
use cloud_llm_client::predict_edits_v3::{RawCompletionRequest, RawCompletionResponse};
use edit_prediction::{DebugEvent, EditPredictionStore, Zeta2RawConfig};
use futures::{AsyncReadExt as _, FutureExt as _, StreamExt as _, future::Shared};
use gpui::{AppContext as _, AsyncApp, Task};
use http_client::{AsyncBody, HttpClient, Method};
use reqwest_client::ReqwestClient;
use std::{
    fs,
    sync::{
        Arc, Mutex, OnceLock,
        atomic::{AtomicUsize, Ordering::SeqCst},
    },
};
use zeta_prompt::ZetaFormat;

static ANTHROPIC_CLIENT: OnceLock<AnthropicClient> = OnceLock::new();
static OPENAI_CLIENT: OnceLock<OpenAiClient> = OnceLock::new();

pub async fn run_prediction(
    example: &mut Example,
    args: &PredictArgs,
    app_state: Arc<EpAppState>,
    example_progress: &ExampleProgress,
    mut cx: AsyncApp,
) -> anyhow::Result<()> {
    let repetition_count = args.repetitions;

    if let Some(existing_prediction) = example.predictions.first() {
        let has_prediction = existing_prediction.actual_patch.is_some()
            || !existing_prediction.actual_output.is_empty();
        if has_prediction {
            match args.provider {
                None => return Ok(()),
                Some(provider) if existing_prediction.provider == provider => return Ok(()),
                Some(_) => example.predictions.clear(),
            }
        }
    }

    let Some(provider) = args.provider else {
        anyhow::bail!(
            "No existing predictions found. Use --provider to specify which model to use for prediction."
        );
    };

    if matches!(
        provider,
        PredictionProvider::TeacherMultiRegion(..)
            | PredictionProvider::TeacherMultiRegionNonBatching(..)
    ) {
        anyhow::bail!("Teacher multi-region providers are not supported for prediction.");
    }

    if let PredictionProvider::Teacher(backend, _)
    | PredictionProvider::TeacherNonBatching(backend, _) = provider
    {
        run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;
        run_format_prompt(
            example,
            &FormatPromptArgs { provider },
            app_state.clone(),
            example_progress,
            cx,
        )
        .await?;

        let step_progress = example_progress.start(Step::Predict);
        let batched = matches!(
            provider,
            PredictionProvider::Teacher(..) | PredictionProvider::TeacherMultiRegion(..)
        );
        return predict_teacher(
            example,
            backend,
            batched,
            repetition_count,
            args.cache_only,
            &step_progress,
        )
        .await;
    }

    if let PredictionProvider::Baseten(format) = provider {
        run_format_prompt(
            example,
            &FormatPromptArgs {
                provider: PredictionProvider::Zeta2(format),
            },
            app_state.clone(),
            example_progress,
            cx,
        )
        .await?;

        let step_progress = example_progress.start(Step::Predict);
        return predict_baseten(example, format, &step_progress).await;
    }

    run_load_project(example, app_state.clone(), example_progress, cx.clone()).await?;
    run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;

    let step_progress = example_progress.start(Step::Predict);

    if matches!(
        provider,
        PredictionProvider::Zeta1 | PredictionProvider::Zeta2(_)
    ) {
        step_progress.set_substatus("authenticating");
        static AUTHENTICATED: OnceLock<Shared<Task<()>>> = OnceLock::new();
        AUTHENTICATED
            .get_or_init(|| {
                let client = app_state.client.clone();
                cx.spawn(async move |cx| {
                    if let Err(e) = client.sign_in_with_optional_connect(true, cx).await {
                        eprintln!("Authentication failed: {}", e);
                    }
                })
                .shared()
            })
            .clone()
            .await;
    }

    let ep_store = cx
        .update(|cx| EditPredictionStore::try_global(cx))
        .context("EditPredictionStore not initialized")?;

    ep_store.update(&mut cx, |store, _cx| {
        let model = match provider {
            PredictionProvider::Zeta1 => edit_prediction::EditPredictionModel::Zeta,
            PredictionProvider::Zeta2(_) => edit_prediction::EditPredictionModel::Zeta,
            PredictionProvider::Mercury => edit_prediction::EditPredictionModel::Mercury,
            PredictionProvider::Teacher(..)
            | PredictionProvider::TeacherMultiRegion(..)
            | PredictionProvider::TeacherNonBatching(..)
            | PredictionProvider::TeacherMultiRegionNonBatching(..)
            | PredictionProvider::Repair
            | PredictionProvider::Baseten(_) => {
                unreachable!()
            }
        };
        store.set_edit_prediction_model(model);

        // If user specified a non-default Zeta2 version, configure raw endpoint.
        // ZED_ZETA_MODEL env var is optional.
        if let PredictionProvider::Zeta2(format) = provider {
            if format != ZetaFormat::default() {
                let model_id = std::env::var("ZED_ZETA_MODEL").ok();
                let environment = std::env::var("ZED_ZETA_ENVIRONMENT").ok();
                store.set_zeta2_raw_config(Zeta2RawConfig {
                    model_id,
                    environment,
                    format,
                });
            }
        }
    });
    step_progress.set_substatus("configuring model");
    let state = example.state.as_ref().context("state must be set")?;
    let run_dir = RUN_DIR.join(&example.spec.name);

    let updated_example = Arc::new(Mutex::new(example.clone()));
    let current_run_ix = Arc::new(AtomicUsize::new(0));

    let mut debug_rx = ep_store.update(&mut cx, |store, cx| store.debug_info(&state.project, cx));
    let debug_task = cx.background_spawn({
        let updated_example = updated_example.clone();
        let current_run_ix = current_run_ix.clone();
        let run_dir = run_dir.clone();
        async move {
            while let Some(event) = debug_rx.next().await {
                let run_ix = current_run_ix.load(SeqCst);
                let mut updated_example = updated_example.lock().unwrap();

                let run_dir = if repetition_count > 1 {
                    run_dir.join(format!("{:03}", run_ix))
                } else {
                    run_dir.clone()
                };

                match event {
                    DebugEvent::EditPredictionStarted(request) => {
                        assert_eq!(updated_example.predictions.len(), run_ix + 1);

                        if let Some(prompt) = request.prompt {
                            fs::write(run_dir.join("prediction_prompt.md"), &prompt)?;
                            if matches!(provider, PredictionProvider::Zeta2(_)) {
                                updated_example.prompt.get_or_insert(ExamplePrompt {
                                    input: prompt,
                                    expected_output: None,
                                    rejected_output: None,
                                    provider,
                                    prefill: None,
                                });
                            }
                        }
                    }
                    DebugEvent::EditPredictionFinished(request) => {
                        assert_eq!(updated_example.predictions.len(), run_ix + 1);

                        if let Some(output) = request.model_output {
                            fs::write(run_dir.join("prediction_response.md"), &output)?;
                            updated_example
                                .predictions
                                .last_mut()
                                .unwrap()
                                .actual_output = output;
                        }
                        if run_ix >= repetition_count {
                            break;
                        }
                    }
                    _ => {}
                }
            }
            anyhow::Ok(())
        }
    });

    for ix in 0..repetition_count {
        current_run_ix.store(ix, SeqCst);
        let run_dir = if repetition_count > 1 {
            run_dir.join(format!("{:03}", ix))
        } else {
            run_dir.clone()
        };

        if repetition_count > 1 {
            step_progress.set_substatus(format!(
                "running prediction {}/{}",
                ix + 1,
                repetition_count
            ));
        } else {
            step_progress.set_substatus("running prediction");
        }

        fs::create_dir_all(&run_dir)?;
        if LATEST_EXAMPLE_RUN_DIR.is_symlink() {
            fs::remove_file(&*LATEST_EXAMPLE_RUN_DIR)?;
        }
        #[cfg(unix)]
        std::os::unix::fs::symlink(&run_dir, &*LATEST_EXAMPLE_RUN_DIR)?;
        #[cfg(windows)]
        std::os::windows::fs::symlink_dir(&run_dir, &*LATEST_EXAMPLE_RUN_DIR)?;

        updated_example
            .lock()
            .unwrap()
            .predictions
            .push(ExamplePrediction {
                actual_patch: None,
                actual_output: String::new(),
                actual_cursor: None,
                error: None,
                provider,
                cumulative_logprob: None,
                avg_logprob: None,
            });

        step_progress.set_substatus("requesting prediction");
        let prediction = ep_store
            .update(&mut cx, |store, cx| {
                store.request_prediction(
                    &state.project,
                    &state.buffer,
                    state.cursor_position,
                    cloud_llm_client::PredictEditsRequestTrigger::Cli,
                    cx,
                )
            })
            .await?;

        let actual_patch = prediction.and_then(|prediction| {
            let prediction = prediction.prediction.ok()?;
            prediction
                .edit_preview
                .as_unified_diff(prediction.snapshot.file(), &prediction.edits)
        });

        let has_prediction = actual_patch.as_ref().is_some_and(|p| !p.is_empty());

        updated_example
            .lock()
            .unwrap()
            .predictions
            .last_mut()
            .unwrap()
            .actual_patch = actual_patch;

        if ix == repetition_count - 1 {
            let (info, style) = if has_prediction {
                ("predicted", InfoStyle::Normal)
            } else {
                ("no prediction", InfoStyle::Warning)
            };
            step_progress.set_info(info, style);
        }
    }

    ep_store.update(&mut cx, |store, _| {
        store.remove_project(&state.project);
    });
    debug_task.await?;

    *example = Arc::into_inner(updated_example)
        .ok_or_else(|| anyhow::anyhow!("Failed to unwrap Arc"))?
        .into_inner()
        .map_err(|_| anyhow::anyhow!("Failed to unwrap Mutex"))?;
    Ok(())
}

async fn predict_teacher(
    example: &mut Example,
    backend: TeacherBackend,
    batched: bool,
    repetition_count: usize,
    cache_only: bool,
    step_progress: &crate::progress::StepProgress,
) -> anyhow::Result<()> {
    match backend {
        TeacherBackend::Sonnet45 | TeacherBackend::Sonnet46 => {
            predict_anthropic(
                example,
                backend,
                batched,
                repetition_count,
                cache_only,
                step_progress,
            )
            .await
        }
        TeacherBackend::Gpt52 => {
            predict_openai(
                example,
                backend,
                batched,
                repetition_count,
                cache_only,
                step_progress,
            )
            .await
        }
    }
}

async fn predict_anthropic(
    example: &mut Example,
    backend: TeacherBackend,
    batched: bool,
    repetition_count: usize,
    cache_only: bool,
    step_progress: &crate::progress::StepProgress,
) -> anyhow::Result<()> {
    let llm_model_name = backend.model_name();
    let max_tokens = 16384;
    let llm_client = ANTHROPIC_CLIENT.get_or_init(|| {
        let client = if batched {
            AnthropicClient::batch(&crate::paths::LLM_CACHE_DB)
        } else {
            AnthropicClient::plain()
        };
        client.expect("Failed to create Anthropic client")
    });

    let prompt = example.prompt.as_ref().context("Prompt is required")?;

    for ix in 0..repetition_count {
        if repetition_count > 1 {
            step_progress.set_substatus(format!(
                "running prediction {}/{}",
                ix + 1,
                repetition_count
            ));
        } else {
            step_progress.set_substatus("running prediction");
        }

        let messages = vec![anthropic::Message {
            role: anthropic::Role::User,
            content: vec![anthropic::RequestContent::Text {
                text: prompt.input.clone(),
                cache_control: None,
            }],
        }];

        let seed = if repetition_count > 1 { Some(ix) } else { None };
        let Some(response) = llm_client
            .generate(llm_model_name, max_tokens, messages, seed, cache_only)
            .await?
        else {
            // Request stashed for batched processing
            continue;
        };

        let actual_output = response
            .content
            .into_iter()
            .filter_map(|content| match content {
                anthropic::ResponseContent::Text { text } => Some(text),
                _ => None,
            })
            .collect::<Vec<String>>()
            .join("\n");

        let parser_provider = if batched {
            example
                .prompt
                .as_ref()
                .map(|prompt| prompt.provider)
                .unwrap_or(PredictionProvider::Teacher(backend, ZetaFormat::default()))
        } else {
            match example.prompt.as_ref().map(|prompt| prompt.provider) {
                Some(PredictionProvider::TeacherMultiRegion(_))
                | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
                    PredictionProvider::TeacherMultiRegionNonBatching(backend)
                }
                _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
            }
        };

        let (actual_patch, actual_cursor) = match parser_provider {
            PredictionProvider::TeacherMultiRegion(_)
            | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
                TeacherMultiRegionPrompt::parse(example, &actual_output)?
            }
            _ => TeacherPrompt::parse(example, &actual_output)?,
        };

        let prediction = ExamplePrediction {
            actual_patch: Some(actual_patch),
            actual_output,
            actual_cursor,
            error: None,
            provider: if batched {
                match example.prompt.as_ref().map(|prompt| prompt.provider) {
                    Some(PredictionProvider::TeacherMultiRegion(_)) => {
                        PredictionProvider::TeacherMultiRegion(backend)
                    }
                    _ => PredictionProvider::Teacher(backend, ZetaFormat::default()),
                }
            } else {
                match example.prompt.as_ref().map(|prompt| prompt.provider) {
                    Some(PredictionProvider::TeacherMultiRegion(_))
                    | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
                        PredictionProvider::TeacherMultiRegionNonBatching(backend)
                    }
                    _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
                }
            },
            cumulative_logprob: None,
            avg_logprob: None,
        };

        example.predictions.push(prediction);
    }
    Ok(())
}

async fn predict_openai(
    example: &mut Example,
    backend: TeacherBackend,
    batched: bool,
    repetition_count: usize,
    cache_only: bool,
    step_progress: &crate::progress::StepProgress,
) -> anyhow::Result<()> {
    let llm_model_name = backend.model_name();
    let max_tokens = 16384;
    let llm_client = OPENAI_CLIENT.get_or_init(|| {
        let client = if batched {
            OpenAiClient::batch(&crate::paths::LLM_CACHE_DB)
        } else {
            OpenAiClient::plain()
        };
        client.expect("Failed to create OpenAI client")
    });

    let prompt = example.prompt.as_ref().context("Prompt is required")?;

    for ix in 0..repetition_count {
        if repetition_count > 1 {
            step_progress.set_substatus(format!(
                "running prediction {}/{}",
                ix + 1,
                repetition_count
            ));
        } else {
            step_progress.set_substatus("running prediction");
        }

        let messages = vec![open_ai::RequestMessage::User {
            content: open_ai::MessageContent::Plain(prompt.input.clone()),
        }];

        let seed = if repetition_count > 1 { Some(ix) } else { None };
        let Some(response) = llm_client
            .generate(llm_model_name, max_tokens, messages, seed, cache_only)
            .await?
        else {
            // Request stashed for batched processing
            continue;
        };

        let actual_output = response
            .choices
            .into_iter()
            .filter_map(|choice| match choice.message {
                open_ai::RequestMessage::Assistant { content, .. } => content.map(|c| match c {
                    open_ai::MessageContent::Plain(text) => text,
                    open_ai::MessageContent::Multipart(parts) => parts
                        .into_iter()
                        .filter_map(|p| match p {
                            open_ai::MessagePart::Text { text } => Some(text),
                            _ => None,
                        })
                        .collect::<Vec<_>>()
                        .join(""),
                }),
                _ => None,
            })
            .collect::<Vec<String>>()
            .join("\n");

        let parser_provider = if batched {
            example
                .prompt
                .as_ref()
                .map(|prompt| prompt.provider)
                .unwrap_or(PredictionProvider::Teacher(backend, ZetaFormat::default()))
        } else {
            match example.prompt.as_ref().map(|prompt| prompt.provider) {
                Some(PredictionProvider::TeacherMultiRegion(_))
                | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
                    PredictionProvider::TeacherMultiRegionNonBatching(backend)
                }
                _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
            }
        };

        let (actual_patch, actual_cursor) = match parser_provider {
            PredictionProvider::TeacherMultiRegion(_)
            | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
                TeacherMultiRegionPrompt::parse(example, &actual_output)?
            }
            _ => TeacherPrompt::parse(example, &actual_output)?,
        };

        let prediction = ExamplePrediction {
            actual_patch: Some(actual_patch),
            actual_output,
            actual_cursor,
            error: None,
            provider: if batched {
                match example.prompt.as_ref().map(|prompt| prompt.provider) {
                    Some(PredictionProvider::TeacherMultiRegion(_)) => {
                        PredictionProvider::TeacherMultiRegion(backend)
                    }
                    _ => PredictionProvider::Teacher(backend, ZetaFormat::default()),
                }
            } else {
                match example.prompt.as_ref().map(|prompt| prompt.provider) {
                    Some(PredictionProvider::TeacherMultiRegion(_))
                    | Some(PredictionProvider::TeacherMultiRegionNonBatching(_)) => {
                        PredictionProvider::TeacherMultiRegionNonBatching(backend)
                    }
                    _ => PredictionProvider::TeacherNonBatching(backend, ZetaFormat::default()),
                }
            },
            cumulative_logprob: None,
            avg_logprob: None,
        };

        example.predictions.push(prediction);
    }
    Ok(())
}

pub async fn predict_baseten(
    example: &mut Example,
    format: ZetaFormat,
    step_progress: &StepProgress,
) -> anyhow::Result<()> {
    let model_id =
        std::env::var("ZED_ZETA_MODEL").context("ZED_ZETA_MODEL environment variable required")?;

    let api_key =
        std::env::var("BASETEN_API_KEY").context("BASETEN_API_KEY environment variable not set")?;

    let prompt = example.prompt.as_ref().context("Prompt is required")?;
    let prompt_text = prompt.input.clone();
    let prefill = prompt.prefill.clone().unwrap_or_default();

    step_progress.set_substatus("running prediction via baseten");

    let environment: String = <&'static str>::from(&format).to_lowercase();
    let url = format!(
        "https://model-{model_id}.api.baseten.co/environments/{environment}/sync/v1/completions"
    );

    let request_body = RawCompletionRequest {
        model: model_id,
        prompt: prompt_text.clone(),
        max_tokens: Some(2048),
        temperature: Some(0.),
        stop: vec![],
        environment: None,
    };

    let body_bytes =
        serde_json::to_vec(&request_body).context("Failed to serialize request body")?;

    let http_client: Arc<dyn HttpClient> = Arc::new(ReqwestClient::new());
    let request = http_client::Request::builder()
        .method(Method::POST)
        .uri(&url)
        .header("Content-Type", "application/json")
        .header("Authorization", format!("Api-Key {api_key}"))
        .body(AsyncBody::from(body_bytes))?;

    let mut response = http_client.send(request).await?;
    let status = response.status();

    let mut body = String::new();
    response
        .body_mut()
        .read_to_string(&mut body)
        .await
        .context("Failed to read Baseten response body")?;

    if !status.is_success() {
        anyhow::bail!("Baseten API returned {status}: {body}");
    }

    let completion: RawCompletionResponse =
        serde_json::from_str(&body).context("Failed to parse Baseten response")?;

    let actual_output = completion
        .choices
        .into_iter()
        .next()
        .map(|choice| choice.text)
        .unwrap_or_default();

    let actual_output = format!("{prefill}{actual_output}");

    let (actual_patch, actual_cursor) =
        parse_prediction_output(example, &actual_output, PredictionProvider::Zeta2(format))?;

    let prediction = ExamplePrediction {
        actual_patch: Some(actual_patch),
        actual_output,
        actual_cursor,
        error: None,
        provider: PredictionProvider::Baseten(format),
        cumulative_logprob: None,
        avg_logprob: None,
    };

    example.predictions.push(prediction);
    Ok(())
}

pub async fn sync_batches(provider: Option<&PredictionProvider>) -> anyhow::Result<()> {
    match provider {
        Some(PredictionProvider::Teacher(backend, _))
        | Some(PredictionProvider::TeacherMultiRegion(backend)) => match backend {
            TeacherBackend::Sonnet45 | TeacherBackend::Sonnet46 => {
                let llm_client = ANTHROPIC_CLIENT.get_or_init(|| {
                    AnthropicClient::batch(&crate::paths::LLM_CACHE_DB)
                        .expect("Failed to create Anthropic client")
                });
                llm_client
                    .sync_batches()
                    .await
                    .context("Failed to sync Anthropic batches")?;
            }
            TeacherBackend::Gpt52 => {
                let llm_client = OPENAI_CLIENT.get_or_init(|| {
                    OpenAiClient::batch(&crate::paths::LLM_CACHE_DB)
                        .expect("Failed to create OpenAI client")
                });
                llm_client
                    .sync_batches()
                    .await
                    .context("Failed to sync OpenAI batches")?;
            }
        },
        _ => (),
    };
    Ok(())
}

pub async fn reprocess_after_batch_wait(
    examples: &mut [Example],
    args: &PredictArgs,
) -> anyhow::Result<()> {
    let Some(PredictionProvider::Teacher(backend, _)) = args.provider else {
        return Ok(());
    };

    let mut reprocessed = 0;
    for example in examples.iter_mut() {
        let has_prediction = example
            .predictions
            .iter()
            .any(|p| p.actual_patch.is_some() || !p.actual_output.is_empty());
        if has_prediction || example.prompt.is_none() {
            continue;
        }

        let example_progress = Progress::global().start_group(&example.spec.name);
        let step_progress = example_progress.start(Step::Predict);
        predict_teacher(
            example,
            backend,
            true,
            args.repetitions,
            false,
            &step_progress,
        )
        .await?;
        reprocessed += 1;
    }

    if reprocessed > 0 {
        eprintln!("Reprocessed {} example(s) with batch results", reprocessed);
    }

    Ok(())
}

pub async fn wait_for_batches(provider: Option<&PredictionProvider>) -> anyhow::Result<()> {
    let poll_interval = std::time::Duration::from_secs(30);

    loop {
        let pending = pending_batch_count(provider)?;
        if pending == 0 {
            break;
        }

        eprintln!(
            "Waiting for {} pending batch request(s) to complete... (polling every {}s)",
            pending,
            poll_interval.as_secs()
        );
        std::thread::sleep(poll_interval);

        sync_batches(provider).await?;
    }

    Ok(())
}

fn pending_batch_count(provider: Option<&PredictionProvider>) -> anyhow::Result<usize> {
    match provider {
        Some(PredictionProvider::Teacher(backend, _)) => match backend {
            TeacherBackend::Sonnet45 | TeacherBackend::Sonnet46 => {
                let llm_client = ANTHROPIC_CLIENT.get_or_init(|| {
                    AnthropicClient::batch(&crate::paths::LLM_CACHE_DB)
                        .expect("Failed to create Anthropic client")
                });
                llm_client.pending_batch_count()
            }
            TeacherBackend::Gpt52 => {
                let llm_client = OPENAI_CLIENT.get_or_init(|| {
                    OpenAiClient::batch(&crate::paths::LLM_CACHE_DB)
                        .expect("Failed to create OpenAI client")
                });
                llm_client.pending_batch_count()
            }
        },
        _ => Ok(0),
    }
}