diff --git a/lib/ruby_llm/chat.rb b/lib/ruby_llm/chat.rb index c33a0deb..6a23a501 100644 --- a/lib/ruby_llm/chat.rb +++ b/lib/ruby_llm/chat.rb @@ -47,6 +47,16 @@ def with_tools(*tools) self end + def with_google_search + raise UnsupportedFunctionsError, "Model #{@model.id} doesn't support function calling" unless @model.supports_functions + raise UnsupportedFunctionsError, "Google search is only supported with Gemini models" unless @model.provider == 'gemini' + + @tools = [{ + google_search: {} + }] + self + end + def model=(model_id) @model = Models.find model_id @provider = Models.provider_for model_id @@ -109,6 +119,10 @@ def handle_tool_calls(response, &) end def execute_tool(tool_call) + if tool_call.name.to_sym == :google_search + return tool_call.arguments + end + tool = tools[tool_call.name.to_sym] args = tool_call.arguments tool.call(args) diff --git a/lib/ruby_llm/providers/gemini/chat.rb b/lib/ruby_llm/providers/gemini/chat.rb index bc8ec365..7643923b 100644 --- a/lib/ruby_llm/providers/gemini/chat.rb +++ b/lib/ruby_llm/providers/gemini/chat.rb @@ -26,6 +26,12 @@ def complete(messages, tools:, temperature:, model:, &block) # rubocop:disable M end end + def with_google_search + @tools ||= [] + @tools << { google_search: {} } + self + end + # Format methods can be private private diff --git a/lib/ruby_llm/providers/gemini/tools.rb b/lib/ruby_llm/providers/gemini/tools.rb index 6ed53e17..9e674742 100644 --- a/lib/ruby_llm/providers/gemini/tools.rb +++ b/lib/ruby_llm/providers/gemini/tools.rb @@ -9,9 +9,17 @@ module Tools def format_tools(tools) return [] if tools.empty? - [{ - functionDeclarations: tools.values.map { |tool| function_declaration_for(tool) } - }] + formatted_tools = tools.map do |tool| + if tool.is_a?(Hash) && tool.key?(:google_search) + tool + else + { + functionDeclarations: [function_declaration_for(tool)] + } + end + end + + formatted_tools.flatten end # Extract tool calls from response data