riffer 0.28.0 → 0.29.0

data/lib/riffer/providers/amazon_bedrock.rb

@@ -63,7 +63,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
       system: partitioned_messages[:system],
       messages: partitioned_messages[:conversation],
       **options.except(:tools, :structured_output)
-    }
+    } #: Hash[Symbol, untyped]
 
     if tools && !tools.empty?
       params[:tool_config] = {
@@ -92,18 +92,17 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Hash[Symbol, untyped]) -> Aws::BedrockRuntime::Types::ConverseResponse
+  #: (Hash[Symbol, untyped]) -> Aws::BedrockRuntime::Client::_ConverseResponseSuccess
   def execute_generate(params)
     @client.converse(**params)
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ConverseResponse) -> Riffer::TokenUsage?
+  #: (Aws::BedrockRuntime::Client::_ConverseResponseSuccess) -> Riffer::Providers::TokenUsage?
   def extract_token_usage(response)
     usage = response.usage
-    return nil unless usage
 
-    Riffer::TokenUsage.new(
+    Riffer::Providers::TokenUsage.new(
       input_tokens: usage.input_tokens,
       output_tokens: usage.output_tokens,
       cache_creation_tokens: usage.cache_write_input_tokens,
@@ -112,7 +111,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ConverseResponse) -> String
+  #: (Aws::BedrockRuntime::Client::_ConverseResponseSuccess) -> String
   def extract_content(response)
     content_blocks = response.output&.message&.content
     return "" if content_blocks.nil? || content_blocks.empty?
@@ -127,12 +126,12 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ConverseResponse) -> Array[Riffer::Messages::Assistant::ToolCall]
+  #: (Aws::BedrockRuntime::Client::_ConverseResponseSuccess) -> Array[Riffer::Messages::Assistant::ToolCall]
   def extract_tool_calls(response)
     content_blocks = response.output&.message&.content
     return [] if content_blocks.nil? || content_blocks.empty?
 
-    tool_calls = []
+    tool_calls = [] #: Array[Riffer::Messages::Assistant::ToolCall]
 
     content_blocks.each do |block|
       if block.respond_to?(:tool_use) && block.tool_use
@@ -153,7 +152,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
     current_state = {
       text: nil,
       tool_call: nil
-    }
+    } #: Hash[Symbol, untyped]
 
     @client.converse_stream(**params) do |stream|
       stream.on_event do |event|
@@ -196,7 +195,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockStartEvent, state: Hash[Symbol, untyped], yielder: Enumerator[Riffer::StreamEvents::Base, void]) -> void
+  #: (Aws::BedrockRuntime::Types::ContentBlockStartEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_content_block_start_tool_use(event, state:, yielder:)
     state[:tool_call] = {
       id: event.start.tool_use.tool_use_id,
@@ -206,7 +205,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockDeltaEvent, state: Hash[Symbol, untyped], yielder: Enumerator[Riffer::StreamEvents::Base, void]) -> void
+  #: (Aws::BedrockRuntime::Types::ContentBlockDeltaEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_content_block_delta_text_delta(event, state:, yielder:)
     delta_text = event.delta.text
     state[:text] ||= ""
@@ -215,7 +214,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockDeltaEvent, state: Hash[Symbol, untyped], yielder: Enumerator[Riffer::StreamEvents::Base, void]) -> void
+  #: (Aws::BedrockRuntime::Types::ContentBlockDeltaEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_content_block_delta_tool_use(event, state:, yielder:)
     input_delta = event.delta.tool_use.input
 
@@ -229,14 +228,14 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockStopEvent, state: Hash[Symbol, untyped], yielder: Enumerator[Riffer::StreamEvents::Base, void]) -> void
+  #: (Aws::BedrockRuntime::Types::ContentBlockStopEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_content_block_stop_text_delta(_event, state:, yielder:)
     yielder << Riffer::StreamEvents::TextDone.new(state[:text])
     state[:text] = nil
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockStopEvent, state: Hash[Symbol, untyped], yielder: Enumerator[Riffer::StreamEvents::Base, void]) -> void
+  #: (Aws::BedrockRuntime::Types::ContentBlockStopEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_content_block_stop_tool_use(_event, state:, yielder:)
     tool_call = state[:tool_call]
     yielder << Riffer::StreamEvents::ToolCallDone.new(
@@ -249,10 +248,10 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ConverseStreamMetadataEvent, state: Hash[Symbol, untyped], yielder: Enumerator[Riffer::StreamEvents::Base, void]) -> void
+  #: (Aws::BedrockRuntime::Types::ConverseStreamMetadataEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_metadata_usage(event, state:, yielder:)
     yielder << Riffer::StreamEvents::TokenUsageDone.new(
-      token_usage: Riffer::TokenUsage.new(
+      token_usage: Riffer::Providers::TokenUsage.new(
         input_tokens: event.usage.input_tokens,
         output_tokens: event.usage.output_tokens,
         cache_creation_tokens: event.usage.cache_write_input_tokens,
@@ -264,8 +263,8 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   #--
   #: (Array[Riffer::Messages::Base]) -> Hash[Symbol, untyped]
   def partition_messages(messages)
-    system_prompts = []
-    conversation_messages = []
+    system_prompts = [] #: Array[Hash[Symbol, untyped]]
+    conversation_messages = [] #: Array[Hash[Symbol, untyped]]
 
     messages.each do |message|
       case message
@@ -291,7 +290,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   #--
   #: (Riffer::Messages::Assistant) -> Hash[Symbol, untyped]
   def convert_assistant_to_bedrock_format(message)
-    content = []
+    content = [] #: Array[Hash[Symbol, untyped]]
     content << {text: message.content} if message.content && !message.content.empty?
 
     message.tool_calls.each do |tc|
@@ -326,7 +325,7 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Riffer::FilePart) -> Hash[Symbol, untyped]
+  #: (Riffer::Messages::FilePart) -> Hash[Symbol, untyped]
   def convert_file_part_to_bedrock_format(file)
     format = bedrock_format(file.media_type)
 

data/lib/riffer/providers/anthropic.rb

@@ -26,7 +26,7 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
 
     api_key ||= Riffer.config.anthropic.api_key
 
-    @client = Anthropic::Client.new(api_key: api_key, **options)
+    @client = ::Anthropic::Client.new(api_key: api_key, **options)
   end
 
   private
@@ -46,11 +46,11 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
       messages: partitioned_messages[:conversation],
       max_tokens: max_tokens,
       **options.except(:tools, :max_tokens, :structured_output, :web_search)
-    }
+    } #: Hash[Symbol, untyped]
 
     params[:system] = partitioned_messages[:system] if partitioned_messages[:system]
 
-    anthropic_tools = []
+    anthropic_tools = [] #: Array[Hash[Symbol, untyped]]
     anthropic_tools.concat(tools.map { |t| convert_tool_to_anthropic_format(t) }) if tools && !tools.empty?
 
     if web_search
@@ -83,12 +83,11 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   end
 
   #--
-  #: (Anthropic::Models::Message) -> Riffer::TokenUsage?
+  #: (Anthropic::Models::Message) -> Riffer::Providers::TokenUsage?
   def extract_token_usage(response)
     usage = response.usage
-    return nil unless usage
 
-    Riffer::TokenUsage.new(
+    Riffer::Providers::TokenUsage.new(
       input_tokens: usage.input_tokens,
       output_tokens: usage.output_tokens,
       cache_creation_tokens: usage.cache_creation_input_tokens,
@@ -105,7 +104,7 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
     text_content = ""
 
     content_blocks.each do |block|
-      text_content = block.text if block.type.to_s == "text"
+      text_content = block.text if block.is_a?(::Anthropic::Models::TextBlock)
     end
 
     text_content
@@ -117,10 +116,10 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
     content_blocks = response.content
     return [] if content_blocks.nil? || content_blocks.empty?
 
-    tool_calls = []
+    tool_calls = [] #: Array[Riffer::Messages::Assistant::ToolCall]
 
     content_blocks.each do |block|
-      if block.type.to_s == "tool_use"
+      if block.is_a?(::Anthropic::Models::ToolUseBlock)
         tool_calls << Riffer::Messages::Assistant::ToolCall.new(
           call_id: block.id,
           name: decode_tool_name(block.name, tools: @current_tools),
@@ -142,7 +141,7 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
       web_search_index: nil,
       web_search_json: nil,
       web_search_query: nil
-    }
+    } #: Hash[Symbol, untyped]
 
     # Workaround for anthropics/anthropic-sdk-ruby#182: force identity
     # encoding so Net::HTTP/Zlib doesn't buffer SSE chunks until EOF.
@@ -154,24 +153,24 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
     begin
       stream.each do |event|
         case event
-        when Anthropic::Models::RawContentBlockStartEvent
+        when ::Anthropic::Models::RawContentBlockStartEvent
           handle_raw_content_block_start(event, state: current_state)
-        when Anthropic::Models::RawContentBlockDeltaEvent
+        when ::Anthropic::Models::RawContentBlockDeltaEvent
           handle_raw_content_block_delta(event, state: current_state)
-        when Anthropic::Streaming::TextEvent
+        when ::Anthropic::Helpers::Streaming::TextEvent
           handle_text_event(event, state: current_state, yielder: yielder)
-        when Anthropic::Streaming::ThinkingEvent
+        when ::Anthropic::Helpers::Streaming::ThinkingEvent
           handle_thinking_event(event, state: current_state, yielder: yielder)
-        when Anthropic::Streaming::InputJsonEvent
+        when ::Anthropic::Helpers::Streaming::InputJsonEvent
           handle_input_json_event(event, state: current_state, yielder: yielder)
-        when Anthropic::Streaming::ContentBlockStopEvent
-          block_type = event.content_block&.type.to_s
-          handle_content_block_stop_text(event, state: current_state, yielder: yielder) if block_type == "text" && current_state[:text]
-          handle_content_block_stop_tool_use(event, state: current_state, yielder: yielder) if block_type == "tool_use"
-          handle_content_block_stop_thinking(event, state: current_state, yielder: yielder) if block_type == "thinking" && current_state[:reasoning]
-          handle_content_block_stop_server_tool_use(event, state: current_state, yielder: yielder) if block_type == "server_tool_use"
-          handle_content_block_stop_web_search_result(event, state: current_state, yielder: yielder) if block_type == "web_search_tool_result"
-        when Anthropic::Streaming::MessageStopEvent
+        when ::Anthropic::Helpers::Streaming::ContentBlockStopEvent
+          block = event.content_block
+          handle_content_block_stop_text(event, state: current_state, yielder: yielder) if block.is_a?(::Anthropic::Models::TextBlock) && current_state[:text]
+          handle_content_block_stop_tool_use(event, state: current_state, yielder: yielder) if block.is_a?(::Anthropic::Models::ToolUseBlock)
+          handle_content_block_stop_thinking(event, state: current_state, yielder: yielder) if block.is_a?(::Anthropic::Models::ThinkingBlock) && current_state[:reasoning]
+          handle_content_block_stop_server_tool_use(event, state: current_state, yielder: yielder) if block.is_a?(::Anthropic::Models::ServerToolUseBlock)
+          handle_content_block_stop_web_search_result(event, state: current_state, yielder: yielder) if block.is_a?(::Anthropic::Models::WebSearchToolResultBlock)
+        when ::Anthropic::Helpers::Streaming::MessageStopEvent
           handle_message_stop(event, accumulated_message: stream.accumulated_message, yielder: yielder)
         end
       end
@@ -292,7 +291,7 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
     return unless usage
 
     yielder << Riffer::StreamEvents::TokenUsageDone.new(
-      token_usage: Riffer::TokenUsage.new(
+      token_usage: Riffer::Providers::TokenUsage.new(
         input_tokens: usage.input_tokens,
         output_tokens: usage.output_tokens,
         cache_creation_tokens: usage.cache_creation_input_tokens,
@@ -304,8 +303,8 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   #--
   #: (Array[Riffer::Messages::Base]) -> Hash[Symbol, untyped]
   def partition_messages(messages)
-    system_prompts = []
-    conversation_messages = []
+    system_prompts = [] #: Array[Hash[Symbol, untyped]]
+    conversation_messages = [] #: Array[Hash[Symbol, untyped]]
 
     messages.each do |message|
       case message
@@ -342,7 +341,7 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   #--
   #: (Riffer::Messages::Assistant) -> Hash[Symbol, untyped]
   def convert_assistant_to_anthropic_format(message)
-    content = []
+    content = [] #: Array[Hash[Symbol, untyped]]
     content << {type: "text", text: message.content} if message.content && !message.content.empty?
 
     message.tool_calls.each do |tc|
@@ -358,7 +357,7 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   end
 
   #--
-  #: (Riffer::FilePart) -> Hash[Symbol, untyped]
+  #: (Riffer::Messages::FilePart) -> Hash[Symbol, untyped]
   def convert_file_part_to_anthropic_format(file)
     type = file.image? ? "image" : "document"
 

data/lib/riffer/providers/base.rb

@@ -39,10 +39,10 @@ class Riffer::Providers::Base
   # Generates text using the provider.
   #
   #--
-  #: (?prompt: String?, ?system: String?, ?messages: Array[Hash[Symbol, untyped] | Riffer::Messages::Base]?, ?model: String?, ?files: Array[Hash[Symbol, untyped] | Riffer::FilePart]?, **untyped) -> Riffer::Messages::Assistant
+  #: (?prompt: String?, ?system: String?, ?messages: Array[Hash[Symbol, untyped] | Riffer::Messages::Base]?, ?model: String?, ?files: Array[Hash[Symbol, untyped] | Riffer::Messages::FilePart]?, **untyped) -> Riffer::Messages::Assistant
   def generate_text(prompt: nil, system: nil, messages: nil, model: nil, files: nil, **options)
     validate_input!(prompt: prompt, system: system, messages: messages)
-    @current_tools = options[:tools] || []
+    @current_tools = options[:tools] || [] #: Array[singleton(Riffer::Tool)]
     messages = normalize_messages(prompt: prompt, system: system, messages: messages, files: files)
     validate_normalized_messages!(messages)
     messages = merge_consecutive_messages(messages)
@@ -65,10 +65,10 @@ class Riffer::Providers::Base
   # Streams text from the provider.
   #
   #--
-  #: (?prompt: String?, ?system: String?, ?messages: Array[Hash[Symbol, untyped] | Riffer::Messages::Base]?, ?model: String?, ?files: Array[Hash[Symbol, untyped] | Riffer::FilePart]?, **untyped) -> Enumerator[Riffer::StreamEvents::Base, void]
+  #: (?prompt: String?, ?system: String?, ?messages: Array[Hash[Symbol, untyped] | Riffer::Messages::Base]?, ?model: String?, ?files: Array[Hash[Symbol, untyped] | Riffer::Messages::FilePart]?, **untyped) -> Enumerator[Riffer::StreamEvents::Base, void]
   def stream_text(prompt: nil, system: nil, messages: nil, model: nil, files: nil, **options)
     validate_input!(prompt: prompt, system: system, messages: messages)
-    @current_tools = options[:tools] || []
+    @current_tools = options[:tools] || [] #: Array[singleton(Riffer::Tool)]
     messages = normalize_messages(prompt: prompt, system: system, messages: messages, files: files)
     validate_normalized_messages!(messages)
     messages = merge_consecutive_messages(messages)
@@ -87,7 +87,7 @@ class Riffer::Providers::Base
   end
 
   #--
-  #: (String, tools: Array[Riffer::Tool]) -> String
+  #: (String, tools: Array[singleton(Riffer::Tool)]) -> String
   def decode_tool_name(wire_name, tools:)
     tool = tools.find { |t| encode_tool_name(t.name) == wire_name }
     tool ? tool.name : wire_name
@@ -112,7 +112,7 @@ class Riffer::Providers::Base
   end
 
   #--
-  #: (untyped) -> Riffer::TokenUsage?
+  #: (untyped) -> Riffer::Providers::TokenUsage?
   def extract_token_usage(response)
     raise NotImplementedError, "Subclasses must implement #extract_token_usage"
   end
@@ -166,7 +166,7 @@ class Riffer::Providers::Base
   end
 
   #--
-  #: (prompt: String?, system: String?, messages: Array[Hash[Symbol, untyped] | Riffer::Messages::Base]?, ?files: Array[Hash[Symbol, untyped] | Riffer::FilePart]?) -> Array[Riffer::Messages::Base]
+  #: (prompt: String?, system: String?, messages: Array[Hash[Symbol, untyped] | Riffer::Messages::Base]?, ?files: Array[Hash[Symbol, untyped] | Riffer::Messages::FilePart]?) -> Array[Riffer::Messages::Base]
   def normalize_messages(prompt:, system:, messages:, files: nil)
     if messages && files && !files.empty?
       raise Riffer::ArgumentError, "cannot provide both files and messages; attach files to individual messages instead"
@@ -176,10 +176,11 @@ class Riffer::Providers::Base
       return messages.map { |msg| convert_to_message_object(msg) }
     end
 
-    result = []
+    result = [] #: Array[Riffer::Messages::Base]
     result << Riffer::Messages::System.new(system) if system
     file_parts = (files || []).map { |f| convert_to_file_part(f) }
-    result << Riffer::Messages::User.new(prompt, files: file_parts)
+    prompt_text = prompt #: String
+    result << Riffer::Messages::User.new(prompt_text, files: file_parts)
     result
   end
 

data/lib/riffer/providers/gemini.rb

@@ -36,7 +36,7 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
     params = {
       model: model,
       contents: partitioned[:contents]
-    }
+    } #: Hash[Symbol, untyped]
 
     params[:systemInstruction] = partitioned[:system_instruction] if partitioned[:system_instruction]
 
@@ -96,12 +96,12 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
   end
 
   #--
-  #: (Hash[Symbol, untyped]) -> Riffer::TokenUsage?
+  #: (Hash[Symbol, untyped]) -> Riffer::Providers::TokenUsage?
   def extract_token_usage(response)
     usage = response[:usageMetadata]
     return nil unless usage
 
-    Riffer::TokenUsage.new(
+    Riffer::Providers::TokenUsage.new(
       input_tokens: usage[:promptTokenCount] || 0,
       output_tokens: usage[:candidatesTokenCount] || 0
     )
@@ -114,6 +114,7 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
     body = params.except(:model)
 
     uri = URI("#{BASE_URI}/#{api_path(model, "streamGenerateContent")}?alt=sse")
+    host = uri.hostname #: String
     request = Net::HTTP::Post.new(uri)
     request["Content-Type"] = "application/json"
     request["x-goog-api-key"] = @api_key
@@ -126,7 +127,8 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
       buffer << chunk
 
       while (match = buffer.match(/\r?\n\r?\n/))
-        frame = buffer.slice!(0, match.end(0)).strip
+        match_end = match.end(0) #: Integer
+        frame = buffer.slice!(0, match_end).to_s.strip
         next unless frame.start_with?("data: ")
 
         json_str = frame.delete_prefix("data: ").strip
@@ -155,7 +157,7 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
         usage = parsed[:usageMetadata]
         if usage && usage[:candidatesTokenCount]
           yielder << Riffer::StreamEvents::TokenUsageDone.new(
-            token_usage: Riffer::TokenUsage.new(
+            token_usage: Riffer::Providers::TokenUsage.new(
               input_tokens: usage[:promptTokenCount] || 0,
               output_tokens: usage[:candidatesTokenCount] || 0
             )
@@ -164,7 +166,7 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
       end
     end
 
-    Net::HTTP.start(uri.hostname, uri.port, use_ssl: true, open_timeout: @open_timeout, read_timeout: @read_timeout) do |http|
+    Net::HTTP.start(host, uri.port, use_ssl: true, open_timeout: @open_timeout, read_timeout: @read_timeout) do |http|
       http.request(request) do |response|
         handle_api_error!(response) unless response.is_a?(Net::HTTPSuccess)
 
@@ -182,8 +184,8 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
   #--
   #: (Array[Riffer::Messages::Base]) -> Hash[Symbol, untyped]
   def partition_messages(messages)
-    system_parts = []
-    contents = []
+    system_parts = [] #: Array[Hash[Symbol, untyped]]
+    contents = [] #: Array[Hash[Symbol, untyped]]
 
     messages.each do |message|
       case message
@@ -212,7 +214,7 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
       end
     end
 
-    result = {contents: contents}
+    result = {contents: contents} #: Hash[Symbol, untyped]
     result[:system_instruction] = {parts: system_parts} unless system_parts.empty?
     result
   end
@@ -220,7 +222,7 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
   #--
   #: (Riffer::Messages::Assistant) -> Hash[Symbol, untyped]
   def convert_assistant_to_gemini_format(message)
-    parts = []
+    parts = [] #: Array[Hash[Symbol, untyped]]
     parts << {text: message.content} if message.content && !message.content.empty?
 
     message.tool_calls.each do |tc|
@@ -236,7 +238,7 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
   end
 
   #--
-  #: (Riffer::FilePart) -> Hash[Symbol, untyped]
+  #: (Riffer::Messages::FilePart) -> Hash[Symbol, untyped]
   def convert_file_part_to_gemini_format(file)
     if file.url?
       raise Riffer::ArgumentError,
@@ -268,11 +270,12 @@ class Riffer::Providers::Gemini < Riffer::Providers::Base
   #: (String, Hash[Symbol, untyped]) -> Net::HTTPResponse
   def post_request(path, body)
     uri = URI("#{BASE_URI}/#{path}")
+    host = uri.hostname #: String
     request = Net::HTTP::Post.new(uri)
     request["Content-Type"] = "application/json"
     request["x-goog-api-key"] = @api_key
     request.body = body.to_json
-    Net::HTTP.start(uri.hostname, uri.port, use_ssl: true, open_timeout: @open_timeout, read_timeout: @read_timeout) { |http| http.request(request) }
+    Net::HTTP.start(host, uri.port, use_ssl: true, open_timeout: @open_timeout, read_timeout: @read_timeout) { |http| http.request(request) }
   end
 
   #--

data/lib/riffer/providers/mock.rb

@@ -25,10 +25,21 @@ class Riffer::Providers::Mock < Riffer::Providers::Base
 
   # Initializes the mock provider.
   #
+  # +responses:+ accepts an array of response hashes in the same shape
+  # +#stub_response+ takes — raw +tool_calls:+ hashes are normalised to
+  # +Riffer::Messages::Assistant::ToolCall+ instances. This is the canonical
+  # way to pre-configure canned LLM responses on an agent via
+  # +provider_options responses: [...]+.
+  #
+  #   Riffer::Providers::Mock.new(responses: [
+  #     {content: "", tool_calls: [{name: "tool_a", arguments: "{}"}]},
+  #     {content: "Final answer"}
+  #   ])
+  #
   #--
   #: (**untyped) -> void
   def initialize(**options)
-    @responses = options[:responses] || []
+    @responses = (options[:responses] || []).map { |r| normalize_response(r) }
     @current_index = 0
     @calls = []
     @stubbed_responses = []
@@ -40,19 +51,12 @@ class Riffer::Providers::Mock < Riffer::Providers::Base
   #
   #   provider.stub_response("Hello")
   #   provider.stub_response("", tool_calls: [{name: "my_tool", arguments: '{"key":"value"}'}])
-  #   provider.stub_response("Final response", token_usage: Riffer::TokenUsage.new(input_tokens: 10, output_tokens: 5))
+  #   provider.stub_response("Final response", token_usage: Riffer::Providers::TokenUsage.new(input_tokens: 10, output_tokens: 5))
   #
   #--
-  #: (String, ?tool_calls: Array[Hash[Symbol, untyped]], ?token_usage: Riffer::TokenUsage?) -> void
+  #: (String, ?tool_calls: Array[Hash[Symbol, untyped]], ?token_usage: Riffer::Providers::TokenUsage?) -> void
   def stub_response(content, tool_calls: [], token_usage: nil)
-    formatted_tool_calls = tool_calls.map.with_index do |tc, idx|
-      Riffer::Messages::Assistant::ToolCall.new(
-        call_id: tc[:call_id] || tc[:id] || "mock_call_#{idx}",
-        name: tc[:name],
-        arguments: tc[:arguments].is_a?(String) ? tc[:arguments] : tc[:arguments].to_json
-      )
-    end
-    @stubbed_responses << {role: "assistant", content: content, tool_calls: formatted_tool_calls, token_usage: token_usage}
+    @stubbed_responses << normalize_response(content: content, tool_calls: tool_calls, token_usage: token_usage)
   end
 
   # Clears all stubbed responses.
@@ -65,6 +69,30 @@ class Riffer::Providers::Mock < Riffer::Providers::Base
 
   private
 
+  # Normalises a response hash into Mock's internal format. Accepts the
+  # +#stub_response+ kwargs shape (+content:+, +tool_calls:+, +token_usage:+)
+  # or a pre-built hash with already-converted ToolCall instances. Raw
+  # +tool_calls:+ hashes are wrapped in +Riffer::Messages::Assistant::ToolCall+.
+  #
+  #--
+  #: (Hash[Symbol, untyped]) -> Hash[Symbol, untyped]
+  def normalize_response(response)
+    formatted_tool_calls = (response[:tool_calls] || []).map.with_index do |tc, idx|
+      next tc if tc.is_a?(Riffer::Messages::Assistant::ToolCall)
+      Riffer::Messages::Assistant::ToolCall.new(
+        call_id: tc[:call_id] || tc[:id] || "mock_call_#{idx}",
+        name: tc[:name],
+        arguments: tc[:arguments].is_a?(String) ? tc[:arguments] : tc[:arguments].to_json
+      )
+    end
+    {
+      role: response[:role] || "assistant",
+      content: response[:content] || "",
+      tool_calls: formatted_tool_calls,
+      token_usage: response[:token_usage]
+    }
+  end
+
   #--
   #: (Array[Riffer::Messages::Base], String?, Hash[Symbol, untyped]) -> Hash[Symbol, untyped]
   def build_request_params(messages, model, options)
@@ -82,7 +110,7 @@ class Riffer::Providers::Mock < Riffer::Providers::Base
   end
 
   #--
-  #: (untyped) -> Riffer::TokenUsage?
+  #: (untyped) -> Riffer::Providers::TokenUsage?
   def extract_token_usage(response)
     response[:token_usage]
   end
@@ -146,7 +174,7 @@ class Riffer::Providers::Mock < Riffer::Providers::Base
   def next_response
     if @stubbed_responses.any?
       @stubbed_responses.shift
-    elsif @responses[@current_index]
+    elsif @current_index < @responses.size
       response = @responses[@current_index]
       @current_index += 1
       response