Skip to main content

bamboo_engine/runtime/managers/adapters/
llm.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use bamboo_agent_core::tools::ToolSchema;
5use bamboo_agent_core::{AgentError, AgentEvent, Session};
6use bamboo_llm::LLMProvider;
7use tokio::sync::mpsc;
8use tokio_util::sync::CancellationToken;
9
10use crate::runtime::config::AgentLoopConfig;
11use crate::runtime::managers::llm::{LlmManager, LlmRoundOutput};
12
13/// Default LLM manager that delegates to existing runner functions.
14pub struct DefaultLlmManager {
15    llm: Arc<dyn LLMProvider>,
16}
17
18impl DefaultLlmManager {
19    pub fn new(llm: Arc<dyn LLMProvider>) -> Self {
20        Self { llm }
21    }
22}
23
24#[async_trait]
25impl LlmManager for DefaultLlmManager {
26    #[allow(clippy::too_many_arguments)]
27    async fn execute_round(
28        &self,
29        session: &mut Session,
30        config: &AgentLoopConfig,
31        event_tx: &mpsc::Sender<AgentEvent>,
32        cancel_token: &CancellationToken,
33        session_id: &str,
34        model_name: &str,
35        tool_schemas: &[ToolSchema],
36    ) -> Result<LlmRoundOutput, AgentError> {
37        let result = crate::runtime::runner::round_lifecycle::execute_llm_round(
38            session,
39            config,
40            &self.llm,
41            event_tx,
42            cancel_token,
43            session_id,
44            model_name,
45            tool_schemas,
46        )
47        .await?;
48
49        if let Some(error) = result.terminal_validation_error {
50            return Err(error);
51        }
52
53        let (content, reasoning_content, tool_calls) = {
54            let stream = &result.stream_output;
55            (
56                stream.content.clone(),
57                stream.reasoning_content.clone(),
58                stream.tool_calls.clone(),
59            )
60        };
61
62        Ok(LlmRoundOutput {
63            content,
64            reasoning_content,
65            tool_calls,
66            prompt_tokens: result.prompt_tokens,
67            completion_tokens: result.completion_tokens,
68            response_id: None,
69            round_usage: result.attempt_usage,
70        })
71    }
72
73    async fn attempt_overflow_recovery(
74        &self,
75        session: &mut Session,
76        config: &AgentLoopConfig,
77        session_id: &str,
78        event_tx: &mpsc::Sender<AgentEvent>,
79    ) -> Result<bool, AgentError> {
80        let model_name = config.model_name.as_deref().unwrap_or("unknown");
81
82        crate::runtime::runner::round_lifecycle::force_overflow_context_recovery(
83            session,
84            config,
85            model_name,
86            session_id,
87            &self.llm,
88            Some(event_tx),
89        )
90        .await
91    }
92}