1 files changed, 18 insertions, 5 deletions

diff --git a/src/client/ollama.rs b/src/client/ollama.rs
index 6408d2e..24d2dfe 100644
--- a/src/client/ollama.rs
+++ b/src/client/ollama.rs
@@ -1,5 +1,5 @@
 use super::{
-    catch_error, message::*, CompletionDetails, ExtraConfig, Model, ModelConfig, OllamaClient,
+    catch_error, message::*, CompletionOutput, ExtraConfig, Model, ModelData, OllamaClient,
     PromptAction, PromptKind, SendData, SseHandler,
 };
 
@@ -15,7 +15,7 @@ pub struct OllamaConfig {
     pub api_base: Option<String>,
     pub api_auth: Option<String>,
     pub chat_endpoint: Option<String>,
-    pub models: Vec<ModelConfig>,
+    pub models: Vec<ModelData>,
     pub extra: Option<ExtraConfig>,
 }
 
@@ -59,17 +59,18 @@ impl OllamaClient {
 
 impl_client_trait!(OllamaClient, send_message, send_message_streaming);
 
-async fn send_message(builder: RequestBuilder) -> Result<(String, CompletionDetails)> {
+async fn send_message(builder: RequestBuilder) -> Result<CompletionOutput> {
     let res = builder.send().await?;
     let status = res.status();
     let data = res.json().await?;
     if !status.is_success() {
         catch_error(&data, status.as_u16())?;
     }
+    debug!("non-stream-data: {data}");
     let text = data["message"]["content"]
         .as_str()
         .ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
-    Ok((text.to_string(), CompletionDetails::default()))
+    Ok(CompletionOutput::new(text))
 }
 
 async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandler) -> Result<()> {
@@ -86,6 +87,7 @@ async fn send_message_streaming(builder: RequestBuilder, handler: &mut SseHandle
                 continue;
             }
             let data: Value = serde_json::from_slice(&chunk)?;
+            debug!("stream-data: {data}");
             if data["done"].is_boolean() {
                 if let Some(text) = data["message"]["content"].as_str() {
                     handler.text(text)?;
@@ -103,10 +105,13 @@ fn build_body(data: SendData, model: &Model) -> Result<Value> {
         messages,
         temperature,
         top_p,
+        functions: _,
         stream,
     } = data;
 
+    let mut is_tool_call = false;
     let mut network_image_urls = vec![];
+
     let messages: Vec<Value> = messages
         .into_iter()
         .map(|message| {
@@ -141,10 +146,18 @@ fn build_body(data: SendData, model: &Model) -> Result<Value> {
                     let content = content.join("\n\n");
                     json!({ "role": role, "content": content, "images": images })
                 }
+                MessageContent::ToolResults(_) => {
+                    is_tool_call = true;
+                    json!({ "role": role })
+                }
             }
         })
         .collect();
 
+    if is_tool_call {
+        bail!("The client does not support function calling",);
+    }
+
     if !network_image_urls.is_empty() {
         bail!(
             "The model does not support network images: {:?}",
@@ -153,7 +166,7 @@ fn build_body(data: SendData, model: &Model) -> Result<Value> {
     }
 
     let mut body = json!({
-        "model": &model.name,
+        "model": &model.name(),
         "messages": messages,
         "stream": stream,
         "options": {},