daimon 0.16.0

A Rust-native AI agent framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use crate::error::Result;
use crate::model::types::{ChatRequest, ChatResponse};
use crate::tool::ToolCall;

/// Outcome of a middleware hook. `Continue` proceeds to the next layer;
/// `ShortCircuit` stops the pipeline and returns the given response.
#[derive(Debug)]
pub enum MiddlewareAction {
    /// Proceed to the next middleware / model call.
    Continue,
    /// Stop processing and return this response directly.
    ShortCircuit(ChatResponse),
}

/// Composable layer that can inspect and mutate requests, responses, and tool
/// calls flowing through the agent's ReAct loop.
///
/// All methods default to no-op `Continue`. Override only the hooks you need.
pub trait Middleware: Send + Sync {
    /// Called before each model invocation. May mutate the request or short-circuit.
    fn on_request(
        &self,
        _request: &mut ChatRequest,
    ) -> impl Future<Output = Result<MiddlewareAction>> + Send {
        async { Ok(MiddlewareAction::Continue) }
    }

    /// Called after the model responds, before the agent processes tool calls or
    /// returns the final text. May mutate the response or short-circuit.
    fn on_response(
        &self,
        _response: &mut ChatResponse,
    ) -> impl Future<Output = Result<MiddlewareAction>> + Send {
        async { Ok(MiddlewareAction::Continue) }
    }

    /// Called before each tool is executed. May mutate the tool call arguments
    /// or short-circuit to skip execution.
    fn on_tool_call(
        &self,
        _call: &mut ToolCall,
    ) -> impl Future<Output = Result<MiddlewareAction>> + Send {
        async { Ok(MiddlewareAction::Continue) }
    }
}

/// Object-safe wrapper for [`Middleware`].
pub trait ErasedMiddleware: Send + Sync {
    fn on_request_erased<'a>(
        &'a self,
        request: &'a mut ChatRequest,
    ) -> Pin<Box<dyn Future<Output = Result<MiddlewareAction>> + Send + 'a>>;

    fn on_response_erased<'a>(
        &'a self,
        response: &'a mut ChatResponse,
    ) -> Pin<Box<dyn Future<Output = Result<MiddlewareAction>> + Send + 'a>>;

    fn on_tool_call_erased<'a>(
        &'a self,
        call: &'a mut ToolCall,
    ) -> Pin<Box<dyn Future<Output = Result<MiddlewareAction>> + Send + 'a>>;
}

impl<T: Middleware> ErasedMiddleware for T {
    fn on_request_erased<'a>(
        &'a self,
        request: &'a mut ChatRequest,
    ) -> Pin<Box<dyn Future<Output = Result<MiddlewareAction>> + Send + 'a>> {
        Box::pin(self.on_request(request))
    }

    fn on_response_erased<'a>(
        &'a self,
        response: &'a mut ChatResponse,
    ) -> Pin<Box<dyn Future<Output = Result<MiddlewareAction>> + Send + 'a>> {
        Box::pin(self.on_response(response))
    }

    fn on_tool_call_erased<'a>(
        &'a self,
        call: &'a mut ToolCall,
    ) -> Pin<Box<dyn Future<Output = Result<MiddlewareAction>> + Send + 'a>> {
        Box::pin(self.on_tool_call(call))
    }
}

/// Shared ownership of middleware.
pub type SharedMiddleware = Arc<dyn ErasedMiddleware>;