use axum::{
body::Body,
response::{IntoResponse, Response},
};
use bytes::Bytes;
use futures::StreamExt;
use crate::{app_config::AppType, provider::Provider};
use super::{
error::ProxyError,
metrics::estimate_tokens_from_char_count,
response::{PreparedResponse, StreamCompletion},
server::ProxyServerState,
usage::{
log_buffered_response, log_error_request, log_stream_response, RequestLogContext,
StreamLogCollector,
},
};
#[cfg(test)]
mod tests;
pub struct ResponseHandler;
#[derive(Clone)]
pub struct SuccessSyncInfo {
pub app_type: AppType,
pub provider: Provider,
pub current_provider_id_at_start: String,
}
impl ResponseHandler {
pub async fn finish_buffered(
state: &ProxyServerState,
response_result: Result<PreparedResponse, ProxyError>,
status: reqwest::StatusCode,
success_sync: Option<SuccessSyncInfo>,
request_log: Option<RequestLogContext>,
) -> Response {
match response_result {
Ok(response) => {
let PreparedResponse {
response,
estimated_output_tokens,
upstream_error_summary,
body_bytes,
..
} = response;
if let (Some(request_log), Some(body_bytes)) =
(request_log.as_ref(), body_bytes.as_ref())
{
log_buffered_response(state, request_log, status.as_u16(), body_bytes).await;
}
state
.record_estimated_output_tokens(estimated_output_tokens)
.await;
if status.is_success() {
if let Some(success_sync) = success_sync {
state
.sync_successful_provider_selection(
&success_sync.app_type,
&success_sync.provider,
&success_sync.current_provider_id_at_start,
)
.await;
}
state.record_request_success().await;
} else {
state
.record_upstream_failure(status, upstream_error_summary)
.await;
}
response
}
Err(error) => {
if let Some(request_log) = request_log.as_ref() {
log_error_request(state, request_log, &error).await;
}
state.record_request_error(&error).await;
proxy_error_response(error)
}
}
}
pub async fn finish_streaming(
state: &ProxyServerState,
response_result: Result<PreparedResponse, ProxyError>,
status: reqwest::StatusCode,
success_sync: Option<SuccessSyncInfo>,
request_log: Option<RequestLogContext>,
) -> Response {
match response_result {
Ok(response) => {
track_streaming_response(state.clone(), response, status, success_sync, request_log)
}
Err(error) => {
if let Some(request_log) = request_log.as_ref() {
log_error_request(state, request_log, &error).await;
}
state.record_request_error(&error).await;
proxy_error_response(error)
}
}
}
}
fn track_streaming_response(
state: ProxyServerState,
response: PreparedResponse,
status: reqwest::StatusCode,
success_sync: Option<SuccessSyncInfo>,
request_log: Option<RequestLogContext>,
) -> Response {
let PreparedResponse {
response,
stream_completion,
upstream_error_summary,
body_bytes,
..
} = response;
let (parts, body) = response.into_parts();
let mut recorder = StreamingOutcomeRecorder::new(
state,
stream_completion,
status,
upstream_error_summary,
body_bytes,
success_sync,
request_log,
);
let tracked_stream = async_stream::stream! {
let mut stream = body.into_data_stream();
while let Some(next) = stream.next().await {
match next {
Ok(chunk) => {
recorder.record_chunk(&chunk);
yield Ok(chunk)
}
Err(error) => {
recorder.finish();
yield Err(std::io::Error::other(error));
return;
}
}
}
recorder.finish();
};
Response::from_parts(parts, Body::from_stream(tracked_stream))
}
struct StreamingOutcomeRecorder {
state: ProxyServerState,
stream_completion: Option<StreamCompletion>,
status: reqwest::StatusCode,
upstream_error_summary: Option<String>,
body_bytes: Option<Bytes>,
success_sync: Option<SuccessSyncInfo>,
request_log: Option<RequestLogContext>,
log_collector: Option<StreamLogCollector>,
output_char_count: u64,
finished: bool,
}
impl StreamingOutcomeRecorder {
fn new(
state: ProxyServerState,
stream_completion: Option<StreamCompletion>,
status: reqwest::StatusCode,
upstream_error_summary: Option<String>,
body_bytes: Option<Bytes>,
success_sync: Option<SuccessSyncInfo>,
request_log: Option<RequestLogContext>,
) -> Self {
let log_collector = request_log
.as_ref()
.map(|request_log| StreamLogCollector::new(request_log.started_at));
Self {
state,
stream_completion,
status,
upstream_error_summary,
body_bytes,
success_sync,
request_log,
log_collector,
output_char_count: 0,
finished: false,
}
}
fn record_chunk(&mut self, chunk: &bytes::Bytes) {
self.output_char_count = self
.output_char_count
.saturating_add(String::from_utf8_lossy(chunk).chars().count() as u64);
if let Some(log_collector) = self.log_collector.as_mut() {
log_collector.record_chunk(chunk);
}
}
fn finish(&mut self) {
if self.finished {
return;
}
self.finished = true;
let state = self.state.clone();
let estimated_output_tokens = estimate_tokens_from_char_count(self.output_char_count);
let request_log = self.request_log.clone();
let log_collector = self.log_collector.clone();
let body_bytes = self.body_bytes.clone();
if !self.status.is_success() {
let status = self.status;
let upstream_error_summary = self.upstream_error_summary.clone();
tokio::spawn(async move {
if let Some(request_log) = request_log.as_ref() {
if let Some(body_bytes) = body_bytes.as_ref() {
log_buffered_response(&state, request_log, status.as_u16(), body_bytes)
.await;
} else if let Some(log_collector) = log_collector.as_ref() {
log_stream_response(&state, request_log, status.as_u16(), log_collector)
.await;
}
}
state
.record_estimated_output_tokens(estimated_output_tokens)
.await;
state
.record_upstream_failure(status, upstream_error_summary)
.await;
});
return;
}
match self
.stream_completion
.as_ref()
.and_then(StreamCompletion::outcome)
{
Some(Err(message)) => {
tokio::spawn(async move {
if let Some(request_log) = request_log.as_ref() {
log_error_request(
&state,
request_log,
&ProxyError::RequestFailed(message.clone()),
)
.await;
}
state
.record_estimated_output_tokens(estimated_output_tokens)
.await;
state.record_request_error_message(message).await;
});
}
Some(Ok(())) => {
let success_sync = self.success_sync.clone();
tokio::spawn(async move {
if let Some(request_log) = request_log.as_ref() {
if let Some(body_bytes) = body_bytes.as_ref() {
log_buffered_response(
&state,
request_log,
reqwest::StatusCode::OK.as_u16(),
body_bytes,
)
.await;
} else if let Some(log_collector) = log_collector.as_ref() {
log_stream_response(
&state,
request_log,
reqwest::StatusCode::OK.as_u16(),
log_collector,
)
.await;
}
}
state
.record_estimated_output_tokens(estimated_output_tokens)
.await;
if let Some(success_sync) = success_sync {
state
.sync_successful_provider_selection(
&success_sync.app_type,
&success_sync.provider,
&success_sync.current_provider_id_at_start,
)
.await;
}
state.record_request_success().await;
});
}
None if body_bytes.is_some() => {
let success_sync = self.success_sync.clone();
let status = self.status;
tokio::spawn(async move {
if let Some(request_log) = request_log.as_ref() {
if let Some(body_bytes) = body_bytes.as_ref() {
log_buffered_response(&state, request_log, status.as_u16(), body_bytes)
.await;
}
}
state
.record_estimated_output_tokens(estimated_output_tokens)
.await;
if let Some(success_sync) = success_sync {
state
.sync_successful_provider_selection(
&success_sync.app_type,
&success_sync.provider,
&success_sync.current_provider_id_at_start,
)
.await;
}
state.record_request_success().await;
});
}
None => {
tokio::spawn(async move {
if let Some(request_log) = request_log.as_ref() {
log_error_request(
&state,
request_log,
&ProxyError::RequestFailed(
"stream terminated before completion".to_string(),
),
)
.await;
}
state
.record_estimated_output_tokens(estimated_output_tokens)
.await;
state
.record_request_error_message(
"stream terminated before completion".to_string(),
)
.await;
});
}
}
}
}
impl Drop for StreamingOutcomeRecorder {
fn drop(&mut self) {
self.finish();
}
}
pub fn proxy_error_response(error: ProxyError) -> Response {
error.into_response()
}