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}