batch_mode_batch_schema/
batch_response_body.rs

1// ---------------- [ File: src/batch_response_body.rs ]
2crate::ix!();
3
4#[derive(Debug,Serialize)]
5//#[serde(tag = "object", content = "data")]
6pub enum BatchResponseBody {
7
8    #[serde(rename = "chat.completion")]
9    Success(BatchSuccessResponseBody),
10
11    #[serde(rename = "error")]
12    Error(BatchErrorResponseBody),
13}
14
15impl<'de> Deserialize<'de> for BatchResponseBody {
16    fn deserialize<D>(deserializer: D) -> Result<BatchResponseBody, D::Error>
17    where
18        D: Deserializer<'de>,
19    {
20        let value: serde_json::Value = Deserialize::deserialize(deserializer)?;
21
22        if value.get("error").is_some() {
23            let error_body = BatchErrorResponseBody::deserialize(&value)
24                .map_err(serde::de::Error::custom)?;
25            Ok(BatchResponseBody::Error(error_body))
26        } else {
27            let success_body = BatchSuccessResponseBody::deserialize(&value)
28                .map_err(serde::de::Error::custom)?;
29            Ok(BatchResponseBody::Success(success_body))
30        }
31    }
32}
33
34impl BatchResponseBody {
35    /// Returns `Some(&BatchSuccessResponseBody)` if the response is a success.
36    pub fn as_success(&self) -> Option<&BatchSuccessResponseBody> {
37        if let BatchResponseBody::Success(ref success_body) = *self {
38            Some(success_body)
39        } else {
40            None
41        }
42    }
43
44    /// Returns `Some(&BatchErrorResponseBody)` if the response is an error.
45    pub fn as_error(&self) -> Option<&BatchErrorResponseBody> {
46        if let BatchResponseBody::Error(ref error_body) = *self {
47            Some(error_body)
48        } else {
49            None
50        }
51    }
52
53    /// Retrieves the `id` if the response is successful.
54    pub fn id(&self) -> Option<&str> {
55        self.as_success().map(|body| body.id())
56    }
57
58    /// Retrieves the `object` if the response is successful.
59    pub fn object(&self) -> Option<&str> {
60        self.as_success().map(|body| body.object())
61    }
62
63    /// Retrieves the `model` if the response is successful.
64    pub fn model(&self) -> Option<&str> {
65        self.as_success().map(|body| body.model())
66    }
67
68    /// Retrieves the `choices` if the response is successful.
69    pub fn choices(&self) -> Option<&[BatchChoice]> {
70        self.as_success().map(|body| body.choices())
71    }
72
73    /// Retrieves the `usage` if the response is successful.
74    pub fn usage(&self) -> Option<&BatchUsage> {
75        self.as_success().map(|body| body.usage())
76    }
77
78    /// Retrieves the `system_fingerprint` if the response is successful.
79    pub fn system_fingerprint(&self) -> Option<&str> {
80        self.as_success().and_then(|body| body.system_fingerprint())
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87    use serde_json::json;
88
89    #[test]
90    fn test_success_body_deserialization() {
91
92        let json_data = json!({
93            "id":                 "chatcmpl-AVW7Z2Dd49g7Zq5eVExww6dlKA8T9",
94            "object":             "chat.completion",
95            "created":            1732075005,
96            "model":              "gpt-4o-2024-08-06",
97            "choices":            [],
98            "usage":              {
99                "prompt_tokens":      40,
100                "completion_tokens": 360,
101                "total_tokens":      400,
102            },
103            "system_fingerprint": "fp_7f6be3efb0"
104        });
105
106        let body: BatchResponseBody = serde_json::from_value(json_data).unwrap();
107
108        match body {
109            BatchResponseBody::Success(success_body) => {
110                assert_eq!(success_body.id(), "chatcmpl-AVW7Z2Dd49g7Zq5eVExww6dlKA8T9");
111            }
112            _ => panic!("Expected success body"),
113        }
114    }
115
116    #[test]
117    fn test_error_body_deserialization() {
118        let json_data = json!({
119            "error": {
120                "message": "An error occurred",
121                "type": "server_error",
122                "param": null,
123                "code": null
124            }
125        });
126
127        let body: BatchResponseBody = serde_json::from_value(json_data).unwrap();
128        match body {
129            BatchResponseBody::Error(error_body) => {
130                assert_eq!(error_body.error().message(), "An error occurred");
131            }
132            _ => panic!("Expected error body"),
133        }
134    }
135}
batch_mode_batch_schema/batch_response_body.rs

batch_mode_batch_schema/
batch_response_body.rs
