cognis 0.3.1

Cognis umbrella crate: agent builder, multi-agent orchestration, memory, middleware (rate limit, retry, PII, prompt caching), built-in tools, and re-exports of cognis-core, cognis-graph, cognis-llm, and cognis-rag.
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

//! Fall through to a backup `Client` if the primary errors.

use std::sync::Arc;

use async_trait::async_trait;

use cognis_core::Result;
use cognis_llm::chat::ChatResponse;
use cognis_llm::Client;

use super::{Middleware, MiddlewareCtx, Next};

/// On any error from the primary chain, retries the call against `fallback`
/// (a different `Client`). For multi-tier fallback, stack multiple
/// `ModelFallback` layers.
pub struct ModelFallback {
    fallback: Client,
}

impl ModelFallback {
    /// Wrap with a fallback client.
    pub fn new(fallback: Client) -> Self {
        Self { fallback }
    }
}

#[async_trait]
impl Middleware for ModelFallback {
    async fn call(&self, ctx: MiddlewareCtx, next: Arc<dyn Next>) -> Result<ChatResponse> {
        match next.invoke(ctx.clone()).await {
            Ok(r) => Ok(r),
            Err(_) => {
                self.fallback
                    .provider()
                    .chat_completion_with_tools(ctx.messages, ctx.tool_defs, ctx.opts)
                    .await
            }
        }
    }

    fn name(&self) -> &str {
        "ModelFallback"
    }
}

#[cfg(test)]
mod tests {
    use super::super::tests_util::*;
    use super::*;
    use crate::middleware::MiddlewarePipeline;

    use cognis_core::{CognisError, Message};
    use cognis_llm::chat::ChatOptions;

    #[tokio::test]
    async fn falls_through_on_error() {
        let primary_provider = make_flaky_provider(|_| {
            Err(CognisError::Network {
                status_code: Some(500),
                message: "boom".into(),
            })
        });
        let primary = Client::new(primary_provider);
        let backup_provider = make_flaky_provider(|_| Ok("backup".into()));
        let backup = Client::new(backup_provider);
        let pipe = MiddlewarePipeline::new()
            .push(ModelFallback::new(backup))
            .build(primary);
        let r = pipe
            .invoke(
                vec![Message::human("hi")],
                Vec::new(),
                ChatOptions::default(),
            )
            .await
            .unwrap();
        assert_eq!(r.message.content(), "backup");
    }

    #[tokio::test]
    async fn primary_wins_on_success() {
        let primary = Client::new(make_flaky_provider(|_| Ok("primary".into())));
        let backup = Client::new(make_flaky_provider(|_| Ok("backup".into())));
        let pipe = MiddlewarePipeline::new()
            .push(ModelFallback::new(backup))
            .build(primary);
        let r = pipe
            .invoke(
                vec![Message::human("hi")],
                Vec::new(),
                ChatOptions::default(),
            )
            .await
            .unwrap();
        assert_eq!(r.message.content(), "primary");
    }
}