dify_llm 1.6.4

data/lib/ruby_llm/providers/vertexai/models.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI
+      # Models methods for the Vertex AI integration
+      module Models
+        # Gemini and other Google models that aren't returned by the API
+        KNOWN_GOOGLE_MODELS = %w[
+          gemini-2.5-flash-lite
+          gemini-2.5-pro
+          gemini-2.5-flash
+          gemini-2.0-flash-lite-001
+          gemini-2.0-flash-001
+          gemini-2.0-flash
+          gemini-2.0-flash-exp
+          gemini-1.5-pro-002
+          gemini-1.5-pro
+          gemini-1.5-flash-002
+          gemini-1.5-flash
+          gemini-1.5-flash-8b
+          gemini-pro
+          gemini-pro-vision
+          gemini-exp-1206
+          gemini-exp-1121
+          gemini-embedding-001
+          text-embedding-005
+          text-embedding-004
+          text-multilingual-embedding-002
+        ].freeze
+
+        def list_models
+          all_models = []
+          page_token = nil
+
+          all_models.concat(build_known_models)
+
+          loop do
+            response = @connection.get('publishers/google/models') do |req|
+              req.headers['x-goog-user-project'] = @config.vertexai_project_id
+              req.params = { pageSize: 100 }
+              req.params[:pageToken] = page_token if page_token
+            end
+
+            publisher_models = response.body['publisherModels'] || []
+            publisher_models.each do |model_data|
+              next if model_data['launchStage'] == 'DEPRECATED'
+
+              model_id = extract_model_id_from_path(model_data['name'])
+              all_models << build_model_from_api_data(model_data, model_id)
+            end
+
+            page_token = response.body['nextPageToken']
+            break unless page_token
+          end
+
+          all_models
+        rescue StandardError => e
+          RubyLLM.logger.debug "Error fetching Vertex AI models: #{e.message}"
+          build_known_models
+        end
+
+        private
+
+        def build_known_models
+          KNOWN_GOOGLE_MODELS.map do |model_id|
+            Model::Info.new(
+              id: model_id,
+              name: model_id,
+              provider: slug,
+              family: determine_model_family(model_id),
+              created_at: nil,
+              context_window: nil,
+              max_output_tokens: nil,
+              modalities: nil,
+              capabilities: %w[streaming function_calling],
+              pricing: nil,
+              metadata: {
+                source: 'known_models'
+              }
+            )
+          end
+        end
+
+        def build_model_from_api_data(model_data, model_id)
+          Model::Info.new(
+            id: model_id,
+            name: model_id,
+            provider: slug,
+            family: determine_model_family(model_id),
+            created_at: nil,
+            context_window: nil,
+            max_output_tokens: nil,
+            modalities: nil,
+            capabilities: extract_capabilities(model_data),
+            pricing: nil,
+            metadata: {
+              version_id: model_data['versionId'],
+              open_source_category: model_data['openSourceCategory'],
+              launch_stage: model_data['launchStage'],
+              supported_actions: model_data['supportedActions'],
+              publisher_model_template: model_data['publisherModelTemplate']
+            }
+          )
+        end
+
+        def extract_model_id_from_path(path)
+          path.split('/').last
+        end
+
+        def determine_model_family(model_id)
+          case model_id
+          when /^gemini-2\.\d+/ then 'gemini-2'
+          when /^gemini-1\.\d+/ then 'gemini-1.5'
+          when /^text-embedding/ then 'text-embedding'
+          when /bison/ then 'palm'
+          else 'gemini'
+          end
+        end
+
+        def extract_capabilities(model_data)
+          capabilities = ['streaming']
+          model_name = model_data['name']
+          capabilities << 'function_calling' if model_name.include?('gemini')
+          capabilities.uniq
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/streaming.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI
+      # Streaming methods for the Vertex AI implementation
+      module Streaming
+        def stream_url
+          "projects/#{@config.vertexai_project_id}/locations/#{@config.vertexai_location}/publishers/google/models/#{@model}:streamGenerateContent?alt=sse" # rubocop:disable Layout/LineLength
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # Google Vertex AI implementation
+    class VertexAI < Gemini
+      include VertexAI::Chat
+      include VertexAI::Streaming
+      include VertexAI::Embeddings
+      include VertexAI::Models
+
+      def initialize(config)
+        super
+        @authorizer = nil
+      end
+
+      def api_base
+        "https://#{@config.vertexai_location}-aiplatform.googleapis.com/v1beta1"
+      end
+
+      def headers
+        {
+          'Authorization' => "Bearer #{access_token}"
+        }
+      end
+
+      class << self
+        def configuration_requirements
+          %i[vertexai_project_id vertexai_location]
+        end
+      end
+
+      private
+
+      def access_token
+        return 'test-token' if defined?(VCR) && !VCR.current_cassette.recording?
+
+        initialize_authorizer unless @authorizer
+        @authorizer.fetch_access_token!['access_token']
+      end
+
+      def initialize_authorizer
+        require 'googleauth'
+        @authorizer = ::Google::Auth.get_application_default(
+          scope: [
+            'https://www.googleapis.com/auth/cloud-platform',
+            'https://www.googleapis.com/auth/generative-language.retriever'
+          ]
+        )
+      rescue LoadError
+        raise Error, 'The googleauth gem is required for Vertex AI. Please add it to your Gemfile: gem "googleauth"'
+      end
+    end
+  end
+end

data/lib/ruby_llm/railtie.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # Rails integration for RubyLLM
+  class Railtie < Rails::Railtie
+    initializer 'ruby_llm.inflections' do
+      ActiveSupport::Inflector.inflections(:en) do |inflect|
+        inflect.acronym 'RubyLLM'
+      end
+    end
+
+    initializer 'ruby_llm.active_record' do
+      ActiveSupport.on_load :active_record do
+        model_registry_class = RubyLLM.config.model_registry_class
+
+        if model_registry_class
+          require 'ruby_llm/active_record/acts_as'
+          ::ActiveRecord::Base.include RubyLLM::ActiveRecord::ActsAs
+        else
+          require 'ruby_llm/active_record/acts_as_legacy'
+          ::ActiveRecord::Base.include RubyLLM::ActiveRecord::ActsAsLegacy
+
+          Rails.logger.warn(
+            'RubyLLM: String-based model fields are deprecated and will be removed in RubyLLM 2.0.0. ' \
+            'Please migrate to the DB-backed model registry. ' \
+            "Run 'rails generate ruby_llm:migrate_model_fields' to upgrade."
+          )
+        end
+      end
+    end
+
+    generators do
+      require 'generators/ruby_llm/install_generator'
+      require 'generators/ruby_llm/migrate_model_fields_generator'
+    end
+
+    rake_tasks do
+      load 'tasks/ruby_llm.rake'
+    end
+  end
+end

data/lib/ruby_llm/stream_accumulator.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # Assembles streaming responses from LLMs into complete messages.
+  class StreamAccumulator
+    attr_reader :content, :model_id, :conversation_id, :tool_calls
+
+    def initialize
+      @content = +''
+      @tool_calls = {}
+      @input_tokens = 0
+      @output_tokens = 0
+      @latest_tool_call_id = nil
+    end
+
+    def add(chunk)
+      RubyLLM.logger.debug chunk.inspect if RubyLLM.config.log_stream_debug
+      @model_id ||= chunk.model_id
+      @conversation_id ||= chunk.conversation_id
+
+      if chunk.tool_call?
+        accumulate_tool_calls chunk.tool_calls
+      else
+        @content << (chunk.content || '')
+      end
+
+      count_tokens chunk
+      RubyLLM.logger.debug inspect if RubyLLM.config.log_stream_debug
+    end
+
+    def to_message(response)
+      Message.new(
+        role: :assistant,
+        content: content.empty? ? nil : content,
+        model_id: model_id,
+        conversation_id: conversation_id,
+        tool_calls: tool_calls_from_stream,
+        input_tokens: @input_tokens.positive? ? @input_tokens : nil,
+        output_tokens: @output_tokens.positive? ? @output_tokens : nil,
+        raw: response
+      )
+    end
+
+    private
+
+    def tool_calls_from_stream
+      tool_calls.transform_values do |tc|
+        arguments = if tc.arguments.is_a?(String) && !tc.arguments.empty?
+                      JSON.parse(tc.arguments)
+                    elsif tc.arguments.is_a?(String)
+                      {}
+                    else
+                      tc.arguments
+                    end
+
+        ToolCall.new(
+          id: tc.id,
+          name: tc.name,
+          arguments: arguments
+        )
+      end
+    end
+
+    def accumulate_tool_calls(new_tool_calls)
+      RubyLLM.logger.debug "Accumulating tool calls: #{new_tool_calls}" if RubyLLM.config.log_stream_debug
+      new_tool_calls.each_value do |tool_call|
+        if tool_call.id
+          tool_call_id = tool_call.id.empty? ? SecureRandom.uuid : tool_call.id
+          tool_call_arguments = tool_call.arguments.empty? ? +'' : tool_call.arguments
+          @tool_calls[tool_call.id] = ToolCall.new(
+            id: tool_call_id,
+            name: tool_call.name,
+            arguments: tool_call_arguments
+          )
+          @latest_tool_call_id = tool_call.id
+        else
+          existing = @tool_calls[@latest_tool_call_id]
+          existing.arguments << tool_call.arguments if existing
+        end
+      end
+    end
+
+    def find_tool_call(tool_call_id)
+      if tool_call_id.nil?
+        @tool_calls[@latest_tool_call]
+      else
+        @latest_tool_call_id = tool_call_id
+        @tool_calls[tool_call_id]
+      end
+    end
+
+    def count_tokens(chunk)
+      @input_tokens = chunk.input_tokens if chunk.input_tokens
+      @output_tokens = chunk.output_tokens if chunk.output_tokens
+    end
+  end
+end

data/lib/ruby_llm/streaming.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # Handles streaming responses from AI providers.
+  module Streaming
+    module_function
+
+    def stream_response(connection, payload, additional_headers = {}, &block)
+      accumulator = StreamAccumulator.new
+
+      response = connection.post stream_url, payload do |req|
+        req.headers = additional_headers.merge(req.headers) unless additional_headers.empty?
+        if faraday_1?
+          req.options[:on_data] = handle_stream do |chunk|
+            accumulator.add chunk
+            block.call chunk
+          end
+        else
+          req.options.on_data = handle_stream do |chunk|
+            accumulator.add chunk
+            block.call chunk
+          end
+        end
+      end
+
+      message = accumulator.to_message(response)
+      RubyLLM.logger.debug "Stream completed: #{message.content}"
+      message
+    end
+
+    def handle_stream(&block)
+      to_json_stream do |data|
+        block.call(build_chunk(data)) if data
+      end
+    end
+
+    private
+
+    def faraday_1?
+      Faraday::VERSION.start_with?('1')
+    end
+
+    def to_json_stream(&)
+      buffer = +''
+      parser = EventStreamParser::Parser.new
+
+      create_stream_processor(parser, buffer, &)
+    end
+
+    def create_stream_processor(parser, buffer, &)
+      if faraday_1?
+        legacy_stream_processor(parser, &)
+      else
+        stream_processor(parser, buffer, &)
+      end
+    end
+
+    def process_stream_chunk(chunk, parser, env, &)
+      RubyLLM.logger.debug "Received chunk: #{chunk}" if RubyLLM.config.log_stream_debug
+
+      if error_chunk?(chunk)
+        handle_error_chunk(chunk, env)
+      else
+        yield handle_sse(chunk, parser, env, &)
+      end
+    end
+
+    def legacy_stream_processor(parser, &block)
+      proc do |chunk, _size|
+        process_stream_chunk(chunk, parser, nil, &block)
+      end
+    end
+
+    def stream_processor(parser, buffer, &block)
+      proc do |chunk, _bytes, env|
+        if env&.status == 200
+          process_stream_chunk(chunk, parser, env, &block)
+        else
+          handle_failed_response(chunk, buffer, env)
+        end
+      end
+    end
+
+    def error_chunk?(chunk)
+      chunk.start_with?('event: error')
+    end
+
+    def handle_error_chunk(chunk, env)
+      error_data = chunk.split("\n")[1].delete_prefix('data: ')
+      status, _message = parse_streaming_error(error_data)
+      parsed_data = JSON.parse(error_data)
+
+      error_response = if faraday_1?
+                         Struct.new(:body, :status).new(parsed_data, status)
+                       else
+                         env.merge(body: parsed_data, status: status)
+                       end
+
+      ErrorMiddleware.parse_error(provider: self, response: error_response)
+    rescue JSON::ParserError => e
+      RubyLLM.logger.debug "Failed to parse error chunk: #{e.message}"
+    end
+
+    def handle_failed_response(chunk, buffer, env)
+      buffer << chunk
+      error_data = JSON.parse(buffer)
+      error_response = env.merge(body: error_data)
+      ErrorMiddleware.parse_error(provider: self, response: error_response)
+    rescue JSON::ParserError
+      RubyLLM.logger.debug "Accumulating error chunk: #{chunk}"
+    end
+
+    def handle_sse(chunk, parser, env, &block)
+      parser.feed(chunk) do |type, data|
+        case type.to_sym
+        when :error
+          handle_error_event(data, env)
+        else
+          yield handle_data(data, &block) unless data == '[DONE]'
+        end
+      end
+    end
+
+    def handle_data(data)
+      JSON.parse(data)
+    rescue JSON::ParserError => e
+      RubyLLM.logger.debug "Failed to parse data chunk: #{e.message}"
+    end
+
+    def handle_error_event(data, env)
+      status, _message = parse_streaming_error(data)
+      parsed_data = JSON.parse(data)
+
+      error_response = if faraday_1?
+                         Struct.new(:body, :status).new(parsed_data, status)
+                       else
+                         env.merge(body: parsed_data, status: status)
+                       end
+
+      ErrorMiddleware.parse_error(provider: self, response: error_response)
+    rescue JSON::ParserError => e
+      RubyLLM.logger.debug "Failed to parse error event: #{e.message}"
+    end
+
+    def parse_streaming_error(data)
+      error_data = JSON.parse(data)
+      [500, error_data['message'] || 'Unknown streaming error']
+    rescue JSON::ParserError => e
+      RubyLLM.logger.debug "Failed to parse streaming error: #{e.message}"
+      [500, "Failed to parse error: #{data}"]
+    end
+  end
+end

data/lib/ruby_llm/tool.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # Parameter definition for Tool methods.
+  class Parameter
+    attr_reader :name, :type, :description, :required
+
+    def initialize(name, type: 'string', desc: nil, required: true)
+      @name = name
+      @type = type
+      @description = desc
+      @required = required
+    end
+  end
+
+  # Base class for creating tools that AI models can use
+  class Tool
+    # Stops conversation continuation after tool execution
+    class Halt
+      attr_reader :content
+
+      def initialize(content)
+        @content = content
+      end
+
+      def to_s
+        @content.to_s
+      end
+    end
+
+    class << self
+      def description(text = nil)
+        return @description unless text
+
+        @description = text
+      end
+
+      def param(name, **options)
+        parameters[name] = Parameter.new(name, **options)
+      end
+
+      def parameters
+        @parameters ||= {}
+      end
+    end
+
+    def name
+      klass_name = self.class.name
+      normalized = klass_name.to_s.dup.force_encoding('UTF-8').unicode_normalize(:nfkd)
+      normalized.encode('ASCII', replace: '')
+                .gsub(/[^a-zA-Z0-9_-]/, '-')
+                .gsub(/([A-Z]+)([A-Z][a-z])/, '\1_\2')
+                .gsub(/([a-z\d])([A-Z])/, '\1_\2')
+                .downcase
+                .delete_suffix('_tool')
+    end
+
+    def description
+      self.class.description
+    end
+
+    def parameters
+      self.class.parameters
+    end
+
+    def call(args)
+      RubyLLM.logger.debug "Tool #{name} called with: #{args.inspect}"
+      result = execute(**args.transform_keys(&:to_sym))
+      RubyLLM.logger.debug "Tool #{name} returned: #{result.inspect}"
+      result
+    end
+
+    def execute(...)
+      raise NotImplementedError, 'Subclasses must implement #execute'
+    end
+
+    protected
+
+    def halt(message)
+      Halt.new(message)
+    end
+  end
+end

data/lib/ruby_llm/tool_call.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # Represents a function call from an AI model to a Tool.
+  class ToolCall
+    attr_reader :id, :name, :arguments
+
+    def initialize(id:, name:, arguments: {})
+      @id = id
+      @name = name
+      @arguments = arguments
+    end
+
+    def to_h
+      {
+        id: @id,
+        name: @name,
+        arguments: @arguments
+      }
+    end
+  end
+end

data/lib/ruby_llm/utils.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # Provides utility functions for data manipulation within the RubyLLM library
+  module Utils
+    module_function
+
+    def hash_get(hash, key)
+      hash[key.to_sym] || hash[key.to_s]
+    end
+
+    def to_safe_array(item)
+      case item
+      when Array
+        item
+      when Hash
+        [item]
+      else
+        Array(item)
+      end
+    end
+
+    def to_time(value)
+      return unless value
+
+      value.is_a?(Time) ? value : Time.parse(value.to_s)
+    end
+
+    def to_date(value)
+      return unless value
+
+      value.is_a?(Date) ? value : Date.parse(value.to_s)
+    end
+
+    def deep_merge(original, overrides)
+      original.merge(overrides) do |_key, original_value, overrides_value|
+        if original_value.is_a?(Hash) && overrides_value.is_a?(Hash)
+          deep_merge(original_value, overrides_value)
+        else
+          overrides_value
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  VERSION = '1.6.4'
+end