1#[derive(Debug, PartialEq)]
2pub enum CallbackReason {
3 Start,
4 Data,
5 End,
6 Done,
7 Error(String),
8}
9
10#[derive(Debug, PartialEq)]
11pub enum Error {
12 CurlError(curl::Error),
13 FromUtf8Error(std::string::FromUtf8Error),
14 SerdeJsonError(String),
15}
16
17impl From<curl::Error> for Error {
18 fn from(err: curl::Error) -> Self {
19 Self::CurlError(err)
20 }
21}
22
23impl From<std::string::FromUtf8Error> for Error {
24 fn from(err: std::string::FromUtf8Error) -> Self {
25 Self::FromUtf8Error(err)
26 }
27}
28
29impl From<serde_json::Error> for Error {
30 fn from(err: serde_json::Error) -> Self {
31 Self::SerdeJsonError(err.to_string())
32 }
33}
34
35#[derive(Debug, serde::Deserialize)]
36pub struct Permission {
37 pub id: String,
38 pub object: String,
39 pub created: u64,
40 pub allow_create_engine: bool,
41 pub allow_sampling: bool,
42 pub allow_logprobs: bool,
43 pub allow_search_indices: bool,
44 pub allow_view: bool,
45 pub allow_fine_tuning: bool,
46 pub organization: String,
47 pub group: serde_json::Value,
48 pub is_blocking: bool,
49}
50
51#[derive(Debug, serde::Deserialize)]
52pub struct Model {
53 pub id: String,
54 pub object: String,
55 pub created: u64,
56 pub owned_by: String,
57 pub permission: Vec<Permission>,
58 pub root: String,
59 pub parent: Option<String>,
60}
61
62#[derive(Debug, serde::Deserialize)]
63pub struct ModelList {
64 pub object: String,
65 pub data: Vec<Model>,
66}
67
68#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
69pub struct Message {
70 #[serde(skip_serializing_if = "Option::is_none")]
71 pub role: Option<String>,
72 #[serde(skip_serializing_if = "Option::is_none")]
73 pub content: Option<String>,
74}
75
76#[derive(Debug, Clone, serde::Serialize)]
77pub struct RequestBody {
78 pub model: String,
79 pub messages: Vec<Message>,
80 #[serde(skip_serializing_if = "Option::is_none")]
81 pub temperature: Option<f64>,
82 #[serde(skip_serializing_if = "Option::is_none")]
83 pub stream: Option<bool>,
84 #[serde(skip_serializing_if = "Option::is_none")]
85 pub user: Option<String>,
86}
87
88#[derive(Debug, serde::Deserialize, serde::Serialize)]
89pub struct Choice {
90 pub index: u64,
91 pub delta: Message,
92 pub finish_reason: Option<String>,
93}
94
95#[derive(Debug, serde::Deserialize, serde::Serialize)]
96pub struct Completion {
97 pub id: String,
98 pub object: String,
99 pub created: u64,
100 pub model: String,
101 pub choices: Vec<Choice>,
102}
103
104mod internal {
105 use curl;
106
107 pub fn init(api_key: &str, url: &str) -> Result<curl::easy::Easy, super::Error> {
108 let mut easy = curl::easy::Easy::new();
109 easy.url(url)?;
110 let mut headers = curl::easy::List::new();
111 headers.append(&format!("Authorization: Bearer {}", api_key))?;
112 headers.append("Content-Type: application/json")?;
113 headers.append("Accept: text/event-stream")?;
114 easy.http_headers(headers)?;
115 Ok(easy)
116 }
117}
118
119pub mod ll {
120 pub async fn list_models<F>(api_key: &str, f: F) -> Result<(), super::Error>
121 where
122 F: Fn(&[u8]),
123 {
124 let mut easy = super::internal::init(api_key, "https://api.openai.com/v1/models")?;
125 easy.get(true)?;
126 let mut transfer = easy.transfer();
127 transfer.write_function(|data| {
128 f(data);
129 Ok(data.len())
130 })?;
131 transfer.perform()?;
132 Ok(())
133 }
134
135 pub async fn completions<F>(
136 api_key: &str,
137 request_body: &super::RequestBody,
138 f: F,
139 ) -> Result<(), super::Error>
140 where
141 F: Fn(&[u8]),
142 {
143 let string_body = serde_json::to_string(request_body)?;
144 let mut easy =
145 super::internal::init(api_key, "https://api.openai.com/v1/chat/completions")?;
146 easy.post(true)?;
147 easy.post_fields_copy(string_body.as_bytes())?;
148 let mut transfer = easy.transfer();
149 transfer.write_function(|data| {
150 f(data);
151 Ok(data.len())
152 })?;
153 transfer.perform()?;
154 Ok(())
155 }
156}
157
158pub mod hl {
159 pub async fn list_models(api_key: &str) -> Result<String, super::Error> {
160 let amv = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
161 let cloned_amv = amv.clone();
162 super::ll::list_models(&api_key, move |data| {
163 let mut v = cloned_amv.lock().unwrap();
164 v.extend_from_slice(data)
165 })
166 .await?;
167 let v = amv.lock().unwrap().clone();
168 Ok(String::from_utf8(v)?)
169 }
170
171 pub async fn completions<F>(
172 api_key: &str,
173 request_body: &super::RequestBody,
174 f: F,
175 ) -> Result<(), super::Error>
176 where
177 F: Fn(String),
178 {
179 super::ll::completions(&api_key, request_body, |data| {
180 f(String::from_utf8(data.to_vec()).unwrap());
181 })
182 .await?;
183 Ok(())
184 }
185}
186
187pub async fn list_models(api_key: &str) -> Result<Vec<ModelList>, Error> {
188 let json = hl::list_models(api_key).await?;
189 Ok(serde_json::from_str(&json)?)
190}
191
192pub async fn completions<F>(api_key: &str, request_body: &RequestBody, f: F) -> Result<(), Error>
193where
194 F: Fn(CallbackReason, Option<Completion>),
195{
196 hl::completions(&api_key, request_body, |data| {
197 let is_debug = std::env::var("OPENAI_DEBUG").unwrap_or(String::from("")) != "";
198 data.lines().for_each(|line| {
199 let prefix = "data: ";
200 let prefix_len = prefix.len();
201 if line.starts_with(prefix) {
202 if is_debug {
203 eprintln!("{}", line);
204 }
205 if line == "data: [DONE]" {
206 f(CallbackReason::Done, None);
207 } else {
208 let string_json = match line.char_indices().nth(prefix_len) {
209 Some((i, _)) => &line[i..],
210 None => "",
211 };
212 let completion: Completion = serde_json::from_str(string_json).unwrap();
213 if completion.choices[0].delta.role == Some(String::from("assistant")) {
214 f(CallbackReason::Start, Some(completion));
215 } else if completion.choices[0].finish_reason == Some(String::from("stop")) {
216 f(CallbackReason::End, Some(completion));
217 } else {
218 f(CallbackReason::Data, Some(completion));
219 }
220 }
221 } else if line != "" {
222 if is_debug {
223 eprintln!("{}", line);
224 }
225 f(CallbackReason::Error(String::from(line)), None);
226 }
227 });
228 })
229 .await?;
230 Ok(())
231}
232
233#[cfg(test)]
234mod tests {
235 #[test]
236 fn init_handle() {
237 let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
238 let result = crate::internal::init(&api_key, "https://api.openai.com/v1/models");
239 assert!(result.is_ok());
240 }
241
242 #[test]
243 fn ll_list_models() {
244 let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
245 let future = crate::ll::list_models(&api_key, |_| {});
246 let result = futures::executor::block_on(future);
247 assert!(result.is_ok());
248 }
249
250 #[test]
251 fn hl_list_models() {
252 let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
253 let future = crate::hl::list_models(&api_key);
254 let result = futures::executor::block_on(future);
255 assert!(result.is_ok());
256 }
257
258 #[test]
259 fn ll_completions() {
260 let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
261 let request_body = super::RequestBody {
262 model: String::from("gpt-3.5-turbo"),
263 messages: vec![super::Message {
264 role: Some(String::from("user")),
265 content: Some(String::from("Say hello")),
266 }],
267 temperature: None,
268 stream: Some(true),
269 user: None,
270 };
271 let future = super::ll::completions(&api_key, &request_body, |_| {});
272 let result = futures::executor::block_on(future);
273 assert!(result.is_ok());
274 }
275
276 #[test]
277 fn hl_completions() {
278 let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
279 let request_body = super::RequestBody {
280 model: String::from("gpt-3.5-turbo"),
281 messages: vec![super::Message {
282 role: Some(String::from("user")),
283 content: Some(String::from("Say hello")),
284 }],
285 temperature: None,
286 stream: Some(true),
287 user: None,
288 };
289 let count_start = std::cell::Cell::new(0);
290 let count_data = std::cell::Cell::new(0);
291 let count_end = std::cell::Cell::new(0);
292 let future = super::hl::completions(&api_key, &request_body, |data| {
293 data.lines().for_each(|line| {
294 let prefix = "data: {";
295 let prefix_len = prefix.len();
296 if line.starts_with(prefix) {
297 let string_json = match line.char_indices().nth(prefix_len - 1) {
298 Some((i, _)) => &line[i..],
299 None => "",
300 };
301 let completion: super::Completion = serde_json::from_str(string_json).unwrap();
302 if completion.choices[0].delta.role == Some(String::from("assistant")) {
303 count_start.set(count_start.get() + 1);
304 } else if completion.choices[0].finish_reason == Some(String::from("stop")) {
305 count_end.set(count_end.get() + 1);
306 } else {
307 count_data.set(count_data.get() + 1);
308 }
309 }
310 });
311 });
312 let result = futures::executor::block_on(future);
313 assert!(result.is_ok());
314 assert_eq!(count_start.get(), 1);
315 assert!(count_data.get() > 0);
316 assert_eq!(count_end.get(), 1);
317 }
318
319 #[test]
320 fn completions() {
321 let api_key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not defined");
322 let request_body = super::RequestBody {
323 model: String::from("gpt-3.5-turbo"),
324 messages: vec![super::Message {
325 role: Some(String::from("user")),
326 content: Some(String::from("Say hello")),
327 }],
328 temperature: None,
329 stream: Some(true),
330 user: None,
331 };
332 let count_start = std::cell::Cell::new(0);
333 let count_data = std::cell::Cell::new(0);
334 let count_end = std::cell::Cell::new(0);
335 let future = super::completions(&api_key, &request_body, |cr, completion| match cr {
336 super::CallbackReason::Start => {
337 count_start.set(count_start.get() + 1);
338 assert!(completion.is_some());
339 assert_eq!(completion.unwrap().choices.len(), 1);
340 }
341 super::CallbackReason::Data => {
342 count_data.set(count_data.get() + 1);
343 assert!(completion.is_some());
344 assert_eq!(completion.unwrap().choices.len(), 1);
345 }
346 super::CallbackReason::End => {
347 count_end.set(count_end.get() + 1);
348 assert!(completion.is_some());
349 assert_eq!(completion.unwrap().choices.len(), 1);
350 }
351 super::CallbackReason::Done => {
352 assert!(completion.is_none());
353 }
354 super::CallbackReason::Error(_) => {
355 assert!(completion.is_none());
356 }
357 });
358 let result = futures::executor::block_on(future);
359 assert!(result.is_ok());
360 assert_eq!(count_start.get(), 1);
361 assert!(count_data.get() > 0);
362 assert_eq!(count_end.get(), 1);
363 }
364}