use rpi_ai::types::{AssistantMessage, ImageContent, TextContent, ToolResultMessage, Usage};
use std::collections::HashSet;
use std::sync::Arc;
use crate::message::AgentMessage;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ToolExecutionMode {
Sequential,
#[default]
Parallel,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum QueueMode {
All,
#[default]
OneAtATime,
}
#[derive(Debug, Clone, Default)]
pub struct AgentToolResult {
pub content: Vec<TextContentOrImage>,
pub details: serde_json::Value,
pub usage: Option<Usage>,
pub added_tool_names: Vec<String>,
pub terminate: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TextContentOrImage {
Text(TextContent),
Image(ImageContent),
}
impl TextContentOrImage {
pub fn text<S: Into<String>>(s: S) -> Self {
TextContentOrImage::Text(TextContent {
kind: rpi_ai::types::TextContentType,
text: s.into(),
text_signature: None,
})
}
}
impl AgentToolResult {
pub fn text(message: impl Into<String>) -> Self {
Self {
content: vec![TextContentOrImage::text(message)],
details: serde_json::Value::Null,
usage: None,
added_tool_names: Vec::new(),
terminate: false,
}
}
pub fn error_text(message: impl Into<String>) -> Self {
Self::text(message)
}
pub fn into_content(self) -> Vec<rpi_ai::types::Content> {
self.content
.into_iter()
.map(|c| match c {
TextContentOrImage::Text(t) => rpi_ai::types::Content::Text(t),
TextContentOrImage::Image(i) => rpi_ai::types::Content::Image(i),
})
.collect()
}
}
impl From<AgentToolResult> for Result<AgentToolResult, crate::AgentError> {
fn from(r: AgentToolResult) -> Self {
Ok(r)
}
}
pub type ToolResultPartial = AgentToolResult;
#[derive(Debug, Clone, Default)]
pub struct BeforeToolCallResult {
pub block: bool,
pub reason: Option<String>,
pub terminate: bool,
pub args: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Default)]
pub struct AfterToolCallResult {
pub content: Option<Vec<TextContentOrImage>>,
pub details: Option<serde_json::Value>,
pub is_error: Option<bool>,
pub usage: Option<Usage>,
pub terminate: Option<bool>,
}
pub struct BeforeToolCallContext<'a> {
pub assistant_message: &'a AssistantMessage,
pub tool_call: &'a rpi_ai::types::ToolCall,
pub args: &'a serde_json::Value,
pub context: &'a AgentContext,
}
pub struct AfterToolCallContext<'a> {
pub assistant_message: &'a AssistantMessage,
pub tool_call: &'a rpi_ai::types::ToolCall,
pub args: &'a serde_json::Value,
pub result: &'a AgentToolResult,
pub is_error: bool,
pub context: &'a AgentContext,
}
pub struct ShouldStopAfterTurnContext<'a> {
pub message: &'a AssistantMessage,
pub tool_results: &'a [ToolResultMessage],
pub context: &'a AgentContext,
pub new_messages: &'a [AgentMessage],
}
#[derive(Clone, Default)]
pub struct AgentContext {
pub system_prompt: String,
pub messages: Vec<AgentMessage>,
pub tools: Vec<Arc<dyn crate::agent_tool::AgentTool>>,
}
impl std::fmt::Debug for AgentContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AgentContext")
.field("system_prompt", &self.system_prompt)
.field("messages", &self.messages)
.field(
"tools",
&self
.tools
.iter()
.map(|t| t.schema().name.as_str())
.collect::<Vec<_>>(),
)
.finish()
}
}
impl AgentContext {
pub fn new(messages: Vec<AgentMessage>) -> Self {
Self {
system_prompt: String::new(),
messages,
tools: Vec::new(),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct AgentLoopTurnUpdate {
pub context: Option<AgentContext>,
pub model: Option<rpi_ai::model::Model>,
pub thinking_level: Option<rpi_ai::types::ThinkingLevel>,
}
#[derive(Clone)]
pub struct AgentState {
pub system_prompt: String,
pub model: rpi_ai::model::Model,
pub thinking_level: rpi_ai::types::ThinkingLevel,
pub tools: Vec<Arc<dyn crate::agent_tool::AgentTool>>,
pub messages: Vec<AgentMessage>,
pub is_streaming: bool,
pub streaming_message: Option<AgentMessage>,
pub pending_tool_calls: HashSet<String>,
pub error_message: Option<String>,
}
impl std::fmt::Debug for AgentState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AgentState")
.field("system_prompt", &self.system_prompt)
.field("model", &self.model)
.field("thinking_level", &self.thinking_level)
.field(
"tools",
&self
.tools
.iter()
.map(|t| t.schema().name.as_str())
.collect::<Vec<_>>(),
)
.field("messages", &self.messages)
.field("is_streaming", &self.is_streaming)
.field("streaming_message", &self.streaming_message)
.field("pending_tool_calls", &self.pending_tool_calls)
.field("error_message", &self.error_message)
.finish()
}
}
impl Default for AgentState {
fn default() -> Self {
Self {
system_prompt: String::new(),
model: default_model(),
thinking_level: rpi_ai::types::ThinkingLevel::Off,
tools: Vec::new(),
messages: Vec::new(),
is_streaming: false,
streaming_message: None,
pending_tool_calls: HashSet::new(),
error_message: None,
}
}
}
pub(crate) fn default_model() -> rpi_ai::model::Model {
rpi_ai::model::Model::new(
"unknown",
"unknown",
rpi_ai::types::Api::Other("unknown".into()),
"unknown",
"",
)
}