bamboo_engine/runtime/managers/adapters/
llm_mini_loop.rs

1use async_trait::async_trait;
2use bamboo_agent_core::tools::ToolCall;
3use bamboo_agent_core::{AgentError, Session};
4use bamboo_domain::{Message, Role};
5use bamboo_llm::LLMProvider;
6
7use crate::runtime::managers::mini_loop::{MiniLoopDecision, MiniLoopExecutor};
8
9/// Production `MiniLoopExecutor` backed by a real LLM provider.
10///
11/// Uses a fast/cheap model for lightweight decisions such as task complexity
12/// classification, compression decisions, retry classification, and routing.
13pub struct LLMMiniLoopExecutor {
14    provider: std::sync::Arc<dyn LLMProvider>,
15    model: String,
16    timeout_context: crate::runtime::stream::handler::StreamTimeoutContext,
17}
18
19impl LLMMiniLoopExecutor {
20    pub fn new(provider: std::sync::Arc<dyn LLMProvider>, model: String) -> Self {
21        Self {
22            provider,
23            model,
24            timeout_context: crate::runtime::stream::handler::StreamTimeoutContext::default(),
25        }
26    }
27
28    pub fn with_timeout_policy(
29        provider: std::sync::Arc<dyn LLMProvider>,
30        model: String,
31        policy: bamboo_config::StreamTimeoutConfig,
32        provider_name: Option<&str>,
33    ) -> Self {
34        let timeout_context = crate::runtime::stream::handler::StreamTimeoutContext::new(
35            policy,
36            provider_name,
37            Some(&model),
38        );
39        Self {
40            provider,
41            model,
42            timeout_context,
43        }
44    }
45}
46
47#[async_trait]
48impl MiniLoopExecutor for LLMMiniLoopExecutor {
49    async fn decide(
50        &self,
51        _session: &Session,
52        prompt: &str,
53        context: &str,
54    ) -> Result<MiniLoopDecision, AgentError> {
55        let user_content = if context.is_empty() {
56            prompt.to_string()
57        } else {
58            format!("Context:\n{}\n\n{}", context, prompt)
59        };
60
61        let now = chrono::Utc::now();
62        let messages = vec![
63            Message {
64                id: String::new(),
65                role: Role::System,
66                content: "You are a task complexity classifier. Respond with exactly one word: simple, standard, or complex.".to_string(),
67                reasoning: None,
68                reasoning_signature: None,
69                content_parts: None,
70                image_ocr: None,
71                phase: None,
72                tool_calls: None,
73                tool_call_id: None,
74                tool_success: None,
75                compressed: false,
76                compressed_by_event_id: None,
77                never_compress: false,
78                compression_level: 0,
79                created_at: now,
80                metadata: None,
81            },
82            Message {
83                id: String::new(),
84                role: Role::User,
85                content: user_content,
86                reasoning: None,
87                reasoning_signature: None,
88                content_parts: None,
89                image_ocr: None,
90                phase: None,
91                tool_calls: None,
92                tool_call_id: None,
93                tool_success: None,
94                compressed: false,
95                compressed_by_event_id: None,
96                never_compress: false,
97                compression_level: 0,
98                created_at: now,
99                metadata: None,
100            },
101        ];
102
103        let options = bamboo_llm::provider::LLMRequestOptions {
104            session_id: None,
105            reasoning_effort: None,
106            parallel_tool_calls: None,
107            required_tool: None,
108            responses: None,
109            request_purpose: Some("mini_loop".to_string()),
110            cache: None,
111        };
112        let timeout_context = self.timeout_context.clone().begin_request();
113        let cancel_token = tokio_util::sync::CancellationToken::new();
114        let stream = crate::runtime::stream::handler::await_stream_bootstrap(
115            self.provider.chat_stream_with_options(
116                &messages,
117                &[],
118                Some(256),
119                &self.model,
120                Some(&options),
121            ),
122            &cancel_token,
123            "mini-loop",
124            &timeout_context,
125        )
126        .await?
127        .map_err(|e| AgentError::LLM(e.to_string()))?;
128
129        let output = crate::runtime::stream::handler::consume_llm_stream_silent_with_context(
130            stream,
131            &cancel_token,
132            "mini-loop",
133            &timeout_context,
134        )
135        .await?;
136
137        Ok(MiniLoopDecision {
138            answer: output.content.trim().to_string(),
139            prompt_tokens: 0,
140            completion_tokens: 0,
141        })
142    }
143
144    async fn evaluate_task(
145        &self,
146        _session: &Session,
147        _tool_calls: &[ToolCall],
148        _round: usize,
149    ) -> Result<MiniLoopDecision, AgentError> {
150        // Task evaluation is handled by the dedicated task_evaluation module.
151        // This implementation returns an empty decision as a fallback.
152        Ok(MiniLoopDecision {
153            answer: String::new(),
154            prompt_tokens: 0,
155            completion_tokens: 0,
156        })
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163    use bamboo_agent_core::tools::ToolSchema;
164    use bamboo_domain::ReasoningEffort;
165    use bamboo_llm::provider::LLMRequestOptions;
166    use bamboo_llm::{LLMChunk, LLMError, LLMProvider, LLMStream};
167    use futures::stream;
168    use std::sync::{Arc, Mutex};
169
170    #[derive(Default)]
171    struct CaptureProvider {
172        captured_max_tokens: Mutex<Vec<Option<u32>>>,
173        captured_reasoning: Mutex<Vec<Option<ReasoningEffort>>>,
174    }
175
176    #[async_trait]
177    impl LLMProvider for CaptureProvider {
178        async fn chat_stream(
179            &self,
180            _messages: &[Message],
181            _tools: &[ToolSchema],
182            _max_output_tokens: Option<u32>,
183            _model: &str,
184        ) -> Result<LLMStream, LLMError> {
185            Ok(Box::pin(stream::iter(vec![
186                Ok::<LLMChunk, LLMError>(LLMChunk::Token("simple".to_string())),
187                Ok::<LLMChunk, LLMError>(LLMChunk::Done),
188            ])))
189        }
190
191        async fn chat_stream_with_options(
192            &self,
193            messages: &[Message],
194            tools: &[ToolSchema],
195            max_output_tokens: Option<u32>,
196            model: &str,
197            options: Option<&LLMRequestOptions>,
198        ) -> Result<LLMStream, LLMError> {
199            self.captured_max_tokens
200                .lock()
201                .expect("lock should not be poisoned")
202                .push(max_output_tokens);
203            self.captured_reasoning
204                .lock()
205                .expect("lock should not be poisoned")
206                .push(options.and_then(|o| o.reasoning_effort));
207            self.chat_stream(messages, tools, max_output_tokens, model)
208                .await
209        }
210    }
211
212    #[tokio::test]
213    async fn mini_loop_sends_no_reasoning_effort() {
214        let provider = Arc::new(CaptureProvider::default());
215        let executor = LLMMiniLoopExecutor::new(provider.clone(), "fast-model".to_string());
216        let session = Session::new("test", "model");
217
218        let decision = executor
219            .decide(&session, "classify this task", "")
220            .await
221            .expect("decide should succeed");
222        assert_eq!(decision.answer, "simple");
223
224        let captured_reasoning = provider
225            .captured_reasoning
226            .lock()
227            .expect("lock should not be poisoned");
228        assert_eq!(
229            captured_reasoning.as_slice(),
230            [None],
231            "mini_loop should not request reasoning to avoid thinking budget consuming output tokens"
232        );
233    }
234
235    #[tokio::test]
236    async fn mini_loop_max_tokens_accommodates_provider_default_reasoning() {
237        let provider = Arc::new(CaptureProvider::default());
238        let executor = LLMMiniLoopExecutor::new(provider.clone(), "fast-model".to_string());
239        let session = Session::new("test", "model");
240
241        let _ = executor
242            .decide(&session, "classify this task", "")
243            .await
244            .expect("decide should succeed");
245
246        let captured = provider
247            .captured_max_tokens
248            .lock()
249            .expect("lock should not be poisoned");
250        let max_tokens = captured[0].expect("max_output_tokens should be set");
251        // Even if the provider has a default_reasoning_effort of High (4096 thinking budget),
252        // 256 tokens is enough for the one-word response while leaving room for thinking.
253        assert!(
254            max_tokens >= 256,
255            "max_output_tokens ({}) should be at least 256 to accommodate potential provider default reasoning",
256            max_tokens
257        );
258    }
259}
bamboo_engine/runtime/managers/adapters/llm_mini_loop.rs

bamboo_engine/runtime/managers/adapters/
llm_mini_loop.rs
