Skip to content

Commit

Permalink
Build support for arrays / objects on tools
Browse files Browse the repository at this point in the history
This allows much more complex calls with tools
  • Loading branch information
ksylvest committed Nov 4, 2024
1 parent a8c2ffd commit 12d2320
Show file tree
Hide file tree
Showing 14 changed files with 481 additions and 81 deletions.
2 changes: 1 addition & 1 deletion Gemfile.lock
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
PATH
remote: .
specs:
omniai (1.9.0)
omniai (1.9.1)
event_stream_parser
http
zeitwerk
Expand Down
30 changes: 26 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -72,16 +72,38 @@ require 'omniai/google'

CLIENT = OmniAI::Google::Client.new

LOCATION = OmniAI::Tool::Property.object(
properties: {
city: OmniAI::Tool::Property.string(description: 'e.g. "Toronto"'),
country: OmniAI::Tool::Property.string(description: 'e.g. "Canada"'),
},
required: %i[city country]
)

LOCATIONS = OmniAI::Tool::Property.array(
min_items: 1,
max_items: 5,
items: LOCATION
)

UNIT = OmniAI::Tool::Property.string(enum: %w[celcius fahrenheit])

WEATHER = proc do |locations:, unit: 'celsius'|
locations.map do |location|
"#{rand(20..50)}° #{unit} in #{location[:city]}, #{location[:country]}"
end.join("\n")
end

TOOL = OmniAI::Tool.new(
proc { |location:, unit: 'celsius'| "#{rand(20..50)}° #{unit} in #{location}" },
WEATHER,
name: 'Weather',
description: 'Lookup the weather in a location',
parameters: OmniAI::Tool::Parameters.new(
properties: {
location: OmniAI::Tool::Property.string(description: 'e.g. Toronto'),
unit: OmniAI::Tool::Property.string(enum: %w[celcius fahrenheit]),
locations: LOCATIONS,
unit: UNIT,
},
required: %i[location]
required: %i[locations]
)
)

Expand Down
30 changes: 26 additions & 4 deletions examples/tools
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,38 @@ require 'omniai/google'

CLIENT = OmniAI::Google::Client.new

LOCATION = OmniAI::Tool::Property.object(
properties: {
city: OmniAI::Tool::Property.string(description: 'e.g. "Toronto"'),
country: OmniAI::Tool::Property.string(description: 'e.g. "Canada"'),
},
required: %i[city country]
)

LOCATIONS = OmniAI::Tool::Property.array(
min_items: 1,
max_items: 5,
items: LOCATION
)

UNIT = OmniAI::Tool::Property.string(enum: %w[celcius fahrenheit])

WEATHER = proc do |locations:, unit: 'celsius'|
locations.map do |location|
"#{rand(20..50)}° #{unit} in #{location[:city]}, #{location[:country]}"
end.join("\n")
end

TOOL = OmniAI::Tool.new(
proc { |location:, unit: 'celsius'| "#{rand(20..50)}° #{unit} in #{location}" },
WEATHER,
name: 'Weather',
description: 'Lookup the weather in a location',
parameters: OmniAI::Tool::Parameters.new(
properties: {
location: OmniAI::Tool::Property.string(description: 'e.g. Toronto'),
unit: OmniAI::Tool::Property.string(enum: %w[celcius fahrenheit]),
locations: LOCATIONS,
unit: UNIT,
},
required: %i[location]
required: %i[locations]
)
)

Expand Down
74 changes: 74 additions & 0 deletions lib/omniai/tool/array.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
# frozen_string_literal: true

module OmniAI
class Tool
# Represents a schema object.
#
# @example
# array = OmniAI::Tool::Array.new(
# description: 'A list of people.',
# items: OmniAI::Tool::Object.new(
# properties: {
# name: OmniAI::Tool::Property.string(description: 'The name of the person.'),
# age: OmniAI::Tool::Property.integer(description: 'The age of the person.'),
# },
# required: %i[name]
# ),
# min_items: 1,
# max_items: 5,
# })
class Array
TYPE = 'array'

# @!attribute [rw] items
# @return [OmniAI::Tool::Object, OmniAI::Tool::Array, OmniAI::Tool::Property]
attr_accessor :items

# @!attribute [rw] max_items
# @return [Integer, nil]
attr_accessor :max_items

# @!attribute [rw] min_items
# @return [Integer, nil]
attr_accessor :min_items

# @!attribute [rw] description
# @return [String, nil]
attr_accessor :description

# @param items [OmniAI::Tool::Object, OmniAI::Tool::Array, OmniAI::Tool::Property] required
# @param max_items [Integer] optional
# @param min_items [Integer] optional
# @param description [String] optional
def initialize(items:, max_items: nil, min_items: nil, description: nil)
@items = items
@description = description
@max_items = max_items
@min_items = min_items
end

# @example
# array.serialize # => { type: 'array', items: { type: 'string' } }
#
# @return [Hash]
def serialize
{
type: TYPE,
description: @description,
items: @items.serialize,
maxItems: @max_items,
minItems: @min_items,
}.compact
end

# @example
# array.parse(['1', '2', '3']) # => [1, 2, 3]
# @param args [Array]
#
# @return [Array]
def parse(args)
args.map { |arg| @items.parse(arg) }
end
end
end
end
62 changes: 62 additions & 0 deletions lib/omniai/tool/object.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# frozen_string_literal: true

module OmniAI
class Tool
# Represents a schema object.
#
# @example
# object = OmniAI::Tool::Object.new(
# properties: {
# name: OmniAI::Tool::Property.string(description: 'The name of the person.'),
# age: OmniAI::Tool::Property.integer(description: 'The age of the person.'),
# },
# required: %i[name]
# })
class Object
TYPE = 'object'

# @!attribute [rw] properties
# @return [Hash]
attr_accessor :properties

# @!attribute [rw] required
# @return [Array<String>]
attr_accessor :required

# @!attribute [rw] description
# @return [String, nil]
attr_accessor :description

# @param properties [Hash]
# @param required [Array<String>]
# @return [OmniAI::Tool::Parameters]
def initialize(properties: {}, required: [], description: nil)
@properties = properties
@required = required
@description = description
end

# @return [Hash]
def serialize
{
type: TYPE,
description: @description,
properties: @properties.transform_values(&:serialize),
required: @required,
}.compact
end

# @param args [Hash]
#
# @return [Hash]
def parse(args)
result = {}
@properties.each do |name, property|
value = args[String(name)]
result[name.intern] = property.parse(value) unless value.nil?
end
result
end
end
end
end
54 changes: 15 additions & 39 deletions lib/omniai/tool/parameters.rb
Original file line number Diff line number Diff line change
Expand Up @@ -2,46 +2,22 @@

module OmniAI
class Tool
# Usage:
# Parameters are used to define the arguments for a tool.
#
# parameters = OmniAI::Tool::Parameters.new(properties: {
# n: OmniAI::Tool::Parameters.integer(description: 'The nth number to calculate.')
# required: %i[n]
# })
class Parameters
module Type
OBJECT = 'object'
end

# @param type [String]
# @param properties [Hash]
# @param required [Array<String>]
# @return [OmniAI::Tool::Parameters]
def initialize(type: Type::OBJECT, properties: {}, required: [])
@type = type
@properties = properties
@required = required
end

# @return [Hash]
def serialize
{
type: @type,
properties: @properties.transform_values(&:serialize),
required: @required,
}.compact
end

# @param args [Hash]
# @return [Hash]
def parse(args)
result = {}
@properties.each do |name, property|
value = args[String(name)]
result[name.intern] = property.parse(value) if value
end
result
end
# @example
# parameters = OmniAI::Tool::Parameters.new(properties: {
# people: OmniAI::Tool::Parameters.array(
# items: OmniAI::Tool::Parameters.object(
# properties: {
# name: OmniAI::Tool::Parameters.string(description: 'The name of the person.'),
# age: OmniAI::Tool::Parameters.integer(description: 'The age of the person.'),
# employeed: OmniAI::Tool::Parameters.boolean(description: 'Is the person employeed?'),
# }
# n: OmniAI::Tool::Parameters.integer(description: 'The nth number to calculate.')
# required: %i[n]
# })
# tool = OmniAI::Tool.new(fibonacci, parameters: parameters)
class Parameters < Object
end
end
end
Loading

0 comments on commit 12d2320

Please sign in to comment.