lib/gen_ai/chat/base.rb in gen-ai-0.4.1 vs lib/gen_ai/chat/base.rb in gen-ai-0.4.2

@@ -16,17 +16,17 @@
 
       def start(history: [], context: nil, examples: [])
         @history = build_history(history.map(&:deep_symbolize_keys), context, examples.map(&:deep_symbolize_keys))
       end
 
-      def message(message, options = {})
+      def message(message, options = {}, &block)
         if @history.size == 1 && @history.first[:role] != 'system'
           append_to_message(message)
         else
           append_to_history({ role: USER_ROLE, content: message })
         end
 
-        response = @model.chat(@history.dup, default_options.merge(options).compact)
+        response = @model.chat(@history.dup, default_options.merge(options).compact, &block)
         append_to_history({ role: ASSISTANT_ROLE, content: response.value })
         response
       end
 
       private