af_agent_runtime/
compactor.rs

1use std::sync::Arc;
2use std::time::Instant;
3
4use af_agent::ChatModel;
5use af_agent_session::Event;
6use af_llm::{ChatMessage, CompletionRequest, Role};
7use async_trait::async_trait;
8
9use crate::{
10    AgentRuntime, CancellationToken, CompactionResult, Compactor, EventWriter, RuntimeError,
11};
12
13const COMPACTION_PROMPT: &str = "Create a compact conversation checkpoint. Preserve: current objective; confirmed decisions and constraints; concrete facts with citations; completed tool outcomes; unresolved questions; pending actions; and the exact next step. Do not invent facts. Use short labeled sections.";
14const COMPACTION_MAX_TOKENS: u32 = 2_048;
15
16/// Default [`Compactor`] that asks the same model for a structured summary through a durable provider attempt.
17pub struct ModelCompactor {
18    model: Arc<dyn ChatModel>,
19}
20
21impl ModelCompactor {
22    /// Compactor over `model`.
23    pub fn new(model: Arc<dyn ChatModel>) -> Self {
24        Self { model }
25    }
26}
27
28#[async_trait]
29impl Compactor for ModelCompactor {
30    fn provider_name(&self) -> Option<&str> {
31        self.model.provider_name()
32    }
33    fn name(&self) -> &str {
34        "agentfactory.model-compactor"
35    }
36
37    async fn summarize(
38        &self,
39        request: &CompletionRequest,
40        cancellation: CancellationToken,
41        deadline: Instant,
42    ) -> Result<CompactionResult, RuntimeError> {
43        let mut request = request.clone();
44        request
45            .messages
46            .insert(0, ChatMessage::system(COMPACTION_PROMPT));
47        request = request.temperature(0.1).max_tokens(COMPACTION_MAX_TOKENS);
48        af_llm::images::select_images(&mut request.messages, af_llm::images::MAX_INPUT_IMAGES)
49            .map_err(|error| RuntimeError::InvalidInput(error.to_string()))?;
50        let (delta_tx, _delta_rx) = tokio::sync::mpsc::unbounded_channel();
51        let completion = self.model.complete_streaming(&request, delta_tx);
52        tokio::pin!(completion);
53        let timeout = tokio::time::sleep_until(deadline.into());
54        tokio::pin!(timeout);
55        let response = tokio::select! {
56            response = &mut completion => response.map_err(|error| RuntimeError::Model(error.to_string()))?,
57            _ = &mut timeout => return Err(RuntimeError::Model("compaction deadline exceeded".into())),
58            _ = cancellation.cancelled() => return Err(RuntimeError::Cancelled),
59        };
60        let summary = response
61            .first_content()
62            .filter(|value| !value.trim().is_empty())
63            .ok_or_else(|| RuntimeError::Model("compactor returned no summary".into()))?
64            .trim()
65            .to_string();
66        let (prompt_tokens, completion_tokens) = response.usage.map_or_else(
67            || {
68                (
69                    request
70                        .messages
71                        .iter()
72                        .filter_map(|message| message.content.as_deref())
73                        .map(|content| content.chars().count().div_ceil(4) as u64)
74                        .sum(),
75                    summary.chars().count().div_ceil(4).max(1) as u64,
76                )
77            },
78            |usage| {
79                (
80                    u64::from(usage.prompt_tokens),
81                    u64::from(usage.completion_tokens),
82                )
83            },
84        );
85        Ok(CompactionResult {
86            source: if response.usage.is_some() {
87                af_agent_session::MeteringSource::Reported
88            } else {
89                af_agent_session::MeteringSource::Estimated
90            },
91            provider: self.model.provider_name().map(str::to_owned),
92            summary,
93            prompt_tokens,
94            completion_tokens,
95        })
96    }
97}
98
99impl AgentRuntime {
100    pub(super) async fn compact_if_needed(
101        &self,
102        writer: &dyn EventWriter,
103        turn: &crate::TurnRequest,
104        step: u32,
105        transcript: &mut Vec<ChatMessage>,
106        context: &[ChatMessage],
107        cancellation: CancellationToken,
108    ) -> Result<(u64, u64), RuntimeError> {
109        let run_id = &turn.run_id;
110        let count = |transcript: &[ChatMessage]| {
111            self.meter.count(
112                &self.model_name,
113                &std::iter::once(ChatMessage::system(self.prompts.render()))
114                    .chain(transcript.iter().cloned())
115                    .chain(context.iter().cloned())
116                    .collect::<Vec<_>>(),
117            )
118        };
119        if !turn.compact && count(transcript) <= self.limits.max_tokens {
120            return Ok((0, 0));
121        }
122        let mut pruned = Vec::new();
123        for message in transcript
124            .iter_mut()
125            .filter(|message| message.role == Role::Tool)
126        {
127            let Some(content) = &message.content else {
128                continue;
129            };
130            if content.chars().count() <= 2_000 {
131                continue;
132            }
133            if let Some(call_id) = &message.tool_call_id {
134                pruned.push(call_id.clone());
135            }
136            message.content = Some(format!(
137                "{}… [tool result pruned]",
138                content.chars().take(512).collect::<String>()
139            ));
140        }
141        if !pruned.is_empty() {
142            writer
143                .append(vec![Event::ToolResultsPruned {
144                    run_id: run_id.clone(),
145                    call_ids: pruned,
146                }])
147                .await?;
148            if !turn.compact && count(transcript) <= self.limits.max_tokens {
149                return Ok((0, 0));
150            }
151        }
152        let compactor = &self.compactor;
153        let compaction_id = format!("compaction:{step}");
154        let operation_id = format!("{compaction_id}:attempt:1");
155        let provider_attempt_id = format!("{run_id}:{operation_id}");
156        let source_through_seq = writer.last_seq().await?;
157        let mut compactable = transcript
158            .iter()
159            .cloned()
160            .chain(context.iter().cloned())
161            .collect::<Vec<_>>();
162        af_llm::images::select_images(&mut compactable, self.limits.max_images)
163            .map_err(|error| RuntimeError::InvalidInput(error.to_string()))?;
164        let mut compaction_request = CompletionRequest::new(&self.model_name, compactable.clone());
165        compaction_request.context = Some(turn.context.clone());
166        compaction_request.provider_attempt_id = Some(provider_attempt_id.clone());
167        let prepared = writer
168            .append_at(
169                source_through_seq,
170                vec![
171                    Event::CompactionStarted {
172                        run_id: run_id.clone(),
173                        compaction_id: compaction_id.clone(),
174                        source_through_seq,
175                    },
176                    Event::ModelRequestPrepared {
177                        metering: Some(af_agent_session::MeteringDetails::Model {
178                            model: self.model_name.clone(),
179                            provider: compactor.provider_name().map(str::to_owned),
180                            provider_attempt_id: provider_attempt_id.parse().map_err(
181                                |error: af_context::ids::EmptyId| {
182                                    RuntimeError::Invariant(error.to_string())
183                                },
184                            )?,
185                            source: af_agent_session::MeteringSource::Estimated,
186                            outcome: af_agent_session::MeteringOutcome::Unknown,
187                        }),
188                        run_id: run_id.clone(),
189                        step,
190                        attempt: 1,
191                        provider_attempt_id: provider_attempt_id.clone(),
192                        operation_id: operation_id.clone(),
193                        reserved_prompt_tokens: self.meter.count(
194                            &self.model_name,
195                            &std::iter::once(ChatMessage::system(COMPACTION_PROMPT))
196                                .chain(compactable.iter().cloned())
197                                .collect::<Vec<_>>(),
198                        ),
199                        reserved_completion_tokens: u64::from(COMPACTION_MAX_TOKENS),
200                        request: serde_json::to_value(&compactable)
201                            .map_err(|error| RuntimeError::Invariant(error.to_string()))?,
202                        prompt_sections: serde_json::Value::Null,
203                    },
204                ],
205            )
206            .await?;
207        let prepared_through_seq = prepared
208            .last()
209            .map_or(source_through_seq, |event| event.seq);
210        let extension_cancellation = cancellation.child();
211        let deadline = std::time::Instant::now() + self.limits.provider_deadline;
212        let result = match crate::extension::await_extension(
213            "context compactor",
214            extension_cancellation.clone(),
215            deadline,
216            compactor.summarize(&compaction_request, extension_cancellation, deadline),
217        )
218        .await
219        {
220            Ok(Ok(result)) => result,
221            Ok(Err(error)) | Err(error) => {
222                writer
223                    .append(vec![Event::CompactionFinished {
224                        run_id: run_id.clone(),
225                        compaction_id,
226                        status: "failed".into(),
227                        error: Some(error.to_string()),
228                    }])
229                    .await?;
230                return Err(error);
231            }
232        };
233        let attribution = Some(af_agent_session::MeteringDetails::Model {
234            model: self.model_name.clone(),
235            provider: result.provider.clone(),
236            provider_attempt_id: provider_attempt_id.parse().map_err(
237                |error: af_context::ids::EmptyId| RuntimeError::Invariant(error.to_string()),
238            )?,
239            source: result.source,
240            outcome: af_agent_session::MeteringOutcome::Completed,
241        });
242        let events = vec![
243            Event::UsageRecorded {
244                metering: attribution.clone(),
245                run_id: run_id.clone(),
246                operation_id: operation_id.clone(),
247                prompt_tokens: result.prompt_tokens,
248                completion_tokens: result.completion_tokens,
249                cost_units: 0,
250            },
251            Event::SummaryReplaced {
252                run_id: run_id.clone(),
253                through_seq: source_through_seq,
254                summary: result.summary.clone(),
255                compactor: compactor.name().into(),
256                model: self.model_name.clone(),
257            },
258            Event::CompactionFinished {
259                run_id: run_id.clone(),
260                compaction_id: compaction_id.clone(),
261                status: "completed".into(),
262                error: None,
263            },
264        ];
265        match writer.append_at(prepared_through_seq, events).await {
266            Ok(_) => crate::replay::replace_with_summary(transcript, &result.summary),
267            Err(RuntimeError::CompactionConflict) => {
268                writer
269                    .append(vec![
270                        Event::UsageRecorded {
271                            metering: attribution.clone(),
272                            run_id: run_id.clone(),
273                            operation_id,
274                            prompt_tokens: result.prompt_tokens,
275                            completion_tokens: result.completion_tokens,
276                            cost_units: 0,
277                        },
278                        Event::CompactionFinished {
279                            run_id: run_id.clone(),
280                            compaction_id,
281                            status: "failed".into(),
282                            error: Some("surface_changed".into()),
283                        },
284                    ])
285                    .await?;
286                return Err(RuntimeError::CompactionConflict);
287            }
288            Err(error) => return Err(error),
289        }
290        Ok((result.prompt_tokens, result.completion_tokens))
291    }
292}
af_agent_runtime/compactor.rs

af_agent_runtime/
compactor.rs
