ollama-ruby 0.0.1 → 0.1.0

data/lib/ollama/client.rb

@@ -1,6 +1,5 @@
 require 'tins/xt/string_camelize'
 require 'tins/annotate'
-require 'excon'
 
 class Ollama::Client
 end
@@ -31,6 +30,8 @@ class Ollama::Client
 
   attr_accessor :output
 
+  attr_reader :base_url
+
   def ssl_verify_peer?
     !!@ssl_verify_peer
   end
@@ -113,11 +114,15 @@ class Ollama::Client
 
   def headers
     {
-      'User-Agent'   => '%s/%s' % [ self.class, Ollama::VERSION ],
+      'User-Agent'   => self.class.user_agent,
       'Content-Type' => 'application/json; charset=utf-8',
     }
   end
 
+  def self.user_agent
+    '%s/%s' % [ self.class, Ollama::VERSION ]
+  end
+
   def excon(url)
     params = {
       connect_timeout: @connect_timeout,

data/lib/ollama/documents/memory_cache.rb

@@ -0,0 +1,44 @@
+class Ollama::Documents::MemoryCache
+  def initialize(prefix:)
+    @prefix = prefix
+    @data   = {}
+  end
+
+  attr_writer :prefix
+
+  def [](key)
+    @data[pre(key)]
+  end
+
+  def []=(key, value)
+    @data[pre(key)] = value
+  end
+
+  def key?(key)
+    @data.key?(pre(key))
+  end
+
+  def delete(key)
+    @data.delete(pre(key))
+  end
+
+  def size
+    @data.size
+  end
+
+  def clear
+    @data.clear
+    self
+  end
+
+  def each(&block)
+    @data.select { |key,| key.start_with?(@prefix) }.each(&block)
+  end
+  include Enumerable
+
+  private
+
+  def pre(key)
+    [ @prefix, key ].join
+  end
+end

data/lib/ollama/documents/redis_cache.rb

@@ -0,0 +1,57 @@
+require 'redis'
+
+class Ollama::Documents::RedisCache
+  def initialize(prefix:, url: ENV['REDIS_URL'])
+    url or raise ArgumentError, 'require redis url'
+    @prefix, @url = prefix, url
+  end
+
+  attr_writer :prefix
+
+  def redis
+    @redis ||= Redis.new(url: @url)
+  end
+
+  def [](key)
+    JSON(redis.get(pre(key)), object_class: Ollama::Documents::Record)
+  end
+
+  def []=(key, value)
+    redis.set(pre(key), JSON(value))
+  end
+
+  def key?(key)
+    !!redis.exists?(pre(key))
+  end
+
+  def delete(key)
+    redis.del(pre(key)) == 1
+  end
+
+  def size
+    s = 0
+    redis.scan_each(match: "#@prefix*") { |key| s += 1 }
+    s
+  end
+
+  def clear
+    redis.scan_each(match: "#@prefix*") { |key| redis.del(key) }
+    self
+  end
+
+  def each(&block)
+    redis.scan_each(match: "#@prefix*") { |key| block.(key, self[unpre(key)]) }
+    self
+  end
+  include Enumerable
+
+  private
+
+  def pre(key)
+    [ @prefix, key ].join
+  end
+
+  def unpre(key)
+    key.sub(/\A#@prefix/, '')
+  end
+end

data/lib/ollama/documents/splitters/character.rb

@@ -0,0 +1,70 @@
+module Ollama::Documents::Splitters
+  class Character
+    DEFAULT_SEPARATOR = /(?:\r?\n){2,}/
+
+    def initialize(separator: DEFAULT_SEPARATOR, include_separator: false, chunk_size: 4096)
+      @separator, @include_separator, @chunk_size = separator, include_separator, chunk_size
+      if include_separator
+        @separator = Regexp.new("(#@separator)")
+      end
+    end
+
+    def split(text)
+      texts = []
+      text.split(@separator) do |t|
+        if @include_separator && t =~ @separator
+          texts.last&.concat t
+        else
+          texts.push(t)
+        end
+      end
+      result = []
+      current_text = +''
+      texts.each do |t|
+        if current_text.size + t.size < @chunk_size
+          current_text += t
+        else
+          current_text.empty? or result << current_text
+          current_text = t
+        end
+      end
+      current_text.empty? or result << current_text
+      result
+    end
+  end
+
+  class RecursiveCharacter
+    DEFAULT_SEPARATORS = [
+      /(?:\r?\n){2,}/,
+      /\r?\n/,
+      /\b/,
+      //,
+    ].freeze
+
+    def initialize(separators: DEFAULT_SEPARATORS, include_separator: false, chunk_size: 4096)
+      separators.empty? and
+        raise ArgumentError, "non-empty array of separators required"
+      @separators, @include_separator, @chunk_size =
+        separators, include_separator, chunk_size
+    end
+
+    def split(text, separators: @separators)
+      separators.empty? and return [ text ]
+      separators = separators.dup
+      separator = separators.shift
+      texts = Character.new(
+        separator:,
+        include_separator: @include_separator,
+        chunk_size: @chunk_size
+      ).split(text)
+      texts.count == 0 and return [ text ]
+      texts.inject([]) do |r, t|
+        if t.size > @chunk_size
+          r.concat(split(t, separators:))
+        else
+          r.concat([ t ])
+        end
+      end
+    end
+  end
+end

data/lib/ollama/documents/splitters/semantic.rb

@@ -0,0 +1,90 @@
+module Ollama::Documents::Splitters
+  class Semantic
+    include Ollama::Utils::Math
+
+    DEFAULT_SEPARATOR = /[.!?]\s*(?:\b|\z)/
+
+    def initialize(ollama:, model:, model_options: nil, separator: DEFAULT_SEPARATOR, chunk_size: 4096)
+      @ollama, @model, @model_options, @separator, @chunk_size =
+        ollama, model, model_options, separator, chunk_size
+    end
+
+    def split(text, batch_size: 100, breakpoint: :percentile, **opts)
+      sentences  = Ollama::Documents::Splitters::Character.new(
+        separator: @separator,
+        include_separator: true,
+        chunk_size: 1,
+      ).split(text)
+      embeddings = sentences.with_infobar(label: 'Split').each_slice(batch_size).inject([]) do |e, batch|
+        e.concat sentence_embeddings(batch)
+        infobar.progress by: batch.size
+        e
+      end
+      infobar.newline
+      distances = embeddings.each_cons(2).map do |a, b|
+        1.0 - cosine_similarity(a:, b:)
+      end
+      max_distance = calculate_breakpoint_threshold(breakpoint, distances, **opts)
+      gaps = distances.each_with_index.select do |d, i|
+        d > max_distance
+      end.transpose.last
+      gaps or return sentences
+      if gaps.last < distances.size
+        gaps << distances.size
+      end
+      if gaps.last < sentences.size - 1
+        gaps << sentences.size - 1
+      end
+      result = []
+      sg = 0
+      current_text = +''
+      gaps.each do |g|
+        sg.upto(g) do |i|
+          sentence = sentences[i]
+          if current_text.size + sentence.size < @chunk_size
+            current_text += sentence
+          else
+            current_text.empty? or result << current_text
+            current_text = sentence
+          end
+        end
+        unless current_text.empty?
+          result << current_text
+          current_text = +''
+        end
+        sg = g.succ
+      end
+      current_text.empty? or result << current_text
+      result
+    end
+
+    private
+
+    def calculate_breakpoint_threshold(breakpoint_method, distances, **opts)
+      sequence = MoreMath::Sequence.new(distances)
+      case breakpoint_method
+      when :percentile
+        percentile = opts.fetch(:percentile, 95)
+        sequence.percentile(percentile)
+      when :standard_deviation
+        percentage = opts.fetch(:percentage, 100)
+        (
+          sequence.mean + sequence.standard_deviation * (percentage / 100.0)
+        ).clamp(0, sequence.max)
+      when :interquartile
+        percentage = opts.fetch(:percentage, 100)
+        iqr = sequence.interquartile_range
+        max = sequence.max
+        (sequence.mean + iqr * (percentage / 100.0)).clamp(0, max)
+      else
+        raise ArgumentError, "invalid breakpoint method #{breakpoint_method}"
+      end
+    end
+
+    def sentence_embeddings(input)
+      @ollama.embed(model: @model, input:, options: @model_options).embeddings.map! {
+        Numo::NArray[*_1]
+      }
+    end
+  end
+end

data/lib/ollama/documents.rb

@@ -0,0 +1,172 @@
+require 'numo/narray'
+require 'digest'
+
+class Ollama::Documents
+end
+require 'ollama/documents/memory_cache'
+require 'ollama/documents/redis_cache'
+module Ollama::Documents::Splitters
+end
+require 'ollama/documents/splitters/character'
+require 'ollama/documents/splitters/semantic'
+
+class Ollama::Documents
+  include Ollama::Utils::Math
+
+  class Record < JSON::GenericObject
+    def to_s
+      my_tags = Ollama::Utils::Tags.new(tags)
+      my_tags.empty? or my_tags = " #{my_tags}"
+      "#<#{self.class} #{text.inspect}#{my_tags} #{similarity || 'n/a'}>"
+    end
+
+    def ==(other)
+      text == other.text
+    end
+
+    alias inspect to_s
+  end
+
+  def initialize(ollama:, model:, model_options: nil, collection: :default, cache: MemoryCache, redis_url: nil)
+    @ollama, @model, @model_options, @collection = ollama, model, model_options, collection
+    @cache, @redis_url = connect_cache(cache), redis_url
+  end
+
+  attr_reader :ollama, :model, :collection
+
+  def collection=(new_collection)
+    @collection = new_collection
+    @cache.prefix = prefix
+  end
+
+  def add(inputs, batch_size: 10, source: nil, tags: [])
+    inputs = Array(inputs)
+    tags   = Ollama::Utils::Tags.new(tags)
+    source and tags.add File.basename(source)
+    inputs.map! { |i|
+      text = i.respond_to?(:read) ? i.read : i.to_s
+      text
+    }
+    inputs.reject! { |i| exist?(i) }
+    inputs.empty? and return self
+    batches = inputs.each_slice(batch_size).
+      with_infobar(
+        label: "Add #{tags}",
+        total: inputs.size
+      )
+    batches.each do |batch|
+      embeddings = fetch_embeddings(model:, options: @model_options, input: batch)
+      batch.zip(embeddings) do |text, embedding|
+        norm       = norm(embedding)
+        self[text] = Record[text:, embedding:, norm:, source:, tags: tags.to_a]
+      end
+      infobar.progress by: batch.size
+    end
+    infobar.newline
+    self
+  end
+  alias << add
+
+  def [](text)
+    @cache[key(text)]
+  end
+
+  def []=(text, record)
+    @cache[key(text)] = record
+  end
+
+  def exist?(text)
+    @cache.key?(key(text))
+  end
+
+  def delete(text)
+    @cache.delete(key(text))
+  end
+
+  def size
+    @cache.size
+  end
+
+  def clear
+    @cache.clear
+  end
+
+  def find(string, tags: nil, prompt: nil)
+    needle      = convert_to_vector(string, prompt:)
+    needle_norm = norm(needle)
+    records = @cache
+    if tags
+      tags = Ollama::Utils::Tags.new(tags)
+      records = records.select { |_key, record| (tags & record.tags).size >= 1 }
+    end
+    records = records.sort_by { |key, record|
+      record.key        = key
+      record.similarity = cosine_similarity(
+        a: needle,
+        b: record.embedding,
+        a_norm: needle_norm,
+        b_norm: record.norm,
+      )
+    }
+    records.transpose.last&.reverse.to_a
+  end
+
+  def collections
+    case @cache
+    when MemoryCache
+      [ @collection ]
+    when RedisCache
+      prefix = '%s-' % self.class
+      Documents::RedisCache.new(prefix:, url: @redis_url).map { _1[/#{prefix}(.*)-/, 1] }.uniq
+    else
+      []
+    end
+  end
+
+  def tags
+    @cache.inject(Ollama::Utils::Tags.new) { |t, (_, record)| t.merge(record.tags) }
+  end
+
+  private
+
+  def connect_cache(cache_class)
+    cache = nil
+    if cache_class == RedisCache
+      begin
+        cache = cache_class.new(prefix:)
+        cache.size
+      rescue Redis::CannotConnectError
+        STDERR.puts(
+          "Cannot connect to redis URL #{@redis_url.inspect}, "\
+          "falling back to MemoryCache."
+        )
+      end
+    end
+  ensure
+    cache ||= MemoryCache.new(prefix:)
+    return cache
+  end
+
+  def convert_to_vector(input, prompt: nil)
+    if prompt
+      input = prompt % input
+    end
+    if input.is_a?(String)
+      Numo::NArray[*fetch_embeddings(model:, input:).first]
+    else
+      super(input)
+    end
+  end
+
+  def fetch_embeddings(model:, input:, options: nil)
+    @ollama.embed(model:, input:, options:).embeddings
+  end
+
+  def prefix
+    '%s-%s-' % [ self.class, @collection ]
+  end
+
+  def key(input)
+    Digest::SHA256.hexdigest(input)
+  end
+end

data/lib/ollama/handlers/progress.rb

@@ -1,4 +1,5 @@
 require 'infobar'
+require 'tins/unit'
 
 class Ollama::Handlers::Progress
   include Ollama::Handlers::Concern
@@ -13,7 +14,9 @@ class Ollama::Handlers::Progress
 
   def call(response)
     infobar.display.output = @output
-    status = response.status
+    if status = response.status
+      infobar.label = status
+    end
     if response.total && response.completed
       if !@last_status or @last_status != status
         @last_status and infobar.newline
@@ -24,13 +27,23 @@ class Ollama::Handlers::Progress
       end
       infobar.counter.progress(by: response.completed - @current)
       @current = response.completed
+      infobar.update(
+        message: message(response.completed, response.total),
+        force: true
+      )
     end
-    if status
-      infobar.label = status
-      infobar.update(message: '%l %c/%t in %te, ETA %e @%E', force: true)
-    elsif error = response.error
+    if error = response.error
       infobar.puts bold { "Error: " } + red { error }
     end
     self
   end
+
+  private
+
+  def message(current, total)
+    progress = '%s/%s' % [ current, total ].map {
+      Tins::Unit.format(_1, format: '%.2f %U')
+    }
+    '%l ' + progress + ' in %te, ETA %e @%E'
+  end
 end

data/lib/ollama/image.rb

@@ -5,26 +5,35 @@ class Ollama::Image
     @data = data
   end
 
+  attr_accessor :path
+
   class << self
-    def for_base64(data)
-      new(data)
+    def for_base64(data, path: nil)
+      obj = new(data)
+      obj.path = path
+      obj
     end
 
-    def for_string(string)
-      for_base64(Base64.encode64(string))
+    def for_string(string, path: nil)
+      for_base64(Base64.encode64(string), path:)
     end
 
-    def for_io(io)
-      for_string(io.read)
+    def for_io(io, path: nil)
+      path ||= io.path
+      for_string(io.read, path:)
     end
 
     def for_filename(path)
-      File.open(path, 'rb') { |io| for_io(io) }
+      File.open(path, 'rb') { |io| for_io(io, path:) }
     end
 
     private :new
   end
 
+  def ==(other)
+    @data == other..data
+  end
+
   def to_s
     @data
   end

data/lib/ollama/utils/chooser.rb

@@ -0,0 +1,30 @@
+require 'amatch'
+require 'search_ui'
+
+module Ollama::Utils::Chooser
+  include SearchUI
+
+  module_function
+
+  def choose(entries)
+    entry = Search.new(
+      match: -> answer {
+        matcher = Amatch::PairDistance.new(answer.downcase)
+        matches = entries.map { |n| [ n, -matcher.similar(n.to_s.downcase) ] }.
+          select { |_, s| s < 0 }.sort_by(&:last).map(&:first)
+        matches.empty? and matches = entries
+        matches.first(Tins::Terminal.lines - 1)
+      },
+      query: -> _answer, matches, selector {
+        matches.each_with_index.map { |m, i|
+          i == selector ? "#{Term::ANSIColor.blue{?⮕}} #{m.on_blue}" : "  #{m}"
+        } * ?\n
+      },
+      found: -> _answer, matches, selector {
+        matches[selector]
+      },
+      output: STDOUT
+    ).start
+    return entry if entry
+  end
+end

data/lib/ollama/utils/colorize_texts.rb

@@ -0,0 +1,42 @@
+class Ollama::Utils::ColorizeTexts
+  include Math
+  include Term::ANSIColor
+  include Ollama::Utils::Width
+
+  def initialize(*texts)
+    texts = texts.map(&:to_a)
+    @texts = Array(texts.flatten)
+  end
+
+  def to_s
+    result = +''
+    @texts.each_with_index do |t, i|
+      color = colors[(t.hash ^ i.hash) % colors.size]
+      wrap(t, percentage: 90).each_line { |l|
+        result << on_color(color) { color(text_color(color)) { l } }
+      }
+      result << "\n##{bold{t.size.to_s}} \n\n"
+    end
+    result
+  end
+
+  private
+
+  def text_color(color)
+    color = Term::ANSIColor::Attribute[color]
+    [
+      Attribute.nearest_rgb_color('#000'),
+      Attribute.nearest_rgb_color('#fff'),
+    ].max_by { |t| t.distance_to(color) }
+  end
+
+  def colors
+    @colors ||= (0..255).map { |i|
+      [
+        128 + 128 * sin(PI * i / 32.0),
+        128 + 128 * sin(PI * i / 64.0),
+        128 + 128 * sin(PI * i / 128.0),
+      ].map { _1.clamp(0, 255).round }
+    }
+  end
+end