Skip to main content

gproxy_protocol/adapt/
compact.rs

1use crate::{
2    HttpBody, WireResponse,
3    adapt::{JsonInvocation, invoke_json},
4    capability::Upstream,
5    codec::{self, CodecLimits},
6    transform::{
7        TransformError, TransformErrorKind,
8        compact::{self, CompactDialectRequest, CompactDialectRequestBody},
9        memory,
10    },
11    wire::{
12        DeclaredFields,
13        openai::compact::{ClientCompactRequestBody, ClientCompactResponseBody},
14    },
15};
16
17#[derive(Debug, Clone, Copy)]
18pub struct CompactLimits {
19    pub max_bytes: u64,
20    pub codec: CodecLimits,
21}
22
23#[derive(Debug)]
24pub enum CompactFailure {
25    Transform(TransformError),
26    Rejected(Box<WireResponse<HttpBody>>),
27}
28
29#[derive(Debug)]
30pub struct CompactError {
31    pub attempted_calls: usize,
32    pub failure: CompactFailure,
33}
34
35impl CompactError {
36    pub fn kind(&self) -> TransformErrorKind {
37        match &self.failure {
38            CompactFailure::Transform(e) => e.kind(),
39            CompactFailure::Rejected(_) => TransformErrorKind::InvalidResult,
40        }
41    }
42}
43
44impl From<TransformError> for CompactError {
45    fn from(error: TransformError) -> Self {
46        Self {
47            attempted_calls: 0,
48            failure: CompactFailure::Transform(error),
49        }
50    }
51}
52
53fn after(error: TransformError) -> CompactError {
54    CompactError {
55        attempted_calls: 1,
56        failure: CompactFailure::Transform(error),
57    }
58}
59
60#[allow(clippy::too_many_arguments)]
61pub async fn compact<U: Upstream>(
62    upstream: &U,
63    target: &U::Target,
64    input: ClientCompactRequestBody,
65    dialect: CompactDialectRequest,
66    target_model: &str,
67    max_tokens: i64,
68    retain_from: usize,
69    limits: CompactLimits,
70) -> Result<ClientCompactResponseBody, CompactError> {
71    let input = input.into_declared();
72    let mut local = limits.codec;
73    local.max_body_bytes = local.max_body_bytes.min(limits.max_bytes);
74    local.max_value_bytes = local.max_value_bytes.min(limits.max_bytes);
75    local.max_buffer_bytes = local.max_buffer_bytes.min(limits.max_bytes);
76    let prepared =
77        compact::build_request(input, dialect, target_model, max_tokens, retain_from, local)?;
78    let Some(request) = prepared.request else {
79        return compact::replacement_history(None, prepared.context.retained_tail)
80            .map_err(CompactError::from);
81    };
82    let retained = codec::encode_json(&prepared.context.retained_tail, local)
83        .map_err(|e| encoding(e, false))?
84        .len() as u64;
85    let mut write = local;
86    write.max_body_bytes = write.max_body_bytes.min(upstream.limits().write_bytes);
87    write.max_value_bytes = write.max_value_bytes.min(write.max_body_bytes);
88    write.max_buffer_bytes = write.max_buffer_bytes.min(write.max_body_bytes);
89    let request_bytes = match &request {
90        CompactDialectRequestBody::Claude(v) => codec::encode_json(&v.body, write),
91        CompactDialectRequestBody::Gemini(v) => codec::encode_json(&v.body, write),
92        CompactDialectRequestBody::OpenAiChat(v) => codec::encode_json(&v.body, write),
93        CompactDialectRequestBody::OpenAiResponses(v) => codec::encode_json(&v.body, write),
94    }
95    .map_err(|e| encoding(e, false))?
96    .len() as u64;
97    if request_bytes
98        .checked_add(retained)
99        .is_none_or(|sum| sum > limits.max_bytes)
100    {
101        return Err(TransformError::new(
102            TransformErrorKind::Limit,
103            "compact.max_bytes",
104            "prepared request plus retained history exceeds bound",
105        )
106        .into());
107    }
108    let text = match request {
109        CompactDialectRequestBody::Claude(request) => extract(
110            invoke_json(upstream, target, request, local).await,
111            memory::claude_text,
112        )?,
113        CompactDialectRequestBody::Gemini(request) => extract(
114            invoke_json(upstream, target, request, local).await,
115            memory::gemini_text,
116        )?,
117        CompactDialectRequestBody::OpenAiChat(request) => extract(
118            invoke_json(upstream, target, request, local).await,
119            memory::chat_text,
120        )?,
121        CompactDialectRequestBody::OpenAiResponses(request) => extract(
122            invoke_json(upstream, target, request, local).await,
123            memory::responses_text,
124        )?,
125    };
126    let result =
127        compact::replacement_history(Some(text), prepared.context.retained_tail).map_err(after)?;
128    codec::encode_json(&result, local).map_err(|e| after(encoding(e, true)))?;
129    Ok(result)
130}
131
132fn extract<T>(
133    response: Result<JsonInvocation<T>, TransformError>,
134    text: fn(T) -> Result<String, TransformError>,
135) -> Result<String, CompactError> {
136    match response.map_err(after)? {
137        JsonInvocation::Success(response) => text(response.body).map_err(after),
138        JsonInvocation::Rejected(response) => Err(CompactError {
139            attempted_calls: 1,
140            failure: CompactFailure::Rejected(Box::new(response)),
141        }),
142    }
143}
144
145fn encoding(e: codec::CodecError, response: bool) -> TransformError {
146    TransformError::new(
147        if e.kind() == codec::CodecErrorKind::Limit {
148            TransformErrorKind::Limit
149        } else if response {
150            TransformErrorKind::InvalidResult
151        } else {
152            TransformErrorKind::InvalidInput
153        },
154        "compact.bytes",
155        e.to_string(),
156    )
157}