tools

#!/usr/bin/env ruby
# frozen_string_literal: true

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(
  WEATHER,
  name: "Weather",
  description: "Lookup the weather in a location",
  parameters: OmniAI::Tool::Parameters.new(
    properties: {
      locations: LOCATIONS,
      unit: UNIT,
    },
    required: %i[locations]
  )
)

completion = CLIENT.chat(tools: [TOOL]) do |prompt|
  prompt.user do |message|
    message.text('What is the weather in "London" in celcius and "Seattle" in fahrenheit?')
  end
end

puts(completion.text)