Skip to main content

pi_agent/
types.rs

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/// A live result returned from a tool execution.
8#[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/// Permission outcome for a tool call. Returned by a [`PermissionPolicy`].
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub enum PermissionDecision {
28    Allow,
29    /// Allow this call and remember the choice for the rest of the run.
30    AllowSession,
31    /// Deny this call; emit an error tool result with `reason`.
32    Deny {
33        reason: String,
34    },
35}
36
37/// User-supplied permission policy. Implementations may prompt interactively or
38/// consult a static allow-list.
39#[async_trait]
40pub trait PermissionPolicy: Send + Sync {
41    async fn check(&self, tool_name: &str, args: &Value) -> PermissionDecision;
42}
43
44/// Always-allow policy — useful for tests and non-interactive runs.
45pub 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/// Tool execution trait — analog of `AgentTool.execute` in TS.
55#[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    /// Whether the tool requires user permission by default. Read-only tools
64    /// (`read`, `ls`, `grep`, `glob`) return `false`; mutating or side-effecting
65    /// tools (`bash`, `write`, `edit`) return `true`.
66    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/// Agent configuration controlling the loop.
81#[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/// Events emitted by the agent loop, mirroring `AgentEvent` in TS.
127#[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    /// Streaming text chunk while the assistant types.
142    TextDelta {
143        delta: String,
144    },
145    /// Streaming thinking chunk.
146    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    /// Permission denied for a tool call (the loop appended an error tool result).
161    PermissionDenied {
162        tool_name: String,
163        reason: String,
164    },
165}