mistral 0.1.0

data/lib/mistral/client.rb

@@ -0,0 +1,229 @@
+# frozen_string_literal: true
+
+require 'mistral/client_base'
+
+module Mistral
+  # Synchronous wrapper around the async client
+  class Client < ClientBase
+    def initialize(
+      api_key: ENV['MISTRAL_API_KEY'],
+      endpoint: ENDPOINT,
+      max_retries: 5,
+      timeout: 120
+    )
+      super(endpoint: endpoint, api_key: api_key, max_retries: max_retries, timeout: timeout)
+
+      @client = HTTP.persistent(ENDPOINT)
+                    .follow
+                    .timeout(timeout)
+                    .use(:line_iterable_body)
+                    .headers('Accept' => 'application/json',
+                      'User-Agent' => "mistral-client-ruby/#{VERSION}",
+                      'Authorization' => "Bearer #{@api_key}",
+                      'Content-Type' => 'application/json'
+                    )
+    end
+
+    # A chat endpoint that returns a single response.
+    #
+    # @param messages [Array<ChatMessage>] An array of messages to chat with, e.g.
+    #   [{role: 'user', content: 'What is the best French cheese?'}]
+    # @param model [String] The name of the model to chat with, e.g. mistral-tiny
+    # @param tools [Array<Hash>] A list of tools to use.
+    # @param temperature [Float] The temperature to use for sampling, e.g. 0.5.
+    # @param max_tokens [Integer] The maximum number of tokens to generate, e.g. 100.
+    # @param top_p [Float] The cumulative probability of tokens to generate, e.g. 0.9.
+    # @param random_seed [Integer] The random seed to use for sampling, e.g. 42.
+    # @param safe_mode [Boolean] Deprecated, use safe_prompt instead.
+    # @param safe_prompt [Boolean] Whether to use safe prompt, e.g. true.
+    # @param tool_choice [String, ToolChoice] The tool choice.
+    # @param response_format [Hash<String, String>, ResponseFormat] The response format.
+    # @return [ChatCompletionResponse] A response object containing the generated text.
+    #
+    def chat(
+      messages:,
+      model: nil,
+      tools: nil,
+      temperature: nil,
+      max_tokens: nil,
+      top_p: nil,
+      random_seed: nil,
+      safe_mode: false,
+      safe_prompt: false,
+      tool_choice: nil,
+      response_format: nil
+    )
+      request = make_chat_request(
+        messages: messages,
+        model: model,
+        tools: tools,
+        temperature: temperature,
+        max_tokens: max_tokens,
+        top_p: top_p,
+        random_seed: random_seed,
+        stream: false,
+        safe_prompt: safe_mode || safe_prompt,
+        tool_choice: tool_choice,
+        response_format: response_format
+      )
+
+      single_response = request('post', 'v1/chat/completions', json: request)
+
+      single_response.each do |response|
+        return ChatCompletionResponse.new(response)
+      end
+
+      raise Mistral::Error.new(message: 'No response received')
+    end
+
+    # A chat endpoint that streams responses.
+    #
+    # @param messages [Array<Any>] An array of messages to chat with, e.g.
+    #   [{role: 'user', content: 'What is the best French cheese?'}]
+    # @param model [String] The name of the model to chat with, e.g. mistral-tiny
+    # @param tools [Array<Hash>] A list of tools to use.
+    # @param temperature [Float] The temperature to use for sampling, e.g. 0.5.
+    # @param max_tokens [Integer] The maximum number of tokens to generate, e.g. 100.
+    # @param top_p [Float] The cumulative probability of tokens to generate, e.g. 0.9.
+    # @param random_seed [Integer] The random seed to use for sampling, e.g. 42.
+    # @param safe_mode [Boolean] Deprecated, use safe_prompt instead.
+    # @param safe_prompt [Boolean] Whether to use safe prompt, e.g. true.
+    # @param tool_choice [String, ToolChoice] The tool choice.
+    # @param response_format [Hash<String, String>, ResponseFormat] The response format.
+    # @return [Enumerator<ChatCompletionStreamResponse>] A generator that yields ChatCompletionStreamResponse objects.
+    #
+    def chat_stream(
+      messages:,
+      model: nil,
+      tools: nil,
+      temperature: nil,
+      max_tokens: nil,
+      top_p: nil,
+      random_seed: nil,
+      safe_mode: false,
+      safe_prompt: false,
+      tool_choice: nil,
+      response_format: nil
+    )
+      request = make_chat_request(
+        messages: messages,
+        model: model,
+        tools: tools,
+        temperature: temperature,
+        max_tokens: max_tokens,
+        top_p: top_p,
+        random_seed: random_seed,
+        stream: true,
+        safe_prompt: safe_mode || safe_prompt,
+        tool_choice: tool_choice,
+        response_format: response_format
+      )
+
+      Enumerator.new do |yielder|
+        request('post', 'v1/chat/completions', json: request, stream: true).each do |json_response|
+          yielder << ChatCompletionStreamResponse.new(**json_response)
+        end
+      end
+    end
+
+    # An embeddings endpoint that returns embeddings for a single, or batch of inputs
+    #
+    # @param model [String] The embedding model to use, e.g. mistral-embed
+    # @param input [String, Array<String>] The input to embed, e.g. ['What is the best French cheese?']
+    #
+    # @return [EmbeddingResponse] A response object containing the embeddings.
+    #
+    def embeddings(model:, input:)
+      request = { model: model, input: input }
+      singleton_response = request('post', 'v1/embeddings', json: request)
+
+      singleton_response.each do |response|
+        return EmbeddingResponse.new(response)
+      end
+
+      raise Mistral::Error.new(message: 'No response received')
+    end
+
+    # Returns a list of the available models
+    #
+    # @return [ModelList] A response object containing the list of models.
+    #
+    def list_models
+      singleton_response = request('get', 'v1/models')
+
+      singleton_response.each do |response|
+        return ModelList.new(response)
+      end
+
+      raise Mistral::Error.new(message: 'No response received')
+    end
+
+    private
+
+    def request(method, path, json: nil, stream: false, attempt: 1)
+      url = File.join(@endpoint, path)
+      headers = {}
+      headers['Accept'] = 'text/event-stream' if stream
+
+      @logger.debug("Sending request: #{method.upcase} #{url} #{json}")
+
+      Enumerator.new do |yielder|
+        response = @client.headers(headers).request(method.downcase.to_sym, url, json: json)
+        check_response_status_codes(response)
+
+        if stream
+          response.body.each_line do |line|
+            processed_line = process_line(line)
+            next if processed_line.nil?
+
+            yielder << processed_line
+          end
+        else
+          yielder << check_response(response)
+        end
+      rescue HTTP::ConnectionError => e
+        raise Mistral::ConnectionError, e.message
+      rescue HTTP::RequestError => e
+        raise Mistral::Error, "Unexpected exception (#{e.class}): #{e.message}"
+      rescue JSON::ParserError
+        raise Mistral::APIError.from_response(response, message: "Failed to decode json body: #{response.body}")
+      rescue Mistral::APIStatusError => e
+        attempt += 1
+
+        raise Mistral::APIStatusError.from_response(response, message: e.message) if attempt > @max_retries
+
+        backoff = 2.0**attempt # exponential backoff
+        sleep(backoff)
+
+        # Retry and yield the response
+        request(method, path, json: json, stream: stream, attempt: attempt).each do |r|
+          yielder << r
+        end
+      end
+    end
+
+    def check_response(response)
+      check_response_status_codes(response)
+
+      json_response = JSON.parse(response.body.to_s)
+
+      if !json_response.key?('object')
+        raise Mistral::Error, "Unexpected response: #{json_response}"
+      elsif json_response['object'] == 'error' # has errors
+        raise Mistral::APIError.from_response(response, message: json_response['message'])
+      end
+
+      json_response
+    end
+
+    def check_response_status_codes(response)
+      if RETRY_STATUS_CODES.include?(response.code)
+        raise APIStatusError.from_response(response, message: "Status: #{response.code}. Message: #{response.body}")
+      elsif response.code >= 400 && response.code < 500
+        raise APIError.from_response(response, message: "Status: #{response.code}. Message: #{response.body}")
+      elsif response.code >= 500
+        raise Mistral::Error.new(message: "Status: #{response.code}. Message: #{response.body}")
+      end
+    end
+  end
+end

data/lib/mistral/client_base.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module Mistral
+  class ClientBase
+    attr_reader :endpoint, :api_key, :max_retries, :timeout
+
+    def initialize(endpoint:, api_key: nil, max_retries: 5, timeout: 120)
+      @endpoint = endpoint
+      @api_key = api_key
+      @max_retries = max_retries
+      @timeout = timeout
+
+      @logger = config_logger
+
+      # For azure endpoints, we default to the mistral model
+      @default_model = 'mistral' if endpoint.include?('inference.azure.com')
+    end
+
+    protected
+
+    def parse_tools(tools)
+      parsed_tools = []
+
+      tools.each do |tool|
+        next unless tool['type'] == 'function'
+
+        parsed_function = {}
+        parsed_function['type'] = tool['type']
+        parsed_function['function'] = if tool['function'].is_a?(Function)
+                                        tool['function'].to_h
+                                      else
+                                        tool['function']
+                                      end
+
+        parsed_tools << parsed_function
+      end
+
+      parsed_tools
+    end
+
+    def parse_tool_choice(tool_choice)
+      tool_choice.is_a?(ToolChoice) ? tool_choice.to_s : tool_choice
+    end
+
+    def parse_response_format(response_format)
+      if response_format.is_a?(ResponseFormat)
+        response_format.to_h
+      else
+        response_format
+      end
+    end
+
+    def parse_messages(messages)
+      parsed_messages = []
+
+      messages.each do |message|
+        parsed_messages << if message.is_a?(ChatMessage)
+                             message.to_h
+                           else
+                             message
+                           end
+      end
+
+      parsed_messages
+    end
+
+    def make_chat_request(
+      messages:,
+      model: nil,
+      tools: nil,
+      temperature: nil,
+      max_tokens: nil,
+      top_p: nil,
+      random_seed: nil,
+      stream: nil,
+      safe_prompt: false,
+      tool_choice: nil,
+      response_format: nil
+    )
+      request_data = {
+        messages: parse_messages(messages),
+        safe_prompt: safe_prompt
+      }
+
+      request_data[:model] = if model.nil?
+                               raise Mistral::Error.new(message: 'model must be provided') if @default_model.nil?
+
+                               @default_model
+                             else
+                               model
+                             end
+
+      request_data[:tools] = parse_tools(tools) unless tools.nil?
+      request_data[:temperature] = temperature unless temperature.nil?
+      request_data[:max_tokens] = max_tokens unless max_tokens.nil?
+      request_data[:top_p] = top_p unless top_p.nil?
+      request_data[:random_seed] = random_seed unless random_seed.nil?
+      request_data[:stream] = stream unless stream.nil?
+      request_data[:tool_choice] = parse_tool_choice(tool_choice) unless tool_choice.nil?
+      request_data[:response_format] = parse_response_format(response_format) unless response_format.nil?
+
+      @logger.debug("Chat request: #{request_data}")
+
+      request_data
+    end
+
+    def process_line(line)
+      return unless line.start_with?('data: ')
+
+      line = line[6..].to_s.strip
+      return if line == '[DONE]'
+
+      JSON.parse(line)
+    end
+
+    def config_logger
+      Logger.new($stdout).tap do |logger|
+        logger.level = ENV.fetch('LOG_LEVEL', 'ERROR')
+
+        logger.formatter = proc do |severity, datetime, progname, msg|
+          "#{datetime.strftime("%Y-%m-%d %H:%M:%S")} #{severity} #{progname}: #{msg}\n"
+        end
+      end
+    end
+  end
+end

data/lib/mistral/constants.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+module Mistral
+  RETRY_STATUS_CODES = [429, 500, 502, 503, 504].freeze
+  ENDPOINT = 'https://api.mistral.ai'
+end

data/lib/mistral/exceptions.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Mistral
+  # Base error class, returned when nothing more specific applies
+  class Error < StandardError
+  end
+
+  class APIError < Error
+    attr_reader :http_status, :headers
+
+    def initialize(message: nil, http_status: nil, headers: nil)
+      super(message: message)
+
+      @http_status = http_status
+      @headers = headers
+    end
+
+    def self.from_response(response, message: nil)
+      new(
+        message: message || response.to_s,
+        http_status: response.code,
+        headers: response.headers.to_h
+      )
+    end
+
+    def to_s
+      "#{self.class.name}(message=#{super}, http_status=#{http_status})"
+    end
+  end
+
+  # Returned when we receive a non-200 response from the API that we should retry
+  class APIStatusError < APIError
+  end
+
+  # Returned when the SDK can not reach the API server for any reason
+  class ConnectionError < Error
+  end
+end

data/lib/mistral/models/chat_completion.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module Mistral
+  class Function < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :name, Types::Strict::String
+    attribute :description, Types::Strict::String
+    attribute :parameters, Types::Strict::Hash
+  end
+
+  ToolType = Types::Strict::String.default('function').enum('function')
+
+  class FunctionCall < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :name, Types::Strict::String
+    attribute :arguments, Types::Strict::String
+  end
+
+  class ToolCall < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :id, Types::Strict::String.default('null')
+    attribute :type, ToolType
+    attribute :function, FunctionCall
+  end
+
+  ResponseFormats = Types::Strict::String.default('text').enum('text', 'json_object')
+
+  ToolChoice = Types::Strict::String.enum('auto', 'any', 'none')
+
+  class ResponseFormat < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :type, ResponseFormats
+  end
+
+  class ChatMessage < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :role, Types::Strict::String
+    attribute :content, Types::Strict::Array.of(Types::Strict::String) | Types::Strict::String
+    attribute? :name, Types::String.optional
+    attribute? :tool_calls, Types::Strict::Array.of(ToolCall).optional
+  end
+
+  class DeltaMessage < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute? :role, Types::Strict::String.optional
+    attribute? :content, Types::Strict::String.optional
+    attribute? :tool_calls, Types::Strict::Array.of(ToolCall).optional
+  end
+
+  FinishReason = Types::Strict::String.enum('stop', 'length', 'error', 'tool_calls')
+
+  class ChatCompletionResponseStreamChoice < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :index, Types::Strict::Integer
+    attribute :delta, DeltaMessage
+    attribute? :finish_reason, FinishReason.optional
+  end
+
+  class ChatCompletionStreamResponse < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :id, Types::Strict::String
+    attribute :model, Types::Strict::String
+    attribute :choices, Types::Strict::Array.of(ChatCompletionResponseStreamChoice)
+    attribute? :created, Types::Strict::Integer.optional
+    attribute? :object, Types::Strict::String.optional
+    attribute? :usage, UsageInfo.optional
+  end
+
+  class ChatCompletionResponseChoice < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :index, Types::Strict::Integer
+    attribute :message, ChatMessage
+    attribute? :finish_reason, FinishReason.optional
+  end
+
+  class ChatCompletionResponse < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :id, Types::Strict::String
+    attribute :object, Types::Strict::String
+    attribute :created, Types::Strict::Integer
+    attribute :model, Types::Strict::String
+    attribute :choices, Types::Strict::Array.of(ChatCompletionResponseChoice)
+    attribute :usage, UsageInfo
+  end
+end

data/lib/mistral/models/common.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Mistral
+  class UsageInfo < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :prompt_tokens, Types::Strict::Integer
+    attribute :total_tokens, Types::Strict::Integer
+    attribute? :completion_tokens, Types::Strict::Integer.optional
+  end
+end

data/lib/mistral/models/embeddings.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Mistral
+  class EmbeddingObject < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :object, Types::Strict::String
+    attribute :embedding, Types::Strict::Array.of(Types::Strict::Float)
+    attribute :index, Types::Strict::Integer
+  end
+
+  class EmbeddingResponse < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :id, Types::Strict::String
+    attribute :object, Types::Strict::String
+    attribute :data, Types::Strict::Array.of(EmbeddingObject)
+    attribute :model, Types::Strict::String
+    attribute :usage, UsageInfo
+  end
+end

data/lib/mistral/models/models.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Mistral
+  class ModelPermission < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :id, Types::Strict::String
+    attribute :object, Types::Strict::String
+    attribute :created, Types::Strict::Integer
+    attribute :allow_create_engine, Types::Strict::Bool.default(false)
+    attribute :allow_sampling, Types::Strict::Bool.default(true)
+    attribute :allow_logprobs, Types::Strict::Bool.default(true)
+    attribute :allow_search_indices, Types::Strict::Bool.default(false)
+    attribute :allow_view, Types::Strict::Bool.default(true)
+    attribute :allow_fine_tuning, Types::Strict::Bool.default(false)
+    attribute :organization, Types::Strict::String.default('*')
+    attribute? :group, Types::Strict::String.optional
+    attribute :is_blocking, Types::Strict::Bool.default(false)
+  end
+
+  class ModelCard < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :id, Types::Strict::String
+    attribute :object, Types::Strict::String
+    attribute :created, Types::Strict::Integer
+    attribute :owned_by, Types::Strict::String
+    attribute? :root, Types::Strict::String.optional
+    attribute? :parent, Types::Strict::String.optional
+    attribute :permission, Types::Strict::Array.of(ModelPermission).default([].freeze)
+  end
+
+  class ModelList < Dry::Struct
+    transform_keys(&:to_sym)
+
+    attribute :object, Types::Strict::String
+    attribute :data, Types::Strict::Array.of(ModelCard)
+  end
+end

data/lib/mistral/version.rb

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

data/lib/mistral.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require 'dry-struct'
+require 'http'
+require 'json'
+require 'logger'
+require 'time'
+
+require 'http/features/line_iterable_body'
+
+module Mistral
+  module Types
+    include Dry.Types()
+  end
+end
+
+require 'mistral/constants'
+require 'mistral/exceptions'
+require 'mistral/models/models'
+require 'mistral/models/common'
+require 'mistral/models/embeddings'
+require 'mistral/models/chat_completion'
+require 'mistral/version'
+require 'mistral/client'