ai-lib 0.4.0

A unified AI SDK for Rust providing a single interface for multiple AI providers with hybrid architecture
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

use async_trait::async_trait;
use std::time::Duration;
use tokio::time::timeout;

use crate::interceptors::{Interceptor, RequestContext};
use crate::types::{AiLibError, ChatCompletionRequest, ChatCompletionResponse};

/// Timeout interceptor that wraps requests with a timeout
pub struct TimeoutInterceptor {
    timeout_duration: Duration,
}

impl TimeoutInterceptor {
    /// Create a new timeout interceptor
    pub fn new(timeout_duration: Duration) -> Self {
        Self { timeout_duration }
    }
}

impl Default for TimeoutInterceptor {
    /// Create with default timeout (30 seconds)
    fn default() -> Self {
        Self::new(Duration::from_secs(30))
    }
}

#[async_trait]
impl Interceptor for TimeoutInterceptor {
    async fn on_request(&self, ctx: &RequestContext, _req: &ChatCompletionRequest) {
        // Log timeout setting (in a real implementation, you'd use proper logging)
        // Timeout setting: In production, this would be logged to metrics
        let _ = (self.timeout_duration, ctx);
    }
}

/// Wrapper that implements timeout logic around a function call
pub struct TimeoutWrapper {
    interceptor: TimeoutInterceptor,
}

impl TimeoutWrapper {
    pub fn new(interceptor: TimeoutInterceptor) -> Self {
        Self { interceptor }
    }
}

impl Default for TimeoutWrapper {
    fn default() -> Self {
        Self::new(TimeoutInterceptor::default())
    }
}

impl TimeoutWrapper {
    /// Execute a function with timeout
    pub async fn execute<F, Fut>(
        &self,
        ctx: &RequestContext,
        _req: &ChatCompletionRequest,
        f: F,
    ) -> Result<ChatCompletionResponse, AiLibError>
    where
        F: Fn() -> Fut,
        Fut: std::future::Future<Output = Result<ChatCompletionResponse, AiLibError>>,
    {
        match timeout(self.interceptor.timeout_duration, f()).await {
            Ok(result) => result,
            Err(_) => Err(AiLibError::TimeoutError(format!(
                "Request timeout after {:?} for {}:{}",
                self.interceptor.timeout_duration, ctx.provider, ctx.model
            ))),
        }
    }
}