summaryrefslogtreecommitdiffstats
path: root/src/client/cohere.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/client/cohere.rs')
-rw-r--r--src/client/cohere.rs32
1 files changed, 5 insertions, 27 deletions
diff --git a/src/client/cohere.rs b/src/client/cohere.rs
index 27acb6c..3f97ece 100644
--- a/src/client/cohere.rs
+++ b/src/client/cohere.rs
@@ -1,12 +1,11 @@
use super::{
- catch_error, extract_system_message, json_stream, message::*, Client, CohereClient,
+ catch_error, extract_system_message, json_stream, message::*, CohereClient,
ExtraConfig, Model, ModelConfig, PromptType, ReplyHandler, SendData,
};
use crate::utils::PromptKind;
use anyhow::{bail, Result};
-use async_trait::async_trait;
use reqwest::{Client as ReqwestClient, RequestBuilder};
use serde::Deserialize;
use serde_json::{json, Value};
@@ -22,26 +21,6 @@ pub struct CohereConfig {
pub extra: Option<ExtraConfig>,
}
-#[async_trait]
-impl Client for CohereClient {
- client_common_fns!();
-
- async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String> {
- let builder = self.request_builder(client, data)?;
- send_message(builder).await
- }
-
- async fn send_message_streaming_inner(
- &self,
- client: &ReqwestClient,
- handler: &mut ReplyHandler,
- data: SendData,
- ) -> Result<()> {
- let builder = self.request_builder(client, data)?;
- send_message_streaming(builder, handler).await
- }
-}
-
impl CohereClient {
list_models_fn!(
CohereConfig,
@@ -74,7 +53,9 @@ impl CohereClient {
}
}
-pub(crate) async fn send_message(builder: RequestBuilder) -> Result<String> {
+impl_client_trait!(CohereClient, send_message, send_message_streaming);
+
+async fn send_message(builder: RequestBuilder) -> Result<String> {
let res = builder.send().await?;
let status = res.status();
let data: Value = res.json().await?;
@@ -85,10 +66,7 @@ pub(crate) async fn send_message(builder: RequestBuilder) -> Result<String> {
Ok(output.to_string())
}
-pub(crate) async fn send_message_streaming(
- builder: RequestBuilder,
- handler: &mut ReplyHandler,
-) -> Result<()> {
+async fn send_message_streaming(builder: RequestBuilder, handler: &mut ReplyHandler) -> Result<()> {
let res = builder.send().await?;
let status = res.status();
if status != 200 {