From 0be700b0b8112ab74ae61ac5baedf01d5767d84f Mon Sep 17 00:00:00 2001
From: Gabriel Tofvesson <contact@w1zzrd.dev>
Date: Fri, 17 Mar 2023 21:33:05 +0100
Subject: [PATCH] Implement edits

---
 src/context.rs |  6 +++++-
 src/edits.rs   | 38 ++++++++++++++++++++++++++++++++++++++
 src/lib.rs     | 39 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 src/edits.rs

diff --git a/src/context.rs b/src/context.rs
index cb78a16..f891cc0 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -1,6 +1,6 @@
 use reqwest::{Client, RequestBuilder};
 
-use crate::{model::{Model, ModelList}, completion::{CompletionRequest, CompletionResponse}, chat::{ChatCompletionResponse, ChatHistory}};
+use crate::{model::{Model, ModelList}, completion::{CompletionRequest, CompletionResponse}, chat::{ChatCompletionResponse, ChatHistory}, edits::{EditRequest, EditResponse}};
 
 pub struct Context {
     api_key: String,
@@ -49,4 +49,8 @@ impl Context {
     pub async fn create_chat_completion(&self, chat_completion_request: ChatHistory) -> anyhow::Result<ChatCompletionResponse> {
         Ok(self.with_auth(Client::builder().build()?.post(&format!("{API_URL}/v1/chat/completions")).json(&chat_completion_request)).send().await?.json::<ChatCompletionResponse>().await?)
     }
+
+    pub async fn create_edit(&self, edit_request: EditRequest) -> anyhow::Result<crate::edits::EditResponse> {
+        Ok(self.with_auth(Client::builder().build()?.post(&format!("{API_URL}/v1/edits")).json(&edit_request)).send().await?.json::<EditResponse>().await?)
+    }
 }
\ No newline at end of file
diff --git a/src/edits.rs b/src/edits.rs
new file mode 100644
index 0000000..a1574d9
--- /dev/null
+++ b/src/edits.rs
@@ -0,0 +1,38 @@
+use derive_builder::Builder;
+use serde::{Serialize, Deserialize};
+
+use crate::completion::Usage;
+
+#[derive(Debug, Serialize, Builder)]
+pub struct EditRequest {
+    #[builder(setter(into))]
+    pub model: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[builder(setter(into, strip_option), default)]
+    pub input: Option<String>,
+    #[builder(setter(into))]
+    pub instruction: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[builder(setter(into, strip_option), default)]
+    pub temperature: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[builder(setter(into, strip_option), default)]
+    pub top_p: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[builder(setter(into, strip_option), default)]
+    pub n: Option<u32>,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct Edit {
+    pub text: String,
+    pub index: i32,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct EditResponse {
+    /* pub object: "edit", */
+    pub created: u64,
+    pub choices: Vec<Edit>,
+    pub usage: Usage
+}
\ No newline at end of file
diff --git a/src/lib.rs b/src/lib.rs
index d9f8b2e..9dd3d9f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@ pub mod context;
 pub mod model;
 pub mod completion;
 pub mod chat;
+pub mod edits;
 
 #[cfg(test)]
 mod tests {
@@ -57,4 +58,42 @@ mod tests {
         assert!(completion.is_ok(), "Could not get completion: {}", completion.unwrap_err());
         assert!(completion.unwrap().choices.len() == 1, "No completion found");
     }
+
+    #[tokio::test]
+    async fn test_edits() {
+        let ctx = get_api();
+        assert!(ctx.is_ok(), "Could not load context");
+        let ctx = ctx.unwrap();
+
+        // Autocorrect English spelling errors
+        let edit = ctx.create_edit(
+            crate::edits::EditRequestBuilder::default()
+            .model("text-davinci-edit-001")
+            .instruction("Correct all spelling mistakes")
+            .input("What a wnoderful day!")
+            .build()
+            .unwrap()
+        ).await;
+
+        assert!(edit.is_ok(), "Could not get edit: {}", edit.unwrap_err());
+        assert!(edit.as_ref().unwrap().choices.len() == 1, "No edit found");
+        assert!(edit.unwrap().choices[0].text.replace("\n", "").eq("What a wonderful day!"));
+
+        
+        // Autocorrect Dutch spelling errors using an English instructino prompt?
+        let edit = ctx.create_edit(
+            crate::edits::EditRequestBuilder::default()
+            .model("text-davinci-edit-001")
+            .instruction("Correct all spelling mistakes")
+            .input("Ik hou van jouw moederr")
+            .build()
+            .unwrap()
+        ).await;
+
+        assert!(edit.is_ok(), "Could not get edit: {}", edit.unwrap_err());
+        assert!(edit.as_ref().unwrap().choices.len() == 1, "No edit found");
+        
+        // This one might be pushing my luck a bit, but it seems to work
+        //assert!(edit.unwrap().choices[0].text.replace("\n", "").eq("Ik hou van jouw moeder"));
+    }
 }
\ No newline at end of file