grok_api 0.1.6

Rust client library for the Grok AI API (xAI)
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

//! Common types used across the API

use serde::{Deserialize, Serialize};

/// Role of a message in a conversation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// System message (instructions for the model)
    System,

    /// User message
    User,

    /// Assistant/model response
    Assistant,

    /// Tool/function response
    Tool,
}

impl Role {
    /// Convert role to string
    pub fn as_str(&self) -> &'static str {
        match self {
            Role::System => "system",
            Role::User => "user",
            Role::Assistant => "assistant",
            Role::Tool => "tool",
        }
    }
}

impl std::fmt::Display for Role {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

impl From<Role> for String {
    fn from(role: Role) -> Self {
        role.as_str().to_string()
    }
}

/// Reason why a completion finished
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
    /// Natural completion
    Stop,

    /// Maximum tokens reached
    Length,

    /// Tool/function call
    ToolCalls,

    /// Content filtered
    ContentFilter,

    /// Other/unknown reason
    #[serde(other)]
    Other,
}

impl StopReason {
    /// Check if completion finished naturally
    pub fn is_complete(&self) -> bool {
        matches!(self, StopReason::Stop)
    }

    /// Check if more tokens were needed
    pub fn needs_more_tokens(&self) -> bool {
        matches!(self, StopReason::Length)
    }

    /// Check if tool calls were made
    pub fn has_tool_calls(&self) -> bool {
        matches!(self, StopReason::ToolCalls)
    }
}

impl std::fmt::Display for StopReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StopReason::Stop => write!(f, "stop"),
            StopReason::Length => write!(f, "length"),
            StopReason::ToolCalls => write!(f, "tool_calls"),
            StopReason::ContentFilter => write!(f, "content_filter"),
            StopReason::Other => write!(f, "other"),
        }
    }
}

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

    #[test]
    fn test_role_conversion() {
        assert_eq!(Role::System.as_str(), "system");
        assert_eq!(Role::User.as_str(), "user");
        assert_eq!(Role::Assistant.as_str(), "assistant");

        let role_string: String = Role::User.into();
        assert_eq!(role_string, "user");
    }

    #[test]
    fn test_stop_reason_checks() {
        assert!(StopReason::Stop.is_complete());
        assert!(!StopReason::Length.is_complete());

        assert!(StopReason::Length.needs_more_tokens());
        assert!(!StopReason::Stop.needs_more_tokens());

        assert!(StopReason::ToolCalls.has_tool_calls());
        assert!(!StopReason::Stop.has_tool_calls());
    }

    #[test]
    fn test_role_serialization() {
        let json = serde_json::to_string(&Role::System).unwrap();
        assert_eq!(json, r#""system""#);

        let role: Role = serde_json::from_str(r#""user""#).unwrap();
        assert_eq!(role, Role::User);
    }
}