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_infrastructure::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        let (content, reasoning_content, tool_calls) = {
50            let stream = &result.stream_output;
51            (
52                stream.content.clone(),
53                stream.reasoning_content.clone(),
54                stream.tool_calls.clone(),
55            )
56        };
57
58        Ok(LlmRoundOutput {
59            content,
60            reasoning_content,
61            tool_calls,
62            prompt_tokens: result.prompt_tokens,
63            completion_tokens: result.completion_tokens,
64            response_id: None,
65            round_usage: result.round_usage,
66        })
67    }
68
69    async fn attempt_overflow_recovery(
70        &self,
71        session: &mut Session,
72        config: &AgentLoopConfig,
73        session_id: &str,
74        event_tx: &mpsc::Sender<AgentEvent>,
75    ) -> Result<bool, AgentError> {
76        let model_name = config.model_name.as_deref().unwrap_or("unknown");
77
78        crate::runtime::runner::round_lifecycle::force_overflow_context_recovery(
79            session,
80            config,
81            model_name,
82            session_id,
83            &self.llm,
84            Some(event_tx),
85        )
86        .await
87    }
88}