FEATURE: display more places where AI is used / Chat streamer (#1278)

* FEATURE: display more places where AI is used

- Usage was not showing automation or image caption in llm list.
- Also: FIX - reasoning models would time out incorrectly after 60 seconds (raised to 10 minutes)

* correct enum not to enumerate non configured models

* FEATURE: implement chat streamer

This implements a basic chat streamer, it provides 2 things:

1. Gives feedback to the user when LLM is generating
2. Streams stuff much more efficiently to client (given it may take 100ms or so per call to update chat)
This commit is contained in:
Sam 2025-04-24 17:22:19 +11:00 committed by GitHub
parent e1731dc3df
commit 2a5c60db10
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 215 additions and 81 deletions

View File

@ -116,6 +116,10 @@ export default class AiLlmsListEditor extends Component {
return i18n("discourse_ai.llms.usage.ai_persona", {
persona: usage.name,
});
} else if (usage.type === "automation") {
return i18n("discourse_ai.llms.usage.automation", {
name: usage.name,
});
} else {
return i18n("discourse_ai.llms.usage." + usage.type);
}

View File

@ -117,6 +117,7 @@
list-style: none;
margin: 0.5em 0 0 0;
display: flex;
flex-wrap: wrap;
li {
font-size: var(--font-down-2);
@ -125,6 +126,7 @@
border: 1px solid var(--primary-low);
padding: 1px 3px;
margin-right: 0.5em;
margin-bottom: 0.5em;
}
}

View File

@ -439,10 +439,12 @@ en:
usage:
ai_bot: "AI bot"
ai_helper: "Helper"
ai_helper_image_caption: "Image caption"
ai_persona: "Persona (%{persona})"
ai_summarization: "Summarize"
ai_embeddings_semantic_search: "AI search"
ai_spam: "Spam"
automation: "Automation (%{name})"
in_use_warning:
one: "This model is currently used by %{settings}. If misconfigured, the feature won't work as expected."
other: "This model is currently used by the following: %{settings}. If misconfigured, features won't work as expected. "

View File

@ -312,11 +312,6 @@ discourse_ai:
default: "1|2" # 1: admins, 2: moderators
allow_any: false
refresh: true
ai_bot_enabled_chat_bots: # TODO(roman): Deprecated. Remove by Sept 2024
type: list
default: "gpt-3.5-turbo"
hidden: true
choices: "DiscourseAi::Configuration::LlmEnumerator.available_ai_bots"
ai_bot_add_to_header:
default: true
client: true

View File

@ -0,0 +1,13 @@
# frozen_string_literal: true
class RemoveOldSettings < ActiveRecord::Migration[7.2]
def up
execute <<~SQL
DELETE FROM site_settings
WHERE name IN ('ai_bot_enabled_chat_bots')
SQL
end
def down
raise ActiveRecord::IrreversibleMigration
end
end

117
lib/ai_bot/chat_streamer.rb Normal file
View File

@ -0,0 +1,117 @@
# frozen_string_literal: true
#
# Chat streaming APIs are a bit slow, this ensures we properly buffer results
# and stream as quickly as possible.
module DiscourseAi
module AiBot
class ChatStreamer
attr_accessor :cancel
attr_reader :reply,
:guardian,
:thread_id,
:force_thread,
:in_reply_to_id,
:channel,
:cancelled
def initialize(message:, channel:, guardian:, thread_id:, in_reply_to_id:, force_thread:)
@message = message
@channel = channel
@guardian = guardian
@thread_id = thread_id
@force_thread = force_thread
@in_reply_to_id = in_reply_to_id
@queue = Queue.new
db = RailsMultisite::ConnectionManagement.current_db
@worker_thread =
Thread.new { RailsMultisite::ConnectionManagement.with_connection(db) { run } }
@client_id =
ChatSDK::Channel.start_reply(
channel_id: message.chat_channel_id,
guardian: guardian,
thread_id: thread_id,
)
end
def <<(partial)
return if partial.to_s.empty?
if @client_id
ChatSDK::Channel.stop_reply(
channel_id: @message.chat_channel_id,
client_id: @client_id,
guardian: @guardian,
thread_id: @thread_id,
)
@client_id = nil
end
if @reply
@queue << partial
else
create_reply(partial)
end
end
def create_reply(message)
@reply =
ChatSDK::Message.create(
raw: message,
channel_id: channel.id,
guardian: guardian,
force_thread: force_thread,
in_reply_to_id: in_reply_to_id,
enforce_membership: !channel.direct_message_channel?,
)
ChatSDK::Message.start_stream(message_id: @reply.id, guardian: @guardian)
if trailing = message.scan(/\s*\z/).first
@queue << trailing
end
end
def done
@queue << :done
@worker_thread.join
ChatSDK::Message.stop_stream(message_id: @reply.id, guardian: @guardian)
@reply
end
private
def run
done = false
while !done
buffer = +""
popped = @queue.pop
break if popped == :done
buffer << popped
begin
while true
popped = @queue.pop(true)
if popped == :done
done = true
break
end
buffer << popped
end
rescue ThreadError
end
streaming = ChatSDK::Message.stream(message_id: reply.id, raw: buffer, guardian: guardian)
if !streaming
cancel.call
@cancelled = true
end
end
end
end
end
end

View File

@ -5,6 +5,9 @@ module DiscourseAi
class Playground
BYPASS_AI_REPLY_CUSTOM_FIELD = "discourse_ai_bypass_ai_reply"
BOT_USER_PREF_ID_CUSTOM_FIELD = "discourse_ai_bot_user_pref_id"
# 10 minutes is enough for vast majority of cases
# there is a small chance that some reasoning models may take longer
MAX_STREAM_DELAY_SECONDS = 600
attr_reader :bot
@ -334,42 +337,38 @@ module DiscourseAi
force_thread = message.thread_id.nil? && channel.direct_message_channel?
in_reply_to_id = channel.direct_message_channel? ? message.id : nil
streamer =
ChatStreamer.new(
message: message,
channel: channel,
guardian: guardian,
thread_id: message.thread_id,
in_reply_to_id: in_reply_to_id,
force_thread: force_thread,
)
new_prompts =
bot.reply(context) do |partial, cancel, placeholder, type|
# no support for tools or thinking by design
next if type == :thinking || type == :tool_details || type == :partial_tool
if !reply
# just eat all leading spaces we can not create the message
next if partial.blank?
reply =
ChatSDK::Message.create(
raw: partial,
thread_id: message.thread_id,
channel_id: channel.id,
guardian: guardian,
in_reply_to_id: in_reply_to_id,
force_thread: force_thread,
enforce_membership: !channel.direct_message_channel?,
)
ChatSDK::Message.start_stream(message_id: reply.id, guardian: guardian)
else
streaming =
ChatSDK::Message.stream(message_id: reply.id, raw: partial, guardian: guardian)
if !streaming
cancel&.call
break
end
end
streamer.cancel = cancel
streamer << partial
break if streamer.cancelled
end
if new_prompts.length > 1 && reply.id
reply = streamer.reply
if new_prompts.length > 1 && reply
ChatMessageCustomPrompt.create!(message_id: reply.id, custom_prompt: new_prompts)
end
ChatSDK::Message.stop_stream(message_id: reply.id, guardian: guardian) if reply
if streamer
streamer.done
streamer = nil
end
reply
ensure
streamer.done if streamer
end
def reply_to(
@ -464,7 +463,7 @@ module DiscourseAi
publish_update(reply_post, { raw: reply_post.cooked })
redis_stream_key = "gpt_cancel:#{reply_post.id}"
Discourse.redis.setex(redis_stream_key, 60, 1)
Discourse.redis.setex(redis_stream_key, MAX_STREAM_DELAY_SECONDS, 1)
end
context.skip_tool_details ||= !bot.persona.class.tool_details
@ -504,7 +503,7 @@ module DiscourseAi
if post_streamer
post_streamer.run_later do
Discourse.redis.expire(redis_stream_key, 60)
Discourse.redis.expire(redis_stream_key, MAX_STREAM_DELAY_SECONDS)
publish_update(reply_post, { raw: raw })
end
end

View File

@ -13,16 +13,22 @@ module DiscourseAi
.where("enabled_chat_bot = ?", true)
.pluck(:id)
.each { |llm_id| rval[llm_id] << { type: :ai_bot } }
AiPersona
.where("force_default_llm = ?", true)
.pluck(:default_llm_id, :name, :id)
.each { |llm_id, name, id| rval[llm_id] << { type: :ai_persona, name: name, id: id } }
end
# this is unconditional, so it is clear that we always signal configuration
AiPersona
.where("default_llm_id IS NOT NULL")
.pluck(:default_llm_id, :name, :id)
.each { |llm_id, name, id| rval[llm_id] << { type: :ai_persona, name: name, id: id } }
if SiteSetting.ai_helper_enabled
model_id = SiteSetting.ai_helper_model.split(":").last.to_i
rval[model_id] << { type: :ai_helper }
rval[model_id] << { type: :ai_helper } if model_id != 0
end
if SiteSetting.ai_helper_image_caption_model
model_id = SiteSetting.ai_helper_image_caption_model.split(":").last.to_i
rval[model_id] << { type: :ai_helper_image_caption } if model_id != 0
end
if SiteSetting.ai_summarization_enabled
@ -42,6 +48,25 @@ module DiscourseAi
rval[model_id] << { type: :ai_spam }
end
if defined?(DiscourseAutomation::Automation)
DiscourseAutomation::Automation
.joins(:fields)
.where(script: %w[llm_report llm_triage])
.where("discourse_automation_fields.name = ?", "model")
.pluck(
"metadata ->> 'value', discourse_automation_automations.name, discourse_automation_automations.id",
)
.each do |model_text, name, id|
next if model_text.blank?
model_id = model_text.split("custom:").last.to_i
if model_id.present?
if model_text =~ /custom:(\d+)/
rval[model_id] << { type: :automation, name: name, id: id }
end
end
end
end
rval
end
@ -85,45 +110,6 @@ module DiscourseAi
values.each { |value_h| value_h[:value] = "custom:#{value_h[:value]}" }
values
end
# TODO(roman): Deprecated. Remove by Sept 2024
def self.old_summarization_options
%w[
gpt-4
gpt-4-32k
gpt-4-turbo
gpt-4o
gpt-3.5-turbo
gpt-3.5-turbo-16k
gemini-pro
gemini-1.5-pro
gemini-1.5-flash
claude-2
claude-instant-1
claude-3-haiku
claude-3-sonnet
claude-3-opus
mistralai/Mixtral-8x7B-Instruct-v0.1
mistralai/Mixtral-8x7B-Instruct-v0.1
]
end
# TODO(roman): Deprecated. Remove by Sept 2024
def self.available_ai_bots
%w[
gpt-3.5-turbo
gpt-4
gpt-4-turbo
gpt-4o
claude-2
gemini-1.5-pro
mixtral-8x7B-Instruct-V0.1
claude-3-opus
claude-3-sonnet
claude-3-haiku
cohere-command-r-plus
]
end
end
end
end

View File

@ -4,6 +4,9 @@ RSpec.describe DiscourseAi::Configuration::LlmEnumerator do
fab!(:fake_model)
fab!(:llm_model)
fab!(:seeded_model)
fab!(:automation) do
Fabricate(:automation, script: "llm_report", name: "some automation", enabled: true)
end
describe "#values_for_serialization" do
it "returns an array for that can be used for serialization" do
@ -37,13 +40,27 @@ RSpec.describe DiscourseAi::Configuration::LlmEnumerator do
end
describe "#global_usage" do
before do
it "returns a hash of Llm models in use globally" do
SiteSetting.ai_helper_model = "custom:#{fake_model.id}"
SiteSetting.ai_helper_enabled = true
expect(described_class.global_usage).to eq(fake_model.id => [{ type: :ai_helper }])
end
it "returns a hash of Llm models in use globally" do
expect(described_class.global_usage).to eq(fake_model.id => [{ type: :ai_helper }])
it "returns information about automation rules" do
automation.fields.create!(
component: "text",
name: "model",
metadata: {
value: "custom:#{fake_model.id}",
},
target: "script",
)
usage = described_class.global_usage
expect(usage).to eq(
{ fake_model.id => [{ type: :automation, name: "some automation", id: automation.id }] },
)
end
it "doesn't error on spam when spam detection is enabled but moderation setting is missing" do

View File

@ -7,7 +7,6 @@ RSpec.describe "Admin AI persona configuration", type: :system, js: true do
before do
SiteSetting.ai_bot_enabled = true
SiteSetting.ai_bot_enabled_chat_bots = "gpt-4"
sign_in(admin)
end