kie-ruby 0.1.0

data/lib/kie/general_task.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Kie
+  # Represents an asynchronous task for General API endpoints
+  class GeneralTask
+    include Task
+
+    RECORD_INFO_ENDPOINT = "/api/v1/jobs/recordInfo"
+
+    # Status mapping from General API raw states to normalized statuses
+    STATUS_MAP = {
+      waiting: :processing,
+      generating: :processing,
+      success: :success,
+      fail: :failed
+    }.freeze
+
+    attr_reader :task_id, :raw_status, :cost_time, :fail_code, :fail_message
+
+    alias id task_id
+
+    def initialize(client:, task_id:)
+      @client = client
+      @task_id = task_id
+      @raw_status = nil
+      @result_urls = []
+      @cost_time = nil
+      @fail_code = nil
+      @fail_message = nil
+    end
+
+    def refresh!
+      response = @client.connection.get(RECORD_INFO_ENDPOINT, taskId: @task_id)
+      handle_response(response)
+      self
+    end
+
+    # Returns normalized status (:processing, :success, :failed)
+    def status
+      return :processing if @raw_status.nil?
+
+      STATUS_MAP.fetch(@raw_status, :processing)
+    end
+
+    # Returns array of result URLs
+    def urls
+      @result_urls
+    end
+
+    # Alias for backwards compatibility
+    alias result_urls urls
+
+    # Returns the error message from the failed task, or nil if not failed
+    def error_message
+      @fail_message
+    end
+
+    # Returns true when task is in pending/waiting state (not yet processing)
+    def pending?
+      @raw_status.nil? || @raw_status == :waiting
+    end
+
+    # Raises the appropriate exception based on the failure reason
+    def raise_task_error!
+      error_class = @fail_code == "CONTENT_POLICY" ? ContentPolicyError : TaskFailedError
+      raise error_class.new(@fail_message, fail_code: @fail_code)
+    end
+
+    private
+
+    def handle_response(response)
+      data = JSON.parse(response.body)
+      handle_api_error(data) unless data["code"] == 200
+
+      update_from_data(data["data"])
+    end
+
+    def handle_api_error(data)
+      case data["code"]
+      when 404
+        raise NotFoundError.new(data["msg"], status_code: data["code"], response_body: data.to_json)
+      else
+        raise ApiError.new(data["msg"], status_code: data["code"], response_body: data.to_json)
+      end
+    end
+
+    def update_from_data(data)
+      @raw_status = data["state"]&.to_sym
+      @cost_time = data["costTime"]
+      @fail_code = data["failCode"]
+      @fail_message = data["failMsg"]
+      parse_result_json(data["resultJson"])
+    end
+
+    def parse_result_json(result_json)
+      return if result_json.nil? || result_json.empty?
+
+      parsed = JSON.parse(result_json)
+      @result_urls = parsed["resultUrls"] || []
+    rescue JSON::ParserError
+      @result_urls = []
+    end
+  end
+end

data/lib/kie/middleware/raise_error.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "json"
+require "faraday"
+
+module Kie
+  module Middleware
+    # Faraday middleware that raises Kie-specific exceptions based on HTTP status codes
+    class RaiseError < Faraday::Middleware
+      ERROR_MAP = {
+        400 => InvalidParametersError,
+        401 => AuthenticationError,
+        402 => InsufficientCreditsError,
+        404 => NotFoundError,
+        422 => InvalidParametersError,
+        429 => RateLimitError,
+        455 => ServiceUnavailableError
+      }.freeze
+
+      def on_complete(env)
+        return if env.success?
+
+        error_class = determine_error_class(env.status)
+        message = extract_message(env.body, env.status)
+
+        raise error_class.new(message, status_code: env.status, response_body: env.body)
+      end
+
+      private
+
+      def determine_error_class(status)
+        return ERROR_MAP[status] if ERROR_MAP.key?(status)
+        return ServerError if status >= 500
+
+        ApiError
+      end
+
+      def extract_message(body, status)
+        data = JSON.parse(body)
+        data["msg"] || "HTTP #{status}"
+      rescue JSON::ParserError
+        "HTTP #{status}"
+      end
+    end
+  end
+end

data/lib/kie/model_definition.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module Kie
+  ModelDefinition = Struct.new(:name, :category, :input_constraints, :parameters, keyword_init: true) do
+    def initialize(...)
+      super
+      freeze
+    end
+  end
+end

data/lib/kie/model_registry.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Kie
+  # Central registry for model definitions
+  class ModelRegistry
+    class << self
+      def register(definition)
+        registry[definition.name] = definition
+      end
+
+      def find(model_name)
+        registry.fetch(model_name) do
+          raise UnknownModelError, "Unknown model: #{model_name}"
+        end
+      end
+
+      def registered?(model_name)
+        registry.key?(model_name)
+      end
+
+      def all
+        registry.values
+      end
+
+      def clear
+        @registry = {}
+      end
+
+      private
+
+      def registry
+        @registry ||= {}
+      end
+    end
+  end
+end

data/lib/kie/models/nano_banana_pro.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+Kie::ModelRegistry.register(
+  Kie::ModelDefinition.new(
+    name: "nano-banana-pro",
+    category: :general,
+    input_constraints: {
+      prompt: { type: :string, max_length: 20_000, required: true },
+      images: { type: :array, max_items: 8, max_file_size: 30_000_000, formats: %w[jpeg png webp], required: false }
+    },
+    parameters: {
+      aspect_ratio: { type: :enum, values: %w[1:1 2:3 3:2 3:4 4:3 4:5 5:4 9:16 16:9 21:9 auto], default: "1:1" },
+      resolution: { type: :enum, values: %w[1K 2K 4K], default: "1K" },
+      output_format: { type: :enum, values: %w[png jpg], default: "png" }
+    }
+  )
+)

data/lib/kie/models/suno_v4.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+Kie::ModelRegistry.register(
+  Kie::ModelDefinition.new(
+    name: "V4",
+    category: :suno,
+    input_constraints: {
+      prompt: { type: :string, max_length: 3000, max_length_non_custom: 500 },
+      style: { type: :string, max_length: 200 },
+      title: { type: :string, max_length: 80 }
+    },
+    parameters: {}
+  )
+)

data/lib/kie/models/suno_v4_5.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+Kie::ModelRegistry.register(
+  Kie::ModelDefinition.new(
+    name: "V4_5",
+    category: :suno,
+    input_constraints: {
+      prompt: { type: :string, max_length: 5000, max_length_non_custom: 500 },
+      style: { type: :string, max_length: 1000 },
+      title: { type: :string, max_length: 80 }
+    },
+    parameters: {}
+  )
+)

data/lib/kie/models/suno_v4_5_all.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+Kie::ModelRegistry.register(
+  Kie::ModelDefinition.new(
+    name: "V4_5ALL",
+    category: :suno,
+    input_constraints: {
+      prompt: { type: :string, max_length: 5000, max_length_non_custom: 500 },
+      style: { type: :string, max_length: 1000 },
+      title: { type: :string, max_length: 80 }
+    },
+    parameters: {}
+  )
+)

data/lib/kie/models/suno_v4_5_plus.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+Kie::ModelRegistry.register(
+  Kie::ModelDefinition.new(
+    name: "V4_5PLUS",
+    category: :suno,
+    input_constraints: {
+      prompt: { type: :string, max_length: 5000, max_length_non_custom: 500 },
+      style: { type: :string, max_length: 1000 },
+      title: { type: :string, max_length: 80 }
+    },
+    parameters: {}
+  )
+)

data/lib/kie/models/suno_v5.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+Kie::ModelRegistry.register(
+  Kie::ModelDefinition.new(
+    name: "V5",
+    category: :suno,
+    input_constraints: {
+      prompt: { type: :string, max_length: 5000, max_length_non_custom: 500 },
+      style: { type: :string, max_length: 1000 },
+      title: { type: :string, max_length: 80 }
+    },
+    parameters: {}
+  )
+)

data/lib/kie/suno_engine.rb

@@ -0,0 +1,129 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Kie
+  # Engine for Suno API endpoints (/api/v1/generate) - Suno music generation
+  class SunoEngine
+    CREATE_TASK_ENDPOINT = "/api/v1/generate"
+
+    OPTIONAL_FIELD_MAPPING = {
+      callback_url: :callBackUrl,
+      negative_tags: :negativeTags,
+      vocal_gender: :vocalGender,
+      style_weight: :styleWeight,
+      weirdness_constraint: :weirdnessConstraint,
+      audio_weight: :audioWeight,
+      persona_id: :personaId
+    }.freeze
+    private_constant :OPTIONAL_FIELD_MAPPING
+
+    def initialize(client:)
+      @client = client
+    end
+
+    def generate(model:, input:, **)
+      definition = ModelRegistry.find(model)
+      validate_input!(definition, input)
+
+      response = @client.connection.post(CREATE_TASK_ENDPOINT) do |req|
+        req.body = build_request_body(definition, input)
+      end
+
+      handle_response(response)
+    end
+
+    private
+
+    def validate_input!(definition, input)
+      constraints = definition.input_constraints
+      custom_mode = input[:custom_mode] || false
+
+      validate_prompt!(constraints, input, custom_mode)
+      validate_style!(constraints, input) if custom_mode && input[:style]
+      validate_title!(constraints, input) if custom_mode && input[:title]
+    end
+
+    def validate_prompt!(constraints, input, custom_mode)
+      prompt = input[:prompt]
+      return unless prompt && constraints[:prompt]
+
+      max_length = if custom_mode
+                     constraints[:prompt][:max_length]
+                   else
+                     constraints[:prompt][:max_length_non_custom] || constraints[:prompt][:max_length]
+                   end
+
+      return unless prompt.length > max_length
+
+      raise ValidationError, "prompt exceeds maximum length of #{max_length} characters"
+    end
+
+    def validate_style!(constraints, input)
+      style = input[:style]
+      return unless style && constraints[:style]
+
+      max_length = constraints[:style][:max_length]
+      return unless style.length > max_length
+
+      raise ValidationError, "style exceeds maximum length of #{max_length} characters"
+    end
+
+    def validate_title!(constraints, input)
+      title = input[:title]
+      return unless title && constraints[:title]
+
+      max_length = constraints[:title][:max_length]
+      return unless title.length > max_length
+
+      raise ValidationError, "title exceeds maximum length of #{max_length} characters"
+    end
+
+    def build_request_body(definition, input)
+      body = {
+        model: definition.name,
+        customMode: input[:custom_mode] || false,
+        instrumental: input[:instrumental] || false,
+        prompt: input[:prompt],
+        callBackUrl: input[:callback_url] || "https://example.invalid/callback"
+      }
+
+      add_custom_mode_fields(body, input) if input[:custom_mode]
+      add_optional_fields(body, input)
+
+      body.to_json
+    end
+
+    def add_custom_mode_fields(body, input)
+      body[:style] = input[:style] if input[:style]
+      body[:title] = input[:title] if input[:title]
+    end
+
+    def add_optional_fields(body, input)
+      OPTIONAL_FIELD_MAPPING.each do |input_key, api_key|
+        body[api_key] = input[input_key] if input[input_key]
+      end
+    end
+
+    def handle_response(response)
+      data = JSON.parse(response.body)
+      handle_api_error(data) unless data["code"] == 200
+
+      SunoTask.new(client: @client, task_id: data["data"]["taskId"])
+    end
+
+    def handle_api_error(data)
+      error_class = error_class_for_code(data["code"])
+      raise error_class.new(data["msg"], status_code: data["code"], response_body: data.to_json)
+    end
+
+    def error_class_for_code(code)
+      case code
+      when 400, 422 then InvalidParametersError
+      when 402 then InsufficientCreditsError
+      when 429 then RateLimitError
+      else ApiError
+      end
+    end
+  end
+end

data/lib/kie/suno_task.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Kie
+  # Represents an asynchronous task for Suno API endpoints (Suno music generation)
+  class SunoTask
+    include Task
+
+    RECORD_INFO_ENDPOINT = "/api/v1/generate/record-info"
+
+    # Status mapping from Suno API raw states to normalized statuses
+    STATUS_MAP = {
+      PENDING: :processing,
+      TEXT_SUCCESS: :processing,
+      FIRST_SUCCESS: :processing,
+      SUCCESS: :success,
+      CREATE_TASK_FAILED: :failed,
+      GENERATE_AUDIO_FAILED: :failed,
+      CALLBACK_EXCEPTION: :failed,
+      SENSITIVE_WORD_ERROR: :failed
+    }.freeze
+
+    attr_reader :task_id, :raw_status, :error_code, :error_message, :suno_data, :audio_urls, :image_urls
+
+    alias id task_id
+
+    def initialize(client:, task_id:)
+      @client = client
+      @task_id = task_id
+      @raw_status = nil
+      @suno_data = []
+      @audio_urls = []
+      @image_urls = []
+      @error_code = nil
+      @error_message = nil
+    end
+
+    def refresh!
+      response = @client.connection.get(RECORD_INFO_ENDPOINT, taskId: @task_id)
+      handle_response(response)
+      self
+    end
+
+    # Returns normalized status (:processing, :success, :failed)
+    def status
+      return :processing if @raw_status.nil?
+
+      STATUS_MAP.fetch(@raw_status, :processing)
+    end
+
+    # Returns array of audio URLs (implements Task interface)
+    def urls
+      @audio_urls
+    end
+
+    # Raises the appropriate exception based on the failure reason
+    def raise_task_error!
+      error_class = @raw_status == :SENSITIVE_WORD_ERROR ? ContentPolicyError : TaskFailedError
+      raise error_class.new(@error_message, fail_code: @raw_status.to_s)
+    end
+
+    private
+
+    def handle_response(response)
+      data = JSON.parse(response.body)
+      handle_api_error(data) unless data["code"] == 200
+
+      update_from_data(data["data"])
+    end
+
+    def handle_api_error(data)
+      case data["code"]
+      when 404
+        raise NotFoundError.new(data["msg"], status_code: data["code"], response_body: data.to_json)
+      else
+        raise ApiError.new(data["msg"], status_code: data["code"], response_body: data.to_json)
+      end
+    end
+
+    def update_from_data(data)
+      @raw_status = data["status"]&.to_sym
+      @error_code = data["errorCode"]
+      @error_message = data["errorMessage"]
+      parse_response(data["response"])
+    end
+
+    def parse_response(response)
+      return if response.nil?
+
+      suno_data = response["sunoData"]
+      return if suno_data.nil? || suno_data.empty?
+
+      @suno_data = suno_data
+      @audio_urls = suno_data.filter_map { |item| item["audioUrl"] }
+      @image_urls = suno_data.filter_map { |item| item["imageUrl"] }
+    end
+  end
+end

data/lib/kie/task.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+module Kie
+  # Common interface module for all task types (General API, Suno API, etc.)
+  # Classes including this module must implement:
+  # - #status: Returns normalized status symbol (:processing, :success, :failed)
+  # - #urls: Returns array of result URLs
+  module Task
+    NORMALIZED_STATUSES = %i[processing success failed].freeze
+
+    def status
+      raise NotImplementedError, "#{self.class} must implement #status"
+    end
+
+    def urls
+      raise NotImplementedError, "#{self.class} must implement #urls"
+    end
+
+    def error_message
+      raise NotImplementedError, "#{self.class} must implement #error_message"
+    end
+
+    def completed?
+      success? || failed?
+    end
+
+    def success?
+      status == :success
+    end
+
+    def failed?
+      status == :failed
+    end
+
+    def processing?
+      status == :processing
+    end
+
+    # Waits for the task to complete by polling refresh!
+    # @param timeout [Numeric] Maximum time to wait in seconds (default: 300)
+    # @param interval [Numeric] Time between polling attempts in seconds (default: 2)
+    # @param raise_on_failure [Boolean] Whether to raise an exception on task failure (default: true)
+    # @return [self] Returns the task instance for chaining
+    # @raise [TimeoutError] if the timeout is exceeded
+    # @raise [TaskFailedError, ContentPolicyError] if the task fails and raise_on_failure is true
+    def wait(timeout: 300, interval: 2, raise_on_failure: true)
+      deadline = Time.now + timeout
+      poll_until_completed(deadline: deadline, interval: interval, raise_on_failure: raise_on_failure)
+    end
+
+    private
+
+    def poll_until_completed(deadline:, interval:, raise_on_failure:)
+      loop do
+        refresh!
+        return self if success?
+        return handle_failure(raise_on_failure) if failed?
+        raise TimeoutError, "Task did not complete within timeout" if Time.now > deadline
+
+        sleep interval
+      end
+    end
+
+    def handle_failure(raise_on_failure)
+      raise_task_error! if raise_on_failure
+      self
+    end
+  end
+end

data/lib/kie/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Kie
+  VERSION = "0.1.0"
+end

data/lib/kie.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require_relative "kie/version"
+
+module Kie
+  class Error < StandardError; end
+end
+
+require_relative "kie/errors"
+require_relative "kie/middleware/raise_error"
+require_relative "kie/task"
+require_relative "kie/model_definition"
+require_relative "kie/model_registry"
+require_relative "kie/client"
+require_relative "kie/general_task"
+require_relative "kie/general_engine"
+require_relative "kie/suno_task"
+require_relative "kie/suno_engine"
+
+# Model definitions
+require_relative "kie/models/nano_banana_pro"
+require_relative "kie/models/suno_v4"
+require_relative "kie/models/suno_v4_5"
+require_relative "kie/models/suno_v4_5_plus"
+require_relative "kie/models/suno_v4_5_all"
+require_relative "kie/models/suno_v5"

data/sig/kie.rbs

@@ -0,0 +1,6 @@
+module Kie
+  VERSION: String
+
+  class Error < StandardError
+  end
+end