From 57844638ce504f1a69eb770f27bbafb7077ccb19 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 2 Mar 2023 23:11:34 +0800
Subject: feat: add config.temperature

---
 src/config.rs |  2 ++
 src/main.rs   | 16 +++++++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 0feed64..729320e 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -7,6 +7,8 @@ use serde::Deserialize;
 pub struct Config {
     /// Openai api key
     pub api_key: String,
+    /// What sampling temperature to use, between 0 and 2
+    pub temperature: Option<f64>,
     /// Set proxy
     pub proxy: Option<String>,
     /// Used only for debugging
diff --git a/src/main.rs b/src/main.rs
index 53c87b7..eb88226 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -26,7 +26,7 @@ const MODEL: &str = "gpt-3.5-turbo";
 const REPL_COMMANDS: [(&str, &str); 6] = [
     (".clear", "Clear the screen"),
     (".clear-history", "Clear the history"),
-    (".exit", " Exit the REPL"),
+    (".exit", "Exit the REPL"),
     (".help", "Print this help message"),
     (".history", "Print the history"),
     (".role", "Specify the role that the AI will play"),
@@ -319,11 +319,16 @@ async fn acquire(client: &Client, config: &Config, content: &str) -> Result<Stri
     if config.dry_run {
         return Ok(content.to_string());
     }
-    let body = json!({
+    let mut body = json!({
         "model": MODEL,
         "messages": [{"role": "user", "content": content}]
     });
 
+    if let Some(v) = config.temperature {
+        body.as_object_mut()
+            .and_then(|m| m.insert("temperature".into(), json!(v)));
+    }
+
     let data: Value = client
         .post(API_URL)
         .bearer_auth(&config.api_key)
@@ -345,12 +350,17 @@ async fn acquire_stream(
     config: &Config,
     content: &str,
 ) -> Result<EventStream<impl Stream<Item = reqwest::Result<bytes::Bytes>>>> {
-    let body = json!({
+    let mut body = json!({
         "model": MODEL,
         "messages": [{"role": "user", "content": content}],
         "stream": true,
     });
 
+    if let Some(v) = config.temperature {
+        body.as_object_mut()
+            .and_then(|m| m.insert("temperature".into(), json!(v)));
+    }
+
     let stream = client
         .post(API_URL)
         .bearer_auth(&config.api_key)
-- 
cgit v1.2.3