Message Role enum

Author: Guantong

Cargo.lock

@@ -63,10 +63,7 @@ fn main() {
         frequency_penalty: None,
         logit_bias: None,
         user: None,
-        messages: vec![Message {
-            role: "user".to_string(),
-            content: "Hello!".to_string(),
-        }],
+        messages: vec![Message { role: Role::User, content: "Hello!".to_string() }],
     };
     let rs = openai.chat_completion_create(&body);
     let choice = rs.unwrap().choices;

README.md

@@ -9,10 +9,7 @@ use crate::requests::Requests;
 use crate::*;
 use serde::{Deserialize, Serialize};
 
-use super::{
-	completions::{Completion, Message},
-	CHAT_COMPLETION_CREATE,
-};
+use super::{completions::Completion, CHAT_COMPLETION_CREATE};
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ChatBody {
@@ -98,10 +95,7 @@ impl ChatApi for OpenAI {
 
 #[cfg(test)]
 mod tests {
-	use crate::{
-		apis::{chat::ChatBody, completions::Message},
-		openai::new_test_openai,
-	};
+	use crate::{apis::chat::ChatBody, openai::new_test_openai, Message, Role};
 
 	use super::ChatApi;
 
@@ -120,10 +114,7 @@ mod tests {
 			frequency_penalty: None,
 			logit_bias: None,
 			user: None,
-			messages: vec![Message {
-				role: "user".to_string(),
-				content: "Hello!".to_string(),
-			}],
+			messages: vec![Message { role: Role::User, content: "Hello!".to_string() }],
 		};
 		let rs = openai.chat_completion_create(&body);
 		let choice = rs.unwrap().choices;

src/apis/chat.rs

@@ -126,21 +126,6 @@ pub struct CompletionsBody {
 	pub user: Option<String>,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
-pub struct Choice {
-	pub text: Option<String>,
-	pub index: u32,
-	pub logprobs: Option<String>,
-	pub finish_reason: Option<String>,
-	pub message: Option<Message>,
-}
-
-#[derive(Debug, Serialize, Deserialize)]
-pub struct Message {
-	pub role: String,
-	pub content: String,
-}
-
 pub trait CompletionsApi {
 	/// Creates a completion for the provided prompt and parameters
 	fn completion_create(&self, completions_body: &CompletionsBody) -> ApiResult<Completion>;

src/apis/completions.rs

@@ -33,3 +33,26 @@ pub struct Usage {
 	pub completion_tokens: Option<u32>,
 	pub total_tokens: Option<u32>,
 }
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Choice {
+	pub text: Option<String>,
+	pub index: u32,
+	pub logprobs: Option<String>,
+	pub finish_reason: Option<String>,
+	pub message: Option<Message>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Message {
+	pub role: Role,
+	pub content: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum Role {
+	System,
+	Assistant,
+	User,
+}

src/apis/mod.rs

@@ -58,7 +58,7 @@ impl OpenAI {
 			_ => {
 				log::warn!("Missing http_proxy or https_proxy");
 				None
-			}
+			},
 		};
 		if let Some(proxy) = proxy {
 			let proxy = ureq::Proxy::new(&proxy).unwrap();