diff options
| author | sigoden <sigoden@gmail.com> | 2024-11-14 06:03:06 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-11-14 06:03:06 +0800 |
| commit | cfa9217422dfa6cf10bed6e6e3fab0722f58588d (patch) | |
| tree | f004b226caae309c860914e7393e516464eecee8 /src/client/model.rs | |
| parent | ff0ea19b48a18e9a1849bd2a653173a3ed4e4560 (diff) | |
| download | aichat-cfa9217422dfa6cf10bed6e6e3fab0722f58588d.tar.gz | |
feat: save function calls in the session (#994)
Diffstat (limited to 'src/client/model.rs')
| -rw-r--r-- | src/client/model.rs | 26 |
1 files changed, 23 insertions, 3 deletions
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() } |
