refactor: handling of system message (#432)

1 files changed, 6 insertions, 2 deletions

diff --git a/src/client/cohere.rs b/src/client/cohere.rs
index 445c145..db2ad28 100644
--- a/src/client/cohere.rs
+++ b/src/client/cohere.rs
@@ -1,5 +1,5 @@
 use super::{
-    json_stream, message::*, patch_system_message, Client, CohereClient, ExtraConfig, Model,
+    extract_sytem_message, json_stream, message::*, Client, CohereClient, ExtraConfig, Model,
     ModelConfig, PromptType, ReplyHandler, SendData,
 };
 
@@ -129,7 +129,7 @@ pub(crate) fn build_body(data: SendData, model: &Model) -> Result<Value> {
         stream,
     } = data;
 
-    patch_system_message(&mut messages);
+    let system_message = extract_sytem_message(&mut messages);
 
     let mut image_urls = vec![];
     let mut messages: Vec<Value> = messages
@@ -174,6 +174,10 @@ pub(crate) fn build_body(data: SendData, model: &Model) -> Result<Value> {
         "message": message,
     });
 
+    if let Some(preamble) = system_message {
+        body["preamble"] = preamble.into();
+    }
+
     if let Some(max_tokens) = model.max_output_tokens {
         body["max_tokens"] = max_tokens.into();
     }