FIX: make AI helper more robust

- If JSON is broken for structured output then lean on a more forgiving parser
- Gemini 2.5 flash does not support temp, support opting out
- Evals for assistant were broken, fix interface
- Add some missing LLMs
- Translator was not mapped correctly to the feature - fix that
- Don't mix XML in prompt for translator

Author: SamSaffron

app/models/llm_model.rb

@@ -65,6 +65,8 @@ def self.provider_params
       google: {
         disable_native_tools: :checkbox,
         enable_thinking: :checkbox,
+        disable_temperature: :checkbox,
+        disable_top_p: :checkbox,
         thinking_tokens: :number,
       },
       azure: {

config/eval-llms.yml

@@ -1,4 +1,27 @@
 llms:
+  o3:
+    display_name: O3
+    name: o3
+    tokenizer: DiscourseAi::Tokenizer::OpenAiTokenizer
+    api_key_env: OPENAI_API_KEY
+    provider: open_ai
+    url: https://api.openai.com/v1/chat/completions
+    max_prompt_tokens: 131072
+    vision_enabled: true
+    provider_params:
+      disable_top_p: true
+      disable_temperature: true
+
+  gpt-41:
+    display_name: GPT-4.1
+    name: gpt-4.1
+    tokenizer: DiscourseAi::Tokenizer::OpenAiTokenizer
+    api_key_env: OPENAI_API_KEY
+    provider: open_ai
+    url: https://api.openai.com/v1/chat/completions
+    max_prompt_tokens: 131072
+    vision_enabled: true
+
   gpt-4o:
     display_name: GPT-4o
     name: gpt-4o
@@ -74,12 +97,25 @@ llms:
     max_prompt_tokens: 1000000
     vision_enabled: true
 
-  gemini-2.0-pro-exp:
+  gemini-2.5-flash:
+    display_name: Gemini 2.5 Flash
+    name: gemini-2-5-flash
+    tokenizer: DiscourseAi::Tokenizer::GeminiTokenizer
+    api_key_env: GEMINI_API_KEY
+    provider: google
+    url: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash
+    max_prompt_tokens: 1000000
+    vision_enabled: true
+    provider_params:
+      disable_top_p: true
+      disable_temperature: true
+
+  gemini-2.0-pro:
     display_name: Gemini 2.0 pro
-    name: gemini-2-0-pro-exp
+    name: gemini-2-0-pro
     tokenizer: DiscourseAi::Tokenizer::GeminiTokenizer
     api_key_env: GEMINI_API_KEY
     provider: google
-    url: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-pro-exp
+    url: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-pro
     max_prompt_tokens: 1000000
     vision_enabled: true

config/locales/client.en.yml

@@ -249,6 +249,7 @@ en:
           markdown_tables: "Generate Markdown table"
           custom_prompt: "Custom prompt"
           image_caption: "Caption images"
+          translator: "Translator"
 
         translation:
           name: "Translation"
@@ -257,7 +258,7 @@ en:
           post_raw_translator: "Post raw translator"
           topic_title_translator: "Topic title translator"
           short_text_translator: "Short text translator"
-        
+
         spam:
           name: "Spam"
           description: "Identifies potential spam using the selected LLM and flags it for site moderators to inspect in the review queue"

evals/lib/eval.rb

@@ -200,12 +200,7 @@ class << user
       user.admin = true
     end
     result =
-      helper.generate_and_send_prompt(
-        name,
-        input,
-        current_user = user,
-        _force_default_locale = false,
-      )
+      helper.generate_and_send_prompt(name, input, current_user = user, force_default_locale: false)
 
     result[:suggestions].first
   end

lib/ai_helper/assistant.rb

@@ -82,7 +82,7 @@ def attach_user_context(context, user = nil, force_default_locale: false)
         context.user_language = "#{locale_hash["name"]}"
 
         if user
-          timezone = user.user_option.timezone || "UTC"
+          timezone = user&.user_option&.timezone || "UTC"
           current_time = Time.now.in_time_zone(timezone)
 
           temporal_context = {
@@ -126,21 +126,34 @@ def generate_prompt(
           )
         context = attach_user_context(context, user, force_default_locale: force_default_locale)
 
-        helper_response = +""
+        bad_json = false
+        json_summary_schema_key = bot.persona.response_format&.first.to_h
+
+        schema_key = json_summary_schema_key["key"]&.to_sym
+        schema_type = json_summary_schema_key["type"]
+
+        if schema_type == "array"
+          helper_response = []
+        else
+          helper_response = +""
+        end
 
         buffer_blk =
           Proc.new do |partial, _, type|
-            json_summary_schema_key = bot.persona.response_format&.first.to_h
-            helper_response = [] if json_summary_schema_key["type"] == "array"
-            if type == :structured_output
-              helper_chunk = partial.read_buffered_property(json_summary_schema_key["key"]&.to_sym)
+            if type == :structured_output && schema_type
+              bad_json ||= partial.broken?
+              helper_chunk = partial.read_buffered_property(schema_key)
               if !helper_chunk.nil? && !helper_chunk.empty?
-                if json_summary_schema_key["type"] != "array"
-                  helper_response = helper_chunk
-                else
+                if !bad_json
                   helper_response << helper_chunk
+                else
+                  if schema_type == "string" || schema_type == "array"
+                    helper_response << helper_chunk
+                  else
+                    helper_response = helper_chunk
+                  end
+                  block.call(helper_chunk) if block
                 end
-                block.call(helper_chunk) if block
               end
             elsif type.blank?
               # Assume response is a regular completion.
@@ -151,6 +164,17 @@ def generate_prompt(
 
         bot.reply(context, &buffer_blk)
 
+        # handle edge cases where structured output is all over the place
+        if bad_json
+          helper_response = helper_response.join if helper_response.is_a?(Array)
+          helper_response =
+            DiscourseAi::Utils::BestEffortJsonParser.extract_key(
+              helper_response,
+              schema_type,
+              schema_key,
+            )
+          block.call(helper_response) if block
+        end
         helper_response
       end
 
@@ -255,7 +279,7 @@ def generate_image_caption(upload, user)
           Proc.new do |partial, _, type|
             if type == :structured_output
               structured_output = partial
-              json_summary_schema_key = bot.persona.response_format&.first.to_h
+              _json_summary_schema_key = bot.persona.response_format&.first.to_h
             end
           end
 
@@ -287,6 +311,11 @@ def build_bot(helper_mode, user)
       end
 
       def find_ai_helper_model(helper_mode, persona_klass)
+        if helper_mode == IMAGE_CAPTION && @image_caption_llm.is_a?(LlmModel)
+          return @image_caption_llm
+        end
+
+        return @helper_llm if helper_mode != IMAGE_CAPTION && @helper_llm.is_a?(LlmModel)
         self.class.find_ai_helper_model(helper_mode, persona_klass)
       end
 
@@ -299,9 +328,9 @@ def self.find_ai_helper_model(helper_mode, persona_klass)
 
         if !model_id
           if helper_mode == IMAGE_CAPTION
-            model_id = @helper_llm || SiteSetting.ai_helper_image_caption_model&.split(":")&.last
+            model_id = SiteSetting.ai_helper_image_caption_model&.split(":")&.last
           else
-            model_id = @image_caption_llm || SiteSetting.ai_helper_model&.split(":")&.last
+            model_id = SiteSetting.ai_helper_model&.split(":")&.last
           end
         end
 

lib/completions/endpoints/gemini.rb

@@ -33,7 +33,8 @@ def normalize_model_params(model_params)
 
           model_params[:topP] = model_params.delete(:top_p) if model_params[:top_p]
 
-          # temperature already supported
+          model_params.delete(:temperature) if llm_model.lookup_custom_param("disable_temperature")
+          model_params.delete(:topP) if llm_model.lookup_custom_param("disable_top_p")
 
           model_params
         end

lib/completions/structured_output.rb

@@ -26,6 +26,10 @@ def <<(raw)
         @partial_json_tracker << raw
       end
 
+      def broken?
+        @partial_json_tracker.broken?
+      end
+
       def read_buffered_property(prop_name)
         # Safeguard: If the model is misbehaving and generating something that's not a JSON,
         # treat response as a normal string.

lib/configuration/feature.rb

@@ -103,6 +103,12 @@ def ai_helper_features
               DiscourseAi::Configuration::Module::AI_HELPER_ID,
               DiscourseAi::Configuration::Module::AI_HELPER,
             ),
+            new(
+              "translator",
+              "ai_helper_translator_persona",
+              DiscourseAi::Configuration::Module::AI_HELPER_ID,
+              DiscourseAi::Configuration::Module::AI_HELPER,
+            ),
             new(
               "custom_prompt",
               "ai_helper_custom_prompt_persona",

lib/personas/translator.rb

@@ -19,11 +19,12 @@ def system_prompt
 
           Format your response as a JSON object with a single key named "output", which has the translation as the value.
           Your output should be in the following format:
-            <output>
-              {"output": "xx"}
-            </output>
+
+          {"output": "xx"}
 
           Where "xx" is replaced by the translation.
+
+          reply with valid JSON only
         PROMPT
       end