v0.4.1: Integrate WebSocket with standard chat.ask interface

Add transport: :websocket support via with_params so chat.ask works
over WebSocket without API changes. Add WebSocket#call for pre-built
payloads. Fix token count preservation in WebSocket responses.

Author: khasinski

CHANGELOG.md

@@ -5,6 +5,17 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.4.1] - 2026-02-24
+
+### Added
+
+- `chat.with_params(transport: :websocket)` integration with standard `chat.ask` interface
+- `WebSocket#call` for accepting pre-built payloads from the provider
+
+### Fixed
+
+- WebSocket responses now preserve token counts from `StreamAccumulator`
+
 ## [0.4.0] - 2026-02-24
 
 ### Added

README.md

@@ -269,59 +269,40 @@ Requires the `websocket-client-simple` gem:
 gem 'websocket-client-simple'
 ```
 
-### Basic usage
+### Usage
 
-```ruby
-ws = RubyLLM::ResponsesAPI::WebSocket.new(api_key: ENV['OPENAI_API_KEY'])
-ws.connect
+Just add `transport: :websocket` to your params -- the standard `chat.ask` API works as-is:
 
-# Stream a response
-message = ws.create_response(
-  model: 'gpt-4o',
-  input: [{ type: 'message', role: 'user', content: 'Hello!' }]
-) do |chunk|
-  print chunk.content if chunk.content
-end
+```ruby
+chat = RubyLLM.chat(model: 'gpt-4o', provider: :openai_responses)
+chat.with_params(transport: :websocket)
 
-puts "\n#{message.content}"
+chat.ask("Hello!")
+chat.ask("What's 2+2?")  # reuses the same WebSocket connection
 ```
 
-### Multi-turn conversations
-
-`previous_response_id` is tracked automatically across turns:
+Streaming works the same way:
 
 ```ruby
-ws.create_response(model: 'gpt-4o', input: [
-  { type: 'message', role: 'user', content: 'My name is Alice.' }
-])
-
-ws.create_response(model: 'gpt-4o', input: [
-  { type: 'message', role: 'user', content: "What's my name?" }
-])
-# => "Alice" (auto-chained via previous_response_id)
+chat.ask("Tell me a story") { |chunk| print chunk.content }
 ```
 
-### With tools
+### Direct WebSocket access
+
+For advanced use cases (raw Responses API format, warmup, explicit connection management):
 
 ```ruby
+ws = RubyLLM::ResponsesAPI::WebSocket.new(api_key: ENV['OPENAI_API_KEY'])
+ws.connect
+
 ws.create_response(
   model: 'gpt-4o',
-  input: [{ type: 'message', role: 'user', content: 'Search for Ruby 3.4 release notes' }],
-  tools: [{ type: 'web_search_preview' }]
-)
-```
-
-### Warmup
-
-Pre-cache model weights without generating output:
+  input: [{ type: 'message', role: 'user', content: 'Hello!' }]
+) { |chunk| print chunk.content }
 
-```ruby
+# Pre-cache model weights
 ws.warmup(model: 'gpt-4o')
-```
-
-### Cleanup
 
-```ruby
 ws.disconnect
 ```
 

lib/ruby_llm/providers/openai_responses.rb

@@ -16,6 +16,16 @@ def api_base
         @config.openai_api_base || 'https://api.openai.com/v1'
       end
 
+      # Override to support WebSocket transport via with_params(transport: :websocket)
+      def complete(messages, tools:, temperature:, model:, params: {}, headers: {}, schema: nil, thinking: nil, &block) # rubocop:disable Metrics/ParameterLists
+        if params[:transport]&.to_sym == :websocket
+          ws_complete(messages, tools: tools, temperature: temperature, model: model,
+                      params: params.except(:transport), schema: schema, thinking: thinking, &block)
+        else
+          super
+        end
+      end
+
       def headers
         {
           'Authorization' => "Bearer #{@config.openai_api_key}",
@@ -137,6 +147,35 @@ def retrieve_container_file_content(container_id, file_id)
 
       private
 
+      def ws_complete(messages, tools:, temperature:, model:, params:, schema:, thinking:, &block)
+        normalized_temperature = maybe_normalize_temperature(temperature, model)
+
+        payload = Utils.deep_merge(
+          render_payload(
+            messages,
+            tools: tools,
+            temperature: normalized_temperature,
+            model: model,
+            stream: true,
+            schema: schema,
+            thinking: thinking
+          ),
+          params
+        )
+
+        ws_connection.connect unless ws_connection.connected?
+        ws_connection.call(payload, &block)
+      end
+
+      def ws_connection
+        @ws_connection ||= WebSocket.new(
+          api_key: @config.openai_api_key,
+          api_base: api_base,
+          organization_id: @config.openai_organization_id,
+          project_id: @config.openai_project_id
+        )
+      end
+
       # DELETE request via the underlying Faraday connection
       # RubyLLM::Connection only exposes get/post, so we use Faraday directly
       def delete_request(url)

lib/ruby_llm/providers/openai_responses/web_socket.rb

@@ -11,14 +11,15 @@ class OpenAIResponses
       #
       # Requires the `websocket-client-simple` gem (soft dependency).
       #
-      # Usage:
+      # Integrated usage (recommended):
+      #   chat = RubyLLM.chat(model: 'gpt-4o', provider: :openai_responses)
+      #   chat.with_params(transport: :websocket)
+      #   chat.ask("Hello!")
+      #
+      # Standalone usage (advanced):
       #   ws = RubyLLM::ResponsesAPI::WebSocket.new(api_key: ENV['OPENAI_API_KEY'])
       #   ws.connect
-      #
-      #   ws.create_response(model: 'gpt-4o', input: [{ type: 'message', role: 'user', content: 'Hi' }]) do |chunk|
-      #     print chunk.content if chunk.content
-      #   end
-      #
+      #   ws.create_response(model: 'gpt-4o', input: [...]) { |chunk| ... }
       #   ws.disconnect
       class WebSocket
         WEBSOCKET_PATH = '/v1/responses'
@@ -73,7 +74,6 @@ def connect(timeout: 10)
             end
           end
 
-          # Route all messages to the current queue (swapped per request)
           @ws.on(:message) do |msg|
             q = @mutex.synchronize { @message_queue }
             q&.push(msg.data)
@@ -89,35 +89,47 @@ def connect(timeout: 10)
           self
         end
 
-        # Send a response.create request and stream chunks via block.
-        # @param model [String] model ID
-        # @param input [Array<Hash>] input items in Responses API format
-        # @param tools [Array<Hash>, nil] tool definitions
-        # @param previous_response_id [String, nil] chain to a prior response
-        # @param instructions [String, nil] system/developer instructions
-        # @param extra [Hash] additional top-level fields forwarded to the API
+        # Send a pre-built payload over WebSocket, streaming chunks via block.
+        # This is the integration point for Provider#complete -- it accepts the
+        # same payload hash that render_payload returns.
+        #
+        # @param payload [Hash] Responses API payload (model, input, tools, etc.)
         # @yield [RubyLLM::Chunk] each streamed chunk
         # @return [RubyLLM::Message] the assembled final message
-        # @raise [ConcurrencyError] if another response is already in flight
-        # @raise [ConnectionError] if not connected
-        def create_response(model:, input:, tools: nil, previous_response_id: nil, instructions: nil, **extra, &block)
+        def call(payload, &block)
           ensure_connected!
           acquire_flight!
 
           queue = Queue.new
           @mutex.synchronize { @message_queue = queue }
 
-          payload = build_payload(
+          envelope = { type: 'response.create', response: payload.except(:stream) }
+          send_json(envelope)
+          accumulate_response(queue, &block)
+        ensure
+          @mutex.synchronize { @message_queue = nil }
+          release_flight!
+        end
+
+        # Send a response.create request using raw Responses API format.
+        # Useful for standalone usage outside the RubyLLM chat interface.
+        #
+        # @param model [String] model ID
+        # @param input [Array<Hash>] input items in Responses API format
+        # @param tools [Array<Hash>, nil] tool definitions
+        # @param previous_response_id [String, nil] chain to a prior response
+        # @param instructions [String, nil] system/developer instructions
+        # @param extra [Hash] additional fields forwarded to the API
+        # @yield [RubyLLM::Chunk] each streamed chunk
+        # @return [RubyLLM::Message] the assembled final message
+        def create_response(model:, input:, tools: nil, previous_response_id: nil, instructions: nil, **extra, &block)
+          payload = build_standalone_payload(
             model: model, input: input, tools: tools,
             previous_response_id: previous_response_id,
             instructions: instructions, **extra
           )
 
-          send_json(payload)
-          accumulate_response(queue, &block)
-        ensure
-          @mutex.synchronize { @message_queue = nil }
-          release_flight!
+          call(payload, &block)
         end
 
         # Warm up the connection by sending a response.create with generate: false.
@@ -209,7 +221,7 @@ def build_headers
           headers
         end
 
-        def build_payload(model:, input:, tools: nil, previous_response_id: nil, instructions: nil, **extra)
+        def build_standalone_payload(model:, input:, tools: nil, previous_response_id: nil, instructions: nil, **extra)
           prev_id = previous_response_id || @last_response_id
           response = { model: model, input: input }
           response[:tools] = tools.map { |t| Tools.tool_for(t) } if tools&.any?
@@ -220,7 +232,7 @@ def build_payload(model:, input:, tools: nil, previous_response_id: nil, instruc
           Compaction.apply_compaction(response, extra)
 
           forwarded = extra.reject { |k, _| KNOWN_PARAMS.include?(k) }
-          { type: 'response.create', response: response.merge(forwarded) }
+          response.merge(forwarded)
         end
 
         def send_json(payload)
@@ -247,24 +259,16 @@ def accumulate_response(queue, &block)
             end
           end
 
-          build_final_message(accumulator)
+          message = accumulator.to_message(nil)
+          message.response_id = @last_response_id
+          message
         end
 
         def track_response_id(data)
           resp_id = data.dig('response', 'id')
           @mutex.synchronize { @last_response_id = resp_id } if resp_id
         end
 
-        def build_final_message(accumulator)
-          Message.new(
-            role: :assistant,
-            content: accumulator.content,
-            tool_calls: accumulator.tool_calls.empty? ? nil : accumulator.tool_calls,
-            model_id: accumulator.model_id,
-            response_id: @last_response_id
-          )
-        end
-
         def ensure_connected!
           raise ConnectionError, 'WebSocket is not connected. Call #connect first.' unless connected?
         end

lib/rubyllm_responses_api.rb

@@ -37,7 +37,7 @@
 module RubyLLM
   # ResponsesAPI namespace for direct access to helpers and version
   module ResponsesAPI
-    VERSION = '0.4.0'
+    VERSION = '0.4.1'
 
     # Shorthand access to built-in tool helpers
     BuiltInTools = Providers::OpenAIResponses::BuiltInTools

ruby_llm-responses_api.gemspec

@@ -2,7 +2,7 @@
 
 Gem::Specification.new do |spec|
   spec.name = 'ruby_llm-responses_api'
-  spec.version = '0.4.0'
+  spec.version = '0.4.1'
   spec.authors = ['Chris Hasinski']
   spec.email = ['krzysztof.hasinski@gmail.com']
 

spec/ruby_llm/providers/openai_responses/web_socket_spec.rb

@@ -279,6 +279,46 @@ def last_sent_payload
     end
   end
 
+  describe '#call' do
+    before { connect_ws! }
+
+    it 'accepts a pre-built payload and wraps in WS envelope' do
+      payload = { model: 'gpt-4o', input: [{ type: 'message', role: 'user', content: 'Hi' }] }
+
+      with_events(standard_events) do
+        ws.call(payload)
+      end
+
+      sent = last_sent_payload
+      expect(sent['type']).to eq('response.create')
+      expect(sent['response']['model']).to eq('gpt-4o')
+    end
+
+    it 'strips :stream from the payload' do
+      payload = { model: 'gpt-4o', input: [], stream: true }
+
+      with_events(standard_events) do
+        ws.call(payload)
+      end
+
+      expect(last_sent_payload['response']).not_to have_key('stream')
+    end
+
+    it 'yields chunks and returns assembled Message' do
+      payload = { model: 'gpt-4o', input: [] }
+      chunks = []
+
+      message = with_events(standard_events(text_deltas: %w[Hello world])) do
+        ws.call(payload) { |chunk| chunks << chunk }
+      end
+
+      expect(chunks.select(&:content).map(&:content)).to eq(%w[Hello world])
+      expect(message).to be_a(RubyLLM::Message)
+      expect(message.content).to eq('Helloworld')
+      expect(message.response_id).to eq('resp_ws_test')
+    end
+  end
+
   describe '#warmup' do
     before { connect_ws! }