riffer 0.29.0 → 0.30.0

data/lib/riffer/agent.rb

@@ -19,6 +19,8 @@ require "json"
 #   agent.generate('Hello!')
 #
 class Riffer::Agent
+  # @rbs self.@config: Riffer::Agent::Config?
+
   include Riffer::Messages::Converter
   extend Riffer::Helpers::ClassNameConverter
 
@@ -100,13 +102,19 @@ class Riffer::Agent
 
   # Gets or sets the maximum number of LLM call steps in the tool-use loop.
   #
-  # Defaults to Riffer::Agent::Config::DEFAULT_MAX_STEPS (16). Set to
-  # +Float::INFINITY+ for unlimited steps.
+  # Defaults to Riffer::Agent::Config::DEFAULT_MAX_STEPS (16). Set to +nil+
+  # for unlimited steps. The splat distinguishes a getter call (no argument)
+  # from setting the limit to +nil+.
+  #
+  #   max_steps        # reads the current limit
+  #   max_steps 8      # cap the loop at 8 steps
+  #   max_steps nil    # unlimited
   #
   #--
-  #: (?Numeric?) -> Numeric
-  def self.max_steps(value = nil)
-    value.nil? ? config.max_steps : (config.max_steps = value)
+  #: (*Numeric?) -> Numeric?
+  def self.max_steps(*value)
+    return config.max_steps if value.empty?
+    config.max_steps = value.first
   end
 
   # Gets or sets the tools used by this agent.
@@ -204,6 +212,30 @@ class Riffer::Agent
     new(context: context).stream(prompt, files: files)
   end
 
+  # Reconstructs a runnable agent from a wire dict produced by +#to_h+.
+  #
+  # Delegates to Riffer::Agent::Serializer.from_h. See it for the
+  # +tool_resolver+ / +tool_runtime+ injection points and what does not
+  # transfer.
+  #
+  #--
+  #: (Hash[Symbol, untyped], ?context: Hash[Symbol, untyped]?, ?tool_resolver: ^(Hash[Symbol, untyped]) -> singleton(Riffer::Tool), ?tool_runtime: (singleton(Riffer::Tools::Runtime) | Riffer::Tools::Runtime | Proc)?) -> Riffer::Agent
+  def self.from_h(hash, context: nil, tool_resolver: Riffer::Agent::Serializer::DEFAULT_TOOL_RESOLVER, tool_runtime: nil)
+    Riffer::Agent::Serializer.from_h(hash, context: context, tool_resolver: tool_resolver, tool_runtime: tool_runtime)
+  end
+
+  # Reconstructs a runnable agent from a JSON string produced by +#to_json+.
+  #
+  # Delegates to Riffer::Agent::Serializer.from_json, which parses the JSON
+  # (with symbol keys) for you. See Riffer::Agent::Serializer.from_h for the
+  # +tool_resolver+ / +tool_runtime+ injection points.
+  #
+  #--
+  #: (String, ?context: Hash[Symbol, untyped]?, ?tool_resolver: ^(Hash[Symbol, untyped]) -> singleton(Riffer::Tool), ?tool_runtime: (singleton(Riffer::Tools::Runtime) | Riffer::Tools::Runtime | Proc)?) -> Riffer::Agent
+  def self.from_json(json, context: nil, tool_resolver: Riffer::Agent::Serializer::DEFAULT_TOOL_RESOLVER, tool_runtime: nil)
+    Riffer::Agent::Serializer.from_json(json, context: context, tool_resolver: tool_resolver, tool_runtime: tool_runtime)
+  end
+
   # Registers a guardrail for input, output, or both phases.
   #
   # [phase] :before, :after, or :around.
@@ -259,6 +291,11 @@ class Riffer::Agent
   #   reserved and cannot be passed by the caller.
   attr_reader :context #: Riffer::Agent::Context
 
+  # The resolved provider name (the part before "provider/"), e.g. +"openai"+.
+  # Resolved eagerly at +Agent.new+ alongside +model_name+; together they
+  # form the provider-neutral model identifier the agent serializes.
+  attr_reader :provider_name #: String
+
   # The resolved model name (the part after "provider/"), used as the model
   # argument on every LLM call. Resolved eagerly at +Agent.new+.
   attr_reader :model_name #: String
@@ -306,10 +343,10 @@ class Riffer::Agent
     @config = config || self.class.config
     @context = Riffer::Agent::Context.new(context || {})
 
-    provider_class, @model_name = resolve_provider_and_model
-    @provider = provider_class.new(**@config.provider_options)
+    @provider_name, @model_name = resolve_provider_and_model
+    @provider = build_provider
 
-    @context.skills = resolve_skills(provider_class)
+    @context.skills = resolve_skills
 
     @structured_output = resolve_structured_output
     @tools = resolve_tools
@@ -374,6 +411,25 @@ class Riffer::Agent
     throw :riffer_interrupt, reason
   end
 
+  # Snapshots this resolved agent into a self-contained, provider-neutral
+  # wire dict. Delegates to Riffer::Agent::Serializer.to_h.
+  #
+  #--
+  #: () -> Hash[Symbol, untyped]
+  def to_h
+    Riffer::Agent::Serializer.to_h(agent: self)
+  end
+
+  # Snapshots this resolved agent into a wire JSON string. Delegates to
+  # Riffer::Agent::Serializer.to_json. The +*+ absorbs the JSON generator
+  # state argument so <tt>JSON.generate(agent)</tt> works too.
+  #
+  #--
+  #: (*untyped) -> String
+  def to_json(*)
+    Riffer::Agent::Serializer.to_json(agent: self)
+  end
+
   private
 
   #--
@@ -393,13 +449,13 @@ class Riffer::Agent
   end
 
   # Resolves +Config#model+ to a "provider/model" string (calling the Proc
-  # form against +@context+), parses it, and looks up the provider class.
+  # form against +@context+) and parses it.
   #
-  # Returns +[provider_class, model_name]+. Raises Riffer::ArgumentError on
-  # an invalid model string or an unregistered provider.
+  # Returns +[provider_name, model_name]+. Raises Riffer::ArgumentError on an
+  # invalid model string.
   #
   #--
-  #: () -> [singleton(Riffer::Providers::Base), String]
+  #: () -> [String, String]
   def resolve_provider_and_model
     model_string = Riffer::Helpers::CallOrValue.resolve(@config.model, context: @context)
     raise Riffer::ArgumentError, "Invalid model string: #{model_string}" unless model_string.is_a?(String)
@@ -410,18 +466,28 @@ class Riffer::Agent
       raise Riffer::ArgumentError, "Invalid model string: #{model_string}"
     end
 
-    provider_class = Riffer::Providers::Repository.find(provider_name)
-    raise Riffer::ArgumentError, "Provider not found: #{provider_name}" unless provider_class
+    [provider_name, model_name]
+  end
 
-    [provider_class, model_name]
+  # Builds the provider client from the resolved +@provider_name+ and the
+  # configured +provider_options+.
+  #
+  # Raises Riffer::ArgumentError on an unregistered provider.
+  #
+  #--
+  #: () -> Riffer::Providers::Base
+  def build_provider
+    provider_class = Riffer::Providers::Repository.find(@provider_name)
+    raise Riffer::ArgumentError, "Provider not found: #{@provider_name}" unless provider_class
+    provider_class.new(**@config.provider_options)
   end
 
   # Resolves the skills backend, lists skills, and selects an adapter.
   # Returns nil if skills are unconfigured or the backend is empty.
   #
   #--
-  #: (singleton(Riffer::Providers::Base)) -> Riffer::Skills::Context?
-  def resolve_skills(provider_class)
+  #: () -> Riffer::Skills::Context?
+  def resolve_skills
     skills_config = @config.skills_config
     return nil unless skills_config
 
@@ -432,7 +498,7 @@ class Riffer::Agent
     return nil if backend.list_skills.empty?
 
     skills = backend.list_skills.to_h { |s| [s.name, s] }
-    adapter_class = skills_config.adapter || provider_class.skills_adapter(@model_name)
+    adapter_class = skills_config.adapter || @provider.class.skills_adapter(@model_name)
     skill_activate_tool_class = skills_config.activate_tool || Riffer.config.skills.default_activate_tool
 
     skills_context = Riffer::Skills::Context.new(

data/lib/riffer/evals/evaluator.rb

@@ -16,6 +16,11 @@
 #   end
 #
 class Riffer::Evals::Evaluator
+  # @rbs self.@instructions: String?
+  # @rbs self.@higher_is_better: bool?
+  # @rbs self.@judge_model: String?
+  # @rbs @judge: Riffer::Evals::Judge?
+
   class << self
     # Gets or sets the evaluation instructions (criteria and scoring rubric).
     #

data/lib/riffer/evals/judge.rb

@@ -19,6 +19,11 @@ require "json"
 #   result[:reason] # => "The response is relevant..."
 #
 class Riffer::Evals::Judge
+  # @rbs @provider_options: Hash[Symbol, untyped]
+  # @rbs @provider_instance: Riffer::Providers::Base?
+  # @rbs @provider_name: String?
+  # @rbs @model_name: String?
+
   # Internal tool for structured evaluation output.
   class EvaluationTool < Riffer::Tool
     identifier "evaluation"

data/lib/riffer/mcp/client.rb

@@ -16,6 +16,8 @@
 class Riffer::Mcp::Client
   include Riffer::Helpers::Dependencies
 
+  # @rbs @client: untyped
+
   #--
   #: (endpoint: String, ?headers: (Hash[String, String] | Proc), ?client: untyped?) -> void
   def initialize(endpoint:, headers: {}, client: nil)

data/lib/riffer/mcp/registration.rb

@@ -7,6 +7,10 @@
 # +tools/list+ call, then generates tool classes.
 #
 class Riffer::Mcp::Registration
+  # @rbs @cancelled: bool
+  # @rbs @tools: Array[singleton(Riffer::Tool)]
+  # @rbs @mutex: Thread::Mutex
+
   # The manifest that describes this server.
   attr_reader :manifest #: Riffer::Mcp::Manifest
 

data/lib/riffer/mcp/registry.rb

@@ -6,6 +6,9 @@
 # Keyed by manifest name. All public methods are mutex-guarded.
 #
 module Riffer::Mcp::Registry
+  # @rbs self.@mutex: Thread::Mutex
+  # @rbs self.@store: Hash[String, Riffer::Mcp::Registration]
+
   @mutex = Mutex.new
   @store = {} #: Hash[String, Riffer::Mcp::Registration]
 

data/lib/riffer/messages/file_part.rb

@@ -15,6 +15,8 @@ require "uri"
 #   file.document?   # => true
 #
 class Riffer::Messages::FilePart
+  # @rbs @url_string: String?
+
   MEDIA_TYPES = {
     ".jpg" => "image/jpeg",
     ".jpeg" => "image/jpeg",

data/lib/riffer/params/param.rb

@@ -18,19 +18,90 @@ class Riffer::Params::Param
   }.freeze #: Hash[Module, String]
 
   # Primitive types allowed for the <tt>of:</tt> keyword on Array params
-  PRIMITIVE_TYPES = (TYPE_MAPPINGS.keys - [Array, Hash]).freeze #: Array[Class]
+  PRIMITIVE_TYPES = (TYPE_MAPPINGS.keys - [Array, Hash]).freeze #: Array[Module]
+
+  # Maps JSON Schema type strings back to Ruby types. The inverse of
+  # TYPE_MAPPINGS, collapsing the three boolean spellings onto
+  # Riffer::Params::Boolean. Used by +from_json_schema+.
+  JSON_TYPE_MAPPINGS = {
+    "string" => String,
+    "integer" => Integer,
+    "number" => Float,
+    "boolean" => Riffer::Params::Boolean,
+    "array" => Array,
+    "object" => Hash
+  }.freeze #: Hash[String, Module]
 
   attr_reader :name #: Symbol
-  attr_reader :type #: Class
+  attr_reader :type #: Module
   attr_reader :required #: bool
   attr_reader :description #: String?
   attr_reader :enum #: Array[untyped]?
   attr_reader :default #: untyped
-  attr_reader :item_type #: Class?
+  attr_reader :item_type #: Module?
   attr_reader :nested_params #: Riffer::Params?
 
   #--
-  #: (name: Symbol, type: Class, required: bool, ?description: String?, ?enum: Array[untyped]?, ?default: untyped, ?item_type: Class?, ?nested_params: Riffer::Params?) -> void
+  # Reconstructs a Param from a single JSON Schema property.
+  #
+  # [name] the parameter name (Symbol).
+  # [schema] the property's JSON Schema (Symbol-keyed).
+  # [required] whether the property appeared in the parent's +required+ list.
+  #
+  # Raises Riffer::ArgumentError on a type outside the Params-expressible subset.
+  #
+  #--
+  #: (Symbol, Hash[Symbol, untyped], required: bool) -> Riffer::Params::Param
+  def self.from_json_schema(name, schema, required:)
+    ruby_type = json_type_to_ruby(schema[:type])
+    item_type, nested = resolve_nesting(ruby_type, schema)
+
+    new(
+      name: name,
+      type: ruby_type,
+      required: required,
+      description: schema[:description],
+      enum: schema[:enum],
+      default: schema[:default],
+      item_type: item_type,
+      nested_params: nested
+    )
+  end
+
+  # Resolves the +[item_type, nested_params]+ pair for a reconstructed Param:
+  # a nested Params for object / array-of-object schemas, an +item_type+ for
+  # typed primitive arrays, and +nil+ for everything else.
+  #
+  #--
+  #: (Module, Hash[Symbol, untyped]) -> [Module?, Riffer::Params?]
+  def self.resolve_nesting(ruby_type, schema)
+    return [nil, Riffer::Params.from_json_schema(schema)] if ruby_type == Hash && schema[:properties]
+    return [nil, nil] unless ruby_type == Array
+
+    items = schema[:items]
+    return [nil, nil] unless items.is_a?(Hash)
+    return [nil, Riffer::Params.from_json_schema(items)] if items[:properties]
+
+    [json_type_to_ruby(items[:type]), nil]
+  end
+  private_class_method :resolve_nesting
+
+  # Resolves a JSON Schema +type+ (a String, or a <tt>[type, "null"]</tt>
+  # union) back to its Ruby type. Returns a Module because
+  # Riffer::Params::Boolean is a Module, not a Class — the same widening the
+  # +type+ attribute uses. Raises Riffer::ArgumentError on a type outside the
+  # Params-expressible subset (the block runs only for an unmapped type).
+  #
+  #--
+  #: (untyped) -> Module
+  def self.json_type_to_ruby(type)
+    key = type.is_a?(Array) ? type.find { |t| t != "null" } : type
+    JSON_TYPE_MAPPINGS.fetch(key) { raise Riffer::ArgumentError, "Unsupported JSON Schema type: #{type.inspect}" }
+  end
+  private_class_method :json_type_to_ruby
+
+  #--
+  #: (name: Symbol, type: Module, required: bool, ?description: String?, ?enum: Array[untyped]?, ?default: untyped, ?item_type: Module?, ?nested_params: Riffer::Params?) -> void
   def initialize(name:, type:, required:, description: nil, enum: nil, default: nil, item_type: nil, nested_params: nil)
     @name = name.to_sym
     @type = type
@@ -74,6 +145,11 @@ class Riffer::Params::Param
   # constraint from the null type, since providers like Anthropic reject
   # <tt>{"type": ["string", "null"], "enum": [...]}</tt>.
   #
+  # In non-strict mode a +default+ is emitted when set (a standard JSON
+  # Schema keyword), making the schema a lossless source for
+  # +Riffer::Params.from_json_schema+. Strict mode omits it, since strict
+  # providers reject the keyword.
+  #
   #--
   #: (?strict: bool) -> Hash[Symbol, untyped]
   def to_json_schema(strict: false)
@@ -91,6 +167,10 @@ class Riffer::Params::Param
     schema = {type: type} #: Hash[Symbol, untyped]
     schema[:description] = description if description
     schema[:enum] = enum if enum
+    # Strict providers reject the +default+ keyword; emit it only in
+    # non-strict mode, where it makes the schema a lossless round-trip
+    # source for +from_json_schema+.
+    schema[:default] = default unless strict || default.nil?
 
     if self.type == Array && nested_params
       schema[:items] = nested_params.to_json_schema(strict: strict)

data/lib/riffer/params.rb

@@ -20,10 +20,41 @@ class Riffer::Params
     @parameters = []
   end
 
+  # Reconstructs a Params from a JSON Schema object (the inverse of
+  # +to_json_schema(strict: false)+).
+  #
+  # Accepts the Symbol-keyed object schema produced by +to_json_schema+
+  # (property-name keys may be String or Symbol — both are normalized).
+  # Reconstructs types, +required+, +description+, +enum+, +default+,
+  # typed-array +item_type+, and nested object/array Params recursively.
+  #
+  # Round-trips losslessly with +to_json_schema(strict: false)+ over the
+  # Params-expressible subset of JSON Schema. Raises Riffer::ArgumentError
+  # on a schema using features outside that subset.
+  #
+  #   schema = params.to_json_schema(strict: false)
+  #   Riffer::Params.from_json_schema(schema) # => equivalent Riffer::Params
+  #
+  #--
+  #: (Hash[Symbol, untyped]) -> Riffer::Params
+  def self.from_json_schema(schema)
+    params = new
+    properties = schema[:properties] || {}
+    required = (schema[:required] || []).map { |key| key.to_s }
+
+    properties.each do |name, property_schema|
+      params.parameters << Riffer::Params::Param.from_json_schema(
+        name.to_sym, property_schema, required: required.include?(name.to_s)
+      )
+    end
+
+    params
+  end
+
   # Defines a required parameter.
   #
   #--
-  #: (Symbol, Class, ?description: String?, ?enum: Array[untyped]?, ?of: Class?) ?{ (Riffer::Params) [self: Riffer::Params] -> void } -> void
+  #: (Symbol, Module, ?description: String?, ?enum: Array[untyped]?, ?of: Module?) ?{ (Riffer::Params) [self: Riffer::Params] -> void } -> void
   def required(name, type, description: nil, enum: nil, of: nil, &block)
     nested = build_nested(type, of, &block)
     @parameters << Riffer::Params::Param.new(
@@ -40,7 +71,7 @@ class Riffer::Params
   # Defines an optional parameter.
   #
   #--
-  #: (Symbol, Class, ?description: String?, ?enum: Array[untyped]?, ?default: untyped, ?of: Class?) ?{ (Riffer::Params) [self: Riffer::Params] -> void } -> void
+  #: (Symbol, Module, ?description: String?, ?enum: Array[untyped]?, ?default: untyped, ?of: Module?) ?{ (Riffer::Params) [self: Riffer::Params] -> void } -> void
   def optional(name, type, description: nil, enum: nil, default: nil, of: nil, &block)
     nested = build_nested(type, of, &block)
     @parameters << Riffer::Params::Param.new(
@@ -126,7 +157,7 @@ class Riffer::Params
   private
 
   #--
-  #: (Class, Class?) ?{ (Riffer::Params) [self: Riffer::Params] -> void } -> Riffer::Params?
+  #: (Module, Module?) ?{ (Riffer::Params) [self: Riffer::Params] -> void } -> Riffer::Params?
   def build_nested(type, of, &block)
     if of && block
       raise Riffer::ArgumentError, "cannot use both of: and a block"

data/lib/riffer/providers/amazon_bedrock.rb

@@ -92,15 +92,16 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Hash[Symbol, untyped]) -> Aws::BedrockRuntime::Client::_ConverseResponseSuccess
+  #: (Hash[Symbol, untyped]) -> untyped
   def execute_generate(params)
     @client.converse(**params)
   end
 
   #--
-  #: (Aws::BedrockRuntime::Client::_ConverseResponseSuccess) -> Riffer::Providers::TokenUsage?
+  #: (untyped) -> Riffer::Providers::TokenUsage?
   def extract_token_usage(response)
-    usage = response.usage
+    typed_response = response #: Aws::BedrockRuntime::Client::_ConverseResponseSuccess
+    usage = typed_response.usage
 
     Riffer::Providers::TokenUsage.new(
       input_tokens: usage.input_tokens,
@@ -111,9 +112,10 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Client::_ConverseResponseSuccess) -> String
+  #: (untyped) -> String
   def extract_content(response)
-    content_blocks = response.output&.message&.content
+    typed_response = response #: Aws::BedrockRuntime::Client::_ConverseResponseSuccess
+    content_blocks = typed_response.output&.message&.content
     return "" if content_blocks.nil? || content_blocks.empty?
 
     text_content = ""
@@ -126,9 +128,10 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Client::_ConverseResponseSuccess) -> Array[Riffer::Messages::Assistant::ToolCall]
+  #: (untyped) -> Array[Riffer::Messages::Assistant::ToolCall]
   def extract_tool_calls(response)
-    content_blocks = response.output&.message&.content
+    typed_response = response #: Aws::BedrockRuntime::Client::_ConverseResponseSuccess
+    content_blocks = typed_response.output&.message&.content
     return [] if content_blocks.nil? || content_blocks.empty?
 
     tool_calls = [] #: Array[Riffer::Messages::Assistant::ToolCall]
@@ -195,28 +198,31 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockStartEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
+  #: (untyped, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_content_block_start_tool_use(event, state:, yielder:)
+    typed_event = event #: Aws::BedrockRuntime::Types::ContentBlockStartEvent
     state[:tool_call] = {
-      id: event.start.tool_use.tool_use_id,
-      name: decode_tool_name(event.start.tool_use.name, tools: @current_tools),
+      id: typed_event.start.tool_use.tool_use_id,
+      name: decode_tool_name(typed_event.start.tool_use.name, tools: @current_tools),
       arguments: ""
     }
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockDeltaEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
+  #: (untyped, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_content_block_delta_text_delta(event, state:, yielder:)
-    delta_text = event.delta.text
+    typed_event = event #: Aws::BedrockRuntime::Types::ContentBlockDeltaEvent
+    delta_text = typed_event.delta.text
     state[:text] ||= ""
     state[:text] += delta_text
     yielder << Riffer::StreamEvents::TextDelta.new(delta_text)
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockDeltaEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
+  #: (untyped, 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
+    typed_event = event #: Aws::BedrockRuntime::Types::ContentBlockDeltaEvent
+    input_delta = typed_event.delta.tool_use.input
 
     state[:tool_call][:arguments] += input_delta
 
@@ -228,14 +234,14 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ContentBlockStopEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
+  #: (untyped, 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::Yielder) -> void
+  #: (untyped, 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(
@@ -248,14 +254,15 @@ class Riffer::Providers::AmazonBedrock < Riffer::Providers::Base
   end
 
   #--
-  #: (Aws::BedrockRuntime::Types::ConverseStreamMetadataEvent, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
+  #: (untyped, state: Hash[Symbol, untyped], yielder: Enumerator::Yielder) -> void
   def handle_metadata_usage(event, state:, yielder:)
+    typed_event = event #: Aws::BedrockRuntime::Types::ConverseStreamMetadataEvent
     yielder << Riffer::StreamEvents::TokenUsageDone.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,
-        cache_read_tokens: event.usage.cache_read_input_tokens
+        input_tokens: typed_event.usage.input_tokens,
+        output_tokens: typed_event.usage.output_tokens,
+        cache_creation_tokens: typed_event.usage.cache_write_input_tokens,
+        cache_read_tokens: typed_event.usage.cache_read_input_tokens
       )
     )
   end

data/lib/riffer/providers/anthropic.rb

@@ -77,15 +77,16 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   end
 
   #--
-  #: (Hash[Symbol, untyped]) -> Anthropic::Models::Message
+  #: (Hash[Symbol, untyped]) -> untyped
   def execute_generate(params)
     @client.messages.create(**params)
   end
 
   #--
-  #: (Anthropic::Models::Message) -> Riffer::Providers::TokenUsage?
+  #: (untyped) -> Riffer::Providers::TokenUsage?
   def extract_token_usage(response)
-    usage = response.usage
+    message = response #: Anthropic::Models::Message
+    usage = message.usage
 
     Riffer::Providers::TokenUsage.new(
       input_tokens: usage.input_tokens,
@@ -96,9 +97,10 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   end
 
   #--
-  #: (Anthropic::Models::Message) -> String
+  #: (untyped) -> String
   def extract_content(response)
-    content_blocks = response.content
+    message = response #: Anthropic::Models::Message
+    content_blocks = message.content
     return "" if content_blocks.nil? || content_blocks.empty?
 
     text_content = ""
@@ -111,9 +113,10 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   end
 
   #--
-  #: (Anthropic::Models::Message) -> Array[Riffer::Messages::Assistant::ToolCall]
+  #: (untyped) -> Array[Riffer::Messages::Assistant::ToolCall]
   def extract_tool_calls(response)
-    content_blocks = response.content
+    message = response #: Anthropic::Models::Message
+    content_blocks = message.content
     return [] if content_blocks.nil? || content_blocks.empty?
 
     tool_calls = [] #: Array[Riffer::Messages::Assistant::ToolCall]
@@ -285,9 +288,10 @@ class Riffer::Providers::Anthropic < Riffer::Providers::Base
   end
 
   #--
-  #: (untyped, accumulated_message: Anthropic::Models::Message?, yielder: Enumerator::Yielder) -> void
+  #: (untyped, accumulated_message: untyped, yielder: Enumerator::Yielder) -> void
   def handle_message_stop(_event, accumulated_message:, yielder:)
-    usage = accumulated_message&.usage
+    message = accumulated_message #: Anthropic::Models::Message?
+    usage = message&.usage
     return unless usage
 
     yielder << Riffer::StreamEvents::TokenUsageDone.new(

data/lib/riffer/providers/base.rb

@@ -17,6 +17,8 @@ require "json"
 # [extract_content] extract text content from the SDK response
 # [extract_tool_calls] extract tool calls from the SDK response
 class Riffer::Providers::Base
+  # @rbs @current_tools: Array[singleton(Riffer::Tool)]
+
   include Riffer::Helpers::Dependencies
   include Riffer::Messages::Converter
 

data/lib/riffer/providers/gemini.rb

@@ -8,6 +8,10 @@ require "uri"
 
 # Google Gemini provider for Gemini models via the Gemini REST API.
 class Riffer::Providers::Gemini < Riffer::Providers::Base
+  # @rbs @api_key: String?
+  # @rbs @open_timeout: Integer
+  # @rbs @read_timeout: Integer
+
   BASE_URI = URI("https://generativelanguage.googleapis.com") #: URI::Generic
   VALID_MODEL_PATTERN = /\A[a-zA-Z0-9._-]+\z/ #: Regexp
   DEFAULT_OPEN_TIMEOUT = 10 #: Integer

data/lib/riffer/providers/mock.rb

@@ -5,6 +5,10 @@
 #
 # No external gems required.
 class Riffer::Providers::Mock < Riffer::Providers::Base
+  # @rbs @responses: Array[Hash[Symbol, untyped]]
+  # @rbs @current_index: Integer
+  # @rbs @stubbed_responses: Array[Hash[Symbol, untyped]]
+
   # Returns the preferred skill adapter for the given mock model.
   #
   # Mock is used to stand in for any real provider in tests, so the model