From 4f8d895154c602fb51a7694911efd96d52129f82 Mon Sep 17 00:00:00 2001 From: sigoden Date: Wed, 24 Apr 2024 07:56:24 +0800 Subject: refactor: handling of system message (#432) --- src/client/cohere.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) (limited to 'src/client/cohere.rs') 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 { stream, } = data; - patch_system_message(&mut messages); + let system_message = extract_sytem_message(&mut messages); let mut image_urls = vec![]; let mut messages: Vec = messages @@ -174,6 +174,10 @@ pub(crate) fn build_body(data: SendData, model: &Model) -> Result { "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(); } -- cgit v1.2.3