use std::time::Instant;
use crate::convert::{ResponseRequestContext, StreamState};
use crate::config::BackendInfo;
use crate::types::chat_api::ChatMessage;
use crate::types::response_api::ResponseRequest;
#[derive(Debug, Default)]
pub struct RouteInfo {
pub selected_backend: Option<BackendInfo>,
pub provider_name: Option<String>,
pub normalized_path: Option<String>,
pub rewritten_path: Option<String>,
}
#[derive(Debug, Default)]
pub struct ConversionFlags {
pub is_conversion_request: bool,
pub is_streaming: bool,
pub is_stream_response: bool,
pub should_convert_stream_response: bool,
}
#[derive(Debug, Default)]
pub struct ConversionBuffers {
pub request_body: Vec<u8>,
pub response_body: Vec<u8>,
pub stream_body_parsed_offset: usize,
}
#[derive(Debug, Default)]
pub struct UpstreamDiagnostics {
pub upstream_status: Option<u16>,
pub upstream_content_type: Option<String>,
pub stream_chunks_parsed: usize,
}
#[derive(Debug, Default)]
pub struct FollowUpContext {
pub pending_conversation_messages: Option<Vec<ChatMessage>>,
pub pending_instructions: Option<String>,
}
#[derive(Debug)]
pub struct ProxyContext {
pub start_time: Instant,
pub model: Option<String>,
pub stream_state: Option<StreamState>,
pub route: RouteInfo,
pub flags: ConversionFlags,
pub buffers: ConversionBuffers,
pub diagnostics: UpstreamDiagnostics,
pub follow_up: FollowUpContext,
}
impl ProxyContext {
pub fn new() -> Self {
Self {
start_time: Instant::now(),
model: None,
stream_state: None,
route: RouteInfo::default(),
flags: ConversionFlags::default(),
buffers: ConversionBuffers::default(),
diagnostics: UpstreamDiagnostics::default(),
follow_up: FollowUpContext::default(),
}
}
pub fn init_from_response_request(&mut self, req: &ResponseRequest) {
if self.model.is_none() {
self.model = Some(req.model.clone());
}
self.flags.is_streaming = req.stream;
if req.stream {
self.flags.is_stream_response = true;
}
if self.flags.is_conversion_request {
let model = self.model.clone().unwrap_or_else(|| "unknown".to_string());
let context = self
.stream_state
.as_ref()
.and_then(|s| s.request_context.clone());
self.stream_state = Some(StreamState::new(
format!("resp_{}", uuid::Uuid::new_v4()),
model,
context,
));
}
}
pub fn init_from_passthrough_json(&mut self, json: &serde_json::Value) {
if self.model.is_none()
&& let Some(model) = json.get("model").and_then(|v| v.as_str())
{
self.model = Some(model.to_string());
}
if let Some(stream) = json.get("stream").and_then(|v| v.as_bool()) {
self.flags.is_streaming = stream;
if stream {
self.flags.is_stream_response = true;
}
}
}
pub fn set_response_request_context(&mut self, context: ResponseRequestContext) {
if let Some(ref mut state) = self.stream_state {
state.request_context = Some(context);
} else {
let model = self.model.clone().unwrap_or_else(|| "unknown".to_string());
self.stream_state = Some(StreamState::new(
format!("resp_{}", uuid::Uuid::new_v4()),
model,
Some(context),
));
}
}
}
impl Default for ProxyContext {
fn default() -> Self {
Self::new()
}
}