discourse-ai/spec/shared/inference/function_list_spec.rb

109 lines
2.6 KiB
Ruby

# frozen_string_literal: true
require "rails_helper"
module DiscourseAi::Inference
describe FunctionList do
let :function_list do
function =
Function.new(name: "get_weather", description: "Get the weather in a city (default to c)")
function.add_parameter(
name: "location",
type: "string",
description: "the city name",
required: true,
)
function.add_parameter(
name: "unit",
type: "string",
description: "the unit of measurement celcius c or fahrenheit f",
enum: %w[c f],
required: false,
)
list = FunctionList.new
list << function
list
end
let :image_function_list do
function = Function.new(name: "image", description: "generates an image")
function.add_parameter(
name: "prompts",
type: "array",
item_type: "string",
required: true,
description: "the prompts",
)
list = FunctionList.new
list << function
list
end
it "can handle function call parsing" do
raw_prompt = <<~PROMPT
<function_calls>
<invoke>
<tool_name>image</tool_name>
<parameters>
<prompts>
[
"an oil painting",
"a cute fluffy orange",
"3 apple's",
"a cat"
]
</prompts>
</parameters>
</invoke>
</function_calls>
PROMPT
parsed = image_function_list.parse_prompt(raw_prompt)
expect(parsed).to eq(
[
{
name: "image",
arguments: {
prompts: ["an oil painting", "a cute fluffy orange", "3 apple's", "a cat"],
},
},
],
)
end
it "can generate a general custom system prompt" do
prompt = function_list.system_prompt
# this is fragile, by design, we need to test something here
#
expected = <<~PROMPT
<tools>
<tool_description>
<tool_name>get_weather</tool_name>
<description>Get the weather in a city (default to c)</description>
<parameters>
<parameter>
<name>location</name>
<type>string</type>
<description>the city name</description>
<required>true</required>
</parameter>
<parameter>
<name>unit</name>
<type>string</type>
<description>the unit of measurement celcius c or fahrenheit f</description>
<required>false</required>
<options>c,f</options>
</parameter>
</parameters>
</tool_description>
</tools>
PROMPT
expect(prompt).to include(expected)
end
end
end