use async_trait::async_trait;
use rucora_core::agent::{Agent, AgentContext, AgentDecision, AgentError, AgentInput, AgentOutput};
use rucora_core::provider::LlmProvider;
use rucora_core::provider::types::{ChatMessage, ChatRequest, LlmParams, Role};
use rucora_core::tool::Tool;
use rucora_core::tool::types::ToolDefinition;
use std::sync::Arc;
use tokio::sync::Mutex;
use crate::agent::ToolRegistry;
use crate::agent::execution::{build_default_execution, DefaultExecution};
use crate::agent::tool_call_config::ToolCallEnhancedConfig;
use crate::conversation::ConversationManager;
pub struct ToolAgent<P> {
provider: Arc<P>,
model: String,
system_prompt: Option<String>,
tools: ToolRegistry,
_max_steps: usize,
conversation_manager: Option<Arc<Mutex<ConversationManager>>>,
llm_params: LlmParams,
execution: DefaultExecution,
}
#[async_trait]
impl<P> Agent for ToolAgent<P>
where
P: LlmProvider + Send + Sync + 'static,
{
async fn think(&self, context: &AgentContext) -> AgentDecision {
if !context.tool_results.is_empty() {
return AgentDecision::Chat {
request: Box::new(self._build_chat_request(context)),
};
}
AgentDecision::Chat {
request: Box::new(self._build_chat_request_with_tools(context)),
}
}
fn name(&self) -> &str {
"tool_agent"
}
fn description(&self) -> Option<&str> {
Some("工具调用 Agent,支持自动决定何时调用工具")
}
async fn run(&self, input: AgentInput) -> Result<AgentOutput, rucora_core::agent::AgentError> {
self.execution.run(self, input).await
}
fn run_stream(
&self,
input: AgentInput,
) -> futures_util::stream::BoxStream<
'static,
Result<rucora_core::channel::types::ChannelEvent, rucora_core::agent::AgentError>,
> {
self.execution.run_stream_simple(input)
}
}
impl<P> ToolAgent<P>
where
P: LlmProvider + Send + Sync + 'static,
{
pub async fn run_stream_text(
&self,
input: impl Into<AgentInput>,
) -> Result<String, rucora_core::agent::AgentError> {
self.execution.run_stream_text(input.into()).await
}
}
impl<P> ToolAgent<P>
where
P: LlmProvider,
{
#[must_use = "构建器必须调用 try_build() 来创建 Agent"]
pub fn builder() -> ToolAgentBuilder<P> {
ToolAgentBuilder::new()
}
fn _build_chat_request(&self, context: &AgentContext) -> ChatRequest {
let mut request = context.default_chat_request_with(&self.llm_params);
self._apply_config(&mut request);
request
}
fn _build_chat_request_with_tools(&self, context: &AgentContext) -> ChatRequest {
let mut request = context.default_chat_request_with(&self.llm_params);
let tool_defs = self._get_tool_definitions();
if !tool_defs.is_empty() {
request.tools = Some(tool_defs);
}
self._apply_config(&mut request);
request
}
fn _apply_config(&self, request: &mut ChatRequest) {
if let Some(ref prompt) = self.system_prompt
&& (request.messages.is_empty()
|| request.messages.first().map(|m| &m.role) != Some(&Role::System))
{
request
.messages
.insert(0, ChatMessage::system(prompt.clone()));
}
request.model = Some(self.model.clone());
}
fn _get_tool_definitions(&self) -> Vec<ToolDefinition> {
self.tools.definitions()
}
pub fn tools(&self) -> Vec<&str> {
self.tools
.tool_names()
.into_iter()
.map(|s| s.as_str())
.collect()
}
pub fn provider(&self) -> &P {
&self.provider
}
pub fn model(&self) -> &str {
&self.model
}
pub fn tool_registry(&self) -> &ToolRegistry {
&self.tools
}
pub async fn get_conversation_history(&self) -> Option<Vec<ChatMessage>> {
match &self.conversation_manager {
Some(conv_arc) => {
let conv = conv_arc.lock().await;
Some(conv.get_messages().to_vec())
}
None => None,
}
}
pub async fn clear_conversation(&self) {
if let Some(ref conv_arc) = self.conversation_manager {
let mut conv = conv_arc.lock().await;
conv.clear();
if let Some(ref prompt) = self.system_prompt {
conv.ensure_system_prompt(prompt);
}
}
}
}
pub struct ToolAgentBuilder<P> {
provider: Option<P>,
system_prompt: Option<String>,
model: Option<String>,
tools: ToolRegistry,
max_steps: usize,
max_tool_concurrency: usize,
with_conversation: bool,
middleware_chain: crate::middleware::MiddlewareChain,
enhanced_config: ToolCallEnhancedConfig,
llm_params: LlmParams,
}
impl<P> ToolAgentBuilder<P> {
pub fn new() -> Self {
Self {
provider: None,
system_prompt: None,
model: None,
tools: ToolRegistry::new(),
max_steps: 10,
max_tool_concurrency: 1,
with_conversation: false,
middleware_chain: crate::middleware::MiddlewareChain::new(),
enhanced_config: ToolCallEnhancedConfig::default(),
llm_params: LlmParams::default(),
}
}
}
impl<P> ToolAgentBuilder<P>
where
P: LlmProvider + Send + Sync + 'static,
{
pub fn provider(mut self, provider: P) -> Self {
self.provider = Some(provider);
self
}
pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.system_prompt = Some(prompt.into());
self
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
self.tools = self.tools.register(tool);
self
}
pub fn tools<I, T>(mut self, tools: I) -> Self
where
I: IntoIterator<Item = T>,
T: Tool + 'static,
{
for tool in tools {
self.tools = self.tools.register(tool);
}
self
}
pub fn tool_registry(mut self, registry: ToolRegistry) -> Self {
self.tools = registry;
self
}
pub fn max_steps(mut self, max: usize) -> Self {
self.max_steps = max;
self
}
pub fn max_tool_concurrency(mut self, max: usize) -> Self {
self.max_tool_concurrency = max.max(1);
self
}
pub fn with_conversation(mut self, enabled: bool) -> Self {
self.with_conversation = enabled;
self
}
pub fn with_middleware_chain(
mut self,
middleware_chain: crate::middleware::MiddlewareChain,
) -> Self {
self.middleware_chain = middleware_chain;
self
}
pub fn with_middleware<M: crate::middleware::Middleware + 'static>(
mut self,
middleware: M,
) -> Self {
self.middleware_chain = self.middleware_chain.with(middleware);
self
}
pub fn with_enhanced_config(mut self, config: ToolCallEnhancedConfig) -> Self {
self.enhanced_config = config;
self
}
pub fn llm_params(mut self, params: LlmParams) -> Self {
self.llm_params = params;
self
}
pub fn temperature(mut self, value: f32) -> Self {
self.llm_params.temperature = Some(value);
self
}
pub fn top_p(mut self, value: f32) -> Self {
self.llm_params.top_p = Some(value);
self
}
pub fn top_k(mut self, value: u32) -> Self {
self.llm_params.top_k = Some(value);
self
}
pub fn max_tokens(mut self, value: u32) -> Self {
self.llm_params.max_tokens = Some(value);
self
}
pub fn frequency_penalty(mut self, value: f32) -> Self {
self.llm_params.frequency_penalty = Some(value);
self
}
pub fn presence_penalty(mut self, value: f32) -> Self {
self.llm_params.presence_penalty = Some(value);
self
}
pub fn stop(mut self, value: Vec<String>) -> Self {
self.llm_params.stop = Some(value);
self
}
pub fn extra_params(mut self, value: serde_json::Value) -> Self {
self.llm_params.extra = Some(value);
self
}
pub fn try_build(self) -> Result<ToolAgent<P>, AgentError> {
let provider = self
.provider
.ok_or_else(|| AgentError::Message("构建 ToolAgent 失败:缺少 provider".to_string()))?;
let model = self
.model
.ok_or_else(|| AgentError::Message("构建 ToolAgent 失败:缺少 model".to_string()))?;
let conversation_manager = if self.with_conversation {
let mut conv = ConversationManager::new();
if let Some(ref prompt) = self.system_prompt {
conv = conv.with_system_prompt(prompt.clone());
}
Some(Arc::new(Mutex::new(conv)))
} else {
None
};
let provider_arc = Arc::new(provider);
let execution = build_default_execution(crate::agent::ExecutionBuildConfig {
provider: provider_arc.clone(),
model: model.clone(),
tools: self.tools.clone(),
system_prompt: self.system_prompt.clone(),
max_steps: self.max_steps,
max_tool_concurrency: self.max_tool_concurrency,
conversation_manager: conversation_manager.clone(),
middleware_chain: self.middleware_chain.clone(),
enhanced_config: self.enhanced_config,
llm_params: self.llm_params.clone(),
});
Ok(ToolAgent {
provider: provider_arc,
model,
system_prompt: self.system_prompt,
tools: self.tools,
_max_steps: self.max_steps,
conversation_manager,
llm_params: self.llm_params,
execution,
})
}
pub fn build(self) -> ToolAgent<P> {
self.try_build()
.unwrap_or_else(|err| panic!("ToolAgentBuilder::build 失败:{err}"))
}
}
impl<P> Default for ToolAgentBuilder<P> {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use rucora_core::test_utils::MockProvider;
#[test]
fn test_tool_agent_builder() {
let _agent = ToolAgentBuilder::<MockProvider>::new()
.provider(MockProvider)
.model("gpt-4o-mini")
.system_prompt("test")
.max_steps(10)
.build();
}
}