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

use serde::{Deserialize, Serialize};
use std::fmt;

use crate::message::Message;

#[derive(Debug, Serialize, Deserialize)]
pub struct Choice {
    index: u64,
    pub message: Message,
    finish_reason: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Usage {
    prompt_tokens: u32,
    completion_tokens: u32,
    total_tokens: u32,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ChatResponse {
    id: String,
    object: String,
    created: u64,
    pub choices: Vec<Choice>,
    usage: Usage,
}

impl fmt::Display for Choice {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{{\"index\":{},\"message\":{},\"finish_reason\":\"{}\"}}",
            self.index, self.message, self.finish_reason
        )
    }
}

impl fmt::Display for ChatResponse {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{{\"id\":\"{}\",\"object\":\"{}\",\"created\":{},\"choices\":[",
            self.id, self.object, self.created
        )?;
        for (i, choice) in self.choices.iter().enumerate() {
            write!(
                f,
                "{}{}",
                choice,
                if i == self.choices.len() - 1 { "" } else { "," }
            )?;
        }
        write!(f, "],\"usage\":{}}}", self.usage)
    }
}

impl fmt::Display for Usage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{{\"prompt_tokens\":{},\"completion_tokens\":{},\"total_tokens\":{}}}",
            self.prompt_tokens, self.completion_tokens, self.total_tokens
        )
    }
}

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

    #[test]
    fn test_display_for_choice() {
        let choice = Choice {
            index: 0,
            message: Message {
                role: "role".to_string(),
                content: Some("content".to_string()),
                name: Some("name".to_string()),
                function_call: Some(FunctionCall {
                    name: "name".to_string(),
                    arguments: "{\"example\":\"this\"}".to_string(),
                }),
            },
            finish_reason: "finish_reason".to_string(),
        };
        assert_eq!(
            format!("{}", choice),
            "{\"index\":0,\"message\":{\"role\":\"role\",\"content\":\"content\",\"name\":\"name\",\"function_call\":{\"name\":\"name\",\"arguments\":{\"example\":\"this\"}}},\"finish_reason\":\"finish_reason\"}"
        );
    }

    #[test]
    fn test_display_chat_response() {
        let chat_response = ChatResponse {
            id: "id".to_string(),
            object: "object".to_string(),
            created: 0,
            choices: vec![Choice {
                index: 0,
                message: Message {
                    role: "role".to_string(),
                    content: Some("content".to_string()),
                    name: Some("name".to_string()),
                    function_call: Some(FunctionCall {
                        name: "name".to_string(),
                        arguments: "{\"example\":\"this\"}".to_string(),
                    }),
                },
                finish_reason: "finish_reason".to_string(),
            }],
            usage: Usage {
                prompt_tokens: 0,
                completion_tokens: 0,
                total_tokens: 0,
            },
        };
        assert_eq!(
            format!("{}", chat_response),
            "{\"id\":\"id\",\"object\":\"object\",\"created\":0,\"choices\":[{\"index\":0,\"message\":{\"role\":\"role\",\"content\":\"content\",\"name\":\"name\",\"function_call\":{\"name\":\"name\",\"arguments\":{\"example\":\"this\"}}},\"finish_reason\":\"finish_reason\"}],\"usage\":{\"prompt_tokens\":0,\"completion_tokens\":0,\"total_tokens\":0}}"
        );
    }

    #[test]
    fn test_display_usage() {
        let usage = Usage {
            prompt_tokens: 0,
            completion_tokens: 0,
            total_tokens: 0,
        };
        assert_eq!(
            format!("{}", usage),
            "{\"prompt_tokens\":0,\"completion_tokens\":0,\"total_tokens\":0}"
        );
    }
}