1use std::sync::Arc;
9
10use futures::future::BoxFuture;
11use futures::stream::BoxStream;
12use oino_ai::{
13 AgentMessage, AssistantMessage, AssistantMessageEvent, Context, JsonObject, JsonValue, Message,
14 Model, StreamOptions, ThinkingLevel, ToolCall, ToolResultMessage, Usage, UserContent,
15};
16use serde::{Deserialize, Serialize};
17use tokio_util::sync::CancellationToken;
18
19use crate::tool::{AgentTool, AgentToolResult};
20
21#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
23#[error("{0}")]
24pub struct HookError(pub String);
25
26impl HookError {
27 #[must_use]
29 pub fn new(message: impl Into<String>) -> Self {
30 Self(message.into())
31 }
32}
33
34#[derive(Clone, Default)]
36pub struct AgentContext {
37 pub messages: Vec<AgentMessage>,
39 pub tools: Vec<Arc<dyn AgentTool>>,
41}
42
43impl std::fmt::Debug for AgentContext {
44 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45 f.debug_struct("AgentContext")
46 .field("messages", &self.messages.len())
47 .field(
48 "tools",
49 &self
50 .tools
51 .iter()
52 .map(|t| t.meta().name.clone())
53 .collect::<Vec<_>>(),
54 )
55 .finish()
56 }
57}
58
59impl AgentContext {
60 #[must_use]
62 pub fn tool(&self, name: &str) -> Option<&Arc<dyn AgentTool>> {
63 self.tools.iter().find(|tool| tool.meta().name == name)
64 }
65}
66
67#[derive(Clone, Debug)]
69pub struct StreamRequest {
70 pub model: Model,
72 pub context: Context,
74 pub options: StreamOptions,
76 pub api_key: Option<String>,
78 pub cancel: CancellationToken,
80}
81
82pub trait StreamFn: Send + Sync {
84 fn stream(&self, request: StreamRequest) -> BoxStream<'static, AssistantMessageEvent>;
86}
87
88pub type ConvertToLlm = Arc<dyn Fn(&[AgentMessage]) -> Vec<Message> + Send + Sync>;
90
91pub trait TransformContext: Send + Sync {
93 fn transform<'a>(
95 &'a self,
96 messages: Vec<AgentMessage>,
97 cancel: &'a CancellationToken,
98 ) -> BoxFuture<'a, Result<Vec<AgentMessage>, HookError>>;
99}
100
101pub trait GetApiKey: Send + Sync {
103 fn api_key<'a>(&'a self, provider: &'a str)
105 -> BoxFuture<'a, Result<Option<String>, HookError>>;
106}
107
108#[derive(Clone, Copy, Debug)]
110pub struct BeforeToolCallContext<'a> {
111 pub assistant_message: &'a AssistantMessage,
113 pub tool_call: &'a ToolCall,
115 pub args: &'a JsonObject,
117 pub context: &'a AgentContext,
119}
120
121#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
123#[serde(rename_all = "camelCase")]
124pub struct BeforeToolCallResult {
125 #[serde(default)]
127 pub block: bool,
128 #[serde(default, skip_serializing_if = "Option::is_none")]
130 pub reason: Option<String>,
131 #[serde(default)]
133 pub terminate: bool,
134 #[serde(default, skip_serializing_if = "Option::is_none")]
136 pub input: Option<JsonObject>,
137}
138
139pub trait BeforeToolCall: Send + Sync {
141 fn before<'a>(
143 &'a self,
144 cx: BeforeToolCallContext<'a>,
145 cancel: &'a CancellationToken,
146 ) -> BoxFuture<'a, Result<Option<BeforeToolCallResult>, HookError>>;
147}
148
149#[derive(Clone, Copy, Debug)]
151pub struct AfterToolCallContext<'a> {
152 pub assistant_message: &'a AssistantMessage,
154 pub tool_call: &'a ToolCall,
156 pub args: &'a JsonObject,
158 pub result: &'a AgentToolResult,
160 pub is_error: bool,
162 pub context: &'a AgentContext,
164}
165
166#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
168#[serde(rename_all = "camelCase")]
169pub struct AfterToolCallResult {
170 #[serde(default, skip_serializing_if = "Option::is_none")]
172 pub content: Option<Vec<UserContent>>,
173 #[serde(default, skip_serializing_if = "Option::is_none")]
175 pub details: Option<JsonValue>,
176 #[serde(default, skip_serializing_if = "Option::is_none")]
178 pub structured_content: Option<JsonValue>,
179 #[serde(default, skip_serializing_if = "Option::is_none")]
181 pub is_error: Option<bool>,
182 #[serde(default, skip_serializing_if = "Option::is_none")]
184 pub usage: Option<Usage>,
185 #[serde(default, skip_serializing_if = "Option::is_none")]
187 pub terminate: Option<bool>,
188}
189
190pub trait AfterToolCall: Send + Sync {
192 fn after<'a>(
194 &'a self,
195 cx: AfterToolCallContext<'a>,
196 cancel: &'a CancellationToken,
197 ) -> BoxFuture<'a, Result<Option<AfterToolCallResult>, HookError>>;
198}
199
200#[derive(Clone, Copy, Debug)]
202pub struct AgentTurnContext<'a> {
203 pub message: &'a AssistantMessage,
205 pub tool_results: &'a [ToolResultMessage],
207 pub context: &'a AgentContext,
209 pub new_messages: &'a [AgentMessage],
211}
212
213#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
215#[serde(tag = "action", rename_all = "lowercase")]
216pub enum TurnDecision {
217 Continue,
219 End,
221}
222
223pub trait FinishTurn: Send + Sync {
225 fn finish<'a>(
227 &'a self,
228 turn: AgentTurnContext<'a>,
229 cancel: &'a CancellationToken,
230 ) -> BoxFuture<'a, Result<Option<TurnDecision>, HookError>>;
231}
232
233#[derive(Clone, Debug, Default)]
235pub struct AgentLoopTurnUpdate {
236 pub context: Option<AgentContext>,
238 pub messages: Option<Vec<AgentMessage>>,
240 pub model: Option<Model>,
242 pub thinking_level: Option<ThinkingLevel>,
244}
245
246pub trait PrepareNextTurn: Send + Sync {
248 fn prepare<'a>(
250 &'a self,
251 turn: AgentTurnContext<'a>,
252 cancel: &'a CancellationToken,
253 ) -> BoxFuture<'a, Result<Option<AgentLoopTurnUpdate>, HookError>>;
254}
255
256#[derive(Clone, Copy, Debug)]
258pub struct PrepareRequestContext<'a> {
259 pub context: &'a AgentContext,
261 pub model: &'a Model,
263 pub thinking_level: ThinkingLevel,
265}
266
267#[derive(Clone, Debug, Default)]
269pub struct AgentRequestUpdate {
270 pub context: Option<AgentContext>,
272 pub model: Option<Model>,
274 pub thinking_level: Option<ThinkingLevel>,
276}
277
278pub trait PrepareRequest: Send + Sync {
280 fn prepare<'a>(
282 &'a self,
283 request: PrepareRequestContext<'a>,
284 cancel: &'a CancellationToken,
285 ) -> BoxFuture<'a, Result<Option<AgentRequestUpdate>, HookError>>;
286}
287
288pub trait QueueSource: Send + Sync {
290 fn drain(&self) -> BoxFuture<'_, Vec<AgentMessage>>;
292}
293
294#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
296#[serde(rename_all = "lowercase")]
297pub enum ToolExecution {
298 #[default]
300 Parallel,
301 Sequential,
303}
304
305#[derive(Clone)]
310pub struct AgentLoopConfig {
311 pub model: Model,
313 pub thinking_level: ThinkingLevel,
315 pub stream_options: StreamOptions,
317 pub stream: Arc<dyn StreamFn>,
319 pub convert_to_llm: ConvertToLlm,
321 pub transform_context: Option<Arc<dyn TransformContext>>,
323 pub get_api_key: Option<Arc<dyn GetApiKey>>,
325 pub api_key: Option<String>,
327 pub before_tool_call: Option<Arc<dyn BeforeToolCall>>,
329 pub after_tool_call: Option<Arc<dyn AfterToolCall>>,
331 pub finish_turn: Option<Arc<dyn FinishTurn>>,
333 pub prepare_request: Option<Arc<dyn PrepareRequest>>,
335 pub prepare_next_turn: Option<Arc<dyn PrepareNextTurn>>,
337 pub get_steering_messages: Option<Arc<dyn QueueSource>>,
339 pub get_follow_up_messages: Option<Arc<dyn QueueSource>>,
341 pub tool_execution: ToolExecution,
343}
344
345impl AgentLoopConfig {
346 #[must_use]
348 pub fn new(model: Model, stream: Arc<dyn StreamFn>) -> Self {
349 Self {
350 model,
351 thinking_level: ThinkingLevel::Off,
352 stream_options: StreamOptions::default(),
353 stream,
354 convert_to_llm: Arc::new(default_convert_to_llm),
355 transform_context: None,
356 get_api_key: None,
357 api_key: None,
358 before_tool_call: None,
359 after_tool_call: None,
360 finish_turn: None,
361 prepare_request: None,
362 prepare_next_turn: None,
363 get_steering_messages: None,
364 get_follow_up_messages: None,
365 tool_execution: ToolExecution::Parallel,
366 }
367 }
368}
369
370#[must_use]
372pub fn default_convert_to_llm(messages: &[AgentMessage]) -> Vec<Message> {
373 messages.iter().filter_map(AgentMessage::as_llm).collect()
374}