From 2d2eda072cd557066ff0c52387b94c5160ae3d7d Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 00:40:04 -0400 Subject: [PATCH 01/11] Usage struct added and add usage attribute to ChatCompletionResponse struct --- src/api/chat_completions.rs | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/api/chat_completions.rs b/src/api/chat_completions.rs index 9306b34..bd81955 100644 --- a/src/api/chat_completions.rs +++ b/src/api/chat_completions.rs @@ -1,4 +1,4 @@ -use super::json_models::chat_completion::ChatCompletionResponse; +use super::json_models::chat_completion::{ChatCompletionResponse, Usage}; use super::Instance; use crate::clue::ClueCollection; use serde_json::json; @@ -25,7 +25,7 @@ impl Instance { &self, link_words: Vec, avoid_words: Vec, - ) -> Result> { + ) -> Result<(ClueCollection, Usage), Box> { let request_body = self.build_request_body(link_words, avoid_words); // Get response from API endpoint @@ -38,12 +38,19 @@ impl Instance { .await .map_err(|_| "Failed to fetch clue collection from API server")?; - // Deserialize the response - let clue_strings = response + let parsed_response = response .json::() .await - .map_err(|_| "Failed to parse clues from API server")? - .choices[0] + .map_err(|_| "Failed to parse clues from API server")?; + + // Extract usage information from the parsed response + let usage = parsed_response.usage.clone(); + + // Extract clue strings from the parsed response + let clue_strings = parsed_response + .choices + .get(0) + .ok_or("No choices returned from API")? .message .content .lines() @@ -53,7 +60,7 @@ impl Instance { // Build clues let clue_collection = ClueCollection::new(clue_strings); - Ok(clue_collection) + Ok((clue_collection, usage)) } fn build_request_body( From 16a9809c3d23e82fe2f0796170381be0186f9782 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 00:41:07 -0400 Subject: [PATCH 02/11] -t && --token-usage flags added --- src/lib.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index d66cda0..14e64f0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,6 +21,10 @@ pub struct Args { #[arg(short, long, value_name = "FILE")] pub output: Option, + /// Print all token usage information + #[arg(short, long = "token-usage")] + pub token: bool, + /// File containing words to link together - the words from your team #[arg(required_unless_present = "get")] pub to_link: Option, From 203bb9dad8d5ec8d64ce29c58e9381fc67048134 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 00:44:58 -0400 Subject: [PATCH 03/11] Additional object returned (Usage) fetching token's usage data also added to fetch_clue_collection() --- src/api/json_models/chat_completion.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/api/json_models/chat_completion.rs b/src/api/json_models/chat_completion.rs index 47fd4b2..9f1b362 100644 --- a/src/api/json_models/chat_completion.rs +++ b/src/api/json_models/chat_completion.rs @@ -10,7 +10,15 @@ pub struct Choice { pub message: Message, } +#[derive(Deserialize, Clone)] +pub struct Usage { + pub prompt_tokens: usize, + pub completion_tokens: usize, + pub total_tokens: usize, +} + #[derive(Deserialize)] pub struct ChatCompletionResponse { pub choices: Vec, + pub usage: Usage, } From a5a58bcdae25f96563c68fd025621339d4874574 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 00:46:47 -0400 Subject: [PATCH 04/11] New feature integrated to main file and output in stderr added --- src/main.rs | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/main.rs b/src/main.rs index fbdf4f0..efe48b2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,8 @@ use clap::Parser; use mastermind::api::Instance; use mastermind::*; +use std::io::{self, Write}; + #[tokio::main] async fn main() -> Result<(), Box> { // Read arguments and environment variables @@ -27,10 +29,11 @@ async fn main() -> Result<(), Box> { let avoid_words = read_words_from_file(args.to_avoid.unwrap())?; // Get clues from API - let clue_collection = api_instance + let (clue_collection, usage) = api_instance .fetch_clue_collection(link_words, avoid_words) .await?; + // Output if clue_collection.is_empty() { println!("The language model didn't return any useful clues. Maybe try again?"); @@ -41,5 +44,15 @@ async fn main() -> Result<(), Box> { clue_collection.display(); } + // If -t is set, output the token usage information + if args.token { + // Write to stderr in the format: prompt_tokens, completion_tokens, total_tokens + writeln!( + io::stderr(), + "\nTokens Usage\n----------------------\nPrompt Tokens: {}\nCompletion Tokens: {}\n----------------------\nTotal Tokens: {}", + usage.prompt_tokens, usage.completion_tokens, usage.total_tokens + )?; + } + Ok(()) } From 0bfd7a7c447426036757649933f7d2fab2a605a3 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 00:48:59 -0400 Subject: [PATCH 05/11] Documentations updated, information about -t & --token-usage provided. --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 14d8383..ff28d59 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,7 @@ Feel free to run the program multiple times to get the best result! - `-o`, `--output` : Specify an output file - `-h`, `--help` : Print help - `-V`, `--version` : Print version +- `-t`, `--token-usage` : Print token usage ## 🛠️ Building From eda931105bbd65cf54d54bffc895492c690ff042 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 17:42:10 -0400 Subject: [PATCH 06/11] Changed error output, api instance doesn't return tuple anymore --- src/main.rs | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/main.rs b/src/main.rs index efe48b2..983c8f3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,8 +3,6 @@ use clap::Parser; use mastermind::api::Instance; use mastermind::*; -use std::io::{self, Write}; - #[tokio::main] async fn main() -> Result<(), Box> { // Read arguments and environment variables @@ -29,7 +27,7 @@ async fn main() -> Result<(), Box> { let avoid_words = read_words_from_file(args.to_avoid.unwrap())?; // Get clues from API - let (clue_collection, usage) = api_instance + let clue_collection = api_instance .fetch_clue_collection(link_words, avoid_words) .await?; @@ -47,11 +45,8 @@ async fn main() -> Result<(), Box> { // If -t is set, output the token usage information if args.token { // Write to stderr in the format: prompt_tokens, completion_tokens, total_tokens - writeln!( - io::stderr(), - "\nTokens Usage\n----------------------\nPrompt Tokens: {}\nCompletion Tokens: {}\n----------------------\nTotal Tokens: {}", - usage.prompt_tokens, usage.completion_tokens, usage.total_tokens - )?; + eprintln!("\nTokens Usage\n----------------------\nPrompt Tokens: {}\nCompletion Tokens: {}\n----------------------\nTotal Tokens: {}", + clue_collection.usage.prompt_tokens, clue_collection.usage.completion_tokens, clue_collection.usage.total_tokens); } Ok(()) From 171d6705e2b411a7f0fcb7d86e8a9819100c4f0a Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 17:44:40 -0400 Subject: [PATCH 07/11] usage import from chat_completion.rs added as atribute to ClueCollection struct, ClueCollection Constructor updated --- src/clue.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/clue.rs b/src/clue.rs index 3f53ce9..4736355 100644 --- a/src/clue.rs +++ b/src/clue.rs @@ -1,7 +1,7 @@ use comfy_table::modifiers::UTF8_ROUND_CORNERS; use comfy_table::presets::UTF8_FULL; use comfy_table::{Attribute, Cell, CellAlignment, ContentArrangement, Table}; - +use crate::api::json_models::chat_completion::Usage; struct Clue { clue_word: String, count: usize, @@ -10,6 +10,7 @@ struct Clue { pub struct ClueCollection { clues: Vec, + pub usage: Usage, } impl Clue { @@ -17,6 +18,8 @@ impl Clue { pub fn new(clue_line: &str) -> Option { let chunks: Vec<&str> = clue_line.split(", ").collect(); + + // Discard empty lines as well as clues with only one word linked if chunks.len() < 4 { return None; @@ -44,13 +47,13 @@ impl Clue { impl ClueCollection { /// Create an instance of `ClueCollection` from `Vec`, which contains lines of clue response from the API - pub fn new(clue_strings: Vec) -> Self { + pub fn new(clue_strings: Vec, usage: Usage) -> Self { let mut clues: Vec = clue_strings.iter().filter_map(|s| Clue::new(s)).collect(); // Sort the clues by the number of words they link together clues.sort_by(|a, b| b.count.cmp(&a.count)); - Self { clues } + Self { clues, usage } } pub fn is_empty(&self) -> bool { From 8c9058140a9260922feb58b3e8473c39250b1b66 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 17:47:21 -0400 Subject: [PATCH 08/11] fetch_clue_collection doesn't return tuple anymore, usage gets returned via ClueCollection --- src/api/chat_completions.rs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/api/chat_completions.rs b/src/api/chat_completions.rs index bd81955..faa7e43 100644 --- a/src/api/chat_completions.rs +++ b/src/api/chat_completions.rs @@ -1,8 +1,10 @@ -use super::json_models::chat_completion::{ChatCompletionResponse, Usage}; +use super::json_models::chat_completion::{ChatCompletionResponse}; use super::Instance; use crate::clue::ClueCollection; use serde_json::json; + + const SYSTEM_PROMPT: &str = r#" You are the spymaster in Codenames. I will give you a list of [agent word], followed by a list of [avoid word]. @@ -25,7 +27,7 @@ impl Instance { &self, link_words: Vec, avoid_words: Vec, - ) -> Result<(ClueCollection, Usage), Box> { + ) -> Result> { let request_body = self.build_request_body(link_words, avoid_words); // Get response from API endpoint @@ -44,7 +46,8 @@ impl Instance { .map_err(|_| "Failed to parse clues from API server")?; // Extract usage information from the parsed response - let usage = parsed_response.usage.clone(); + let token_usage = parsed_response.usage; + // Extract clue strings from the parsed response let clue_strings = parsed_response @@ -58,9 +61,9 @@ impl Instance { .collect::>(); // Build clues - let clue_collection = ClueCollection::new(clue_strings); + let clue_collection = ClueCollection::new(clue_strings, token_usage); - Ok((clue_collection, usage)) + Ok(clue_collection) } fn build_request_body( From e2a69f62fff79bdf170f5dacdc8270881d0ebd46 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 17:49:10 -0400 Subject: [PATCH 09/11] added: chat_completion module public --- src/api/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/api/mod.rs b/src/api/mod.rs index b95b552..8b9fa6f 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,5 +1,5 @@ -mod chat_completions; -mod json_models; +pub mod chat_completions; +pub mod json_models; mod language_models; use dotenv::dotenv; From d423f05cb04ccbda35bcc1f31b9c7b2fe24db569 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 17:50:19 -0400 Subject: [PATCH 10/11] chat completion module is public now --- src/api/json_models/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/api/json_models/mod.rs b/src/api/json_models/mod.rs index 7381405..c19183e 100644 --- a/src/api/json_models/mod.rs +++ b/src/api/json_models/mod.rs @@ -1,2 +1,2 @@ -pub mod chat_completion; pub mod language_model; +pub mod chat_completion; From 8237b32625b842131a99b8526d4241a2c0ef1617 Mon Sep 17 00:00:00 2001 From: mulla028 Date: Thu, 19 Sep 2024 17:51:19 -0400 Subject: [PATCH 11/11] Clone ability of usage discarded --- src/api/json_models/chat_completion.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/api/json_models/chat_completion.rs b/src/api/json_models/chat_completion.rs index 9f1b362..08e1469 100644 --- a/src/api/json_models/chat_completion.rs +++ b/src/api/json_models/chat_completion.rs @@ -10,7 +10,7 @@ pub struct Choice { pub message: Message, } -#[derive(Deserialize, Clone)] +#[derive(Deserialize)] pub struct Usage { pub prompt_tokens: usize, pub completion_tokens: usize,