1use std::sync::Arc;
2
3use async_trait::async_trait;
4use pi_ai::{Content, Message, Model, StreamOptions, ThinkingLevel, Tool};
5use serde_json::Value;
6
7#[derive(Debug, Clone, Default)]
9pub struct AgentToolResult {
10 pub content: Vec<Content>,
11 pub details: Value,
12 pub terminate: bool,
13}
14
15impl AgentToolResult {
16 pub fn text(s: impl Into<String>) -> Self {
17 Self {
18 content: vec![Content::text(s)],
19 details: Value::Null,
20 terminate: false,
21 }
22 }
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
27pub enum PermissionDecision {
28 Allow,
29 AllowSession,
31 Deny {
33 reason: String,
34 },
35}
36
37#[async_trait]
40pub trait PermissionPolicy: Send + Sync {
41 async fn check(&self, tool_name: &str, args: &Value) -> PermissionDecision;
42}
43
44pub struct AllowAllPolicy;
46
47#[async_trait]
48impl PermissionPolicy for AllowAllPolicy {
49 async fn check(&self, _tool_name: &str, _args: &Value) -> PermissionDecision {
50 PermissionDecision::Allow
51 }
52}
53
54#[async_trait]
56pub trait AgentTool: Send + Sync {
57 fn name(&self) -> &str;
58 fn label(&self) -> &str {
59 self.name()
60 }
61 fn description(&self) -> &str;
62 fn parameters(&self) -> Value;
63 fn requires_permission(&self) -> bool {
67 false
68 }
69 async fn execute(&self, tool_call_id: &str, args: Value) -> Result<AgentToolResult, String>;
70}
71
72pub fn tool_def(t: &dyn AgentTool) -> Tool {
73 Tool {
74 name: t.name().to_string(),
75 description: t.description().to_string(),
76 parameters: t.parameters(),
77 }
78}
79
80#[derive(Clone)]
82pub struct AgentConfig {
83 pub model: Model,
84 pub thinking_level: ThinkingLevel,
85 pub stream_options: StreamOptions,
86 pub max_turns: u32,
87 pub tools: Vec<Arc<dyn AgentTool>>,
88 pub system_prompt: String,
89 pub permission: Arc<dyn PermissionPolicy>,
90}
91
92impl AgentConfig {
93 pub fn new(model: Model, system_prompt: impl Into<String>) -> Self {
94 Self {
95 model,
96 thinking_level: ThinkingLevel::Off,
97 stream_options: StreamOptions::default(),
98 max_turns: 32,
99 tools: Vec::new(),
100 system_prompt: system_prompt.into(),
101 permission: Arc::new(AllowAllPolicy),
102 }
103 }
104
105 pub fn with_tools(mut self, tools: Vec<Arc<dyn AgentTool>>) -> Self {
106 self.tools = tools;
107 self
108 }
109
110 pub fn with_max_turns(mut self, n: u32) -> Self {
111 self.max_turns = n;
112 self
113 }
114
115 pub fn with_permission(mut self, p: Arc<dyn PermissionPolicy>) -> Self {
116 self.permission = p;
117 self
118 }
119
120 pub fn with_thinking(mut self, level: ThinkingLevel) -> Self {
121 self.thinking_level = level;
122 self
123 }
124}
125
126#[derive(Debug, Clone)]
128pub enum AgentEvent {
129 AgentStart,
130 AgentEnd {
131 messages: Vec<Message>,
132 },
133 TurnStart,
134 TurnEnd,
135 AssistantMessage {
136 message: Message,
137 },
138 UserMessage {
139 message: Message,
140 },
141 TextDelta {
143 delta: String,
144 },
145 ThinkingDelta {
147 delta: String,
148 },
149 ToolExecutionStart {
150 tool_call_id: String,
151 tool_name: String,
152 args: Value,
153 },
154 ToolExecutionEnd {
155 tool_call_id: String,
156 tool_name: String,
157 is_error: bool,
158 content: Vec<Content>,
159 },
160 PermissionDenied {
162 tool_name: String,
163 reason: String,
164 },
165}