cognis 0.2.1

LLM application framework built on cognis-core
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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

//! Model retry middleware — automatic retry on model failures.
//!
//! Mirrors Python `langchain.agents.middleware.model_retry`.

use async_trait::async_trait;

use cognis_core::error::{CognisError, Result};
use cognis_core::messages::Message;

use super::retry::{should_retry, OnFailure, RetryCondition, RetryConfig};
use super::types::{
    AgentMiddleware, AsyncModelHandler, ModelCallResult, ModelRequest, ModelResponse,
};

/// Middleware that retries model calls on failure with exponential backoff.
#[derive(Default)]
pub struct ModelRetryMiddleware {
    config: RetryConfig,
}

impl ModelRetryMiddleware {
    pub fn new(config: RetryConfig) -> Self {
        Self { config }
    }

    /// Create with default config and specified max retries.
    pub fn with_max_retries(max_retries: usize) -> Self {
        Self {
            config: RetryConfig::new(max_retries),
        }
    }

    /// Set the retry condition.
    pub fn with_retry_on(mut self, retry_on: RetryCondition) -> Self {
        self.config.retry_on = retry_on;
        self
    }

    /// Set the on_failure strategy.
    pub fn with_on_failure(mut self, on_failure: OnFailure) -> Self {
        self.config.on_failure = on_failure;
        self
    }
}

#[async_trait]
impl AgentMiddleware for ModelRetryMiddleware {
    fn name(&self) -> &str {
        "ModelRetryMiddleware"
    }

    async fn wrap_model_call(
        &self,
        request: &ModelRequest,
        handler: &AsyncModelHandler,
    ) -> Result<ModelCallResult> {
        let mut last_error: Option<CognisError> = None;

        for attempt in 0..=self.config.max_retries {
            match handler(request).await {
                Ok(response) => return Ok(ModelCallResult::Response(response)),
                Err(e) => {
                    if !should_retry(&e, &self.config.retry_on)
                        || attempt == self.config.max_retries
                    {
                        last_error = Some(e);
                        break;
                    }
                    last_error = Some(e);
                    let delay = self.config.calculate_delay(attempt);
                    tokio::time::sleep(delay).await;
                }
            }
        }

        let error = last_error.unwrap_or_else(|| CognisError::Other("Unknown error".into()));

        match &self.config.on_failure {
            OnFailure::Error => Err(error),
            OnFailure::Continue => {
                let error_msg = Message::ai(format!(
                    "Model call failed after {} retries: {}",
                    self.config.max_retries, error
                ));
                Ok(ModelCallResult::Response(ModelResponse::new(vec![
                    error_msg,
                ])))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_model_retry_default() {
        let mw = ModelRetryMiddleware::default();
        assert_eq!(mw.name(), "ModelRetryMiddleware");
        assert_eq!(mw.config.max_retries, 2);
        assert!(matches!(mw.config.on_failure, OnFailure::Continue));
        assert!(matches!(mw.config.retry_on, RetryCondition::AnyError));
    }

    #[test]
    fn test_model_retry_with_max_retries() {
        let mw = ModelRetryMiddleware::with_max_retries(5);
        assert_eq!(mw.config.max_retries, 5);
    }

    #[test]
    fn test_model_retry_with_retry_on() {
        let mw = ModelRetryMiddleware::default()
            .with_retry_on(RetryCondition::ErrorContains(vec!["timeout".into()]));
        assert!(matches!(
            mw.config.retry_on,
            RetryCondition::ErrorContains(_)
        ));
    }

    #[test]
    fn test_model_retry_with_on_failure() {
        let mw = ModelRetryMiddleware::default().with_on_failure(OnFailure::Error);
        assert!(matches!(mw.config.on_failure, OnFailure::Error));
    }
}