mojentic/examples/react/
summarization_agent.rs

1//! Summarization agent for the ReAct pattern.
2//!
3//! This agent generates the final answer based on accumulated context.
4
5use crate::agents::BaseAsyncAgent;
6use crate::event::Event;
7use crate::llm::{LlmBroker, LlmMessage, MessageRole};
8use crate::Result;
9use async_trait::async_trait;
10use std::sync::Arc;
11
12use super::events::{FailureOccurred, FinishAndSummarize};
13use super::formatters::format_current_context;
14
15/// Agent responsible for generating the final answer.
16///
17/// This agent reviews the context, plan, and history to synthesize
18/// a complete answer to the user's original query.
19pub struct SummarizationAgent {
20    llm: Arc<LlmBroker>,
21}
22
23impl SummarizationAgent {
24    /// Initialize the summarization agent.
25    ///
26    /// # Arguments
27    ///
28    /// * `llm` - The LLM broker to use for generating summaries.
29    pub fn new(llm: Arc<LlmBroker>) -> Self {
30        Self { llm }
31    }
32
33    /// Generate the prompt for the summarization LLM.
34    fn prompt(&self, event: &FinishAndSummarize) -> String {
35        format!(
36            "Based on the following context, provide a clear and concise answer to the user's query.
37
38{}
39
40Your task:
41Review what we've learned and provide a direct answer to: \"{}\"
42
43Be specific and use the information gathered during our process.",
44            format_current_context(&event.context),
45            event.context.user_query
46        )
47    }
48}
49
50#[async_trait]
51impl BaseAsyncAgent for SummarizationAgent {
52    async fn receive_event_async(&self, event: Box<dyn Event>) -> Result<Vec<Box<dyn Event>>> {
53        // Downcast to FinishAndSummarize
54        let finish_event = match event.as_any().downcast_ref::<FinishAndSummarize>() {
55            Some(e) => e,
56            None => return Ok(vec![]),
57        };
58
59        let prompt = self.prompt(finish_event);
60        println!("\n{}\n{}\n{}\n", "=".repeat(80), prompt, "=".repeat(80));
61
62        // Generate final response
63        let response = match self
64            .llm
65            .generate(
66                &[LlmMessage {
67                    role: MessageRole::User,
68                    content: Some(prompt),
69                    tool_calls: None,
70                    image_paths: None,
71                }],
72                None,
73                None,
74                finish_event.correlation_id.clone(),
75            )
76            .await
77        {
78            Ok(r) => r,
79            Err(e) => {
80                return Ok(vec![Box::new(FailureOccurred {
81                    source: "SummarizationAgent".to_string(),
82                    correlation_id: finish_event.correlation_id.clone(),
83                    context: finish_event.context.clone(),
84                    reason: format!("Error during summarization: {}", e),
85                }) as Box<dyn Event>]);
86            }
87        };
88
89        println!("\n{}", "=".repeat(80));
90        println!("FINAL ANSWER:");
91        println!("{}", "=".repeat(80));
92        println!("{}", response);
93        println!("{}\n", "=".repeat(80));
94
95        // This is a terminal event - return empty list to stop the loop
96        Ok(vec![])
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::llm::gateways::OllamaGateway;
104
105    use super::super::models::{CurrentContext, Plan, ThoughtActionObservation};
106
107    #[test]
108    fn test_summarization_agent_prompt_generation() {
109        let gateway = Arc::new(OllamaGateway::new());
110        let llm = Arc::new(LlmBroker::new("qwen3:32b", gateway, None));
111        let agent = SummarizationAgent::new(llm);
112
113        let mut context = CurrentContext::new("What is the date tomorrow?");
114        context.history.push(ThoughtActionObservation {
115            thought: "Need to resolve date".to_string(),
116            action: "Called resolve_date".to_string(),
117            observation: "2025-11-30".to_string(),
118        });
119
120        let event = FinishAndSummarize {
121            source: "TestSource".to_string(),
122            correlation_id: Some("test-123".to_string()),
123            context,
124            thought: "I have the information needed".to_string(),
125        };
126
127        let prompt = agent.prompt(&event);
128
129        assert!(prompt.contains("What is the date tomorrow?"));
130        assert!(prompt.contains("provide a direct answer"));
131        assert!(prompt.contains("2025-11-30"));
132    }
133
134    #[test]
135    fn test_summarization_agent_with_plan() {
136        let gateway = Arc::new(OllamaGateway::new());
137        let llm = Arc::new(LlmBroker::new("qwen3:32b", gateway, None));
138        let agent = SummarizationAgent::new(llm);
139
140        let mut context = CurrentContext::new("What day is next Friday?");
141        context.plan = Plan {
142            steps: vec!["Resolve date for next Friday".to_string()],
143        };
144        context.history.push(ThoughtActionObservation {
145            thought: "Calculate next Friday".to_string(),
146            action: "Called resolve_date".to_string(),
147            observation: "2025-11-29".to_string(),
148        });
149
150        let event = FinishAndSummarize {
151            source: "TestSource".to_string(),
152            correlation_id: Some("test-456".to_string()),
153            context,
154            thought: "Ready to summarize".to_string(),
155        };
156
157        let prompt = agent.prompt(&event);
158
159        assert!(prompt.contains("What day is next Friday?"));
160        assert!(prompt.contains("Current plan:"));
161        assert!(prompt.contains("2025-11-29"));
162    }
163
164    #[tokio::test]
165    async fn test_summarization_agent_ignores_wrong_event_type() {
166        let gateway = Arc::new(OllamaGateway::new());
167        let llm = Arc::new(LlmBroker::new("qwen3:32b", gateway, None));
168        let agent = SummarizationAgent::new(llm);
169
170        use super::super::events::InvokeThinking;
171
172        let wrong_event = Box::new(InvokeThinking {
173            source: "Wrong".to_string(),
174            correlation_id: None,
175            context: CurrentContext::new("Test"),
176        }) as Box<dyn Event>;
177
178        let result = agent.receive_event_async(wrong_event).await.unwrap();
179        assert!(result.is_empty());
180    }
181}
mojentic/examples/react/summarization_agent.rs

mojentic/examples/react/
summarization_agent.rs
