google-genai 0.1.0

data/lib/google/genai/files.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require_relative 'pagers'
+
+module Google
+  module Genai
+    class Files
+      def initialize(api_client)
+        @api_client = api_client
+      end
+
+      def upload(file:, config: nil)
+        file_path = file.is_a?(String) ? file : file.path
+        raise ArgumentError, "File not found: #{file_path}" unless File.exist?(file_path)
+
+        config ||= {}
+        mime_type = config[:mime_type] || MimeMagic.by_path(file_path)&.type
+        raise ArgumentError, "Could not determine MIME type for file: #{file_path}" unless mime_type
+
+        display_name = config[:display_name] || File.basename(file_path)
+        file_size = File.size(file_path)
+
+        response_data = @api_client.upload_file(file_path, file_size, mime_type, display_name: display_name)
+        file_info = Types::File.new(response_data['file'])
+
+        # Poll for file processing to complete with exponential backoff
+        timeout = 120 # 2 minutes
+        start_time = Time.now
+        delay = 0.1 # Initial delay of 100ms
+
+        loop do
+          file_info = get(name: file_info.name)
+          case file_info.state
+          when 'ACTIVE'
+            return file_info
+          when 'FAILED'
+            raise "File processing failed: #{file_info.error}"
+          end
+
+          raise "File processing timed out" if Time.now - start_time > timeout
+
+          sleep delay
+          delay = [delay * 1.5, 5].min # Increase delay, but cap at 5 seconds
+        end
+      end
+
+      def get(name:, config: nil)
+        file_id = name.start_with?('files/') ? name.split('/').last : name
+        response = @api_client.get("v1beta/files/#{file_id}")
+        Types::File.new(JSON.parse(response.body))
+      end
+
+      def delete(name:, config: nil)
+        file_id = name.start_with?('files/') ? name.split('/').last : name
+        @api_client.delete("v1beta/files/#{file_id}")
+        nil
+      end
+
+      def list(config: nil)
+        list_request = ->(options) do
+          path = "v1beta/files"
+          params = {}
+          params[:pageToken] = options[:page_token] if options&.key?(:page_token)
+          params[:pageSize] = options[:page_size] if options&.key?(:page_size)
+          path += "?#{URI.encode_www_form(params)}" unless params.empty?
+          @api_client.get(path)
+        end
+
+        response = list_request.call(config)
+
+        Pager.new(
+          name: :files,
+          request: list_request,
+          response: response,
+          config: config,
+          item_class: Types::File
+        )
+      end
+
+      def download(file:, config: nil)
+        # TODO: Implement file download logic
+      end
+    end
+  end
+end

data/lib/google/genai/live.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require 'websocket-client-simple'
+require 'json'
+require_relative 'live_music'
+
+module Google
+  module Genai
+    class Live
+      def initialize(api_client)
+        @api_client = api_client
+      end
+
+      def music
+        @music ||= LiveMusic.new(@api_client)
+      end
+
+      def connect(model:, config: nil)
+        raise "Live API is not supported for Vertex AI yet" if @api_client.vertexai
+
+        base_url = "wss://generativelanguage.googleapis.com"
+        version = @api_client.instance_variable_get(:@http_options)&.[](:api_version) || 'v1beta'
+        api_key = @api_client.api_key
+        
+        uri = "#{base_url}/ws/google.ai.generativelanguage.#{version}.GenerativeService.BidiGenerateContent?key=#{api_key}"
+        
+        ws = WebSocket::Client::Simple.connect(uri)
+        
+        session = Session.new(ws)
+        
+        setup_message = {
+          setup: {
+            model: "models/#{model}",
+            generationConfig: config
+          }
+        }
+        ws.send(setup_message.to_json)
+        
+        # Block until the connection is open and initial setup is confirmed.
+        # This is a simplified way to handle the async nature of websockets in a sync method.
+        # A more robust solution would involve a proper event loop.
+        sleep 0.1 until ws.open?
+        
+        yield session
+      ensure
+        ws.close if ws&.open?
+      end
+    end
+
+    class Session
+      def initialize(websocket)
+        @ws = websocket
+        @message_queue = Queue.new
+        
+        @ws.on :message do |msg|
+          @message_queue.push(JSON.parse(msg.data))
+        end
+
+        @ws.on :error do |err|
+          # For now, just print the error. A more robust error handling can be added.
+          puts "WebSocket Error: #{err.message}"
+        end
+      end
+
+      def send_client_content(turns:, turn_complete: true)
+        message = {
+          clientContent: {
+            turns: turns,
+            turnComplete: turn_complete
+          }
+        }
+        @ws.send(message.to_json)
+      end
+
+      def receive
+        @message_queue.pop
+      end
+    end
+  end
+end

data/lib/google/genai/live_music.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+require 'websocket-client-simple'
+require 'json'
+
+module Google
+  module Genai
+    class LiveMusic
+      def initialize(api_client)
+        @api_client = api_client
+      end
+
+      def connect(model:)
+        raise "Live Music API is not supported for Vertex AI yet" if @api_client.vertexai
+
+        base_url = "wss://generativelanguage.googleapis.com"
+        version = @api_client.instance_variable_get(:@http_options)&.[](:api_version) || 'v1beta'
+        api_key = @api_client.api_key
+        
+        uri = "#{base_url}/ws/google.ai.generativelanguage.#{version}.GenerativeService.BidiGenerateMusic?key=#{api_key}"
+        
+        ws = WebSocket::Client::Simple.connect(uri)
+        
+        session = MusicSession.new(ws)
+        
+        setup_message = {
+          setup: {
+            model: "models/#{model}"
+          }
+        }
+        ws.send(setup_message.to_json)
+        
+        sleep 0.1 until ws.open?
+        
+        yield session
+      ensure
+        ws.close if ws&.open?
+      end
+    end
+
+    class MusicSession
+      def initialize(websocket)
+        @ws = websocket
+        @message_queue = Queue.new
+        
+        @ws.on :message do |msg|
+          @message_queue.push(JSON.parse(msg.data))
+        end
+
+        @ws.on :error do |err|
+          puts "WebSocket Error: #{err.message}"
+        end
+      end
+
+      def set_weighted_prompts(prompts:)
+        message = {
+          clientContent: {
+            weightedPrompts: prompts.map(&:to_h)
+          }
+        }
+        @ws.send(message.to_json)
+      end
+
+      def set_music_generation_config(config:)
+        @ws.send({ musicGenerationConfig: config }.to_json)
+      end
+
+      def play
+        _send_control_signal('PLAY')
+      end
+
+      def pause
+        _send_control_signal('PAUSE')
+      end
+
+      def stop
+        _send_control_signal('STOP')
+      end
+
+      def reset_context
+        _send_control_signal('RESET_CONTEXT')
+      end
+
+      def receive
+        @message_queue.pop
+      end
+
+      private
+
+      def _send_control_signal(control_signal)
+        message = {
+          playbackControl: control_signal
+        }
+        @ws.send(message.to_json)
+      end
+    end
+  end
+end

data/lib/google/genai/models.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+require_relative 'types'
+require_relative 'pagers'
+
+module Google
+  module Genai
+    class Models
+      def initialize(api_client)
+        @api_client = api_client
+      end
+
+      def generate_content(model:, contents:, config: nil)
+        path = "v1beta/models/#{model}:generateContent"
+        body = {
+          contents: normalize_contents(contents)
+        }
+        body[:generationConfig] = config if config
+
+        response = @api_client.request(:post, path, body: body)
+        Types::GenerateContentResponse.new(JSON.parse(response.body))
+      end
+
+      def list(config: nil)
+        query_base = config&.dig(:query_base) != false
+        
+        list_request = ->(options) do
+          path = "v1beta/#{query_base ? 'models' : 'tunedModels'}"
+          params = {}
+          params[:pageToken] = options[:page_token] if options&.key?(:page_token)
+          params[:pageSize] = options[:page_size] if options&.key?(:page_size)
+          path += "?#{URI.encode_www_form(params)}" unless params.empty?
+          @api_client.get(path)
+        end
+
+        response = list_request.call(config)
+
+        Pager.new(
+          name: query_base ? :models : :tunedModels,
+          request: list_request,
+          response: response,
+          config: config,
+          item_class: Types::Model
+        )
+      end
+
+      private
+
+      def normalize_contents(contents)
+        contents = [contents] unless contents.is_a?(Array)
+
+        contents.map do |item|
+          case item
+          when String
+            { role: 'user', parts: [{ text: item }] }
+          when Hash
+            item # Assumes it's already in the correct format
+          when Types::Content
+            item.to_h
+          when Types::File
+            { role: 'user', parts: [{ file_data: { mime_type: item.mime_type, file_uri: item.uri } }] }
+          else
+            raise ArgumentError, "Unsupported content type: #{item.class}"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/google/genai/operations.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require_relative 'types'
+
+module Google
+  module Genai
+    class Operations
+      def initialize(api_client)
+        @api_client = api_client
+      end
+
+      def get(name:, config: nil)
+        response = @api_client.get("v1beta/#{name}")
+        Types::Operation.new(JSON.parse(response.body))
+      end
+    end
+  end
+end

data/lib/google/genai/pagers.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Google
+  module Genai
+    class Pager
+      include Enumerable
+
+      attr_reader :name, :page_size, :config, :sdk_http_response
+
+      def initialize(name:, request:, response:, config:, item_class:)
+        @item_class = item_class
+        _init_page(name: name, request: request, response: response, config: config)
+      end
+
+      def each(&block)
+        loop do
+          @page.each(&block)
+          break unless next_page_token
+          next_page
+        end
+      end
+
+      def page
+        @page
+      end
+
+      def next_page
+        raise IndexError, 'No more pages to fetch.' unless next_page_token
+
+        response = @request.call(config: @config)
+        _init_page(name: @name, request: @request, response: response, config: @config)
+        @page
+      end
+
+      private
+
+      def _init_page(name:, request:, response:, config:)
+        @name = name
+        @request = request
+        
+        response_body = JSON.parse(response.body)
+        @page = (response_body[name.to_s] || []).map { |item_data| @item_class.new(item_data) }
+        
+        @sdk_http_response = response
+
+        @config = config ? config.dup : {}
+        @config[:page_token] = response_body['nextPageToken']
+        
+        @page_size = @config[:page_size] || @page.length
+      end
+
+      def next_page_token
+        @config[:page_token]
+      end
+    end
+  end
+end

data/lib/google/genai/tokens.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+require_relative 'types'
+
+module Google
+  module Genai
+    class Tokens
+      def initialize(api_client)
+        @api_client = api_client
+      end
+
+      def create(config: nil)
+        raise "Tokens is only supported for Gemini API" if @api_client.vertexai
+
+        # This is a simplified port of the Python SDK's logic.
+        # The full logic for field masks is complex and will be implemented later.
+        
+        body = {}
+        body[:expireTime] = config[:expire_time] if config&.key?(:expire_time)
+        body[:uses] = config[:uses] if config&.key?(:uses)
+
+        if config&.key?(:live_connect_constraints)
+          # This part of the conversion is complex and will require more detailed mapping.
+          # For now, we'll pass a simplified version.
+          body[:bidiGenerateContentSetup] = {
+            setup: {
+              model: config[:live_connect_constraints][:model]
+            }
+          }
+          if config[:live_connect_constraints][:config]
+            body[:bidiGenerateContentSetup][:setup][:generationConfig] = config[:live_connect_constraints][:config]
+          end
+        end
+
+        response = @api_client.request(:post, "v1alpha/authTokens", body: body)
+        Types::AuthToken.new(JSON.parse(response.body))
+      end
+    end
+  end
+end

data/lib/google/genai/tunings.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require_relative 'types'
+require_relative 'pagers'
+
+module Google
+  module Genai
+    class Tunings
+      def initialize(api_client)
+        @api_client = api_client
+      end
+
+      def create(base_model:, training_dataset:, config: nil)
+        raise "Tuning is only supported for Vertex AI" unless @api_client.vertexai
+
+        body = {
+          baseModel: base_model,
+          supervisedTuningSpec: {
+            trainingDatasetUri: training_dataset[:gcs_uri] || training_dataset[:vertex_dataset_resource]
+          }
+        }
+        # TODO: Add other config options
+
+        response = @api_client.request(:post, "v1beta/tuningJobs", body: body)
+        Types::TuningJob.new(JSON.parse(response.body))
+      end
+
+      def get(name:, config: nil)
+        raise "Tuning is only supported for Vertex AI" unless @api_client.vertexai
+        response = @api_client.get("v1beta/#{name}")
+        Types::TuningJob.new(JSON.parse(response.body))
+      end
+
+      def list(config: nil)
+        raise "Tuning is only supported for Vertex AI" unless @api_client.vertexai
+        
+        list_request = ->(options) do
+          path = "v1beta/tuningJobs"
+          params = {}
+          params[:pageToken] = options[:page_token] if options&.key?(:page_token)
+          params[:pageSize] = options[:page_size] if options&.key?(:page_size)
+          path += "?#{URI.encode_www_form(params)}" unless params.empty?
+          @api_client.get(path)
+        end
+
+        response = list_request.call(config)
+
+        Pager.new(
+          name: :tuningJobs,
+          request: list_request,
+          response: response,
+          config: config,
+          item_class: Types::TuningJob
+        )
+      end
+
+      def cancel(name:, config: nil)
+        raise "Tuning is only supported for Vertex AI" unless @api_client.vertexai
+        @api_client.request(:post, "v1beta/#{name}:cancel", body: {})
+        nil
+      end
+    end
+  end
+end

data/lib/google/genai/types.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+require "base64"
+
+module Google
+  module Genai
+    module Types
+      class Base
+        def initialize(attributes = {})
+          attributes.each do |key, value|
+            setter = "#{key}="
+            public_send(setter, value) if respond_to?(setter)
+          end
+        end
+      end
+
+      class Blob < Base
+        attr_accessor :mime_type, :data
+
+        def to_h
+          { mimeType: mime_type, data: Base64.strict_encode64(data) }
+        end
+      end
+
+      class FileData < Base
+        attr_accessor :mime_type, :file_uri
+
+        def to_h
+          { mimeType: mime_type, fileUri: file_uri }
+        end
+      end
+
+      class Part < Base
+        attr_accessor :text, :inline_data, :file_data
+
+        def to_h
+          data = {}
+          data[:text] = text if text
+          data[:inlineData] = inline_data.to_h if inline_data
+          data[:fileData] = file_data.to_h if file_data
+          data
+        end
+      end
+
+      class Content < Base
+        attr_accessor :role, :parts
+
+        def initialize(attributes = {})
+          super
+          self.parts = Array(self.parts).map { |p| p.is_a?(Part) ? p : Part.new(p) }
+        end
+
+        def to_h
+          {
+            role: role,
+            parts: parts.map(&:to_h)
+          }
+        end
+      end
+
+      class GenerateContentResponse < Base
+        attr_accessor :candidates
+
+        def initialize(attributes = {})
+          super
+          self.candidates = Array(self.candidates).map { |c| c.is_a?(Candidate) ? c : Candidate.new(c) }
+        end
+
+        def text
+          candidates&.first&.content&.parts&.map(&:text)&.join
+        end
+      end
+
+      class Candidate < Base
+        attr_accessor :content, :finish_reason, :safety_ratings
+
+        def initialize(attributes = {})
+          super
+          self.content = Content.new(self.content) if self.content.is_a?(Hash)
+        end
+      end
+
+      class File < Base
+        attr_accessor :name, :display_name, :mime_type, :size_bytes, :create_time, :update_time, :expiration_time, :sha256_hash, :uri, :state, :error
+      end
+
+      class TuningJob < Base
+        attr_accessor :name, :state, :create_time, :end_time, :tuned_model
+      end
+
+      class CachedContent < Base
+        attr_accessor :name, :display_name, :model, :create_time, :update_time, :expire_time, :usage_metadata
+      end
+
+      class BatchJob < Base
+        attr_accessor :name, :model, :state, :create_time, :end_time, :update_time, :error
+      end
+
+      class Operation < Base
+        attr_accessor :name, :metadata, :done, :error, :response
+      end
+    end
+  end
+end

data/lib/google/genai/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Google
+  module Genai
+    VERSION = "0.1.0"
+  end
+end

data/lib/google/genai.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require "zeitwerk"
+loader = Zeitwerk::Loader.for_gem
+loader.setup
+
+require_relative "genai/version"
+require_relative "genai/api_client"
+require_relative "genai/models"
+require_relative "genai/chats"
+require_relative "genai/files"
+require_relative "genai/tunings"
+require_relative "genai/caches"
+require_relative "genai/batches"
+require_relative "genai/operations"
+require_relative "genai/tokens"
+require_relative "genai/live"
+require_relative "genai/client"
+
+module Google
+  module Genai
+    # Your code goes here...
+  end
+end