aidale-layer 0.1.0

Built-in layers for Aidale (logging, retry, caching, etc.)
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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

//! Logging layer for provider operations.

use aidale_core::error::AiError;
use aidale_core::layer::{Layer, LayeredProvider};
use aidale_core::provider::{ChatCompletionStream, Provider};
use aidale_core::types::*;
use async_trait::async_trait;
use std::fmt::Debug;
use std::sync::Arc;

/// Logging layer that logs provider operations.
#[derive(Debug, Clone)]
pub struct LoggingLayer {
    prefix: String,
}

impl LoggingLayer {
    /// Create a new logging layer
    pub fn new() -> Self {
        Self {
            prefix: "[AI Core]".to_string(),
        }
    }

    /// Create a logging layer with custom prefix
    pub fn with_prefix(prefix: impl Into<String>) -> Self {
        Self {
            prefix: prefix.into(),
        }
    }
}

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

impl<P: Provider> Layer<P> for LoggingLayer {
    type LayeredProvider = LoggingProvider<P>;

    fn layer(&self, inner: P) -> Self::LayeredProvider {
        LoggingProvider {
            inner,
            prefix: self.prefix.clone(),
        }
    }
}

/// Provider wrapped with logging
#[derive(Debug)]
pub struct LoggingProvider<P> {
    inner: P,
    prefix: String,
}

#[async_trait]
impl<P: Provider> LayeredProvider for LoggingProvider<P> {
    type Inner = P;

    fn inner(&self) -> &Self::Inner {
        &self.inner
    }

    async fn layered_chat_completion(
        &self,
        req: ChatCompletionRequest,
    ) -> Result<ChatCompletionResponse, AiError> {
        tracing::debug!(
            "{} chat_completion request: model={}, messages={}",
            self.prefix,
            req.model,
            req.messages.len()
        );

        let start = std::time::Instant::now();
        let result = self.inner.chat_completion(req).await;
        let elapsed = start.elapsed();

        match &result {
            Ok(response) => {
                tracing::debug!(
                    "{} chat_completion success: id={}, tokens={}, elapsed={:?}",
                    self.prefix,
                    response.id,
                    response.usage.total_tokens,
                    elapsed
                );
            }
            Err(e) => {
                tracing::error!(
                    "{} chat_completion error: {:?}, elapsed={:?}",
                    self.prefix,
                    e,
                    elapsed
                );
            }
        }

        result
    }

    async fn layered_stream_chat_completion(
        &self,
        req: ChatCompletionRequest,
    ) -> Result<Box<ChatCompletionStream>, AiError> {
        tracing::debug!(
            "{} stream_chat_completion request: model={}, messages={}",
            self.prefix,
            req.model,
            req.messages.len()
        );

        let start = std::time::Instant::now();
        let result = self.inner.stream_chat_completion(req).await;
        let elapsed = start.elapsed();

        match &result {
            Ok(_) => {
                tracing::debug!(
                    "{} stream_chat_completion success, elapsed={:?}",
                    self.prefix,
                    elapsed
                );
            }
            Err(e) => {
                tracing::error!(
                    "{} stream_chat_completion error: {:?}, elapsed={:?}",
                    self.prefix,
                    e,
                    elapsed
                );
            }
        }

        result
    }
}

#[async_trait]
impl<P: Provider> Provider for LoggingProvider<P> {
    fn info(&self) -> Arc<ProviderInfo> {
        LayeredProvider::layered_info(self)
    }

    async fn chat_completion(
        &self,
        req: ChatCompletionRequest,
    ) -> Result<ChatCompletionResponse, AiError> {
        LayeredProvider::layered_chat_completion(self, req).await
    }

    async fn stream_chat_completion(
        &self,
        req: ChatCompletionRequest,
    ) -> Result<Box<ChatCompletionStream>, AiError> {
        LayeredProvider::layered_stream_chat_completion(self, req).await
    }
}