From 865be2bf75bb62b6aeee059f684400b4b9938a15 Mon Sep 17 00:00:00 2001 From: sigoden Date: Mon, 29 Apr 2024 06:51:03 +0800 Subject: feat: non-streaming returns completion stats (#456) --- src/client/ollama.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) (limited to 'src/client/ollama.rs') diff --git a/src/client/ollama.rs b/src/client/ollama.rs index 5ec50a1..ec83cbd 100644 --- a/src/client/ollama.rs +++ b/src/client/ollama.rs @@ -1,6 +1,6 @@ use super::{ - catch_error, message::*, ExtraConfig, Model, ModelConfig, OllamaClient, PromptType, - ReplyHandler, SendData, + catch_error, message::*, CompletionStats, ExtraConfig, Model, ModelConfig, OllamaClient, + PromptType, ReplyHandler, SendData, }; use crate::utils::PromptKind; @@ -59,17 +59,17 @@ impl OllamaClient { impl_client_trait!(OllamaClient, send_message, send_message_streaming); -async fn send_message(builder: RequestBuilder) -> Result { +async fn send_message(builder: RequestBuilder) -> Result<(String, CompletionStats)> { let res = builder.send().await?; let status = res.status(); let data = res.json().await?; if status != 200 { catch_error(&data, status.as_u16())?; } - let output = data["message"]["content"] + let text = data["message"]["content"] .as_str() .ok_or_else(|| anyhow!("Invalid response data: {data}"))?; - Ok(output.to_string()) + Ok((text.to_string(), CompletionStats::default())) } async fn send_message_streaming(builder: RequestBuilder, handler: &mut ReplyHandler) -> Result<()> { -- cgit v1.2.3