1 files changed, 14 insertions, 9 deletions

diff --git a/src/client/ernie.rs b/src/client/ernie.rs
index 28cb857..1d79138 100644
--- a/src/client/ernie.rs
+++ b/src/client/ernie.rs
@@ -1,7 +1,7 @@
 use super::access_token::*;
 use super::{
-    maybe_catch_error, patch_system_message, sse_stream, Client, CompletionDetails, ErnieClient,
-    ExtraConfig, Model, ModelConfig, PromptAction, PromptKind, SendData, SsMmessage, SseHandler,
+    maybe_catch_error, patch_system_message, sse_stream, Client, CompletionOutput, ErnieClient,
+    ExtraConfig, Model, ModelData, PromptAction, PromptKind, SendData, SsMmessage, SseHandler,
 };
 
 use anyhow::{anyhow, Context, Result};
@@ -20,7 +20,7 @@ pub struct ErnieConfig {
     pub api_key: Option<String>,
     pub secret_key: Option<String>,
     #[serde(default)]
-    pub models: Vec<ModelConfig>,
+    pub models: Vec<ModelData>,
     pub extra: Option<ExtraConfig>,
 }
 
@@ -36,7 +36,7 @@ impl ErnieClient {
 
         let url = format!(
             "{API_BASE}/wenxinworkshop/chat/{}?access_token={access_token}",
-            &self.model.name,
+            &self.model.name(),
         );
 
         debug!("Ernie Request: {url} {body}");
@@ -78,7 +78,7 @@ impl Client for ErnieClient {
         &self,
         client: &ReqwestClient,
         data: SendData,
-    ) -> Result<(String, CompletionDetails)> {
+    ) -> Result<CompletionOutput> {
         self.prepare_access_token().await?;
         let builder = self.request_builder(client, data)?;
         send_message(builder).await
@@ -96,15 +96,17 @@ impl Client for ErnieClient {
     }
 }
 
-async fn send_message(builder: RequestBuilder) -> Result<(String, CompletionDetails)> {
+async fn send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
     let data: Value = builder.send().await?.json().await?;
     maybe_catch_error(&data)?;
+    debug!("non-stream-data: {data}");
     extract_completion_text(&data)
 }
 
 async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandler) -> Result<()> {
     let handle = |message: SsMmessage| -> Result<bool> {
         let data: Value = serde_json::from_str(&message.data)?;
+        debug!("stream-data: {data}");
         if let Some(text) = data["result"].as_str() {
             handler.text(text)?;
         }
@@ -119,6 +121,7 @@ fn build_body(data: SendData, model: &Model) -> Value {
         mut messages,
         temperature,
         top_p,
+        functions: _,
         stream,
     } = data;
 
@@ -145,16 +148,18 @@ fn build_body(data: SendData, model: &Model) -> Value {
     body
 }
 
-fn extract_completion_text(data: &Value) -> Result<(String, CompletionDetails)> {
+fn extract_completion_text(data: &Value) -> Result<CompletionOutput> {
     let text = data["result"]
         .as_str()
         .ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
-    let details = CompletionDetails {
+    let output = CompletionOutput {
+        text: text.to_string(),
+        tool_calls: vec![],
         id: data["id"].as_str().map(|v| v.to_string()),
         input_tokens: data["usage"]["prompt_tokens"].as_u64(),
         output_tokens: data["usage"]["completion_tokens"].as_u64(),
     };
-    Ok((text.to_string(), details))
+    Ok(output)
 }
 
 async fn fetch_access_token(