rustic-ai 0.1.0

A Rust-native agent framework with tool calling, streaming, and multi-provider support for OpenAI, Anthropic, Gemini, and Grok
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

use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use thiserror::Error;

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RequestUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_write_tokens: u64,
    pub cache_read_tokens: u64,
    pub input_audio_tokens: u64,
    pub output_audio_tokens: u64,
    pub details: HashMap<String, u64>,
}

impl RequestUsage {
    pub fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens
    }
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RunUsage {
    pub requests: u64,
    pub tool_calls: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_write_tokens: u64,
    pub cache_read_tokens: u64,
    pub input_audio_tokens: u64,
    pub output_audio_tokens: u64,
    pub details: HashMap<String, u64>,
}

impl RunUsage {
    pub fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens
    }

    pub fn incr_request(&mut self, request: &RequestUsage) {
        self.requests += 1;
        self.input_tokens += request.input_tokens;
        self.output_tokens += request.output_tokens;
        self.cache_write_tokens += request.cache_write_tokens;
        self.cache_read_tokens += request.cache_read_tokens;
        self.input_audio_tokens += request.input_audio_tokens;
        self.output_audio_tokens += request.output_audio_tokens;
        for (k, v) in &request.details {
            *self.details.entry(k.clone()).or_insert(0) += v;
        }
    }

    pub fn incr_tool_call(&mut self) {
        self.tool_calls += 1;
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageLimits {
    pub request_limit: Option<u64>,
    pub tool_calls_limit: Option<u64>,
    pub input_tokens_limit: Option<u64>,
    pub output_tokens_limit: Option<u64>,
    pub total_tokens_limit: Option<u64>,
}

impl Default for UsageLimits {
    fn default() -> Self {
        Self {
            request_limit: Some(50),
            tool_calls_limit: None,
            input_tokens_limit: None,
            output_tokens_limit: None,
            total_tokens_limit: None,
        }
    }
}

impl UsageLimits {
    pub fn check_request(&self, current_requests: u64) -> Result<(), UsageError> {
        if let Some(limit) = self.request_limit
            && current_requests >= limit
        {
            return Err(UsageError::RequestLimitExceeded { limit });
        }
        Ok(())
    }

    pub fn check_tool_call(&self, current_calls: u64) -> Result<(), UsageError> {
        if let Some(limit) = self.tool_calls_limit
            && current_calls >= limit
        {
            return Err(UsageError::ToolCallsLimitExceeded { limit });
        }
        Ok(())
    }

    pub fn check_after_response(&self, usage: &RunUsage) -> Result<(), UsageError> {
        if let Some(limit) = self.input_tokens_limit
            && usage.input_tokens > limit
        {
            return Err(UsageError::InputTokensLimitExceeded { limit });
        }
        if let Some(limit) = self.output_tokens_limit
            && usage.output_tokens > limit
        {
            return Err(UsageError::OutputTokensLimitExceeded { limit });
        }
        if let Some(limit) = self.total_tokens_limit
            && usage.total_tokens() > limit
        {
            return Err(UsageError::TotalTokensLimitExceeded { limit });
        }
        Ok(())
    }
}

#[derive(Debug, Error)]
pub enum UsageError {
    #[error("request limit exceeded (limit {limit})")]
    RequestLimitExceeded { limit: u64 },
    #[error("tool call limit exceeded (limit {limit})")]
    ToolCallsLimitExceeded { limit: u64 },
    #[error("input token limit exceeded (limit {limit})")]
    InputTokensLimitExceeded { limit: u64 },
    #[error("output token limit exceeded (limit {limit})")]
    OutputTokensLimitExceeded { limit: u64 },
    #[error("total token limit exceeded (limit {limit})")]
    TotalTokensLimitExceeded { limit: u64 },
}