Add more metrics for Fireworks Completion Requested (#23062)

Release Notes:

- N/A

Co-authored-by: Thorsten <thorsten@zed.dev>
This commit is contained in:
Antonio Scandurra 2025-01-13 13:04:28 +01:00
parent a917a894bf
commit 7fe8a4449e

View file

@ -470,6 +470,8 @@ async fn predict_edits(
.replace("<outline>", &outline_prefix)
.replace("<events>", &params.input_events)
.replace("<excerpt>", &params.input_excerpt);
let request_start = std::time::Instant::now();
let mut response = fireworks::complete(
&state.http_client,
api_url,
@ -486,13 +488,18 @@ async fn predict_edits(
},
)
.await?;
let duration = request_start.elapsed();
let choice = response
.completion
.choices
.pop()
.context("no output from completion response")?;
state.executor.spawn_detached({
let kinesis_client = state.kinesis_client.clone();
let kinesis_stream = state.config.kinesis_stream.clone();
let headers = response.headers.clone();
let model = model.clone();
async move {
SnowflakeRow::new(
"Fireworks Completion Requested",
@ -501,7 +508,9 @@ async fn predict_edits(
claims.system_id.clone(),
json!({
"model": model.to_string(),
"headers": headers,
"headers": response.headers,
"usage": response.completion.usage,
"duration": duration.as_secs_f64(),
}),
)
.write(&kinesis_client, &kinesis_stream)
@ -510,11 +519,6 @@ async fn predict_edits(
}
});
let choice = response
.completion
.choices
.pop()
.context("no output from completion response")?;
Ok(Json(PredictEditsResponse {
output_excerpt: choice.text,
}))