From d24a19505278ae56cc8bdd29c23de8f171cf7dd8 Mon Sep 17 00:00:00 2001
From: Claas Augner <caugner@mozilla.com>
Date: Tue, 14 May 2024 21:56:18 -0700
Subject: [PATCH] feat(ai-help): upgrade to GPT-4o model

---
 src/ai/constants.rs | 5 ++++-
 src/ai/helpers.rs   | 6 +++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/ai/constants.rs b/src/ai/constants.rs
index f5154ff0..7a6fc4da 100644
--- a/src/ai/constants.rs
+++ b/src/ai/constants.rs
@@ -7,6 +7,7 @@ use crate::ai::embeddings::RelatedDoc;
 pub struct AIHelpConfig {
     pub name: &'static str,
     pub model: &'static str,
+    pub token_model: &'static str,
     pub full_doc: bool,
     pub system_prompt: &'static str,
     pub user_prompt: Option<&'static str>,
@@ -27,6 +28,7 @@ fn join_with_tags(related_docs: Vec<RelatedDoc>) -> String {
 pub const AI_HELP_GPT3_5_FULL_DOC_NEW_PROMPT: AIHelpConfig = AIHelpConfig {
     name: "20230901-full_doc-new_prompt",
     model: "gpt-3.5-turbo-0125",
+    token_model: "gpt-3.5-turbo-0125",
     full_doc: true,
     system_prompt: include_str!("prompts/new_prompt/system.md"),
     user_prompt: None,
@@ -39,7 +41,8 @@ pub const AI_HELP_GPT3_5_FULL_DOC_NEW_PROMPT: AIHelpConfig = AIHelpConfig {
 
 pub const AI_HELP_GPT4_FULL_DOC_NEW_PROMPT: AIHelpConfig = AIHelpConfig {
     name: "20240125-gpt4-full_doc-new_prompt",
-    model: "gpt-4-0125-preview",
+    model: "gpt-4o-2024-05-13",
+    token_model: "gpt-4-0125-preview",
     full_doc: true,
     system_prompt: include_str!("prompts/new_prompt/system.md"),
     user_prompt: None,
diff --git a/src/ai/helpers.rs b/src/ai/helpers.rs
index 6188c29e..15ec600e 100644
--- a/src/ai/helpers.rs
+++ b/src/ai/helpers.rs
@@ -26,11 +26,11 @@ pub fn cap_messages(
     mut init_messages: Vec<ChatCompletionRequestMessage>,
     context_messages: Vec<ChatCompletionRequestMessage>,
 ) -> Result<Vec<ChatCompletionRequestMessage>, AIError> {
-    let init_tokens = num_tokens_from_messages(config.model, &init_messages)?;
+    let init_tokens = num_tokens_from_messages(config.token_model, &init_messages)?;
     if init_tokens + config.max_completion_tokens > config.token_limit {
         return Err(AIError::TokenLimit);
     }
-    let mut context_tokens = num_tokens_from_messages(config.model, &context_messages)?;
+    let mut context_tokens = num_tokens_from_messages(config.token_model, &context_messages)?;
 
     let mut skip = 0;
     while context_tokens + init_tokens + config.max_completion_tokens > config.token_limit {
@@ -38,7 +38,7 @@ pub fn cap_messages(
         if skip >= context_messages.len() {
             return Err(AIError::TokenLimit);
         }
-        context_tokens = num_tokens_from_messages(config.model, &context_messages[skip..])?;
+        context_tokens = num_tokens_from_messages(config.token_model, &context_messages[skip..])?;
     }
     init_messages.extend(context_messages.into_iter().skip(skip));
     Ok(init_messages)