diff options
| author | sigoden <sigoden@gmail.com> | 2024-04-29 08:33:17 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-04-29 08:33:17 +0800 |
| commit | 37a0cd08a92f07ef24e39bdab7c8bced5b59c146 (patch) | |
| tree | 1082ed49fa584a7f8e07763146ab0e73a475b383 /src/client/cohere.rs | |
| parent | 865be2bf75bb62b6aeee059f684400b4b9938a15 (diff) | |
| download | aichat-37a0cd08a92f07ef24e39bdab7c8bced5b59c146.tar.gz | |
refactor: rename some structs (#457)
Diffstat (limited to 'src/client/cohere.rs')
| -rw-r--r-- | src/client/cohere.rs | 14 |
1 files changed, 7 insertions, 7 deletions
diff --git a/src/client/cohere.rs b/src/client/cohere.rs index 6186718..de2a11c 100644 --- a/src/client/cohere.rs +++ b/src/client/cohere.rs @@ -1,6 +1,6 @@ use super::{ - catch_error, extract_system_message, json_stream, message::*, CohereClient, CompletionStats, - ExtraConfig, Model, ModelConfig, PromptType, ReplyHandler, SendData, + catch_error, extract_system_message, json_stream, message::*, CohereClient, CompletionDetails, + ExtraConfig, Model, ModelConfig, PromptType, SendData, SseHandler, }; use crate::utils::PromptKind; @@ -47,7 +47,7 @@ impl CohereClient { impl_client_trait!(CohereClient, send_message, send_message_streaming); -async fn send_message(builder: RequestBuilder) -> Result<(String, CompletionStats)> { +async fn send_message(builder: RequestBuilder) -> Result<(String, CompletionDetails)> { let res = builder.send().await?; let status = res.status(); let data: Value = res.json().await?; @@ -58,7 +58,7 @@ async fn send_message(builder: RequestBuilder) -> Result<(String, CompletionStat cohere_extract_completion(&data) } -async fn send_message_streaming(builder: RequestBuilder, handler: &mut ReplyHandler) -> Result<()> { +async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandler) -> Result<()> { let res = builder.send().await?; let status = res.status(); if status != 200 { @@ -156,15 +156,15 @@ fn build_body(data: SendData, model: &Model) -> Result<Value> { Ok(body) } -fn cohere_extract_completion(data: &Value) -> Result<(String, CompletionStats)> { +fn cohere_extract_completion(data: &Value) -> Result<(String, CompletionDetails)> { let text = data["text"] .as_str() .ok_or_else(|| anyhow!("Invalid response data: {data}"))?; - let stats = CompletionStats { + let details = CompletionDetails { id: data["generation_id"].as_str().map(|v| v.to_string()), input_tokens: data["meta"]["billed_units"]["input_tokens"].as_u64(), output_tokens: data["meta"]["billed_units"]["output_tokens"].as_u64(), }; - Ok((text.to_string(), stats)) + Ok((text.to_string(), details)) } |
