inst_llm 0.1.0 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 33b83be8e20977ac085278558163be88d0acc3b35436e2360bb7e8d0102cb907
-  data.tar.gz: a47ac4068509bef5454dc37904ff2752203194fd3d855a06d126c2d0d9f425e8
+  metadata.gz: 105c04ef5d12358f4ff663bc8c1278e731c0664a40ac016e29ec2defc999b6a3
+  data.tar.gz: 4873e34b6e2677822a6f8cabb3cac06367a6062d04301810856b33b751b260fb
 SHA512:
-  metadata.gz: 313914f2b440882aab196834e08f9ca9eed991380fd17174c74a97656a9a00eb34809a4e2b03ad87201a1c7521698288dd57b71cfacc32ee6b5b265b0737033e
-  data.tar.gz: 26467fb44454136ccb202219d88991c2c3f8a5a304cbda5707929ce98f059fdb231400b1f356cc2037df62f8d85a3eea7d1d79220575bc4679ede4a4141b3914
+  metadata.gz: 6d7c85ea828972e507b8c680e340fb5ad76b5784277a441ecad7fb1ee215a2bc839c426daf76207cacf63b1ca9fe54ff4d7bf1bf3b9d1f78ef413a28e64ddc33
+  data.tar.gz: 7ffed97282a142a6adb4227334ddc333825b2de7225ec61b2f8d1796681a68b693be4799f73aee09cc4f5b482c2c8722a7633f9973d47bb42a8546880d75b997

data/lib/inst_llm/client.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+require "aws-sdk-bedrockruntime"
+require "json"
+
+require_relative "parameter/all"
+require_relative "response/all"
+
+module InstLLM
+  class Client
+    MODELS = {
+      "anthropic.claude-3-sonnet-20240229-v1:0": { format: :claude, provider: :bedrock, type: :chat },
+      "anthropic.claude-3-haiku-20240307-v1:0": { format: :claude, provider: :bedrock, type: :chat },
+
+      "mistral.mistral-7b-instruct-v0:2": { format: :mistral, provider: :bedrock, type: :chat },
+      "mistral.mixtral-8x7b-instruct-v0:1": { format: :mistral, provider: :bedrock, type: :chat },
+      "mistral.mistral-large-2402-v1:0": { format: :mistral, provider: :bedrock, type: :chat },
+
+      "cohere.embed-english-v3": { format: :cohere_embed, provider: :bedrock, type: :embedding },
+      "cohere.embed-multilingual-v3": { format: :cohere_embed, provider: :bedrock, type: :embedding },
+    }.freeze
+
+    def initialize(model, **options)
+      model = model.to_sym
+      raise UnknownArgumentError unless MODELS.key?(model)
+
+      @model = model
+      @options = options
+    end
+
+    def chat(messages, **options)
+      model = (options[:model] || options[:model_id] || @model).to_sym
+      raise ArgumentError, "Model #{model} is not a chat model" unless chat_model?(model)
+
+      response_factory(model, call(model, messages, **options))
+    end
+
+    def embedding(message, **options)
+      model = (options[:model] || options[:model_id] || @model).to_sym
+      raise ArgumentError, "Model #{model} is not an embedding model" unless embedding_model?(model)
+
+      embedding_response_factory(model, call(model, message, **options))
+    end
+
+    private
+
+    def call(model, messages, **options)
+      params = params_factory(model, messages, **options)
+
+      begin
+        res = client.invoke_model(**params)
+      rescue => error
+        raise map_error_type(error)
+      end
+
+      JSON.parse(res.body.read)
+    end
+
+    def chat_model?(model)
+      MODELS[model][:type] == :chat
+    end
+
+    def embedding_model?(model)
+      MODELS[model][:type] == :embedding
+    end
+
+    def client
+      return @client if @client
+
+      case MODELS[@model][:provider]
+      when :bedrock
+        @client = Aws::BedrockRuntime::Client.new(**@options)
+      else
+        raise UnknownArgumentError
+      end
+
+      @client
+    end
+
+    def map_error_type(error)
+      mapped_error_type = nil
+
+      case MODELS[@model][:provider]
+      when :bedrock
+        case error
+        when Aws::BedrockRuntime::Errors::ServiceQuotaExceededException
+          mapped_error_type = ServiceQuotaExceededError
+        when Aws::BedrockRuntime::Errors::ThrottlingException
+          mapped_error_type = ThrottlingError
+        when Aws::BedrockRuntime::Errors::ValidationException
+          if error.message.include?("too long")
+            mapped_error_type = ValidationTooLongError
+          else
+            mapped_error_type = ValidationError
+          end
+        else
+          mapped_error_type = Error
+        end
+      else
+        raise UnknownArgumentError
+      end
+
+      mapped_error_type.new(error.message)
+    end
+
+    def params_factory(model, messages, **options)
+      params_table = {
+        claude: Parameter::ClaudeParameters,
+        cohere_embed: Parameter::CohereEmbedParameters,
+        mistral: Parameter::MistralParameters
+      }
+      params_table[MODELS[model][:format]].new(model: model, messages: messages, **options)
+    end
+
+    def embedding_response_factory(model, response)
+      Response::EmbeddingResponse.send(:"from_#{MODELS[model][:format]}",  model: model, response: response)
+    end
+
+    def response_factory(model, response)
+      Response::ChatResponse.send(:"from_#{MODELS[model][:format]}", model: model, response: response)
+    end
+  end
+end

data/lib/inst_llm/parameter/all.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+require_relative 'claude_parameters'
+require_relative 'cohere_embed_parameters'
+require_relative 'mistral_parameters'

data/lib/inst_llm/parameter/claude_parameters.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module InstLLM
+  module Parameter
+    class ClaudeParameters
+      DEFAULT_OPTIONS = {
+        anthropic_version: "bedrock-2023-05-31",
+        max_tokens: 2000,
+        stop_sequences: nil,
+        temperature: nil,
+        top_k: nil,
+        top_p: nil,
+        system: nil,
+      }.freeze
+
+      def initialize(model:, messages: [], **options)
+        @messages = messages
+        @model = model
+        @options = DEFAULT_OPTIONS.merge(options.slice(*DEFAULT_OPTIONS.keys)).compact
+      end
+
+      def to_hash
+        { model_id: @model, body: { messages: @messages }.merge(@options).to_json }
+      end
+    end
+  end
+end

data/lib/inst_llm/parameter/cohere_embed_parameters.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module InstLLM
+  module Parameter
+    class CohereEmbedParameters
+      DEFAULT_OPTIONS = {
+        input_type: nil,
+        truncate: nil
+      }.freeze
+
+      def initialize(model:, texts: [], **options)
+        @model = model
+        @texts = texts
+        @options = DEFAULT_OPTIONS.merge(options.slice(*DEFAULT_OPTIONS.keys)).compact
+      end
+
+      def to_hash
+        { model_id: @model, body: { texts: @texts }.merge(@options).to_json }
+      end
+    end
+  end
+end

data/lib/inst_llm/parameter/mistral_parameters.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module InstLLM
+  module Parameter
+    class MistralParameters
+      DEFAULT_OPTIONS = {
+        max_tokens: nil,
+        stop: nil,
+        temperature: nil,
+        top_p: nil,
+        top_k: nil
+      }.freeze
+
+      def initialize(model:, messages:, **options)
+        @model = model
+        @messages = messages
+        @options = DEFAULT_OPTIONS.merge(options.slice(*DEFAULT_OPTIONS.keys)).compact
+      end
+
+      def to_hash
+        { model_id: @model, body: { prompt: prompt }.merge(@options).to_json }
+      end
+
+      private
+
+      def prompt
+        system_message = nil
+        prompt = @messages.map do |message|
+          case message[:role].to_sym
+          when :assistant
+            "#{message[:content]}"
+          when :system
+            system_message = message[:content]
+          when :user
+            "[INST] #{message[:content]} [/INST]"
+          else
+            raise UnknownArgumentError
+          end
+        end
+
+        prompt = "<s>" + prompt.join("\n\n")
+
+        if system_message
+          prompt.sub("\[INST\]", "[INST] #{system_message}\n")
+        end
+
+        prompt
+      end
+    end
+  end
+end

data/lib/inst_llm/response/all.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+require_relative "chat_response"
+require_relative "embedding_response"

data/lib/inst_llm/response/chat_response.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+module InstLLM
+  module Response
+    class ChatResponse
+      attr_reader :created, :fingerprint, :stop_reason, :message, :model, :usage
+
+      def initialize(model:, message:, stop_reason:, usage:)
+        @created = Time.now.to_i
+        @fingerprint = SecureRandom.uuid
+        @message = message
+        @model = model
+        @stop_reason = stop_reason
+        @usage = usage
+      end
+
+      class << self
+        def from_claude(model:, response:)
+          new(
+            model: model,
+            message: { role: :assistant, content: response["content"][0]["text"] },
+            stop_reason: response["stop_reason"],
+            usage: {
+              input_tokens: response["usage"]["input_tokens"],
+              output_tokens: response["usage"]["output_tokens"]
+            }
+          )
+        end
+
+        def from_mistral(model:, response:)
+          new(
+            model: model,
+            message: { role: :assistant, content: response["outputs"][0]["text"] },
+            stop_reason: response["outputs"][0]["stop_reason"],
+            usage: {
+              input_tokens: -1,
+              output_tokens: -1
+            }
+          )
+        end
+      end
+    end
+  end
+end

data/lib/inst_llm/response/embedding_response.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module InstLLM
+  module Response
+    class EmbeddingResponse
+      attr_reader :model, :embeddings
+
+      def initialize(model, embeddings)
+        @model = model
+        @embeddings = embeddings
+      end
+
+      class << self
+        def from_cohere_embed(model:, response:)
+          embeddings = response["embeddings"].map.with_index do |embedding, i|
+            { object: "embedding", embedding: embedding, index: i }
+          end
+          new(model, embeddings)
+        end
+      end
+    end
+  end
+end

data/lib/inst_llm/version.rb

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

data/lib/inst_llm.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require_relative "inst_llm/version"
+require_relative "inst_llm/client"
+
+module InstLLM
+  class Error < StandardError; end
+  class UnknownArgumentError < StandardError; end
+  class ServiceQuotaExceededError < StandardError; end
+  class ThrottlingError < StandardError; end
+  class ValidationTooLongError < StandardError; end
+  class ValidationError < StandardError; end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: inst_llm
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
 platform: ruby
 authors:
 - Zach Pendleton
@@ -30,7 +30,17 @@ email:
 executables: []
 extensions: []
 extra_rdoc_files: []
-files: []
+files:
+- lib/inst_llm.rb
+- lib/inst_llm/client.rb
+- lib/inst_llm/parameter/all.rb
+- lib/inst_llm/parameter/claude_parameters.rb
+- lib/inst_llm/parameter/cohere_embed_parameters.rb
+- lib/inst_llm/parameter/mistral_parameters.rb
+- lib/inst_llm/response/all.rb
+- lib/inst_llm/response/chat_response.rb
+- lib/inst_llm/response/embedding_response.rb
+- lib/inst_llm/version.rb
 homepage: https://instructure.com
 licenses:
 - MIT