Skip to main content

adk_agent/
compaction.rs

1//! LLM-based event summarizer for context compaction.
2//!
3//! This module provides [`LlmEventSummarizer`], which uses an LLM to summarize
4//! a window of conversation events into a single compacted event. This is the
5//! Rust equivalent of ADK Python's `LlmEventSummarizer`.
6
7use adk_core::{
8    BaseEventsSummarizer, Content, Event, EventActions, EventCompaction, Llm, LlmRequest, Part,
9    Result,
10};
11use async_trait::async_trait;
12use std::sync::Arc;
13
14const DEFAULT_PROMPT_TEMPLATE: &str = "\
15The following is a conversation history between a user and an AI agent. \
16Please summarize the conversation, focusing on key information and decisions made, \
17as well as any unresolved questions or tasks. \
18The summary should be concise and capture the essence of the interaction.\n\n\
19{conversation_history}";
20
21/// An LLM-based event summarizer for sliding window compaction.
22///
23/// When called with a list of events, this formats the events, generates a
24/// summary using an LLM, and returns a new [`Event`] containing the summary
25/// within an [`EventCompaction`].
26pub struct LlmEventSummarizer {
27    llm: Arc<dyn Llm>,
28    prompt_template: String,
29}
30
31impl LlmEventSummarizer {
32    /// Create a new summarizer using the given LLM.
33    pub fn new(llm: Arc<dyn Llm>) -> Self {
34        Self { llm, prompt_template: DEFAULT_PROMPT_TEMPLATE.to_string() }
35    }
36
37    /// Create a new summarizer with a custom prompt template.
38    /// The template must contain `{conversation_history}` as a placeholder.
39    pub fn with_prompt_template(mut self, template: impl Into<String>) -> Self {
40        self.prompt_template = template.into();
41        self
42    }
43
44    fn format_events_for_prompt(events: &[Event]) -> String {
45        let mut lines = Vec::new();
46        for event in events {
47            if let Some(content) = &event.llm_response.content {
48                for part in &content.parts {
49                    if let Part::Text { text } = part {
50                        lines.push(format!("{}: {}", event.author, text));
51                    }
52                }
53            }
54        }
55        lines.join("\n")
56    }
57}
58
59#[async_trait]
60impl BaseEventsSummarizer for LlmEventSummarizer {
61    async fn summarize_events(&self, events: &[Event]) -> Result<Option<Event>> {
62        if events.is_empty() {
63            return Ok(None);
64        }
65
66        let conversation_history = Self::format_events_for_prompt(events);
67        let prompt = self.prompt_template.replace("{conversation_history}", &conversation_history);
68
69        let request = LlmRequest {
70            model: self.llm.name().to_string(),
71            contents: vec![Content {
72                role: "user".to_string(),
73                parts: vec![Part::Text { text: prompt }],
74            }],
75            tools: Default::default(),
76            config: None,
77            previous_response_id: None,
78        };
79
80        // Generate summary (non-streaming)
81        let mut response_stream = self.llm.generate_content(request, false).await?;
82
83        use futures::StreamExt;
84        let mut summary_content: Option<Content> = None;
85        while let Some(chunk_result) = response_stream.next().await {
86            if let Ok(chunk) = chunk_result
87                && chunk.content.is_some()
88            {
89                summary_content = chunk.content;
90                break;
91            }
92        }
93
94        let Some(mut summary) = summary_content else {
95            return Ok(None);
96        };
97
98        // Ensure the compacted content has the role 'model'
99        summary.role = "model".to_string();
100
101        let start_timestamp = events.first().map(|e| e.timestamp).unwrap_or_default();
102        let end_timestamp = events.last().map(|e| e.timestamp).unwrap_or_default();
103
104        let compaction =
105            EventCompaction { start_timestamp, end_timestamp, compacted_content: summary };
106
107        let actions = EventActions { compaction: Some(compaction), ..Default::default() };
108
109        let mut event = Event::new(Event::new("compaction").invocation_id);
110        event.author = "system".to_string();
111        event.actions = actions;
112
113        Ok(Some(event))
114    }
115}