Implement edits

This commit is contained in:
Gabriel Tofvesson 2023-03-17 21:33:05 +01:00
parent de0a109a4a
commit 0be700b0b8
No known key found for this signature in database
GPG Key ID: 6F1345DF28EDA13E
3 changed files with 82 additions and 1 deletions

View File

@ -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?)
}
}

38
src/edits.rs Normal file
View File

@ -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
}

View File

@ -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"));
}
}