Skip to main content

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