use super::client::{ChatSession, LLMClient};
use super::provider::{ChatStream, LLMProvider};
use super::tool_executor::ToolExecutor;
use super::types::{ChatMessage, LLMError, LLMResult, Tool};
use crate::llm::{
AnthropicConfig, AnthropicProvider, GeminiConfig, GeminiProvider, OllamaConfig, OllamaProvider,
};
use crate::prompt;
use futures::{Stream, StreamExt};
use mofa_kernel::agent::AgentMetadata;
use mofa_kernel::agent::AgentState;
use mofa_kernel::plugin::{AgentPlugin, PluginType};
use mofa_plugins::tts::TTSPlugin;
use std::collections::HashMap;
use std::io::Write;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::{Mutex, RwLock};
pub type TtsAudioStream = Pin<Box<dyn Stream<Item = (Vec<f32>, Duration)> + Send>>;
struct CancellationToken {
cancel: Arc<AtomicBool>,
}
impl CancellationToken {
fn new() -> Self {
Self {
cancel: Arc::new(AtomicBool::new(false)),
}
}
fn is_cancelled(&self) -> bool {
self.cancel.load(Ordering::Relaxed)
}
fn cancel(&self) {
self.cancel.store(true, Ordering::Relaxed);
}
fn clone_token(&self) -> CancellationToken {
CancellationToken {
cancel: Arc::clone(&self.cancel),
}
}
}
pub type TextStream = Pin<Box<dyn Stream<Item = LLMResult<String>> + Send>>;
#[cfg(feature = "kokoro")]
struct TTSStreamHandle {
sink: mofa_plugins::tts::kokoro_wrapper::SynthSink<String>,
_stream_handle: tokio::task::JoinHandle<()>,
}
struct TTSSession {
cancellation_token: CancellationToken,
is_active: Arc<AtomicBool>,
}
impl TTSSession {
fn new(token: CancellationToken) -> Self {
let is_active = Arc::new(AtomicBool::new(true));
TTSSession {
cancellation_token: token,
is_active,
}
}
fn cancel(&self) {
self.cancellation_token.cancel();
self.is_active.store(false, Ordering::Relaxed);
}
fn is_active(&self) -> bool {
self.is_active.load(Ordering::Relaxed)
}
}
struct SentenceBuffer {
buffer: String,
}
impl SentenceBuffer {
fn new() -> Self {
Self {
buffer: String::new(),
}
}
fn push(&mut self, text: &str) -> Option<String> {
for ch in text.chars() {
self.buffer.push(ch);
if matches!(ch, '。' | '!' | '?' | '!' | '?') {
let sentence = self.buffer.trim().to_string();
if !sentence.is_empty() {
self.buffer.clear();
return Some(sentence);
}
}
}
None
}
fn flush(&mut self) -> Option<String> {
if self.buffer.trim().is_empty() {
None
} else {
let remaining = self.buffer.trim().to_string();
self.buffer.clear();
Some(remaining)
}
}
}
#[derive(Debug, Clone)]
pub enum StreamEvent {
Text(String),
ToolCallStart { id: String, name: String },
ToolCallDelta { id: String, arguments_delta: String },
Done(Option<String>),
}
#[derive(Clone)]
pub struct LLMAgentConfig {
pub agent_id: String,
pub name: String,
pub system_prompt: Option<String>,
pub temperature: Option<f32>,
pub max_tokens: Option<u32>,
pub custom_config: HashMap<String, String>,
pub user_id: Option<String>,
pub tenant_id: Option<String>,
pub context_window_size: Option<usize>,
}
impl Default for LLMAgentConfig {
fn default() -> Self {
Self {
agent_id: "llm-agent".to_string(),
name: "LLM Agent".to_string(),
system_prompt: None,
temperature: Some(0.7),
max_tokens: Some(4096),
custom_config: HashMap::new(),
user_id: None,
tenant_id: None,
context_window_size: None,
}
}
}
pub struct LLMAgent {
config: LLMAgentConfig,
metadata: AgentMetadata,
client: LLMClient,
sessions: Arc<RwLock<HashMap<String, Arc<RwLock<ChatSession>>>>>,
active_session_id: Arc<RwLock<String>>,
tools: Vec<Tool>,
tool_executor: Option<Arc<dyn ToolExecutor>>,
event_handler: Option<Box<dyn LLMAgentEventHandler>>,
plugins: Vec<Box<dyn AgentPlugin>>,
state: AgentState,
provider: Arc<dyn LLMProvider>,
prompt_plugin: Option<Box<dyn prompt::PromptTemplatePlugin>>,
tts_plugin: Option<Arc<Mutex<TTSPlugin>>>,
#[cfg(feature = "kokoro")]
cached_kokoro_engine: Arc<Mutex<Option<Arc<mofa_plugins::tts::kokoro_wrapper::KokoroTTS>>>>,
active_tts_session: Arc<Mutex<Option<TTSSession>>>,
message_store: Option<Arc<dyn crate::persistence::MessageStore + Send + Sync>>,
session_store: Option<Arc<dyn crate::persistence::SessionStore + Send + Sync>>,
persistence_user_id: Option<uuid::Uuid>,
persistence_agent_id: Option<uuid::Uuid>,
}
#[async_trait::async_trait]
pub trait LLMAgentEventHandler: Send + Sync {
fn clone_box(&self) -> Box<dyn LLMAgentEventHandler>;
fn as_any(&self) -> &dyn std::any::Any;
async fn before_chat(&self, message: &str) -> LLMResult<Option<String>> {
Ok(Some(message.to_string()))
}
async fn before_chat_with_model(
&self,
message: &str,
_model: &str,
) -> LLMResult<Option<String>> {
self.before_chat(message).await
}
async fn after_chat(&self, response: &str) -> LLMResult<Option<String>> {
Ok(Some(response.to_string()))
}
async fn after_chat_with_metadata(
&self,
response: &str,
_metadata: &super::types::LLMResponseMetadata,
) -> LLMResult<Option<String>> {
self.after_chat(response).await
}
async fn on_tool_call(&self, name: &str, arguments: &str) -> LLMResult<Option<String>> {
let _ = (name, arguments);
Ok(None)
}
async fn on_error(&self, error: &LLMError) -> LLMResult<Option<String>> {
let _ = error;
Ok(None)
}
}
impl Clone for Box<dyn LLMAgentEventHandler> {
fn clone(&self) -> Self {
self.clone_box()
}
}
impl LLMAgent {
pub fn new(config: LLMAgentConfig, provider: Arc<dyn LLMProvider>) -> Self {
Self::with_initial_session(config, provider, None)
}
pub fn with_initial_session(
config: LLMAgentConfig,
provider: Arc<dyn LLMProvider>,
initial_session_id: Option<String>,
) -> Self {
let client = LLMClient::new(provider.clone());
let mut session = if let Some(sid) = initial_session_id {
ChatSession::with_id_str(&sid, LLMClient::new(provider.clone()))
} else {
ChatSession::new(LLMClient::new(provider.clone()))
};
if let Some(ref prompt) = config.system_prompt {
session = session.with_system(prompt.clone());
}
session = session.with_context_window_size(config.context_window_size);
let session_id = session.session_id().to_string();
let session_arc = Arc::new(RwLock::new(session));
let mut sessions = HashMap::new();
sessions.insert(session_id.clone(), session_arc);
let agent_id = config.agent_id.clone();
let name = config.name.clone();
let capabilities = mofa_kernel::agent::AgentCapabilities::builder()
.tags(vec![
"llm".to_string(),
"chat".to_string(),
"text-generation".to_string(),
"multi-session".to_string(),
])
.build();
Self {
config,
metadata: AgentMetadata {
id: agent_id,
name,
description: None,
version: None,
capabilities,
state: AgentState::Created,
},
client,
sessions: Arc::new(RwLock::new(sessions)),
active_session_id: Arc::new(RwLock::new(session_id)),
tools: Vec::new(),
tool_executor: None,
event_handler: None,
plugins: Vec::new(),
state: AgentState::Created,
provider,
prompt_plugin: None,
tts_plugin: None,
#[cfg(feature = "kokoro")]
cached_kokoro_engine: Arc::new(Mutex::new(None)),
active_tts_session: Arc::new(Mutex::new(None)),
message_store: None,
session_store: None,
persistence_user_id: None,
persistence_agent_id: None,
}
}
pub async fn with_initial_session_async(
config: LLMAgentConfig,
provider: Arc<dyn LLMProvider>,
initial_session_id: Option<String>,
message_store: Option<Arc<dyn crate::persistence::MessageStore + Send + Sync>>,
session_store: Option<Arc<dyn crate::persistence::SessionStore + Send + Sync>>,
persistence_user_id: Option<uuid::Uuid>,
persistence_tenant_id: Option<uuid::Uuid>,
persistence_agent_id: Option<uuid::Uuid>,
) -> Self {
let client = LLMClient::new(provider.clone());
let initial_session_id_clone = initial_session_id.clone();
let session = if let (
Some(sid),
Some(msg_store),
Some(sess_store),
Some(user_id),
Some(tenant_id),
Some(agent_id),
) = (
initial_session_id_clone,
message_store.clone(),
session_store.clone(),
persistence_user_id,
persistence_tenant_id,
persistence_agent_id,
) {
let msg_store_clone = msg_store.clone();
let sess_store_clone = sess_store.clone();
let session_uuid = uuid::Uuid::parse_str(&sid).unwrap_or_else(|_| {
tracing::warn!("⚠️ 无效的 session_id 格式 '{}', 将生成新的 UUID", sid);
uuid::Uuid::now_v7()
});
match ChatSession::load(
session_uuid,
LLMClient::new(provider.clone()),
user_id,
agent_id,
tenant_id,
msg_store,
sess_store,
config.context_window_size,
)
.await
{
Ok(loaded_session) => {
tracing::info!(
"✅ 从数据库加载会话: {} ({} 条消息)",
sid,
loaded_session.messages().len()
);
loaded_session
}
Err(e) => {
tracing::info!("📝 创建新会话并持久化: {} (数据库中不存在: {})", sid, e);
let msg_store_clone2 = msg_store_clone.clone();
let sess_store_clone2 = sess_store_clone.clone();
match ChatSession::with_id_and_stores_and_persist(
session_uuid,
LLMClient::new(provider.clone()),
user_id,
agent_id,
tenant_id,
msg_store_clone,
sess_store_clone,
config.context_window_size,
)
.await
{
Ok(mut new_session) => {
if let Some(ref prompt) = config.system_prompt {
new_session = new_session.with_system(prompt.clone());
}
new_session
}
Err(persist_err) => {
tracing::error!("❌ 持久化会话失败: {}, 降级为内存会话", persist_err);
let new_session = ChatSession::with_id_and_stores(
session_uuid,
LLMClient::new(provider.clone()),
user_id,
agent_id,
tenant_id,
msg_store_clone2,
sess_store_clone2,
config.context_window_size,
);
if let Some(ref prompt) = config.system_prompt {
new_session.with_system(prompt.clone())
} else {
new_session
}
}
}
}
}
} else {
let mut session = if let Some(sid) = initial_session_id {
ChatSession::with_id_str(&sid, LLMClient::new(provider.clone()))
} else {
ChatSession::new(LLMClient::new(provider.clone()))
};
if let Some(ref prompt) = config.system_prompt {
session = session.with_system(prompt.clone());
}
session.with_context_window_size(config.context_window_size)
};
let session_id = session.session_id().to_string();
let session_arc = Arc::new(RwLock::new(session));
let mut sessions = HashMap::new();
sessions.insert(session_id.clone(), session_arc);
let agent_id = config.agent_id.clone();
let name = config.name.clone();
let capabilities = mofa_kernel::agent::AgentCapabilities::builder()
.tags(vec![
"llm".to_string(),
"chat".to_string(),
"text-generation".to_string(),
"multi-session".to_string(),
])
.build();
Self {
config,
metadata: AgentMetadata {
id: agent_id,
name,
description: None,
version: None,
capabilities,
state: AgentState::Created,
},
client,
sessions: Arc::new(RwLock::new(sessions)),
active_session_id: Arc::new(RwLock::new(session_id)),
tools: Vec::new(),
tool_executor: None,
event_handler: None,
plugins: Vec::new(),
state: AgentState::Created,
provider,
prompt_plugin: None,
tts_plugin: None,
#[cfg(feature = "kokoro")]
cached_kokoro_engine: Arc::new(Mutex::new(None)),
active_tts_session: Arc::new(Mutex::new(None)),
message_store,
session_store,
persistence_user_id,
persistence_agent_id,
}
}
pub fn config(&self) -> &LLMAgentConfig {
&self.config
}
pub fn client(&self) -> &LLMClient {
&self.client
}
pub async fn current_session_id(&self) -> String {
self.active_session_id.read().await.clone()
}
pub async fn create_session(&self) -> String {
let mut session = ChatSession::new(LLMClient::new(self.provider.clone()));
let mut system_prompt = self.config.system_prompt.clone();
if let Some(ref plugin) = self.prompt_plugin
&& let Some(template) = plugin.get_current_template().await
{
system_prompt = match template.render(&[]) {
Ok(prompt) => Some(prompt),
Err(_) => self.config.system_prompt.clone(),
};
}
if let Some(ref prompt) = system_prompt {
session = session.with_system(prompt.clone());
}
session = session.with_context_window_size(self.config.context_window_size);
let session_id = session.session_id().to_string();
let session_arc = Arc::new(RwLock::new(session));
let mut sessions = self.sessions.write().await;
sessions.insert(session_id.clone(), session_arc);
session_id
}
pub async fn create_session_with_id(&self, session_id: impl Into<String>) -> LLMResult<String> {
let session_id = session_id.into();
{
let sessions = self.sessions.read().await;
if sessions.contains_key(&session_id) {
return Err(LLMError::Other(format!(
"Session with id '{}' already exists",
session_id
)));
}
}
let mut session =
ChatSession::with_id_str(&session_id, LLMClient::new(self.provider.clone()));
let mut system_prompt = self.config.system_prompt.clone();
if let Some(ref plugin) = self.prompt_plugin
&& let Some(template) = plugin.get_current_template().await
{
system_prompt = match template.render(&[]) {
Ok(prompt) => Some(prompt),
Err(_) => self.config.system_prompt.clone(),
};
}
if let Some(ref prompt) = system_prompt {
session = session.with_system(prompt.clone());
}
session = session.with_context_window_size(self.config.context_window_size);
let session_arc = Arc::new(RwLock::new(session));
let mut sessions = self.sessions.write().await;
sessions.insert(session_id.clone(), session_arc);
Ok(session_id)
}
pub async fn switch_session(&self, session_id: &str) -> LLMResult<()> {
let sessions = self.sessions.read().await;
if !sessions.contains_key(session_id) {
return Err(LLMError::Other(format!(
"Session '{}' not found",
session_id
)));
}
drop(sessions);
let mut active = self.active_session_id.write().await;
*active = session_id.to_string();
Ok(())
}
pub async fn get_or_create_session(&self, session_id: impl Into<String>) -> String {
let session_id = session_id.into();
{
let sessions = self.sessions.read().await;
if sessions.contains_key(&session_id) {
return session_id;
}
}
let _ = self.create_session_with_id(&session_id).await;
session_id
}
pub async fn remove_session(&self, session_id: &str) -> LLMResult<()> {
let active = self.active_session_id.read().await.clone();
if active == session_id {
return Err(LLMError::Other(
"Cannot remove active session. Switch to another session first.".to_string(),
));
}
let mut sessions = self.sessions.write().await;
if sessions.remove(session_id).is_none() {
return Err(LLMError::Other(format!(
"Session '{}' not found",
session_id
)));
}
Ok(())
}
pub async fn list_sessions(&self) -> Vec<String> {
let sessions = self.sessions.read().await;
sessions.keys().cloned().collect()
}
pub async fn session_count(&self) -> usize {
let sessions = self.sessions.read().await;
sessions.len()
}
pub async fn has_session(&self, session_id: &str) -> bool {
let sessions = self.sessions.read().await;
sessions.contains_key(session_id)
}
pub async fn tts_speak(&self, text: &str) -> LLMResult<()> {
let tts = self
.tts_plugin
.as_ref()
.ok_or_else(|| LLMError::Other("TTS plugin not configured".to_string()))?;
let mut tts_guard = tts.lock().await;
tts_guard
.synthesize_and_play(text)
.await
.map_err(|e| LLMError::Other(format!("TTS synthesis failed: {}", e)))
}
pub async fn tts_speak_streaming(
&self,
text: &str,
callback: Box<dyn Fn(Vec<u8>) + Send + Sync>,
) -> LLMResult<()> {
let tts = self
.tts_plugin
.as_ref()
.ok_or_else(|| LLMError::Other("TTS plugin not configured".to_string()))?;
let mut tts_guard = tts.lock().await;
tts_guard
.synthesize_streaming(text, callback)
.await
.map_err(|e| LLMError::Other(format!("TTS streaming failed: {}", e)))
}
pub async fn tts_speak_f32_stream(
&self,
text: &str,
callback: Box<dyn Fn(Vec<f32>) + Send + Sync>,
) -> LLMResult<()> {
let tts = self
.tts_plugin
.as_ref()
.ok_or_else(|| LLMError::Other("TTS plugin not configured".to_string()))?;
let mut tts_guard = tts.lock().await;
tts_guard
.synthesize_streaming_f32(text, callback)
.await
.map_err(|e| LLMError::Other(format!("TTS f32 streaming failed: {}", e)))
}
pub async fn tts_create_stream(&self, text: &str) -> LLMResult<TtsAudioStream> {
#[cfg(feature = "kokoro")]
{
use mofa_plugins::tts::kokoro_wrapper::KokoroTTS;
let cached_engine = {
let cache_guard = self.cached_kokoro_engine.lock().await;
cache_guard.clone()
};
let kokoro = if let Some(engine) = cached_engine {
engine
} else {
let tts = self
.tts_plugin
.as_ref()
.ok_or_else(|| LLMError::Other("TTS plugin not configured".to_string()))?;
let tts_guard = tts.lock().await;
let engine = tts_guard
.engine()
.ok_or_else(|| LLMError::Other("TTS engine not initialized".to_string()))?;
if let Some(kokoro_ref) = engine.as_any().downcast_ref::<KokoroTTS>() {
let cloned = kokoro_ref.clone();
let cloned_arc = Arc::new(cloned);
let voice = tts_guard
.stats()
.get("default_voice")
.and_then(|v| v.as_str())
.unwrap_or("default");
{
let mut cache_guard = self.cached_kokoro_engine.lock().await;
*cache_guard = Some(cloned_arc.clone());
}
cloned_arc
} else {
return Err(LLMError::Other("TTS engine is not KokoroTTS".to_string()));
}
};
let voice = "default"; let (mut sink, stream) = kokoro
.create_stream(voice)
.await
.map_err(|e| LLMError::Other(format!("Failed to create TTS stream: {}", e)))?;
sink.synth(text.to_string()).await.map_err(|e| {
LLMError::Other(format!("Failed to submit text for synthesis: {}", e))
})?;
return Ok(Box::pin(stream));
}
#[cfg(not(feature = "kokoro"))]
{
Err(LLMError::Other("Kokoro feature not enabled".to_string()))
}
}
pub async fn tts_speak_f32_stream_batch(
&self,
sentences: Vec<String>,
callback: Box<dyn Fn(Vec<f32>) + Send + Sync>,
) -> LLMResult<()> {
let tts = self
.tts_plugin
.as_ref()
.ok_or_else(|| LLMError::Other("TTS plugin not configured".to_string()))?;
let tts_guard = tts.lock().await;
#[cfg(feature = "kokoro")]
{
use mofa_plugins::tts::kokoro_wrapper::KokoroTTS;
let engine = tts_guard
.engine()
.ok_or_else(|| LLMError::Other("TTS engine not initialized".to_string()))?;
if let Some(kokoro) = engine.as_any().downcast_ref::<KokoroTTS>() {
let voice = tts_guard
.stats()
.get("default_voice")
.and_then(|v| v.as_str())
.unwrap_or("default")
.to_string();
let (mut sink, mut stream) = kokoro
.create_stream(&voice)
.await
.map_err(|e| LLMError::Other(format!("Failed to create TTS stream: {}", e)))?;
tokio::spawn(async move {
while let Some((audio, _took)) = stream.next().await {
callback(audio);
}
});
for sentence in sentences {
sink.synth(sentence)
.await
.map_err(|e| LLMError::Other(format!("Failed to submit text: {}", e)))?;
}
return Ok(());
}
return Err(LLMError::Other("TTS engine is not KokoroTTS".to_string()));
}
#[cfg(not(feature = "kokoro"))]
{
Err(LLMError::Other("Kokoro feature not enabled".to_string()))
}
}
pub fn has_tts(&self) -> bool {
self.tts_plugin.is_some()
}
pub async fn interrupt_tts(&self) -> LLMResult<()> {
let mut session_guard = self.active_tts_session.lock().await;
if let Some(session) = session_guard.take() {
session.cancel();
}
Ok(())
}
pub async fn chat_with_tts(
&self,
session_id: &str,
message: impl Into<String>,
) -> LLMResult<()> {
self.chat_with_tts_internal(session_id, message, None).await
}
pub async fn chat_with_tts_callback(
&self,
session_id: &str,
message: impl Into<String>,
callback: impl Fn(Vec<f32>) + Send + Sync + 'static,
) -> LLMResult<()> {
self.chat_with_tts_internal(session_id, message, Some(Box::new(callback)))
.await
}
#[cfg(feature = "kokoro")]
async fn create_tts_stream_handle(
&self,
callback: Box<dyn Fn(Vec<f32>) + Send + Sync>,
cancellation_token: Option<CancellationToken>,
) -> LLMResult<TTSStreamHandle> {
use mofa_plugins::tts::kokoro_wrapper::KokoroTTS;
let tts = self
.tts_plugin
.as_ref()
.ok_or_else(|| LLMError::Other("TTS plugin not configured".to_string()))?;
let tts_guard = tts.lock().await;
let engine = tts_guard
.engine()
.ok_or_else(|| LLMError::Other("TTS engine not initialized".to_string()))?;
let kokoro = engine
.as_any()
.downcast_ref::<KokoroTTS>()
.ok_or_else(|| LLMError::Other("TTS engine is not KokoroTTS".to_string()))?;
let voice = tts_guard
.stats()
.get("default_voice")
.and_then(|v| v.as_str())
.unwrap_or("default")
.to_string();
let (sink, mut stream) = kokoro
.create_stream(&voice)
.await
.map_err(|e| LLMError::Other(format!("Failed to create TTS stream: {}", e)))?;
let token_clone = cancellation_token.as_ref().map(|t| t.clone_token());
let stream_handle = tokio::spawn(async move {
while let Some((audio, _took)) = stream.next().await {
if let Some(ref token) = token_clone {
if token.is_cancelled() {
break; }
}
callback(audio);
}
});
Ok(TTSStreamHandle {
sink,
_stream_handle: stream_handle,
})
}
async fn chat_with_tts_internal(
&self,
session_id: &str,
message: impl Into<String>,
callback: Option<Box<dyn Fn(Vec<f32>) + Send + Sync>>,
) -> LLMResult<()> {
#[cfg(feature = "kokoro")]
{
use mofa_plugins::tts::kokoro_wrapper::KokoroTTS;
let callback = match callback {
Some(cb) => cb,
None => {
let mut text_stream =
self.chat_stream_with_session(session_id, message).await?;
while let Some(result) = text_stream.next().await {
match result {
Ok(text_chunk) => {
print!("{}", text_chunk);
std::io::stdout().flush().map_err(|e| {
LLMError::Other(format!("Failed to flush stdout: {}", e))
})?;
}
Err(e) if e.to_string().contains("__stream_end__") => break,
Err(e) => return Err(e),
}
}
println!();
return Ok(());
}
};
self.interrupt_tts().await?;
let cancellation_token = CancellationToken::new();
let mut tts_handle = self
.create_tts_stream_handle(callback, Some(cancellation_token.clone_token()))
.await?;
let session = TTSSession::new(cancellation_token);
{
let mut active_session = self.active_tts_session.lock().await;
*active_session = Some(session);
}
let mut buffer = SentenceBuffer::new();
let mut text_stream = self.chat_stream_with_session(session_id, message).await?;
while let Some(result) = text_stream.next().await {
match result {
Ok(text_chunk) => {
{
let active_session = self.active_tts_session.lock().await;
if let Some(ref session) = *active_session {
if !session.is_active() {
return Ok(()); }
}
}
print!("{}", text_chunk);
std::io::stdout().flush().map_err(|e| {
LLMError::Other(format!("Failed to flush stdout: {}", e))
})?;
if let Some(sentence) = buffer.push(&text_chunk) {
if let Err(e) = tts_handle.sink.synth(sentence).await {
eprintln!("[TTS Error] Failed to submit sentence: {}", e);
}
}
}
Err(e) if e.to_string().contains("__stream_end__") => break,
Err(e) => return Err(e),
}
}
if let Some(remaining) = buffer.flush() {
if let Err(e) = tts_handle.sink.synth(remaining).await {
eprintln!("[TTS Error] Failed to submit final sentence: {}", e);
}
}
{
let mut active_session = self.active_tts_session.lock().await;
*active_session = None;
}
let _ = tokio::time::timeout(
tokio::time::Duration::from_secs(30),
tts_handle._stream_handle,
)
.await
.map_err(|_| LLMError::Other("TTS stream processing timeout".to_string()))
.and_then(|r| r.map_err(|e| LLMError::Other(format!("TTS stream task failed: {}", e))));
Ok(())
}
#[cfg(not(feature = "kokoro"))]
{
let mut text_stream = self.chat_stream_with_session(session_id, message).await?;
let mut buffer = SentenceBuffer::new();
let mut sentences = Vec::new();
while let Some(result) = text_stream.next().await {
match result {
Ok(text_chunk) => {
print!("{}", text_chunk);
std::io::stdout().flush().map_err(|e| {
LLMError::Other(format!("Failed to flush stdout: {}", e))
})?;
if let Some(sentence) = buffer.push(&text_chunk) {
sentences.push(sentence);
}
}
Err(e) if e.to_string().contains("__stream_end__") => break,
Err(e) => return Err(e),
}
}
if let Some(remaining) = buffer.flush() {
sentences.push(remaining);
}
if !sentences.is_empty()
&& let Some(cb) = callback
{
for sentence in &sentences {
println!("\n[TTS] {}", sentence);
}
let _ = cb;
}
Ok(())
}
}
async fn get_session_arc(&self, session_id: &str) -> LLMResult<Arc<RwLock<ChatSession>>> {
let sessions = self.sessions.read().await;
sessions
.get(session_id)
.cloned()
.ok_or_else(|| LLMError::Other(format!("Session '{}' not found", session_id)))
}
pub async fn chat(&self, message: impl Into<String>) -> LLMResult<String> {
let session_id = self.active_session_id.read().await.clone();
self.chat_with_session(&session_id, message).await
}
pub async fn chat_with_session(
&self,
session_id: &str,
message: impl Into<String>,
) -> LLMResult<String> {
let message = message.into();
let model = self.provider.default_model();
let processed_message = if let Some(ref handler) = self.event_handler {
match handler.before_chat_with_model(&message, model).await? {
Some(msg) => msg,
None => return Ok(String::new()),
}
} else {
message
};
let session = self.get_session_arc(session_id).await?;
let mut session_guard = session.write().await;
let response = match session_guard.send(&processed_message).await {
Ok(resp) => resp,
Err(e) => {
if let Some(ref handler) = self.event_handler
&& let Some(fallback) = handler.on_error(&e).await?
{
return Ok(fallback);
}
return Err(e);
}
};
let final_response = if let Some(ref handler) = self.event_handler {
let metadata = session_guard.last_response_metadata();
if let Some(meta) = metadata {
match handler.after_chat_with_metadata(&response, meta).await? {
Some(resp) => resp,
None => response,
}
} else {
match handler.after_chat(&response).await? {
Some(resp) => resp,
None => response,
}
}
} else {
response
};
Ok(final_response)
}
pub async fn ask(&self, question: impl Into<String>) -> LLMResult<String> {
let question = question.into();
let mut builder = self.client.chat();
let mut system_prompt = self.config.system_prompt.clone();
if let Some(ref plugin) = self.prompt_plugin
&& let Some(template) = plugin.get_current_template().await
{
match template.render(&[]) {
Ok(prompt) => system_prompt = Some(prompt),
Err(_) => {
system_prompt = self.config.system_prompt.clone();
}
}
}
if let Some(ref system) = system_prompt {
builder = builder.system(system.clone());
}
if let Some(temp) = self.config.temperature {
builder = builder.temperature(temp);
}
if let Some(tokens) = self.config.max_tokens {
builder = builder.max_tokens(tokens);
}
builder = builder.user(question);
if let Some(ref executor) = self.tool_executor {
let tools = if self.tools.is_empty() {
executor.available_tools().await?
} else {
self.tools.clone()
};
if !tools.is_empty() {
builder = builder.tools(tools);
}
builder = builder.with_tool_executor(executor.clone());
let response = builder.send_with_tools().await?;
return response
.content()
.map(|s| s.to_string())
.ok_or_else(|| LLMError::Other("No content in response".to_string()));
}
let response = builder.send().await?;
response
.content()
.map(|s| s.to_string())
.ok_or_else(|| LLMError::Other("No content in response".to_string()))
}
pub async fn set_prompt_scenario(&self, scenario: impl Into<String>) {
let scenario = scenario.into();
if let Some(ref plugin) = self.prompt_plugin {
plugin.set_active_scenario(&scenario).await;
}
}
pub async fn clear_history(&self) {
let session_id = self.active_session_id.read().await.clone();
let _ = self.clear_session_history(&session_id).await;
}
pub async fn clear_session_history(&self, session_id: &str) -> LLMResult<()> {
let session = self.get_session_arc(session_id).await?;
let mut session_guard = session.write().await;
session_guard.clear();
Ok(())
}
pub async fn history(&self) -> Vec<ChatMessage> {
let session_id = self.active_session_id.read().await.clone();
self.get_session_history(&session_id)
.await
.unwrap_or_default()
}
pub async fn get_session_history(&self, session_id: &str) -> LLMResult<Vec<ChatMessage>> {
let session = self.get_session_arc(session_id).await?;
let session_guard = session.read().await;
Ok(session_guard.messages().to_vec())
}
pub fn set_tools(&mut self, tools: Vec<Tool>, executor: Arc<dyn ToolExecutor>) {
self.tools = tools;
self.tool_executor = Some(executor);
}
pub fn set_event_handler(&mut self, handler: Box<dyn LLMAgentEventHandler>) {
self.event_handler = Some(handler);
}
pub fn add_plugin<P: AgentPlugin + 'static>(&mut self, plugin: P) {
self.plugins.push(Box::new(plugin));
}
pub fn add_plugins(&mut self, plugins: Vec<Box<dyn AgentPlugin>>) {
self.plugins.extend(plugins);
}
pub async fn ask_stream(&self, question: impl Into<String>) -> LLMResult<TextStream> {
let question = question.into();
let mut builder = self.client.chat();
if let Some(ref system) = self.config.system_prompt {
builder = builder.system(system.clone());
}
if let Some(temp) = self.config.temperature {
builder = builder.temperature(temp);
}
if let Some(tokens) = self.config.max_tokens {
builder = builder.max_tokens(tokens);
}
builder = builder.user(question);
let chunk_stream = builder.send_stream().await?;
Ok(Self::chunk_stream_to_text_stream(chunk_stream))
}
pub async fn chat_stream(&self, message: impl Into<String>) -> LLMResult<TextStream> {
let session_id = self.active_session_id.read().await.clone();
self.chat_stream_with_session(&session_id, message).await
}
pub async fn chat_stream_with_session(
&self,
session_id: &str,
message: impl Into<String>,
) -> LLMResult<TextStream> {
let message = message.into();
let model = self.provider.default_model();
let processed_message = if let Some(ref handler) = self.event_handler {
match handler.before_chat_with_model(&message, model).await? {
Some(msg) => msg,
None => return Ok(Box::pin(futures::stream::empty())),
}
} else {
message
};
let session = self.get_session_arc(session_id).await?;
let history = {
let session_guard = session.read().await;
session_guard.messages().to_vec()
};
let mut builder = self.client.chat();
if let Some(ref system) = self.config.system_prompt {
builder = builder.system(system.clone());
}
if let Some(temp) = self.config.temperature {
builder = builder.temperature(temp);
}
if let Some(tokens) = self.config.max_tokens {
builder = builder.max_tokens(tokens);
}
builder = builder.messages(history);
builder = builder.user(processed_message.clone());
let chunk_stream = builder.send_stream().await?;
{
let mut session_guard = session.write().await;
session_guard
.messages_mut()
.push(ChatMessage::user(&processed_message));
}
let event_handler = self.event_handler.clone().map(Arc::new);
let wrapped_stream =
Self::create_history_updating_stream(chunk_stream, session, event_handler);
Ok(wrapped_stream)
}
pub async fn ask_stream_raw(&self, question: impl Into<String>) -> LLMResult<ChatStream> {
let question = question.into();
let mut builder = self.client.chat();
if let Some(ref system) = self.config.system_prompt {
builder = builder.system(system.clone());
}
if let Some(temp) = self.config.temperature {
builder = builder.temperature(temp);
}
if let Some(tokens) = self.config.max_tokens {
builder = builder.max_tokens(tokens);
}
builder = builder.user(question);
builder.send_stream().await
}
pub async fn chat_stream_with_full(
&self,
message: impl Into<String>,
) -> LLMResult<(TextStream, tokio::sync::oneshot::Receiver<String>)> {
let session_id = self.active_session_id.read().await.clone();
self.chat_stream_with_full_session(&session_id, message)
.await
}
pub async fn chat_stream_with_full_session(
&self,
session_id: &str,
message: impl Into<String>,
) -> LLMResult<(TextStream, tokio::sync::oneshot::Receiver<String>)> {
let message = message.into();
let model = self.provider.default_model();
let processed_message = if let Some(ref handler) = self.event_handler {
match handler.before_chat_with_model(&message, model).await? {
Some(msg) => msg,
None => {
let (tx, rx) = tokio::sync::oneshot::channel();
let _ = tx.send(String::new());
return Ok((Box::pin(futures::stream::empty()), rx));
}
}
} else {
message
};
let session = self.get_session_arc(session_id).await?;
let history = {
let session_guard = session.read().await;
session_guard.messages().to_vec()
};
let mut builder = self.client.chat();
if let Some(ref system) = self.config.system_prompt {
builder = builder.system(system.clone());
}
if let Some(temp) = self.config.temperature {
builder = builder.temperature(temp);
}
if let Some(tokens) = self.config.max_tokens {
builder = builder.max_tokens(tokens);
}
builder = builder.messages(history);
builder = builder.user(processed_message.clone());
let chunk_stream = builder.send_stream().await?;
{
let mut session_guard = session.write().await;
session_guard
.messages_mut()
.push(ChatMessage::user(&processed_message));
}
let (tx, rx) = tokio::sync::oneshot::channel();
let event_handler = self.event_handler.clone().map(Arc::new);
let wrapped_stream =
Self::create_collecting_stream(chunk_stream, session, tx, event_handler);
Ok((wrapped_stream, rx))
}
fn chunk_stream_to_text_stream(chunk_stream: ChatStream) -> TextStream {
use futures::StreamExt;
let text_stream = chunk_stream.filter_map(|result| async move {
match result {
Ok(chunk) => {
if let Some(choice) = chunk.choices.first()
&& let Some(ref content) = choice.delta.content
&& !content.is_empty()
{
return Some(Ok(content.clone()));
}
None
}
Err(e) => Some(Err(e)),
}
});
Box::pin(text_stream)
}
fn create_history_updating_stream(
chunk_stream: ChatStream,
session: Arc<RwLock<ChatSession>>,
event_handler: Option<Arc<Box<dyn LLMAgentEventHandler>>>,
) -> TextStream {
use super::types::LLMResponseMetadata;
let collected = Arc::new(tokio::sync::Mutex::new(String::new()));
let collected_clone = collected.clone();
let event_handler_clone = event_handler.clone();
let metadata_collected = Arc::new(tokio::sync::Mutex::new(None::<LLMResponseMetadata>));
let metadata_collected_clone = metadata_collected.clone();
let stream = chunk_stream.filter_map(move |result| {
let collected = collected.clone();
let event_handler = event_handler.clone();
let metadata_collected = metadata_collected.clone();
async move {
match result {
Ok(chunk) => {
if let Some(choice) = chunk.choices.first() {
if choice.finish_reason.is_some() {
let metadata = LLMResponseMetadata::from(&chunk);
*metadata_collected.lock().await = Some(metadata);
return None;
}
if let Some(ref content) = choice.delta.content
&& !content.is_empty()
{
let mut collected = collected.lock().await;
collected.push_str(content);
return Some(Ok(content.clone()));
}
}
None
}
Err(e) => {
if let Some(handler) = event_handler {
let _ = handler.on_error(&e).await;
}
Some(Err(e))
}
}
}
});
let stream = stream
.chain(futures::stream::once(async move {
let full_response = collected_clone.lock().await.clone();
let metadata = metadata_collected_clone.lock().await.clone();
if !full_response.is_empty() {
let mut session = session.write().await;
session
.messages_mut()
.push(ChatMessage::assistant(&full_response));
let window_size = session.context_window_size();
if window_size.is_some() {
let current_messages = session.messages().to_vec();
*session.messages_mut() = ChatSession::apply_sliding_window_static(
¤t_messages,
window_size,
);
}
if let Some(handler) = event_handler_clone {
if let Some(meta) = &metadata {
let _ = handler.after_chat_with_metadata(&full_response, meta).await;
} else {
let _ = handler.after_chat(&full_response).await;
}
}
}
Err(LLMError::Other("__stream_end__".to_string()))
}))
.filter_map(|result| async move {
match result {
Ok(s) => Some(Ok(s)),
Err(e) if e.to_string() == "__stream_end__" => None,
Err(e) => Some(Err(e)),
}
});
Box::pin(stream)
}
fn create_collecting_stream(
chunk_stream: ChatStream,
session: Arc<RwLock<ChatSession>>,
tx: tokio::sync::oneshot::Sender<String>,
event_handler: Option<Arc<Box<dyn LLMAgentEventHandler>>>,
) -> TextStream {
use super::types::LLMResponseMetadata;
use futures::StreamExt;
let collected = Arc::new(tokio::sync::Mutex::new(String::new()));
let collected_clone = collected.clone();
let event_handler_clone = event_handler.clone();
let metadata_collected = Arc::new(tokio::sync::Mutex::new(None::<LLMResponseMetadata>));
let metadata_collected_clone = metadata_collected.clone();
let stream = chunk_stream.filter_map(move |result| {
let collected = collected.clone();
let event_handler = event_handler.clone();
let metadata_collected = metadata_collected.clone();
async move {
match result {
Ok(chunk) => {
if let Some(choice) = chunk.choices.first() {
if choice.finish_reason.is_some() {
let metadata = LLMResponseMetadata::from(&chunk);
*metadata_collected.lock().await = Some(metadata);
return None;
}
if let Some(ref content) = choice.delta.content
&& !content.is_empty()
{
let mut collected = collected.lock().await;
collected.push_str(content);
return Some(Ok(content.clone()));
}
}
None
}
Err(e) => {
if let Some(handler) = event_handler {
let _ = handler.on_error(&e).await;
}
Some(Err(e))
}
}
}
});
let stream = stream
.chain(futures::stream::once(async move {
let full_response = collected_clone.lock().await.clone();
let mut processed_response = full_response.clone();
let metadata = metadata_collected_clone.lock().await.clone();
if !full_response.is_empty() {
let mut session = session.write().await;
session
.messages_mut()
.push(ChatMessage::assistant(&processed_response));
let window_size = session.context_window_size();
if window_size.is_some() {
let current_messages = session.messages().to_vec();
*session.messages_mut() = ChatSession::apply_sliding_window_static(
¤t_messages,
window_size,
);
}
if let Some(handler) = event_handler_clone {
if let Some(meta) = &metadata {
if let Ok(Some(resp)) = handler
.after_chat_with_metadata(&processed_response, meta)
.await
{
processed_response = resp;
}
} else if let Ok(Some(resp)) = handler.after_chat(&processed_response).await
{
processed_response = resp;
}
}
}
let _ = tx.send(processed_response);
Err(LLMError::Other("__stream_end__".to_string()))
}))
.filter_map(|result| async move {
match result {
Ok(s) => Some(Ok(s)),
Err(e) if e.to_string() == "__stream_end__" => None,
Err(e) => Some(Err(e)),
}
});
Box::pin(stream)
}
}
pub struct LLMAgentBuilder {
agent_id: String,
name: Option<String>,
provider: Option<Arc<dyn LLMProvider>>,
system_prompt: Option<String>,
temperature: Option<f32>,
max_tokens: Option<u32>,
tools: Vec<Tool>,
tool_executor: Option<Arc<dyn ToolExecutor>>,
event_handler: Option<Box<dyn LLMAgentEventHandler>>,
plugins: Vec<Box<dyn AgentPlugin>>,
custom_config: HashMap<String, String>,
prompt_plugin: Option<Box<dyn prompt::PromptTemplatePlugin>>,
session_id: Option<String>,
user_id: Option<String>,
tenant_id: Option<String>,
context_window_size: Option<usize>,
message_store: Option<Arc<dyn crate::persistence::MessageStore + Send + Sync>>,
session_store: Option<Arc<dyn crate::persistence::SessionStore + Send + Sync>>,
persistence_user_id: Option<uuid::Uuid>,
persistence_tenant_id: Option<uuid::Uuid>,
persistence_agent_id: Option<uuid::Uuid>,
}
impl LLMAgentBuilder {
pub fn new() -> Self {
Self {
agent_id: uuid::Uuid::now_v7().to_string(),
name: None,
provider: None,
system_prompt: None,
temperature: None,
max_tokens: None,
tools: Vec::new(),
tool_executor: None,
event_handler: None,
plugins: Vec::new(),
custom_config: HashMap::new(),
prompt_plugin: None,
session_id: None,
user_id: None,
tenant_id: None,
context_window_size: None,
message_store: None,
session_store: None,
persistence_user_id: None,
persistence_tenant_id: None,
persistence_agent_id: None,
}
}
pub fn with_id(mut self, id: impl Into<String>) -> Self {
self.agent_id = id.into();
self
}
pub fn with_name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
pub fn with_provider(mut self, provider: Arc<dyn LLMProvider>) -> Self {
self.provider = Some(provider);
self
}
pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.system_prompt = Some(prompt.into());
self
}
pub fn with_temperature(mut self, temperature: f32) -> Self {
self.temperature = Some(temperature);
self
}
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
self.max_tokens = Some(max_tokens);
self
}
pub fn with_tool(mut self, tool: Tool) -> Self {
self.tools.push(tool);
self
}
pub fn with_tools(mut self, tools: Vec<Tool>) -> Self {
self.tools = tools;
self
}
pub fn with_tool_executor(mut self, executor: Arc<dyn ToolExecutor>) -> Self {
self.tool_executor = Some(executor);
self
}
pub fn with_event_handler(mut self, handler: Box<dyn LLMAgentEventHandler>) -> Self {
self.event_handler = Some(handler);
self
}
pub fn with_plugin(mut self, plugin: impl AgentPlugin + 'static) -> Self {
self.plugins.push(Box::new(plugin));
self
}
pub fn with_plugins(mut self, plugins: Vec<Box<dyn AgentPlugin>>) -> Self {
self.plugins.extend(plugins);
self
}
pub fn with_persistence_plugin(
mut self,
plugin: crate::persistence::PersistencePlugin,
) -> Self {
self.message_store = Some(plugin.message_store());
self.session_store = plugin.session_store();
self.persistence_user_id = Some(plugin.user_id());
self.persistence_tenant_id = Some(plugin.tenant_id());
self.persistence_agent_id = Some(plugin.agent_id());
let plugin_box: Box<dyn AgentPlugin> = Box::new(plugin.clone());
let event_handler: Box<dyn LLMAgentEventHandler> = Box::new(plugin);
self.plugins.push(plugin_box);
self.event_handler = Some(event_handler);
self
}
pub fn with_prompt_plugin(
mut self,
plugin: impl prompt::PromptTemplatePlugin + 'static,
) -> Self {
self.prompt_plugin = Some(Box::new(plugin));
self
}
pub fn with_hot_reload_prompt_plugin(
mut self,
plugin: prompt::HotReloadableRhaiPromptPlugin,
) -> Self {
self.prompt_plugin = Some(Box::new(plugin));
self
}
pub fn with_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.custom_config.insert(key.into(), value.into());
self
}
pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
self.session_id = Some(session_id.into());
self
}
pub fn with_user(mut self, user_id: impl Into<String>) -> Self {
self.user_id = Some(user_id.into());
self
}
pub fn with_tenant(mut self, tenant_id: impl Into<String>) -> Self {
self.tenant_id = Some(tenant_id.into());
self
}
pub fn with_sliding_window(mut self, size: usize) -> Self {
self.context_window_size = Some(size);
self
}
pub fn from_env() -> LLMResult<Self> {
use super::openai::{OpenAIConfig, OpenAIProvider};
let api_key = std::env::var("OPENAI_API_KEY").map_err(|_| {
LLMError::ConfigError("OPENAI_API_KEY environment variable not set".to_string())
})?;
let mut config = OpenAIConfig::new(api_key);
if let Ok(base_url) = std::env::var("OPENAI_BASE_URL") {
config = config.with_base_url(&base_url);
}
if let Ok(model) = std::env::var("OPENAI_MODEL") {
config = config.with_model(&model);
}
Ok(Self::new()
.with_provider(Arc::new(OpenAIProvider::with_config(config)))
.with_temperature(0.7)
.with_max_tokens(4096))
}
pub fn build(self) -> LLMAgent {
let provider = self
.provider
.expect("LLM provider must be set before building");
let config = LLMAgentConfig {
agent_id: self.agent_id.clone(),
name: self.name.unwrap_or_else(|| self.agent_id.clone()),
system_prompt: self.system_prompt,
temperature: self.temperature,
max_tokens: self.max_tokens,
custom_config: self.custom_config,
user_id: self.user_id,
tenant_id: self.tenant_id,
context_window_size: self.context_window_size,
};
let mut agent = LLMAgent::with_initial_session(config, provider, self.session_id);
agent.prompt_plugin = self.prompt_plugin;
if let Some(executor) = self.tool_executor {
agent.set_tools(self.tools, executor);
}
if let Some(handler) = self.event_handler {
agent.set_event_handler(handler);
}
let mut plugins = self.plugins;
let mut tts_plugin = None;
for i in (0..plugins.len()).rev() {
if plugins[i].as_any().is::<mofa_plugins::tts::TTSPlugin>() {
let plugin = plugins.remove(i);
if let Ok(tts) = plugin.into_any().downcast::<mofa_plugins::tts::TTSPlugin>() {
tts_plugin = Some(Arc::new(Mutex::new(*tts)));
}
}
}
agent.add_plugins(plugins);
agent.tts_plugin = tts_plugin;
agent
}
pub fn try_build(self) -> LLMResult<LLMAgent> {
let provider = self
.provider
.ok_or_else(|| LLMError::ConfigError("LLM provider not set".to_string()))?;
let config = LLMAgentConfig {
agent_id: self.agent_id.clone(),
name: self.name.unwrap_or_else(|| self.agent_id.clone()),
system_prompt: self.system_prompt,
temperature: self.temperature,
max_tokens: self.max_tokens,
custom_config: self.custom_config,
user_id: self.user_id,
tenant_id: self.tenant_id,
context_window_size: self.context_window_size,
};
let mut agent = LLMAgent::with_initial_session(config, provider, self.session_id);
if let Some(executor) = self.tool_executor {
agent.set_tools(self.tools, executor);
}
if let Some(handler) = self.event_handler {
agent.set_event_handler(handler);
}
let mut plugins = self.plugins;
let mut tts_plugin = None;
for i in (0..plugins.len()).rev() {
if plugins[i].as_any().is::<mofa_plugins::tts::TTSPlugin>() {
let plugin = plugins.remove(i);
if let Ok(tts) = plugin.into_any().downcast::<mofa_plugins::tts::TTSPlugin>() {
tts_plugin = Some(Arc::new(Mutex::new(*tts)));
}
}
}
agent.add_plugins(plugins);
agent.tts_plugin = tts_plugin;
Ok(agent)
}
pub async fn build_async(mut self) -> LLMAgent {
let provider = self
.provider
.expect("LLM provider must be set before building");
let tenant_id_for_persistence = self.tenant_id.clone();
let config = LLMAgentConfig {
agent_id: self.agent_id.clone(),
name: self.name.unwrap_or_else(|| self.agent_id.clone()),
system_prompt: self.system_prompt,
temperature: self.temperature,
max_tokens: self.max_tokens,
custom_config: self.custom_config,
user_id: self.user_id,
tenant_id: self.tenant_id,
context_window_size: self.context_window_size,
};
let persistence_tenant_id = if self.session_store.is_some()
&& self.persistence_tenant_id.is_none()
&& tenant_id_for_persistence.is_some()
{
uuid::Uuid::parse_str(&tenant_id_for_persistence.unwrap()).ok()
} else {
self.persistence_tenant_id
};
let mut agent = LLMAgent::with_initial_session_async(
config,
provider,
self.session_id,
self.message_store,
self.session_store,
self.persistence_user_id,
persistence_tenant_id,
self.persistence_agent_id,
)
.await;
agent.prompt_plugin = self.prompt_plugin;
if self.tools.is_empty() {
if let Some(executor) = self.tool_executor.as_ref() {
if let Ok(tools) = executor.available_tools().await {
self.tools = tools;
}
}
}
if let Some(executor) = self.tool_executor {
agent.set_tools(self.tools, executor);
}
let mut plugins = self.plugins;
let mut tts_plugin = None;
let history_loaded_from_plugin = false;
for i in (0..plugins.len()).rev() {
if plugins[i].as_any().is::<mofa_plugins::tts::TTSPlugin>() {
let plugin = plugins.remove(i);
if let Ok(tts) = plugin.into_any().downcast::<mofa_plugins::tts::TTSPlugin>() {
tts_plugin = Some(Arc::new(Mutex::new(*tts)));
}
}
}
if !history_loaded_from_plugin {
for plugin in &plugins {
if plugin.metadata().plugin_type == PluginType::Storage
&& plugin
.metadata()
.capabilities
.contains(&"message_persistence".to_string())
{
tracing::info!("📦 检测到持久化插件,将在 agent 初始化后加载历史");
break;
}
}
}
agent.add_plugins(plugins);
agent.tts_plugin = tts_plugin;
if let Some(handler) = self.event_handler {
agent.set_event_handler(handler);
}
agent
}
}
impl LLMAgentBuilder {
pub fn from_config_file(path: impl AsRef<std::path::Path>) -> LLMResult<Self> {
let config = crate::config::AgentYamlConfig::from_file(path)
.map_err(|e| LLMError::ConfigError(e.to_string()))?;
Self::from_yaml_config(config)
}
pub fn from_yaml_config(config: crate::config::AgentYamlConfig) -> LLMResult<Self> {
let mut builder = Self::new()
.with_id(&config.agent.id)
.with_name(&config.agent.name);
if let Some(llm_config) = config.llm {
let provider = create_provider_from_config(&llm_config)?;
builder = builder.with_provider(Arc::new(provider));
if let Some(temp) = llm_config.temperature {
builder = builder.with_temperature(temp);
}
if let Some(tokens) = llm_config.max_tokens {
builder = builder.with_max_tokens(tokens);
}
if let Some(prompt) = llm_config.system_prompt {
builder = builder.with_system_prompt(prompt);
}
}
Ok(builder)
}
#[cfg(feature = "persistence-postgres")]
pub async fn from_database<S>(store: &S, agent_code: &str) -> LLMResult<Self>
where
S: crate::persistence::AgentStore + Send + Sync,
{
let config = store
.get_agent_by_code_with_provider(agent_code)
.await
.map_err(|e| LLMError::Other(format!("Failed to load agent from database: {}", e)))?
.ok_or_else(|| {
LLMError::Other(format!(
"Agent with code '{}' not found in database",
agent_code
))
})?;
Self::from_agent_config(&config)
}
#[cfg(feature = "persistence-postgres")]
pub async fn from_database_with_tenant<S>(
store: &S,
tenant_id: uuid::Uuid,
agent_code: &str,
) -> LLMResult<Self>
where
S: crate::persistence::AgentStore + Send + Sync,
{
let config = store
.get_agent_by_code_and_tenant_with_provider(tenant_id, agent_code)
.await
.map_err(|e| LLMError::Other(format!("Failed to load agent from database: {}", e)))?
.ok_or_else(|| {
LLMError::Other(format!(
"Agent with code '{}' not found for tenant {}",
agent_code, tenant_id
))
})?;
Self::from_agent_config(&config)
}
#[cfg(feature = "persistence-postgres")]
pub async fn with_database_agent<S>(store: &S, agent_code: &str) -> LLMResult<Self>
where
S: crate::persistence::AgentStore + Send + Sync,
{
Self::from_database(store, agent_code).await
}
#[cfg(feature = "persistence-postgres")]
pub fn from_agent_config(config: &crate::persistence::AgentConfig) -> LLMResult<Self> {
use super::openai::{OpenAIConfig, OpenAIProvider};
let agent = &config.agent;
let provider = &config.provider;
if !agent.agent_status {
return Err(LLMError::Other(format!(
"Agent '{}' is disabled (agent_status = false)",
agent.agent_code
)));
}
if !provider.enabled {
return Err(LLMError::Other(format!(
"Provider '{}' is disabled (enabled = false)",
provider.provider_name
)));
}
let llm_provider: Arc<dyn super::LLMProvider> = match provider.provider_type.as_str() {
"openai" | "azure" | "compatible" | "local" => {
let mut openai_config = OpenAIConfig::new(provider.api_key.clone());
openai_config = openai_config.with_base_url(&provider.api_base);
openai_config = openai_config.with_model(&agent.model_name);
if let Some(temp) = agent.temperature {
openai_config = openai_config.with_temperature(temp);
}
if let Some(max_tokens) = agent.max_completion_tokens {
openai_config = openai_config.with_max_tokens(max_tokens as u32);
}
Arc::new(OpenAIProvider::with_config(openai_config))
}
"anthropic" => {
let mut cfg = AnthropicConfig::new(provider.api_key.clone());
cfg = cfg.with_base_url(&provider.api_base);
cfg = cfg.with_model(&agent.model_name);
if let Some(temp) = agent.temperature {
cfg = cfg.with_temperature(temp);
}
if let Some(tokens) = agent.max_completion_tokens {
cfg = cfg.with_max_tokens(tokens as u32);
}
Arc::new(AnthropicProvider::with_config(cfg))
}
"gemini" => {
let mut cfg = GeminiConfig::new(provider.api_key.clone());
cfg = cfg.with_base_url(&provider.api_base);
cfg = cfg.with_model(&agent.model_name);
if let Some(temp) = agent.temperature {
cfg = cfg.with_temperature(temp);
}
if let Some(tokens) = agent.max_completion_tokens {
cfg = cfg.with_max_tokens(tokens as u32);
}
Arc::new(GeminiProvider::with_config(cfg))
}
"ollama" => {
let mut ollama_config = OllamaConfig::new();
ollama_config = ollama_config.with_base_url(&provider.api_base);
ollama_config = ollama_config.with_model(&agent.model_name);
if let Some(temp) = agent.temperature {
ollama_config = ollama_config.with_temperature(temp);
}
if let Some(max_tokens) = agent.max_completion_tokens {
ollama_config = ollama_config.with_max_tokens(max_tokens as u32);
}
Arc::new(OllamaProvider::with_config(ollama_config))
}
other => {
return Err(LLMError::Other(format!(
"Unsupported provider type: {}",
other
)));
}
};
let mut builder = Self::new()
.with_id(agent.id.clone())
.with_name(agent.agent_name.clone())
.with_provider(llm_provider)
.with_system_prompt(agent.system_prompt.clone())
.with_tenant(agent.tenant_id.to_string());
if let Some(temp) = agent.temperature {
builder = builder.with_temperature(temp);
}
if let Some(tokens) = agent.max_completion_tokens {
builder = builder.with_max_tokens(tokens as u32);
}
if let Some(limit) = agent.context_limit {
builder = builder.with_sliding_window(limit as usize);
}
if let Some(ref params) = agent.custom_params {
if let Some(obj) = params.as_object() {
for (key, value) in obj.iter() {
let value_str: String = match value {
serde_json::Value::String(s) => s.clone(),
serde_json::Value::Bool(b) => b.to_string(),
serde_json::Value::Number(n) => n.to_string(),
_ => value.to_string(),
};
builder = builder.with_config(key.as_str(), value_str);
}
}
}
if let Some(ref format) = agent.response_format {
builder = builder.with_config("response_format", format);
}
if let Some(stream) = agent.stream {
builder = builder.with_config("stream", if stream { "true" } else { "false" });
}
Ok(builder)
}
}
fn create_provider_from_config(
config: &crate::config::LLMYamlConfig,
) -> LLMResult<super::openai::OpenAIProvider> {
use super::openai::{OpenAIConfig, OpenAIProvider};
match config.provider.as_str() {
"openai" => {
let api_key = config
.api_key
.clone()
.or_else(|| std::env::var("OPENAI_API_KEY").ok())
.ok_or_else(|| LLMError::ConfigError("OpenAI API key not set".to_string()))?;
let mut openai_config = OpenAIConfig::new(api_key);
if let Some(ref model) = config.model {
openai_config = openai_config.with_model(model);
}
if let Some(ref base_url) = config.base_url {
openai_config = openai_config.with_base_url(base_url);
}
if let Some(temp) = config.temperature {
openai_config = openai_config.with_temperature(temp);
}
if let Some(tokens) = config.max_tokens {
openai_config = openai_config.with_max_tokens(tokens);
}
Ok(OpenAIProvider::with_config(openai_config))
}
"azure" => {
let endpoint = config.base_url.clone().ok_or_else(|| {
LLMError::ConfigError("Azure endpoint (base_url) not set".to_string())
})?;
let api_key = config
.api_key
.clone()
.or_else(|| std::env::var("AZURE_OPENAI_API_KEY").ok())
.ok_or_else(|| LLMError::ConfigError("Azure API key not set".to_string()))?;
let deployment = config
.deployment
.clone()
.or_else(|| config.model.clone())
.ok_or_else(|| {
LLMError::ConfigError("Azure deployment name not set".to_string())
})?;
Ok(OpenAIProvider::azure(endpoint, api_key, deployment))
}
"compatible" | "local" => {
let base_url = config.base_url.clone().ok_or_else(|| {
LLMError::ConfigError("base_url not set for compatible provider".to_string())
})?;
let model = config
.model
.clone()
.unwrap_or_else(|| "default".to_string());
Ok(OpenAIProvider::local(base_url, model))
}
other => Err(LLMError::ConfigError(format!(
"Unknown provider: {}",
other
))),
}
}
#[async_trait::async_trait]
impl mofa_kernel::agent::MoFAAgent for LLMAgent {
fn id(&self) -> &str {
&self.metadata.id
}
fn name(&self) -> &str {
&self.metadata.name
}
fn capabilities(&self) -> &mofa_kernel::agent::AgentCapabilities {
use mofa_kernel::agent::AgentCapabilities;
static CAPABILITIES: std::sync::OnceLock<AgentCapabilities> = std::sync::OnceLock::new();
CAPABILITIES.get_or_init(|| {
AgentCapabilities::builder()
.tag("llm")
.tag("chat")
.tag("text-generation")
.input_type(mofa_kernel::agent::InputType::Text)
.output_type(mofa_kernel::agent::OutputType::Text)
.supports_streaming(true)
.supports_tools(true)
.build()
})
}
async fn initialize(
&mut self,
ctx: &mofa_kernel::agent::AgentContext,
) -> mofa_kernel::agent::AgentResult<()> {
let mut plugin_config = mofa_kernel::plugin::PluginConfig::new();
for (k, v) in &self.config.custom_config {
plugin_config.set(k, v);
}
if let Some(user_id) = &self.config.user_id {
plugin_config.set("user_id", user_id);
}
if let Some(tenant_id) = &self.config.tenant_id {
plugin_config.set("tenant_id", tenant_id);
}
let session_id = self.active_session_id.read().await.clone();
plugin_config.set("session_id", session_id);
let plugin_ctx =
mofa_kernel::plugin::PluginContext::new(self.id()).with_config(plugin_config);
for plugin in &mut self.plugins {
plugin
.load(&plugin_ctx)
.await
.map_err(|e| mofa_kernel::agent::AgentError::InitializationFailed(e.to_string()))?;
plugin
.init_plugin()
.await
.map_err(|e| mofa_kernel::agent::AgentError::InitializationFailed(e.to_string()))?;
}
self.state = mofa_kernel::agent::AgentState::Ready;
let _ = ctx;
Ok(())
}
async fn execute(
&mut self,
input: mofa_kernel::agent::AgentInput,
_ctx: &mofa_kernel::agent::AgentContext,
) -> mofa_kernel::agent::AgentResult<mofa_kernel::agent::AgentOutput> {
use mofa_kernel::agent::{AgentError, AgentInput, AgentOutput};
let message = match input {
AgentInput::Text(text) => text,
AgentInput::Json(json) => json.to_string(),
_ => {
return Err(AgentError::ValidationFailed(
"Unsupported input type for LLMAgent".to_string(),
));
}
};
let response = self
.chat(&message)
.await
.map_err(|e| AgentError::ExecutionFailed(format!("LLM chat failed: {}", e)))?;
Ok(AgentOutput::text(response))
}
async fn shutdown(&mut self) -> mofa_kernel::agent::AgentResult<()> {
for plugin in &mut self.plugins {
plugin
.unload()
.await
.map_err(|e| mofa_kernel::agent::AgentError::ShutdownFailed(e.to_string()))?;
}
self.state = mofa_kernel::agent::AgentState::Shutdown;
Ok(())
}
fn state(&self) -> mofa_kernel::agent::AgentState {
self.state.clone()
}
}
pub fn simple_llm_agent(
agent_id: impl Into<String>,
provider: Arc<dyn LLMProvider>,
system_prompt: impl Into<String>,
) -> LLMAgent {
LLMAgentBuilder::new()
.with_id(agent_id)
.with_provider(provider)
.with_system_prompt(system_prompt)
.build()
}
pub fn agent_from_config(path: impl AsRef<std::path::Path>) -> LLMResult<LLMAgent> {
LLMAgentBuilder::from_config_file(path)?.try_build()
}