From cfa9217422dfa6cf10bed6e6e3fab0722f58588d Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 14 Nov 2024 06:03:06 +0800
Subject: feat: save function calls in the session (#994)

---
 src/client/model.rs | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

(limited to 'src/client/model.rs')

diff --git a/src/client/model.rs b/src/client/model.rs
index f8a23ff..af864e2 100644
--- a/src/client/model.rs
+++ b/src/client/model.rs
@@ -1,6 +1,7 @@
 use super::{
     list_chat_models, list_embedding_models, list_reranker_models,
-    message::{Message, MessageContent},
+    message::{Message, MessageContent, MessageContentPart},
+    ToolResults,
 };
 
 use crate::config::Config;
@@ -229,8 +230,27 @@ impl Model {
             .iter()
             .map(|v| match &v.content {
                 MessageContent::Text(text) => estimate_token_length(text),
-                MessageContent::Array(_) => 0,
-                MessageContent::ToolResults(_) => 0,
+                MessageContent::Array(list) => list
+                    .iter()
+                    .map(|v| match v {
+                        MessageContentPart::Text { text } => estimate_token_length(text),
+                        MessageContentPart::ImageUrl { .. } => 0,
+                    })
+                    .sum(),
+                MessageContent::ToolResults(results) => {
+                    let ToolResults {
+                        tool_results, text, ..
+                    } = results;
+                    estimate_token_length(text)
+                        + tool_results
+                            .iter()
+                            .map(|v| {
+                                serde_json::to_string(v)
+                                    .map(|v| estimate_token_length(&v))
+                                    .unwrap_or_default()
+                            })
+                            .sum::<usize>()
+                }
             })
             .sum()
     }
-- 
cgit v1.2.3