From 8c95869a530bc3fe48d1c79952b8e3c68c527aaf Mon Sep 17 00:00:00 2001
From: Kevin Sylvestre <kevin@ksylvest.com>
Date: Thu, 18 Jul 2024 13:24:00 -0700
Subject: [PATCH] Fix missing block in chat API

---
 Gemfile.lock          | 2 +-
 lib/omniai/chat.rb    | 4 ++++
 lib/omniai/client.rb  | 5 ++++-
 lib/omniai/version.rb | 2 +-
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/Gemfile.lock b/Gemfile.lock
index dbb87db..e7e0cbc 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    omniai (1.6.0)
+    omniai (1.6.1)
       event_stream_parser
       http
       zeitwerk
diff --git a/lib/omniai/chat.rb b/lib/omniai/chat.rb
index beb85ad..4ab8b45 100644
--- a/lib/omniai/chat.rb
+++ b/lib/omniai/chat.rb
@@ -57,7 +57,11 @@ def self.process!(...)
     # @param stream [Proc, IO, nil] optional
     # @param tools [Array<OmniAI::Tool>] optional
     # @param format [Symbol, nil] optional - :json
+    #
     # @yield [prompt] optional
+    # @yieldparam prompt [OmniAI::Chat::Prompt]
+    #
+    # @return [OmniAi::Chat]
     def initialize(prompt = nil, client:, model:, temperature: nil, stream: nil, tools: nil, format: nil, &block)
       raise ArgumentError, 'prompt or block is required' if !prompt && !block
 
diff --git a/lib/omniai/client.rb b/lib/omniai/client.rb
index fafe85a..22867c9 100644
--- a/lib/omniai/client.rb
+++ b/lib/omniai/client.rb
@@ -131,8 +131,11 @@ def connection
     # @param stream [Proc, nil] optional
     # @param tools [Array<OmniAI::Tool>] optional
     #
+    # @yield [prompt] optional
+    # @yieldparam prompt [OmniAI::Chat::Prompt]
+    #
     # @return [OmniAI::Chat::Completion]
-    def chat(messages, model:, temperature: nil, format: nil, stream: nil, tools: nil)
+    def chat(messages, model:, temperature: nil, format: nil, stream: nil, tools: nil, &)
       raise NotImplementedError, "#{self.class.name}#chat undefined"
     end
 
diff --git a/lib/omniai/version.rb b/lib/omniai/version.rb
index 89f71d5..02732ab 100644
--- a/lib/omniai/version.rb
+++ b/lib/omniai/version.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module OmniAI
-  VERSION = '1.6.0'
+  VERSION = '1.6.1'
 end