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        };
78
79        // Generate summary (non-streaming)
80        let mut response_stream = self.llm.generate_content(request, false).await?;
81
82        use futures::StreamExt;
83        let mut summary_content: Option<Content> = None;
84        while let Some(chunk_result) = response_stream.next().await {
85            if let Ok(chunk) = chunk_result {
86                if chunk.content.is_some() {
87                    summary_content = chunk.content;
88                    break;
89                }
90            }
91        }
92
93        let Some(mut summary) = summary_content else {
94            return Ok(None);
95        };
96
97        // Ensure the compacted content has the role 'model'
98        summary.role = "model".to_string();
99
100        let start_timestamp = events.first().map(|e| e.timestamp).unwrap_or_default();
101        let end_timestamp = events.last().map(|e| e.timestamp).unwrap_or_default();
102
103        let compaction =
104            EventCompaction { start_timestamp, end_timestamp, compacted_content: summary };
105
106        let actions = EventActions { compaction: Some(compaction), ..Default::default() };
107
108        let mut event = Event::new(Event::new("compaction").invocation_id);
109        event.author = "system".to_string();
110        event.actions = actions;
111
112        Ok(Some(event))
113    }
114}