From 571d1022f628cb7d2a3125664bd3293bac4471b5 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sat, 1 Jun 2024 17:47:49 +0800 Subject: refactor: rename some client structs and methods (#555) * rename `Completeion*` to `ChatCompletions*` * rename `send_message*` to `chat_completions*` * rename `request_builder` to `chat_completions_builder` * rename `build_body` to `build_chat_completions_body` * rename `extract_completion` to `extract_chat_completions` * format * remove unused config fields --- src/client/cohere.rs | 31 +++++++++++++++++-------------- 1 file changed, 17 insertions(+), 14 deletions(-) (limited to 'src/client/cohere.rs') diff --git a/src/client/cohere.rs b/src/client/cohere.rs index e093778..e0a5eec 100644 --- a/src/client/cohere.rs +++ b/src/client/cohere.rs @@ -1,7 +1,7 @@ use super::{ - catch_error, extract_system_message, json_stream, message::*, Client, CohereClient, - CompletionData, CompletionOutput, ExtraConfig, Model, ModelData, ModelPatches, PromptAction, - PromptKind, SseHandler, ToolCall, + catch_error, extract_system_message, json_stream, message::*, ChatCompletionsData, + ChatCompletionsOutput, Client, CohereClient, ExtraConfig, Model, ModelData, ModelPatches, + PromptAction, PromptKind, SseHandler, ToolCall, }; use anyhow::{bail, Result}; @@ -27,14 +27,14 @@ impl CohereClient { pub const PROMPTS: [PromptAction<'static>; 1] = [("api_key", "API Key:", true, PromptKind::String)]; - fn request_builder( + fn chat_completions_builder( &self, client: &ReqwestClient, - data: CompletionData, + data: ChatCompletionsData, ) -> Result { let api_key = self.get_api_key()?; - let mut body = build_body(data, &self.model)?; + let mut body = build_chat_completions_body(data, &self.model)?; self.patch_request_body(&mut body); let url = API_URL; @@ -47,9 +47,9 @@ impl CohereClient { } } -impl_client_trait!(CohereClient, send_message, send_message_streaming); +impl_client_trait!(CohereClient, chat_completions, chat_completions_streaming); -async fn send_message(builder: RequestBuilder) -> Result { +async fn chat_completions(builder: RequestBuilder) -> Result { let res = builder.send().await?; let status = res.status(); let data: Value = res.json().await?; @@ -58,10 +58,13 @@ async fn send_message(builder: RequestBuilder) -> Result { } debug!("non-stream-data: {data}"); - extract_completion(&data) + extract_chat_completions(&data) } -async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandler) -> Result<()> { +async fn chat_completions_streaming( + builder: RequestBuilder, + handler: &mut SseHandler, +) -> Result<()> { let res = builder.send().await?; let status = res.status(); if !status.is_success() { @@ -97,8 +100,8 @@ async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandle Ok(()) } -fn build_body(data: CompletionData, model: &Model) -> Result { - let CompletionData { +fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Result { + let ChatCompletionsData { mut messages, temperature, top_p, @@ -224,7 +227,7 @@ fn build_body(data: CompletionData, model: &Model) -> Result { Ok(body) } -fn extract_completion(data: &Value) -> Result { +fn extract_chat_completions(data: &Value) -> Result { let text = data["text"].as_str().unwrap_or_default(); let mut tool_calls = vec![]; @@ -246,7 +249,7 @@ fn extract_completion(data: &Value) -> Result { if text.is_empty() && tool_calls.is_empty() { bail!("Invalid response data: {data}"); } - let output = CompletionOutput { + let output = ChatCompletionsOutput { text: text.to_string(), tool_calls, id: data["generation_id"].as_str().map(|v| v.to_string()), -- cgit v1.2.3