butterfly_bot/
client.rs

1use std::sync::Arc;
2
3use futures::stream::BoxStream;
4
5use crate::config::Config;
6use crate::error::{ButterflyBotError, Result};
7use crate::factories::agent_factory::ButterflyBotFactory;
8use crate::interfaces::plugins::Tool;
9use crate::services::agent::UiEvent;
10use crate::services::query::{ProcessOptions, ProcessResult, QueryService, UserInput};
11use tokio::sync::broadcast;
12
13pub struct ButterflyBot {
14    query_service: QueryService,
15}
16
17impl ButterflyBot {
18    pub async fn from_config(config: Config) -> Result<Self> {
19        let query_service = ButterflyBotFactory::create_from_config(config).await?;
20        Ok(Self { query_service })
21    }
22
23    pub async fn from_config_with_events(
24        config: Config,
25        ui_event_tx: Option<broadcast::Sender<UiEvent>>,
26    ) -> Result<Self> {
27        let query_service =
28            ButterflyBotFactory::create_from_config_with_events(config, ui_event_tx).await?;
29        Ok(Self { query_service })
30    }
31
32    pub async fn from_store(db_path: &str) -> Result<Self> {
33        let config = Config::from_store(db_path)?.resolve_vault()?;
34        let agent = Self::from_config(config).await?;
35        Ok(agent)
36    }
37
38    pub async fn from_store_with_events(
39        db_path: &str,
40        ui_event_tx: Option<broadcast::Sender<UiEvent>>,
41    ) -> Result<Self> {
42        let config = Config::from_store(db_path)?.resolve_vault()?;
43        let agent = Self::from_config_with_events(config, ui_event_tx).await?;
44        Ok(agent)
45    }
46
47    pub fn process_text_stream<'a>(
48        &'a self,
49        user_id: &'a str,
50        message: &'a str,
51        prompt: Option<&'a str>,
52    ) -> BoxStream<'a, Result<String>> {
53        let service = &self.query_service;
54        service.process_text_stream(user_id, message, prompt)
55    }
56
57    pub async fn process(
58        &self,
59        user_id: &str,
60        input: UserInput,
61        options: ProcessOptions,
62    ) -> Result<ProcessResult> {
63        self.query_service.process(user_id, input, options).await
64    }
65
66    pub async fn delete_user_history(&self, user_id: &str) -> Result<()> {
67        self.query_service.delete_user_history(user_id).await
68    }
69
70    pub async fn get_user_history(&self, user_id: &str, limit: usize) -> Result<Vec<String>> {
71        self.query_service.get_user_history(user_id, limit).await
72    }
73
74    pub async fn search_memory(
75        &self,
76        user_id: &str,
77        query: &str,
78        limit: usize,
79    ) -> Result<Vec<String>> {
80        self.query_service
81            .search_memory(user_id, query, limit)
82            .await
83    }
84
85    pub async fn preload_context(&self, user_id: &str) -> Result<()> {
86        self.query_service.preload_context(user_id).await
87    }
88
89    pub async fn set_heartbeat_markdown(&self, heartbeat_markdown: Option<String>) {
90        let agent_service = self.query_service.agent_service();
91        agent_service
92            .set_heartbeat_markdown(heartbeat_markdown)
93            .await;
94    }
95
96    pub async fn set_prompt_markdown(&self, prompt_markdown: Option<String>) {
97        let agent_service = self.query_service.agent_service();
98        agent_service.set_prompt_markdown(prompt_markdown).await;
99    }
100
101    pub async fn register_tool(&self, tool: Arc<dyn Tool>) -> Result<bool> {
102        let agent_service = self.query_service.agent_service();
103        let registry = agent_service.tool_registry.clone();
104        if !registry.register_tool(tool.clone()).await {
105            return Ok(false);
106        }
107        let assigned = registry
108            .assign_tool_to_agent(agent_service.agent_name(), tool.name())
109            .await;
110        if !assigned {
111            return Err(ButterflyBotError::Runtime(
112                "Tool registered but could not assign to agent".to_string(),
113            ));
114        }
115        Ok(true)
116    }
117
118    pub async fn brain_tick(&self) {
119        let agent_service = self.query_service.agent_service();
120        agent_service.dispatch_brain_tick().await;
121    }
122}
butterfly_bot/client.rs

butterfly_bot/
client.rs
