From 6f8abb252d60e21d58c22ac37936d8338d6584f0 Mon Sep 17 00:00:00 2001
From: Claas Augner <caugner@mozilla.com>
Date: Tue, 19 Dec 2023 11:15:52 +0100
Subject: [PATCH] fix(ai-help): use GPT-3.5 for free users

---
 ai-test/src/ai_help.rs |  4 +++-
 src/ai/constants.rs    | 12 ++++++++++++
 src/ai/help.rs         | 10 ++++++++--
 src/api/ai_help.rs     |  2 +-
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/ai-test/src/ai_help.rs b/ai-test/src/ai_help.rs
index cd4411be..c07bc4aa 100644
--- a/ai-test/src/ai_help.rs
+++ b/ai-test/src/ai_help.rs
@@ -95,7 +95,9 @@ pub async fn ai_help_all(
                     function_call: None,
                 })
                 .collect();
-            if let Some(req) = prepare_ai_help_req(openai_client, supabase_pool, messages).await? {
+            if let Some(req) =
+                prepare_ai_help_req(openai_client, supabase_pool, true, messages).await?
+            {
                 let mut res = openai_client.chat().create(req.req.clone()).await?;
                 let res = res.choices.pop().map(|res| res.message);
                 let storage = Storage { req, res };
diff --git a/src/ai/constants.rs b/src/ai/constants.rs
index ec45aec5..7d93d4a6 100644
--- a/src/ai/constants.rs
+++ b/src/ai/constants.rs
@@ -16,6 +16,18 @@ pub struct AIHelpConfig {
     pub make_context: fn(Vec<RelatedDoc>) -> String,
 }
 
+pub const AI_HELP_GPT3_5_FULL_DOC_NEW_PROMPT: AIHelpConfig = AIHelpConfig {
+    name: "20230901-gpt4-full_doc-new_pormpt",
+    model: "gpt-3.5-turbo-1106",
+    full_doc: true,
+    system_prompt: include_str!("prompts/new_prompt/system.md"),
+    user_prompt: None,
+    token_limit: 16_384,
+    context_limit: 12_000,
+    max_completion_tokens: 2_048,
+    make_context: |related_docs| related_docs.into_iter().map(|d| d.content).join("\n"),
+};
+
 pub const AI_HELP_GPT4_FULL_DOC_NEW_PROMPT: AIHelpConfig = AIHelpConfig {
     name: "20230901-gpt4-full_doc-new_pormpt",
     model: "gpt-4-1106-preview",
diff --git a/src/ai/help.rs b/src/ai/help.rs
index 5b40e5a8..1714a8d3 100644
--- a/src/ai/help.rs
+++ b/src/ai/help.rs
@@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     ai::{
-        constants::AI_HELP_GPT4_FULL_DOC_NEW_PROMPT,
+        constants::{AI_HELP_GPT3_5_FULL_DOC_NEW_PROMPT, AI_HELP_GPT4_FULL_DOC_NEW_PROMPT},
         embeddings::{get_related_docs, get_related_macro_docs},
         error::AIError,
         helpers::{cap_messages, into_user_messages, sanitize_messages},
@@ -35,9 +35,15 @@ pub struct AIHelpRequest {
 pub async fn prepare_ai_help_req(
     client: &Client<OpenAIConfig>,
     pool: &SupaPool,
+    is_subscriber: bool,
     messages: Vec<ChatCompletionRequestMessage>,
 ) -> Result<Option<AIHelpRequest>, AIError> {
-    let config = AI_HELP_GPT4_FULL_DOC_NEW_PROMPT;
+    let config = if is_subscriber {
+        AI_HELP_GPT4_FULL_DOC_NEW_PROMPT
+    } else {
+        AI_HELP_GPT3_5_FULL_DOC_NEW_PROMPT
+    };
+
     let open_ai_messages = sanitize_messages(messages);
 
     // TODO: sign messages os we don't check again
diff --git a/src/api/ai_help.rs b/src/api/ai_help.rs
index 12ffb54b..71f46625 100644
--- a/src/api/ai_help.rs
+++ b/src/api/ai_help.rs
@@ -398,7 +398,7 @@ pub async fn ai_help(
             )?;
         }
 
-        match prepare_ai_help_req(client, pool, messages).await? {
+        match prepare_ai_help_req(client, pool, user.is_subscriber(), messages).await? {
             Some(ai_help_req) => {
                 let sources = ai_help_req.refs;
                 let created_at = match record_sources(