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

af_agent_runtime/
compactor.rs
