From 7ca21cc3299c7501f8ad9cca2c10c12a55dbf3c4 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Tue, 10 Dec 2024 05:59:19 +1100
Subject: [PATCH] FEATURE: first class support for OpenRouter (#1011)

* FEATURE: first class support for OpenRouter

This new implementation supports picking quantization and provider pref

Also:

- Improve logging for summary generation
- Improve error message when contacting LLMs fails

* Better support for full screen artifacts on iPad

Support back button to close full screen
---
 app/models/ai_api_audit_log.rb                |  1 +
 app/models/llm_model.rb                       |  5 ++
 .../discourse/components/ai-artifact.gjs      | 46 ++++++++++--------
 .../modules/ai-bot/common/ai-artifact.scss    | 25 +++++-----
 config/locales/client.en.yml                  |  5 +-
 config/locales/server.en.yml                  |  1 +
 lib/ai_bot/playground.rb                      | 20 ++++++--
 lib/completions/dialects/chat_gpt.rb          |  3 +-
 lib/completions/endpoints/base.rb             |  1 +
 lib/completions/endpoints/open_router.rb      | 42 +++++++++++++++++
 lib/completions/llm.rb                        | 19 ++++++++
 .../strategies/hot_topic_gists.rb             |  8 ++--
 lib/summarization/strategies/topic_summary.rb |  4 +-
 spec/fabricators/llm_model_fabricator.rb      |  9 ++++
 .../completions/endpoints/open_router_spec.rb | 47 +++++++++++++++++++
 15 files changed, 192 insertions(+), 44 deletions(-)
 create mode 100644 lib/completions/endpoints/open_router.rb
 create mode 100644 spec/lib/completions/endpoints/open_router_spec.rb

diff --git a/app/models/ai_api_audit_log.rb b/app/models/ai_api_audit_log.rb
index ca33a2de..075f1444 100644
--- a/app/models/ai_api_audit_log.rb
+++ b/app/models/ai_api_audit_log.rb
@@ -15,6 +15,7 @@ class AiApiAuditLog < ActiveRecord::Base
     Ollama = 7
     SambaNova = 8
     Mistral = 9
+    OpenRouter = 10
   end
 
   def next_log_id
diff --git a/app/models/llm_model.rb b/app/models/llm_model.rb
index e8c21f21..98d39623 100644
--- a/app/models/llm_model.rb
+++ b/app/models/llm_model.rb
@@ -47,6 +47,11 @@ class LlmModel < ActiveRecord::Base
         disable_system_prompt: :checkbox,
         enable_native_tool: :checkbox,
       },
+      open_router: {
+        disable_native_tools: :checkbox,
+        provider_order: :text,
+        provider_quantizations: :text,
+      },
     }
   end
 
diff --git a/assets/javascripts/discourse/components/ai-artifact.gjs b/assets/javascripts/discourse/components/ai-artifact.gjs
index e8ba7090..e1958392 100644
--- a/assets/javascripts/discourse/components/ai-artifact.gjs
+++ b/assets/javascripts/discourse/components/ai-artifact.gjs
@@ -1,12 +1,16 @@
 import Component from "@glimmer/component";
 import { tracked } from "@glimmer/tracking";
-import { on } from "@ember/modifier";
 import { action } from "@ember/object";
 import { service } from "@ember/service";
 import DButton from "discourse/components/d-button";
 import htmlClass from "discourse/helpers/html-class";
 import getURL from "discourse-common/lib/get-url";
 
+// note the panel for artifact full screen can not be at position 0,0
+// otherwise this hack will not activate.
+// https://github.com/discourse/discourse/blob/b8325f2190a8c0a9022405c219faeac6f0f98ca5/app/assets/javascripts/discourse/app/components/scrolling-post-stream.js#L77-L77
+// this will cause post stream to navigate to a different post
+
 export default class AiArtifactComponent extends Component {
   @service siteSettings;
   @tracked expanded = false;
@@ -15,17 +19,29 @@ export default class AiArtifactComponent extends Component {
   constructor() {
     super(...arguments);
     this.keydownHandler = this.handleKeydown.bind(this);
+    this.popStateHandler = this.handlePopState.bind(this);
+    window.addEventListener("popstate", this.popStateHandler);
   }
 
   willDestroy() {
     super.willDestroy(...arguments);
     window.removeEventListener("keydown", this.keydownHandler);
+    window.removeEventListener("popstate", this.popStateHandler);
   }
 
   @action
   handleKeydown(event) {
     if (event.key === "Escape" || event.key === "Esc") {
-      this.expanded = false;
+      history.back();
+    }
+  }
+
+  @action
+  handlePopState(event) {
+    const state = event.state;
+    this.expanded = state?.artifactId === this.args.artifactId;
+    if (!this.expanded) {
+      window.removeEventListener("keydown", this.keydownHandler);
     }
   }
 
@@ -52,12 +68,17 @@ export default class AiArtifactComponent extends Component {
 
   @action
   toggleView() {
-    this.expanded = !this.expanded;
-    if (this.expanded) {
+    if (!this.expanded) {
+      window.history.pushState(
+        { artifactId: this.args.artifactId },
+        "",
+        window.location.href + "#artifact-fullscreen"
+      );
       window.addEventListener("keydown", this.keydownHandler);
     } else {
-      window.removeEventListener("keydown", this.keydownHandler);
+      history.back();
     }
+    this.expanded = !this.expanded;
   }
 
   get wrapperClasses() {
@@ -66,25 +87,12 @@ export default class AiArtifactComponent extends Component {
     }`;
   }
 
-  @action
-  artifactPanelHover() {
-    // retrrigger animation
-    const panel = document.querySelector(".ai-artifact__panel");
-    panel.style.animation = "none"; // Stop the animation
-    setTimeout(() => {
-      panel.style.animation = ""; // Re-trigger the animation by removing the none style
-    }, 0);
-  }
-
   <template>
     {{#if this.expanded}}
       {{htmlClass "ai-artifact-expanded"}}
     {{/if}}
     <div class={{this.wrapperClasses}}>
-      <div
-        class="ai-artifact__panel--wrapper"
-        {{on "mouseleave" this.artifactPanelHover}}
-      >
+      <div class="ai-artifact__panel--wrapper">
         <div class="ai-artifact__panel">
           <DButton
             class="btn-flat btn-icon-text"
diff --git a/assets/stylesheets/modules/ai-bot/common/ai-artifact.scss b/assets/stylesheets/modules/ai-bot/common/ai-artifact.scss
index cd32601a..7c2260cc 100644
--- a/assets/stylesheets/modules/ai-bot/common/ai-artifact.scss
+++ b/assets/stylesheets/modules/ai-bot/common/ai-artifact.scss
@@ -37,27 +37,21 @@ html.ai-artifact-expanded {
   }
 
   .ai-artifact__panel--wrapper {
-    display: block;
     position: fixed;
     top: 0;
-    left: 0;
-    right: 0;
-    height: 4em;
+    left: 2em;
+    right: 2em;
+    height: 2em;
     z-index: 1000000;
-    &:hover {
-      .ai-artifact__panel {
-        transform: translateY(0) !important;
-        animation: none;
-      }
-    }
+    animation: vanishing 0.5s 3s forwards;
   }
 
   .ai-artifact__panel {
     display: block;
     position: fixed;
     top: 0;
-    left: 0;
-    right: 0;
+    left: 2em;
+    right: 2em;
     height: 2em;
     transition: transform 0.5s ease-in-out;
     animation: slideUp 0.5s 3s forwards;
@@ -75,7 +69,6 @@ html.ai-artifact-expanded {
         .d-icon {
           color: var(--secondary-high);
         }
-        //color: var(--secondary-vary-low);
       }
     }
   }
@@ -85,6 +78,12 @@ html.ai-artifact-expanded {
     }
   }
 
+  @keyframes vanishing {
+    to {
+      display: none;
+    }
+  }
+
   iframe {
     position: fixed;
     top: 0;
diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
index a4348df3..0e403243 100644
--- a/config/locales/client.en.yml
+++ b/config/locales/client.en.yml
@@ -351,6 +351,7 @@ en:
           CDCK: "CDCK"
           samba_nova: "SambaNova"
           mistral: "Mistral"
+          open_router: "OpenRouter"
           fake: "Custom"
 
         provider_fields:
@@ -360,6 +361,8 @@ en:
           disable_system_prompt: "Disable system message in prompts"
           enable_native_tool: "Enable native tool support"
           disable_native_tools: "Disable native tool support (use XML based tools)"
+          provider_order: "Provider order (comma delimited list)"
+          provider_quantizations: "Order of provider quantizations (comma delimited list eg: fp16,fp8)"
 
       related_topics:
         title: "Related topics"
@@ -436,7 +439,7 @@ en:
 
       ai_artifact:
         expand_view_label: "Expand view"
-        collapse_view_label: "Exit Fullscreen (ESC)"
+        collapse_view_label: "Exit Fullscreen (ESC or Back button)"
         click_to_run_label: "Run Artifact"
 
       ai_bot:
diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml
index cfb30941..dca9e309 100644
--- a/config/locales/server.en.yml
+++ b/config/locales/server.en.yml
@@ -252,6 +252,7 @@ en:
         failed_to_share: "Failed to share the conversation"
         conversation_deleted: "Conversation share deleted successfully"
     ai_bot:
+      reply_error: "Sorry, it looks like our system encountered an unexpected issue while trying to reply.\n\n[details='Error details']\n%{details}\n[/details]"
       default_pm_prefix: "[Untitled AI bot PM]"
       personas:
         default_llm_required: "Default LLM model is required prior to enabling Chat"
diff --git a/lib/ai_bot/playground.rb b/lib/ai_bot/playground.rb
index f0ed1940..8f95cf57 100644
--- a/lib/ai_bot/playground.rb
+++ b/lib/ai_bot/playground.rb
@@ -533,14 +533,26 @@ module DiscourseAi
           reply_post.post_custom_prompt.update!(custom_prompt: prompt)
         end
 
+        reply_post
+      rescue => e
+        if reply_post
+          details = e.message.to_s
+          reply = "#{reply}\n\n#{I18n.t("discourse_ai.ai_bot.reply_error", details: details)}"
+          reply_post.revise(
+            bot.bot_user,
+            { raw: reply },
+            skip_validations: true,
+            skip_revision: true,
+          )
+        end
+        raise e
+      ensure
         # since we are skipping validations and jobs we
         # may need to fix participant count
-        if reply_post.topic.private_message? && reply_post.topic.participant_count < 2
+        if reply_post && reply_post.topic && reply_post.topic.private_message? &&
+             reply_post.topic.participant_count < 2
           reply_post.topic.update!(participant_count: 2)
         end
-
-        reply_post
-      ensure
         post_streamer&.finish(skip_callback: true)
         publish_final_update(reply_post) if stream_reply
         if reply_post && post.post_number == 1 && post.topic.private_message?
diff --git a/lib/completions/dialects/chat_gpt.rb b/lib/completions/dialects/chat_gpt.rb
index 5cc75d6e..9882adcf 100644
--- a/lib/completions/dialects/chat_gpt.rb
+++ b/lib/completions/dialects/chat_gpt.rb
@@ -6,7 +6,8 @@ module DiscourseAi
       class ChatGpt < Dialect
         class << self
           def can_translate?(llm_model)
-            llm_model.provider == "open_ai" || llm_model.provider == "azure"
+            llm_model.provider == "open_router" || llm_model.provider == "open_ai" ||
+              llm_model.provider == "azure"
           end
         end
 
diff --git a/lib/completions/endpoints/base.rb b/lib/completions/endpoints/base.rb
index 7bc8c076..732e8169 100644
--- a/lib/completions/endpoints/base.rb
+++ b/lib/completions/endpoints/base.rb
@@ -21,6 +21,7 @@ module DiscourseAi
               DiscourseAi::Completions::Endpoints::Cohere,
               DiscourseAi::Completions::Endpoints::SambaNova,
               DiscourseAi::Completions::Endpoints::Mistral,
+              DiscourseAi::Completions::Endpoints::OpenRouter,
             ]
 
             endpoints << DiscourseAi::Completions::Endpoints::Ollama if Rails.env.development?
diff --git a/lib/completions/endpoints/open_router.rb b/lib/completions/endpoints/open_router.rb
new file mode 100644
index 00000000..32d84986
--- /dev/null
+++ b/lib/completions/endpoints/open_router.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Completions
+    module Endpoints
+      class OpenRouter < OpenAi
+        def self.can_contact?(model_provider)
+          %w[open_router].include?(model_provider)
+        end
+
+        def prepare_request(payload)
+          headers = { "Content-Type" => "application/json" }
+          api_key = llm_model.api_key
+
+          headers["Authorization"] = "Bearer #{api_key}"
+          headers["X-Title"] = "Discourse AI"
+          headers["HTTP-Referer"] = "https://www.discourse.org/ai"
+
+          Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
+        end
+
+        def prepare_payload(prompt, model_params, dialect)
+          payload = super
+
+          if quantizations = llm_model.provider_params["provider_quantizations"].presence
+            options = quantizations.split(",").map(&:strip)
+
+            payload[:provider] = { quantizations: options }
+          end
+
+          if order = llm_model.provider_params["provider_order"].presence
+            options = order.split(",").map(&:strip)
+            payload[:provider] ||= {}
+            payload[:provider][:order] = options
+          end
+
+          payload
+        end
+      end
+    end
+  end
+end
diff --git a/lib/completions/llm.rb b/lib/completions/llm.rb
index 9de8f9fe..ea103331 100644
--- a/lib/completions/llm.rb
+++ b/lib/completions/llm.rb
@@ -108,6 +108,24 @@ module DiscourseAi
                   endpoint: "https://api.mistral.ai/v1/chat/completions",
                   provider: "mistral",
                 },
+                {
+                  id: "open_router",
+                  models: [
+                    {
+                      name: "meta-llama/llama-3.3-70b-instruct",
+                      tokens: 128_000,
+                      display_name: "Llama 3.3 70B",
+                    },
+                    {
+                      name: "google/gemini-flash-1.5-exp",
+                      tokens: 1_000_000,
+                      display_name: "Gemini Flash 1.5 Exp",
+                    },
+                  ],
+                  tokenizer: DiscourseAi::Tokenizer::OpenAiTokenizer,
+                  endpoint: "https://openrouter.ai/api/v1/chat/completions",
+                  provider: "open_router",
+                },
               ]
             end
         end
@@ -124,6 +142,7 @@ module DiscourseAi
             azure
             samba_nova
             mistral
+            open_router
           ]
           if !Rails.env.production?
             providers << "fake"
diff --git a/lib/summarization/strategies/hot_topic_gists.rb b/lib/summarization/strategies/hot_topic_gists.rb
index 144507fc..1708f2b8 100644
--- a/lib/summarization/strategies/hot_topic_gists.rb
+++ b/lib/summarization/strategies/hot_topic_gists.rb
@@ -64,7 +64,7 @@ module DiscourseAi
               .map { |item| "(#{item[:id]} #{item[:poster]} said: #{item[:text]} " }
               .join("\n")
 
-          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip)
+          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip, topic_id: target.id)
             You are an advanced summarization bot. Your task is to update an existing single-sentence summary by integrating new developments from a conversation.
             Analyze the most recent messages to identify key updates or shifts in the main topic and reflect these in the updated summary.
             Emphasize new significant information or developments within the context of the initial conversation theme.
@@ -103,7 +103,7 @@ module DiscourseAi
           statements =
             contents.to_a.map { |item| "(#{item[:id]} #{item[:poster]} said: #{item[:text]} " }
 
-          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip)
+          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip, topic_id: target.id)
             You are an advanced summarization bot. Analyze a given conversation and produce a concise,
             single-sentence summary that conveys the main topic and current developments to someone with no prior context.
 
@@ -124,9 +124,9 @@ module DiscourseAi
             ### Context:
 
             #{content_title.present? ? "The discussion title is: " + content_title + ". (DO NOT REPEAT THIS IN THE SUMMARY)\n" : ""}
-            
+
             The conversation began with the following statement:
-        
+
             #{statements.shift}\n
           TEXT
 
diff --git a/lib/summarization/strategies/topic_summary.rb b/lib/summarization/strategies/topic_summary.rb
index 546a9d25..c252fe33 100644
--- a/lib/summarization/strategies/topic_summary.rb
+++ b/lib/summarization/strategies/topic_summary.rb
@@ -33,7 +33,7 @@ module DiscourseAi
           input =
             contents.map { |item| "(#{item[:id]} #{item[:poster]} said: #{item[:text]})" }.join
 
-          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT)
+          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT, topic_id: target.id)
             You are an advanced summarization bot tasked with enhancing an existing summary by incorporating additional posts.
 
             ### Guidelines:
@@ -76,7 +76,7 @@ module DiscourseAi
           input =
             contents.map { |item| "(#{item[:id]} #{item[:poster]} said: #{item[:text]} " }.join
 
-          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip)
+          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip, topic_id: target.id)
             You are an advanced summarization bot that generates concise, coherent summaries of provided text.
 
             - Only include the summary, without any additional commentary.
diff --git a/spec/fabricators/llm_model_fabricator.rb b/spec/fabricators/llm_model_fabricator.rb
index 37c0fcbb..87b1ef25 100644
--- a/spec/fabricators/llm_model_fabricator.rb
+++ b/spec/fabricators/llm_model_fabricator.rb
@@ -28,6 +28,15 @@ Fabricator(:hf_model, from: :llm_model) do
   provider "hugging_face"
 end
 
+Fabricator(:open_router_model, from: :llm_model) do
+  display_name "OpenRouter"
+  name "openrouter-1.0"
+  provider "open_router"
+  tokenizer "DiscourseAi::Tokenizer::OpenAiTokenizer"
+  max_prompt_tokens 64_000
+  url "https://openrouter.ai/api/v1/chat/completions"
+end
+
 Fabricator(:vllm_model, from: :llm_model) do
   display_name "Llama 3.1 vLLM"
   name "meta-llama/Meta-Llama-3.1-70B-Instruct"
diff --git a/spec/lib/completions/endpoints/open_router_spec.rb b/spec/lib/completions/endpoints/open_router_spec.rb
new file mode 100644
index 00000000..e21dd695
--- /dev/null
+++ b/spec/lib/completions/endpoints/open_router_spec.rb
@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+RSpec.describe DiscourseAi::Completions::Endpoints::OpenRouter do
+  fab!(:user)
+  fab!(:open_router_model)
+
+  subject(:endpoint) { described_class.new(open_router_model) }
+
+  it "supports provider quantization and order selection" do
+    open_router_model.provider_params["provider_quantizations"] = "int8,int16"
+    open_router_model.provider_params["provider_order"] = "Google, Amazon Bedrock"
+    open_router_model.save!
+
+    parsed_body = nil
+    stub_request(:post, open_router_model.url).with(
+      body: proc { |body| parsed_body = JSON.parse(body, symbolize_names: true) },
+      headers: {
+        "Content-Type" => "application/json",
+        "X-Title" => "Discourse AI",
+        "HTTP-Referer" => "https://www.discourse.org/ai",
+        "Authorization" => "Bearer 123",
+      },
+    ).to_return(
+      status: 200,
+      body: { "choices" => [message: { role: "assistant", content: "world" }] }.to_json,
+    )
+
+    proxy = DiscourseAi::Completions::Llm.proxy("custom:#{open_router_model.id}")
+    result = proxy.generate("hello", user: user)
+
+    expect(result).to eq("world")
+
+    expected = {
+      model: "openrouter-1.0",
+      messages: [
+        { role: "system", content: "You are a helpful bot" },
+        { role: "user", content: "hello" },
+      ],
+      provider: {
+        quantizations: %w[int8 int16],
+        order: ["Google", "Amazon Bedrock"],
+      },
+    }
+
+    expect(parsed_body).to eq(expected)
+  end
+end