Skip to content

Commit

Permalink
Add more metrics for Fireworks Completion Requested (#23062)
Browse files Browse the repository at this point in the history
Release Notes:

- N/A

Co-authored-by: Thorsten <[email protected]>
  • Loading branch information
as-cii and mrnugget authored Jan 13, 2025
1 parent f2ab00c commit c26553d
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions crates/collab/src/llm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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",
Expand All @@ -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)
Expand All @@ -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,
}))
Expand Down

0 comments on commit c26553d

Please sign in to comment.