Skip to content

Commit 56b42a9

Browse files
romanrizziromanrizzi
authored
FIX: Switch to DiscourseAi's StructuredOutput (#288)
1 parent 39fc55f commit 56b42a9

File tree

5 files changed

+47
-59
lines changed

5 files changed

+47
-59
lines changed

‎app/services/discourse_ai/base_translator.rb‎

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,15 @@ def translate
1414
messages: [{ type: :user, content: formatted_content, id: "user" }],
1515
)
1616

17-
response =
17+
structured_output =
1818
DiscourseAi::Completions::Llm.proxy(SiteSetting.ai_helper_model).generate(
1919
prompt,
2020
user: Discourse.system_user,
2121
feature_name: "translator-translate",
22-
extra_model_params: response_format,
22+
response_format: response_format,
2323
)
2424

25-
JSON.parse(response)&.dig("translation")
25+
structured_output&.read_latest_buffered_chunk&.dig(:translation)
2626
end
2727

2828
def formatted_content
@@ -31,22 +31,20 @@ def formatted_content
3131

3232
def response_format
3333
{
34-
response_format: {
35-
type: "json_schema",
36-
json_schema: {
37-
name: "reply",
38-
schema: {
39-
type: "object",
40-
properties: {
41-
translation: {
42-
type: "string",
43-
},
34+
type: "json_schema",
35+
json_schema: {
36+
name: "reply",
37+
schema: {
38+
type: "object",
39+
properties: {
40+
translation: {
41+
type: "string",
4442
},
45-
required: ["translation"],
46-
additionalProperties: false,
4743
},
48-
strict: true,
44+
required: ["translation"],
45+
additionalProperties: false,
4946
},
47+
strict: true,
5048
},
5149
}
5250
end

‎app/services/discourse_ai/language_detector.rb‎

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -49,35 +49,33 @@ def detect
4949
messages: [{ type: :user, content: @text, id: "user" }],
5050
)
5151

52-
response =
52+
structured_output =
5353
DiscourseAi::Completions::Llm.proxy(SiteSetting.ai_helper_model).generate(
5454
prompt,
5555
user: Discourse.system_user,
5656
feature_name: "translator-language-detect",
57-
extra_model_params: response_format,
57+
response_format: response_format,
5858
)
5959

60-
locale = JSON.parse(response)&.dig("locale")
60+
structured_output&.read_latest_buffered_chunk&.dig(:locale)
6161
end
6262

6363
def response_format
6464
{
65-
response_format: {
66-
type: "json_schema",
67-
json_schema: {
68-
name: "reply",
69-
schema: {
70-
type: "object",
71-
properties: {
72-
locale: {
73-
type: "string",
74-
},
65+
type: "json_schema",
66+
json_schema: {
67+
name: "reply",
68+
schema: {
69+
type: "object",
70+
properties: {
71+
locale: {
72+
type: "string",
7573
},
76-
required: ["locale"],
77-
additionalProperties: false,
7874
},
79-
strict: true,
75+
required: ["locale"],
76+
additionalProperties: false,
8077
},
78+
strict: true,
8179
},
8280
}
8381
end

‎spec/services/discourse_ai/base_translator_spec.rb‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
describe ".translate" do
1414
let(:text_to_translate) { "cats are great" }
1515
let(:target_language) { "de" }
16-
let(:llm_response) { "{\"translation\":\"hur dur hur dur!\"}" }
16+
let(:llm_response) { "hur dur hur dur!" }
1717

1818
it "creates the correct prompt" do
1919
post_translator = DiscourseAi::PostTranslator.new(text_to_translate, target_language)
@@ -32,6 +32,10 @@
3232
mock_llm = instance_double(DiscourseAi::Completions::Llm)
3333
post_translator = DiscourseAi::PostTranslator.new(text_to_translate, target_language)
3434

35+
structured_output =
36+
DiscourseAi::Completions::StructuredOutput.new({ translation: { type: "string" } })
37+
structured_output << { translation: llm_response }.to_json
38+
3539
allow(DiscourseAi::Completions::Prompt).to receive(:new).and_return(mock_prompt)
3640
allow(DiscourseAi::Completions::Llm).to receive(:proxy).with(
3741
SiteSetting.ai_helper_model,
@@ -40,8 +44,8 @@
4044
mock_prompt,
4145
user: Discourse.system_user,
4246
feature_name: "translator-translate",
43-
extra_model_params: post_translator.response_format,
44-
).and_return(llm_response)
47+
response_format: post_translator.response_format,
48+
).and_return(structured_output)
4549

4650
post_translator.translate
4751
end

‎spec/services/discourse_ai/language_detector_spec.rb‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
describe ".detect" do
1414
let(:locale_detector) { described_class.new("meow") }
15-
let(:llm_response) { "{\"translation\":\"hur dur hur dur!\"}" }
15+
let(:llm_response) { "hur dur hur dur!" }
1616

1717
it "creates the correct prompt" do
1818
allow(DiscourseAi::Completions::Prompt).to receive(:new).with(
@@ -29,6 +29,10 @@
2929
mock_prompt = instance_double(DiscourseAi::Completions::Prompt)
3030
mock_llm = instance_double(DiscourseAi::Completions::Llm)
3131

32+
structured_output =
33+
DiscourseAi::Completions::StructuredOutput.new({ locale: { type: "string" } })
34+
structured_output << { locale: llm_response }.to_json
35+
3236
allow(DiscourseAi::Completions::Prompt).to receive(:new).and_return(mock_prompt)
3337
allow(DiscourseAi::Completions::Llm).to receive(:proxy).with(
3438
SiteSetting.ai_helper_model,
@@ -37,8 +41,8 @@
3741
mock_prompt,
3842
user: Discourse.system_user,
3943
feature_name: "translator-language-detect",
40-
extra_model_params: locale_detector.response_format,
41-
).and_return(llm_response)
44+
response_format: locale_detector.response_format,
45+
).and_return(structured_output)
4246

4347
locale_detector.detect
4448
end

‎spec/services/discourse_ai_spec.rb‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
describe ".detect!" do
2929
it "returns the detected language" do
3030
locale = "de"
31-
DiscourseAi::Completions::Llm.with_prepared_responses([locale_json(locale)]) do
31+
DiscourseAi::Completions::Llm.with_prepared_responses([locale]) do
3232
expect(DiscourseTranslator::Provider::DiscourseAi.detect!(post)).to eq locale
3333
end
3434
end
@@ -41,29 +41,23 @@
4141
end
4242

4343
it "translates the post and returns [locale, translated_text]" do
44-
DiscourseAi::Completions::Llm.with_prepared_responses(
45-
[translation_json("some translated text")],
46-
) do
44+
DiscourseAi::Completions::Llm.with_prepared_responses(["some translated text"]) do
4745
translated_text = DiscourseTranslator::Provider::DiscourseAi.translate_translatable!(post)
4846
expect(translated_text).to eq "<p>some translated text</p>"
4947
end
5048
end
5149

5250
it "translates the topic" do
5351
allow(::DiscourseAi::TopicTranslator).to receive(:new).and_call_original
54-
DiscourseAi::Completions::Llm.with_prepared_responses(
55-
[translation_json("some translated text")],
56-
) do
52+
DiscourseAi::Completions::Llm.with_prepared_responses(["some translated text"]) do
5753
translated_text = DiscourseTranslator::Provider::DiscourseAi.translate_translatable!(topic)
5854
expect(translated_text).to eq "some translated text"
5955
end
6056
end
6157

6258
it "sends the content for splitting and the split content for translation" do
6359
post.update(raw: "#{"a" * 3000} #{"b" * 3000}")
64-
DiscourseAi::Completions::Llm.with_prepared_responses(
65-
%w[lol wut].map { |content| translation_json(content) },
66-
) do
60+
DiscourseAi::Completions::Llm.with_prepared_responses(%w[lol wut]) do
6761
expect(
6862
DiscourseTranslator::Provider::DiscourseAi.translate_translatable!(post),
6963
).to eq "<p>lolwut</p>"
@@ -73,20 +67,10 @@
7367

7468
describe ".translate_text!" do
7569
it "returns the translated text" do
76-
DiscourseAi::Completions::Llm.with_prepared_responses(
77-
[translation_json("some translated text")],
78-
) do
70+
DiscourseAi::Completions::Llm.with_prepared_responses(["some translated text"]) do
7971
translated_text = DiscourseTranslator::Provider::DiscourseAi.translate_text!("derp")
8072
expect(translated_text).to eq "some translated text"
8173
end
8274
end
8375
end
84-
85-
def locale_json(content)
86-
{ locale: content }.to_json
87-
end
88-
89-
def translation_json(content)
90-
{ translation: content }.to_json
91-
end
9276
end

0 commit comments

Comments
 (0)