Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions crates/edit_prediction_cli/src/anthropic_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use reqwest_client::ReqwestClient;
use sqlez::bindable::Bind;
use sqlez::bindable::StaticColumnCount;
use sqlez_macros::sql;
use std::collections::HashSet;
use std::hash::Hash;
use std::hash::Hasher;
use std::path::Path;
Expand Down Expand Up @@ -517,6 +518,7 @@ impl BatchingLlmClient {

let mut current_batch_rows = Vec::new();
let mut current_batch_size = 0usize;
let mut pending_hashes: HashSet<String> = HashSet::new();
loop {
let rows: Vec<(String, String)> = {
let connection = self.connection.lock().unwrap();
Expand All @@ -534,9 +536,15 @@ impl BatchingLlmClient {

// Split rows into sub-batches based on size
let mut batches_to_upload = Vec::new();
let mut new_rows_added = 0;

for row in rows {
let (hash, request_str) = row;

// Skip rows already added to current_batch_rows but not yet uploaded
if pending_hashes.contains(&hash) {
continue;
}
let serializable_request: SerializableRequest = serde_json::from_str(&request_str)?;

let messages: Vec<Message> = serializable_request
Expand Down Expand Up @@ -586,15 +594,28 @@ impl BatchingLlmClient {
current_batch_size = 0;
}

pending_hashes.insert(hash.clone());
current_batch_rows.push((hash, batch_request));
current_batch_size += estimated_size;
new_rows_added += 1;
}

// If no new rows were added this iteration, all pending requests are already
// in current_batch_rows, so we should break to avoid an infinite loop
if new_rows_added == 0 {
break;
}

// Only upload full batches, keep the partial batch for the next iteration
// Upload each sub-batch
for (batch_rows, batch_size) in batches_to_upload {
let request_hashes: Vec<String> =
batch_rows.iter().map(|(hash, _)| hash.clone()).collect();

// Remove uploaded hashes from pending set
for hash in &request_hashes {
pending_hashes.remove(hash);
}
let batch_requests: Vec<anthropic::batches::BatchRequest> =
batch_rows.into_iter().map(|(_, req)| req).collect();

Expand Down
2 changes: 2 additions & 0 deletions crates/edit_prediction_cli/src/prompts/repair.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ Your previous prediction has quality issues that need to be addressed. Please ge

{quality_feedback}

{token_change_info}

## Your Previous Prediction (word-diff format)

`````
Expand Down
19 changes: 18 additions & 1 deletion crates/edit_prediction_cli/src/repair.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use crate::{
anthropic_client::AnthropicClient,
example::{ActualCursor, Example, ExamplePrediction},
format_prompt::{TeacherPrompt, extract_last_codeblock},
metrics::count_patch_token_changes,
openai_client::OpenAiClient,
parse_output::run_parse_output,
paths::LLM_CACHE_DB,
Expand Down Expand Up @@ -168,10 +169,26 @@ pub fn build_repair_message(example: &Example) -> Result<String> {

let actual_patch_word_diff = unified_to_word_diff(actual_patch);

let token_counts = count_patch_token_changes(actual_patch);
let mut token_change_info = format!(
"\n## Token Change Statistics\n\n\
- **Deleted tokens**: {}\n\
- **Inserted tokens**: {}",
token_counts.deleted_tokens, token_counts.inserted_tokens,
);
if token_counts.deleted_tokens > 100 || token_counts.inserted_tokens > 100 {
token_change_info.push_str(
"\n\n> **Note:** The token change count is high. \
Consider producing a more scoped edit that targets only the lines \
that truly need to change, rather than rewriting large sections.",
);
}

let prompt_template = crate::prompt_assets::get_prompt("repair.md");
Ok(prompt_template
.replace("{actual_patch_word_diff}", &actual_patch_word_diff)
.replace("{quality_feedback}", &quality_feedback))
.replace("{quality_feedback}", &quality_feedback)
.replace("{token_change_info}", &token_change_info))
}

/// Check if an example needs repair based on QA feedback or computed scores.
Expand Down