119 lines
3.6 KiB
Ruby
119 lines
3.6 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
class TestDialect < DiscourseAi::Completions::Dialects::Dialect
|
|
attr_accessor :max_prompt_tokens
|
|
|
|
def trim(messages)
|
|
trim_messages(messages)
|
|
end
|
|
|
|
def system_msg(msg)
|
|
msg
|
|
end
|
|
|
|
def user_msg(msg)
|
|
msg
|
|
end
|
|
|
|
def model_msg(msg)
|
|
msg
|
|
end
|
|
|
|
def tokenizer
|
|
DiscourseAi::Tokenizer::OpenAiTokenizer
|
|
end
|
|
end
|
|
|
|
RSpec.describe DiscourseAi::Completions::Dialects::Dialect do
|
|
fab!(:llm_model)
|
|
|
|
describe "#translate" do
|
|
let(:five_token_msg) { "This represents five tokens." }
|
|
let(:tools) do
|
|
[
|
|
{
|
|
name: "echo",
|
|
description: "echo a string",
|
|
parameters: [
|
|
{ name: "text", type: "string", description: "string to echo", required: true },
|
|
],
|
|
},
|
|
]
|
|
end
|
|
|
|
it "injects done message when tool_choice is :none and last message follows tool pattern" do
|
|
tool_call_prompt = { name: "echo", arguments: { text: "test message" } }
|
|
|
|
prompt = DiscourseAi::Completions::Prompt.new("System instructions", tools: tools)
|
|
prompt.push(type: :user, content: "echo test message")
|
|
prompt.push(type: :tool_call, content: tool_call_prompt.to_json, id: "123", name: "echo")
|
|
prompt.push(type: :tool, content: "test message".to_json, name: "echo", id: "123")
|
|
prompt.tool_choice = :none
|
|
|
|
dialect = TestDialect.new(prompt, llm_model)
|
|
dialect.max_prompt_tokens = 100 # Set high enough to avoid trimming
|
|
|
|
translated = dialect.translate
|
|
|
|
expect(translated).to eq(
|
|
[
|
|
{ type: :system, content: "System instructions" },
|
|
{ type: :user, content: "echo test message" },
|
|
{
|
|
type: :tool_call,
|
|
content:
|
|
"<function_calls>\n<invoke>\n<tool_name>echo</tool_name>\n<parameters>\n<text>test message</text>\n</parameters>\n</invoke>\n</function_calls>",
|
|
id: "123",
|
|
name: "echo",
|
|
},
|
|
{
|
|
type: :tool,
|
|
id: "123",
|
|
name: "echo",
|
|
content:
|
|
"<function_results>\n<result>\n<tool_name>echo</tool_name>\n<json>\n\"test message\"\n</json>\n</result>\n</function_results>\n\n#{::DiscourseAi::Completions::Dialects::XmlTools::DONE_MESSAGE}",
|
|
},
|
|
],
|
|
)
|
|
end
|
|
end
|
|
|
|
describe "#trim_messages" do
|
|
let(:five_token_msg) { "This represents five tokens." }
|
|
|
|
it "should trim tool messages if tool_calls are trimmed" do
|
|
prompt = DiscourseAi::Completions::Prompt.new(five_token_msg)
|
|
prompt.push(type: :user, content: five_token_msg)
|
|
prompt.push(type: :tool_call, content: five_token_msg, id: 1)
|
|
prompt.push(type: :tool, content: five_token_msg, id: 1)
|
|
prompt.push(type: :user, content: five_token_msg)
|
|
|
|
dialect = TestDialect.new(prompt, llm_model)
|
|
dialect.max_prompt_tokens = 15 # fits the user messages and the tool_call message
|
|
|
|
trimmed = dialect.trim(prompt.messages)
|
|
|
|
expect(trimmed).to eq(
|
|
[{ type: :system, content: five_token_msg }, { type: :user, content: five_token_msg }],
|
|
)
|
|
end
|
|
|
|
it "limits the system message to 60% of available tokens" do
|
|
prompt = DiscourseAi::Completions::Prompt.new("I'm a system message consisting of 10 tokens")
|
|
prompt.push(type: :user, content: five_token_msg)
|
|
|
|
dialect = TestDialect.new(prompt, llm_model)
|
|
dialect.max_prompt_tokens = 15
|
|
|
|
trimmed = dialect.trim(prompt.messages)
|
|
|
|
expect(trimmed).to eq(
|
|
[
|
|
{ type: :system, content: "I'm a system message consisting of 10" },
|
|
{ type: :user, content: five_token_msg },
|
|
],
|
|
)
|
|
end
|
|
end
|
|
end
|