boxcars 0.8.3 → 0.8.5

data/lib/boxcars/engine/groq.rb

@@ -19,10 +19,11 @@ module Boxcars
                           "You should ask targeted questions"
 
     def initialize(name: DEFAULT_NAME, description: DEFAULT_DESCRIPTION, prompts: [], batch_size: 20, **kwargs)
+      user_id = kwargs.delete(:user_id)
       @groq_params = DEFAULT_PARAMS.merge(kwargs) # Corrected typo here
       @prompts = prompts
       @batch_size = batch_size
-      super(description:, name:)
+      super(description:, name:, user_id:)
     end
 
     # Renamed from open_ai_client to groq_client for clarity
@@ -60,7 +61,8 @@ module Boxcars
         request_context = {
           prompt: current_prompt_object,
           inputs:,
-          conversation_for_api: api_request_params&.dig(:messages)
+          conversation_for_api: api_request_params&.dig(:messages),
+          user_id:
         }
         track_ai_generation(
           duration_ms:,
@@ -74,23 +76,23 @@ module Boxcars
       # If there's an error, raise it to maintain backward compatibility with existing tests
       raise response_data[:error] if response_data[:error]
 
-      response_data
+      response_data[:parsed_json]
     end
 
     def run(question, **)
       prompt = Prompt.new(template: question)
-      response_data = client(prompt:, inputs: {}, **)
-      answer = _groq_handle_call_outcome(response_data:)
+      response = client(prompt:, inputs: {}, **)
+      answer = extract_answer(response)
       Boxcars.debug("Answer: #{answer}", :cyan)
       answer
     end
 
+    private
+
     def default_params
-      @groq_params # Use instance variable
+      @groq_params
     end
 
-    private
-
     # Helper methods for the client method
     def _prepare_groq_request_params(prompt_object, inputs, current_params)
       messages_hash_from_prompt = prompt_object.as_messages(inputs)

data/lib/boxcars/engine/intelligence_base.rb

@@ -17,10 +17,11 @@ module Boxcars
     # @param batch_size [Integer] The number of prompts to send to the Engine at a time.
     # @param kwargs [Hash] Additional parameters to pass to the Engine.
     def initialize(provider:, description:, name:, prompts: [], batch_size: 20, **kwargs)
+      user_id = kwargs.delete(:user_id)
       @provider = provider
       # Start with defaults, merge other kwargs, then explicitly set model if provided in initialize
       @all_params = default_model_params.merge(kwargs)
-      super(description:, name:, prompts:, batch_size:)
+      super(description:, name:, prompts:, batch_size:, user_id:)
     end
 
     # can be overridden by provider subclass
@@ -68,7 +69,7 @@ module Boxcars
 
       adapter = adapter(api_key:, params:)
       convo = prompt.as_intelligence_conversation(inputs:)
-      request_context = { prompt: prompt&.as_prompt(inputs:)&.[](:prompt), inputs:, conversation_for_api: convo.to_h }
+      request_context = { user_id:, prompt: prompt&.as_prompt(inputs:)&.[](:prompt), inputs:, conversation_for_api: convo.to_h }
       request = Intelligence::ChatRequest.new(adapter:)
 
       start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
@@ -118,17 +119,6 @@ module Boxcars
 
     private
 
-    def extract_answer(response)
-      # Handle different response formats
-      if response["choices"]
-        response["choices"].map { |c| c.dig("message", "content") || c["text"] }.join("\n").strip
-      elsif response["candidates"]
-        response["candidates"].map { |c| c.dig("content", "parts", 0, "text") }.join("\n").strip
-      else
-        response["output"] || response.to_s
-      end
-    end
-
     def check_response(response)
       return if response.is_a?(Hash) && response.key?("choices")
 

data/lib/boxcars/engine/ollama.rb

@@ -19,10 +19,11 @@ module Boxcars
                           "You should ask targeted questions"
 
     def initialize(name: DEFAULT_NAME, description: DEFAULT_DESCRIPTION, prompts: [], batch_size: 2, **kwargs)
+      user_id = kwargs.delete(:user_id)
       @ollama_params = DEFAULT_PARAMS.merge(kwargs)
       @prompts = prompts
       @batch_size = batch_size # Retain if used by other methods
-      super(description:, name:)
+      super(description:, name:, user_id:)
     end
 
     # Renamed from open_ai_client to ollama_client for clarity
@@ -63,7 +64,8 @@ module Boxcars
         request_context = {
           prompt: current_prompt_object,
           inputs:,
-          conversation_for_api: api_request_params&.dig(:messages)
+          conversation_for_api: api_request_params&.dig(:messages),
+          user_id:
         }
         track_ai_generation(
           duration_ms:,

data/lib/boxcars/engine/openai.rb

@@ -1,244 +1,256 @@
 # frozen_string_literal: true
 
-require 'openai'
-require 'json'
-require 'securerandom'
+require "openai"
+require "json"
+require "securerandom"
 
 module Boxcars
-  # A engine that uses OpenAI's API.
+  # Engine that talks to OpenAI’s REST API.
   class Openai < Engine
     include UnifiedObservability
-    attr_reader :prompts, :open_ai_params, :model_kwargs, :batch_size
+
+    CHAT_MODEL_REGEX = /(^gpt-4)|(-turbo\b)|(^o\d)|(gpt-3\.5-turbo)/
+    O_SERIES_REGEX   = /^o/
 
     DEFAULT_PARAMS = {
       model: "gpt-4o-mini",
       temperature: 0.1,
       max_tokens: 4096
     }.freeze
-    DEFAULT_NAME = "OpenAI engine"
-    DEFAULT_DESCRIPTION = "useful for when you need to use AI to answer questions. " \
-                          "You should ask targeted questions"
-
-    def initialize(name: DEFAULT_NAME, description: DEFAULT_DESCRIPTION, prompts: [], batch_size: 20, **kwargs)
-      @open_ai_params = DEFAULT_PARAMS.merge(kwargs)
-      # Special handling for o1-mini model (deprecated?)
-      if @open_ai_params[:model] =~ /^o/ && @open_ai_params[:max_tokens]
-        @open_ai_params[:max_completion_tokens] = @open_ai_params.delete(:max_tokens)
-        @open_ai_params.delete(:temperature) # o1-mini might not support temperature
-      end
 
-      @prompts = prompts
-      @batch_size = batch_size
-      super(description:, name:)
+    DEFAULT_NAME        = "OpenAI engine"
+    DEFAULT_DESCRIPTION = "Useful when you need AI to answer questions. Ask targeted questions."
+
+    attr_reader :prompts, :open_ai_params, :batch_size
+
+    # --------------------------------------------------------------------------
+    #  Construction
+    # --------------------------------------------------------------------------
+    def initialize(name: DEFAULT_NAME,
+                   description: DEFAULT_DESCRIPTION,
+                   prompts: [],
+                   batch_size: 20,
+                   **kwargs)
+      user_id          = kwargs.delete(:user_id)
+      @open_ai_params  = adjust_for_o_series!(DEFAULT_PARAMS.merge(kwargs))
+      @prompts         = prompts
+      @batch_size      = batch_size
+      super(description:, name:, user_id:)
     end
 
-    def self.open_ai_client(openai_access_token: nil)
-      access_token = Boxcars.configuration.openai_access_token(openai_access_token:)
-      organization_id = Boxcars.configuration.organization_id
-      # log_errors is good for the gem's own logging
-      ::OpenAI::Client.new(access_token:, organization_id:, log_errors: true)
-    end
+    # --------------------------------------------------------------------------
+    #  Public API
+    # --------------------------------------------------------------------------
+    def client(prompt:, inputs: {}, openai_access_token: nil, **kwargs)
+      start_time       = Time.now
+      response_data    = { response_obj: nil, parsed_json: nil,
+                           success: false, error: nil, status_code: nil }
+      current_params   = open_ai_params.merge(kwargs)
+      is_chat_model    = chat_model?(current_params[:model])
+      prompt_object    = prompt.is_a?(Array) ? prompt.first : prompt
+      api_request      = build_api_request(prompt_object, inputs, current_params, chat: is_chat_model)
 
-    def conversation_model?(model_name)
-      !!(model_name =~ /(^gpt-4)|(-turbo\b)|(^o\d)|(gpt-3\.5-turbo)/) # Added gpt-3.5-turbo
+      begin
+        raw_response = execute_api_call(
+          self.class.open_ai_client(openai_access_token:),
+          is_chat_model,
+          api_request
+        )
+        process_response(raw_response, response_data)
+      rescue ::OpenAI::Error, StandardError => e
+        handle_error(e, response_data)
+      ensure
+        track_openai_observability(
+          {
+            start_time:,
+            prompt_object: prompt_object,
+            inputs: inputs,
+            api_request: api_request,
+            current_params: current_params,
+            is_chat_model: is_chat_model
+          },
+          response_data
+        )
+      end
+
+      handle_call_outcome(response_data:)
     end
 
-    def _prepare_openai_chat_request(prompt_object, inputs, current_params)
-      get_params(prompt_object, inputs, current_params.dup)
+    # Convenience one-shot helper used by Engine#generate
+    def run(question, **)
+      prompt      = Prompt.new(template: question)
+      raw_json    = client(prompt:, inputs: {}, **)
+      extract_answer_from_choices(raw_json["choices"]).tap do |ans|
+        Boxcars.debug("Answer: #{ans}", :cyan)
+      end
     end
 
-    def _prepare_openai_completion_request(prompt_object, inputs, current_params)
-      prompt_text_for_api = prompt_object.as_prompt(inputs:)
-      prompt_text_for_api = prompt_text_for_api[:prompt] if prompt_text_for_api.is_a?(Hash) && prompt_text_for_api.key?(:prompt)
-      { prompt: prompt_text_for_api }.merge(current_params).tap { |p| p.delete(:messages) }
+    # Expose the defaults so callers can introspect or dup/merge them
+    def default_params = open_ai_params
+
+    # --------------------------------------------------------------------------
+    #  Class helpers
+    # --------------------------------------------------------------------------
+    def self.open_ai_client(openai_access_token: nil)
+      ::OpenAI::Client.new(
+        access_token: Boxcars.configuration.openai_access_token(openai_access_token:),
+        organization_id: Boxcars.configuration.organization_id,
+        log_errors: true
+      )
     end
 
-    def _execute_openai_api_call(client, is_chat_model, api_request_params)
-      if is_chat_model
-        log_messages_debug(api_request_params[:messages]) if Boxcars.configuration.log_prompts && api_request_params[:messages]
-        client.chat(parameters: api_request_params)
-      else
-        Boxcars.debug("Prompt after formatting:\n#{api_request_params[:prompt]}", :cyan) if Boxcars.configuration.log_prompts
-        client.completions(parameters: api_request_params)
+    # -- Public helper -------------------------------------------------------------
+    # Some callers outside this class still invoke `check_response` directly.
+    # It simply raises if the JSON body contains an "error" payload.
+    def check_response(response) # rubocop:disable Naming/PredicateMethod
+      if (msg = openai_error_message(response))
+        raise Boxcars::Error, msg
       end
+
+      true
     end
 
-    def _process_openai_response(raw_response, response_data)
-      response_data[:response_obj] = raw_response
-      response_data[:parsed_json] = raw_response # Already parsed by OpenAI gem
+    private
 
-      if raw_response && !raw_response["error"]
-        response_data[:success] = true
-        response_data[:status_code] = 200 # Inferred
+    # -- Request construction ---------------------------------------------------
+    def build_api_request(prompt_object, inputs, params, chat:)
+      if chat
+        build_chat_params(prompt_object, inputs, params.dup)
       else
-        response_data[:success] = false
-        err_details = raw_response["error"] if raw_response
-        msg = err_details ? "#{err_details['type']}: #{err_details['message']}" : "Unknown OpenAI API Error"
-        response_data[:error] ||= StandardError.new(msg) # Use ||= to not overwrite existing exception
+        build_completion_params(prompt_object, inputs, params.dup)
       end
     end
 
-    def _handle_openai_api_error(error, response_data)
-      response_data[:error] = error
-      response_data[:success] = false
-      response_data[:status_code] = error.http_status if error.respond_to?(:http_status)
+    def build_chat_params(prompt_object, inputs, params)
+      po        = if prompt_object.is_a?(Boxcars::Prompt)
+                    prompt_object
+                  else
+                    Boxcars::Prompt.new(template: prompt_object.to_s)
+                  end
+      formatted = po.as_messages(inputs).merge(params)
+      adjust_for_o_series!(formatted)
     end
 
-    def _handle_openai_standard_error(error, response_data)
-      response_data[:error] = error
-      response_data[:success] = false
+    def build_completion_params(prompt_object, inputs, params)
+      prompt_txt = prompt_object.as_prompt(inputs:)
+      prompt_txt = prompt_txt[:prompt] if prompt_txt.is_a?(Hash) && prompt_txt.key?(:prompt)
+      { prompt: prompt_txt }.merge(params).tap { |h| h.delete(:messages) }
     end
 
-    def client(prompt:, inputs: {}, openai_access_token: nil, **kwargs)
-      start_time = Time.now
-      response_data = { response_obj: nil, parsed_json: nil, success: false, error: nil, status_code: nil }
-      current_params = open_ai_params.merge(kwargs)
-      current_prompt_object = prompt.is_a?(Array) ? prompt.first : prompt
-      api_request_params = nil
-      is_chat_model = conversation_model?(current_params[:model])
-
-      begin
-        clnt = Openai.open_ai_client(openai_access_token:)
-        api_request_params = if is_chat_model
-                               _prepare_openai_chat_request(current_prompt_object, inputs, current_params)
-                             else
-                               _prepare_openai_completion_request(current_prompt_object, inputs, current_params)
-                             end
-        raw_response = _execute_openai_api_call(clnt, is_chat_model, api_request_params)
-        _process_openai_response(raw_response, response_data)
-      rescue ::OpenAI::Error => e
-        _handle_openai_api_error(e, response_data)
-      rescue StandardError => e
-        _handle_openai_standard_error(e, response_data)
-      ensure
-        call_context = {
-          start_time:,
-          prompt_object: current_prompt_object,
-          inputs:,
-          api_request_params:,
-          current_params:,
-          is_chat_model:
-        }
-        _track_openai_observability(call_context, response_data)
+    # -- API call / response ----------------------------------------------------
+    def execute_api_call(client, chat_mode, api_request)
+      if chat_mode
+        log_messages_debug(api_request[:messages]) if Boxcars.configuration.log_prompts
+        client.chat(parameters: api_request)
+      else
+        Boxcars.debug("Prompt after formatting:\n#{api_request[:prompt]}", :cyan) if Boxcars.configuration.log_prompts
+        client.completions(parameters: api_request)
       end
-
-      _openai_handle_call_outcome(response_data:)
     end
 
-    # Called by Engine#generate to check the response from the client.
-    # @param response [Hash] The parsed JSON response from the OpenAI API.
-    # @raise [Boxcars::Error] if the response contains an error.
-    # rubocop:disable Naming/PredicateMethod
-    def check_response(response)
-      if response.is_a?(Hash) && response["error"]
-        err_details = response["error"]
-        msg = err_details ? "#{err_details['type']}: #{err_details['message']}" : "Unknown OpenAI API Error in check_response"
-        raise Boxcars::Error, msg
+    def process_response(raw, data)
+      data[:response_obj] = raw
+      data[:parsed_json]  = raw
+
+      if (msg = openai_error_message(raw))
+        data[:success]     = false
+        data[:status_code] = raw&.dig("error", "code") || 500
+        data[:error]       = StandardError.new(msg)
+      else
+        data[:success]     = true
+        data[:status_code] = 200
       end
-      true
     end
-    # rubocop:enable Naming/PredicateMethod
 
-    def run(question, **)
-      prompt = Prompt.new(template: question)
-      # client now returns the raw JSON response. We need to extract the answer.
-      raw_response = client(prompt:, inputs: {}, **)
-      answer = _extract_openai_answer_from_choices(raw_response["choices"])
-      Boxcars.debug("Answer: #{answer}", :cyan)
-      answer
+    def handle_error(error, data)
+      data[:error]        = error
+      data[:success]      = false
+      data[:status_code]  = error.respond_to?(:http_status) ? error.http_status : 500
     end
 
-    def default_params
-      open_ai_params
+    def handle_call_outcome(response_data:)
+      return response_data[:parsed_json] if response_data[:success]
+
+      if response_data[:error]
+        raise_api_error(response_data[:error])
+      else
+        raise_body_error(response_data[:response_obj])
+      end
     end
 
-    private
+    # -- Extraction helpers -----------------------------------------------------
+    def extract_answer_from_choices(choices)
+      raise Error, "OpenAI: No choices found in response" unless choices.is_a?(Array) && choices.any?
 
-    def log_messages_debug(messages)
-      return unless messages.is_a?(Array)
+      content = choices.map { |c| c.dig("message", "content") }.compact
+      return content.join("\n").strip unless content.empty?
+
+      text = choices.map { |c| c["text"] }.compact
+      return text.join("\n").strip unless text.empty?
 
-      Boxcars.debug(messages.last(2).map { |p| ">>>>>> Role: #{p[:role]} <<<<<<\n#{p[:content]}" }.join("\n"), :cyan)
+      raise Error, "OpenAI: Could not extract answer from choices"
     end
 
-    def get_params(prompt_object, inputs, params)
-      # Ensure prompt_object is a Boxcars::Prompt
-      current_prompt_object = if prompt_object.is_a?(Boxcars::Prompt)
-                                prompt_object
-                              else
-                                Boxcars::Prompt.new(template: prompt_object.to_s)
-                              end
+    # -- Utility helpers --------------------------------------------------------
+    def chat_model?(model_name) = CHAT_MODEL_REGEX.match?(model_name)
 
-      # Use as_messages for chat models
-      formatted_params = current_prompt_object.as_messages(inputs).merge(params)
+    def openai_error_message(json)
+      err = json&.dig("error")
+      return unless err
 
-      # Handle models like o1-mini that don't support the system role
-      if formatted_params[:model] =~ /^o/ && formatted_params[:messages].first&.fetch(:role)&.to_s == 'system'
-        formatted_params[:messages].first[:role] = :user
-      end
-      # o1-mini specific param adjustments (already in initialize, but good to ensure here if params are rebuilt)
-      if formatted_params[:model] =~ /^o/
-        formatted_params.delete(:response_format)
-        formatted_params.delete(:stop)
-        if formatted_params.key?(:max_tokens) && !formatted_params.key?(:max_completion_tokens)
-          formatted_params[:max_completion_tokens] = formatted_params.delete(:max_tokens)
-        end
-        formatted_params.delete(:temperature)
-      end
-      formatted_params
+      err.is_a?(Hash) ? "#{err['type']}: #{err['message']}" : err.to_s
     end
 
-    def _handle_openai_error_outcome(error_data)
-      detailed_error_message = error_data.message
-      if error_data.respond_to?(:json_body) && error_data.json_body
-        detailed_error_message += " - Details: #{error_data.json_body}"
+    def adjust_for_o_series!(params)
+      return params unless params[:model] =~ O_SERIES_REGEX
+
+      params[:messages][0][:role] = :user if params.dig(:messages, 0, :role).to_s == "system"
+      params.delete(:response_format)
+      params.delete(:stop)
+      if params.key?(:max_tokens) && !params.key?(:max_completion_tokens)
+        params[:max_completion_tokens] =
+          params.delete(:max_tokens)
       end
-      Boxcars.error("OpenAI Error: #{detailed_error_message} (#{error_data.class.name})", :red)
-      raise error_data
+      params.delete(:temperature)
+      params
     end
 
-    def _handle_openai_response_body_error(response_obj)
-      err_details = response_obj&.dig("error")
-      msg = err_details ? "#{err_details['type']}: #{err_details['message']}" : "Unknown error from OpenAI API"
-      raise Error, msg
-    end
+    def log_messages_debug(messages)
+      return unless messages.is_a?(Array)
 
-    def _extract_openai_answer_from_choices(choices)
-      raise Error, "OpenAI: No choices found in response" unless choices.is_a?(Array) && !choices.empty?
+      Boxcars.debug(
+        messages.last(2).map { |m| ">>>>>> Role: #{m[:role]} <<<<<<\n#{m[:content]}" }.join("\n"),
+        :cyan
+      )
+    end
 
-      if choices.first&.dig("message", "content")
-        choices.map { |c| c.dig("message", "content") }.join("\n").strip
-      elsif choices.first&.dig("text")
-        choices.map { |c| c["text"] }.join("\n").strip
-      else
-        raise Error, "OpenAI: Could not extract answer from choices"
-      end
+    # -- Error raising ----------------------------------------------------------
+    def raise_api_error(err)
+      msg = err.message
+      msg += " - Details: #{err.json_body}" if err.respond_to?(:json_body) && err.json_body
+      Boxcars.error("OpenAI Error: #{msg} (#{err.class})", :red)
+      raise err
     end
 
-    def _openai_handle_call_outcome(response_data:)
-      if response_data[:error]
-        _handle_openai_error_outcome(response_data[:error])
-      elsif !response_data[:success] # e.g. raw_response["error"] was present
-        _handle_openai_response_body_error(response_data[:response_obj]) # Raises an error
-      else
-        response_data[:parsed_json] # Return the raw parsed JSON for Engine#generate
-      end
+    def raise_body_error(response_obj)
+      raise Error, openai_error_message(response_obj) || "Unknown error from OpenAI API"
     end
 
-    def _track_openai_observability(call_context, response_data)
-      duration_ms = ((Time.now - call_context[:start_time]) * 1000).round
-      is_chat_model = call_context[:is_chat_model]
-      api_request_params = call_context[:api_request_params] || {}
-      request_context = {
-        prompt: call_context[:prompt_object],
-        inputs: call_context[:inputs],
-        conversation_for_api: is_chat_model ? api_request_params[:messages] : api_request_params[:prompt]
-      }
+    # -- Observability ----------------------------------------------------------
+    def track_openai_observability(call_ctx, response_data)
+      duration_ms = ((Time.now - call_ctx[:start_time]) * 1000).round
+      api_req     = call_ctx[:api_request] || {}
 
       track_ai_generation(
-        duration_ms:,
-        current_params: call_context[:current_params],
-        request_context:,
-        response_data:,
+        duration_ms: duration_ms,
+        current_params: call_ctx[:current_params],
+        request_context: {
+          prompt: call_ctx[:prompt_object],
+          inputs: call_ctx[:inputs],
+          user_id: user_id,
+          conversation_for_api: call_ctx[:is_chat_model] ? api_req[:messages] : api_req[:prompt]
+        },
+        response_data: response_data,
         provider: :openai
       )
     end

data/lib/boxcars/engine/perplexityai.rb

@@ -20,10 +20,11 @@ module Boxcars
                           "You should ask targeted questions"
 
     def initialize(name: DEFAULT_NAME, description: DEFAULT_DESCRIPTION, prompts: [], batch_size: 20, **kwargs)
+      user_id = kwargs.delete(:user_id)
       @perplexity_params = DEFAULT_PARAMS.merge(kwargs)
       @prompts = prompts
       @batch_size = batch_size # Retain if used by generate
-      super(description:, name:)
+      super(description:, name:, user_id:)
     end
 
     # Perplexity models are conversational.
@@ -96,6 +97,7 @@ module Boxcars
         request_context = {
           prompt: current_prompt_object,
           inputs:,
+          user_id:,
           conversation_for_api: api_request_params&.dig(:messages)
         }
         track_ai_generation(

data/lib/boxcars/engine/unified_observability.rb

@@ -52,7 +52,8 @@ module Boxcars
         '$ai_latency': duration_seconds,
         '$ai_http_status': extract_status_code(response_data) || (response_data[:success] ? 200 : 500),
         '$ai_base_url': get_base_url_for_provider(provider),
-        '$ai_is_error': !response_data[:success]
+        '$ai_is_error': !response_data[:success],
+        user_id:
       }
 
       # Add error details if present

data/lib/boxcars/engine.rb

@@ -3,18 +3,20 @@
 module Boxcars
   # @abstract
   class Engine
-    attr_reader :prompts, :batch_size
+    attr_reader :prompts, :batch_size, :user_id
 
     # An Engine is used by Boxcars to generate output from prompts
     # @param name [String] The name of the Engine. Defaults to classname.
     # @param description [String] A description of the Engine.
     # @param prompts [Array<Prompt>] The prompts to use for the Engine.
     # @param batch_size [Integer] The number of prompts to send to the Engine at a time.
-    def initialize(description: 'Engine', name: nil, prompts: [], batch_size: 20)
+    # @param user_id [String, Integer] The ID of the user using this Engine (optional for observability).
+    def initialize(description: 'Engine', name: nil, prompts: [], batch_size: 20, user_id: nil)
       @name = name || self.class.name
       @description = description
       @prompts = prompts
       @batch_size = batch_size
+      @user_id = user_id
     end
 
     # Get an answer from the Engine.
@@ -60,14 +62,8 @@ module Boxcars
         sub_prompts.each do |sprompt, inputs|
           client_response = client(prompt: sprompt, inputs:, **params)
 
-          # Handle different response formats:
-          # - New format: response_data hash with :parsed_json key (Groq, Gemini)
-          # - Legacy format: direct API response hash (OpenAI, others)
-          api_response_hash = if client_response.is_a?(Hash) && client_response.key?(:parsed_json)
-                                client_response[:parsed_json]
-                              else
-                                client_response
-                              end
+          # All engines now return the parsed API response hash directly
+          api_response_hash = client_response
 
           # Ensure we have a hash to work with
           unless api_response_hash.is_a?(Hash)
@@ -101,6 +97,17 @@ module Boxcars
       end
       EngineResult.new(generations:, engine_output: { token_usage: })
     end
+
+    def extract_answer(response)
+      # Handle different response formats
+      if response["choices"]
+        response["choices"].map { |c| c.dig("message", "content") || c["text"] }.join("\n").strip
+      elsif response["candidates"]
+        response["candidates"].map { |c| c.dig("content", "parts", 0, "text") }.join("\n").strip
+      else
+        response["output"] || response.to_s
+      end
+    end
   end
 end