mistral_rb 0.1.1 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 69d426e25d1c19668cbde210ee7269a98e5769cb852c5184220c99460c7c2657
-  data.tar.gz: d0b6828a2313d6bdfd18e0879c8ba8b05a1559a3e4c045d0d26970a2e40fa1c9
+  metadata.gz: 2e43cdb7fa62132f8342d13ec666a4fe48571dc0dbf5b64a2476c23d4da86aa5
+  data.tar.gz: 402b62dfe6f4ca6a2b77d3e13127b673037374fe75ffbe874c4cafef2fbdaeb3
 SHA512:
-  metadata.gz: 647b05e72b0699d8e6c3bf74408cdb3b8cbc633b37221f0b53e074d1a28e7bf21ff74023c921f83b6a3df8882693d726716a328644e9b9bec53fc4d8dbebc333
-  data.tar.gz: ae4b88b5773aa88981220b8cf3c1c301ca262275be1f67c6d37b26e195896df099512632d405afd0ec85a4e985873e48670b9a282614b9e74a20289c7b0e70ef
+  metadata.gz: 921683a278ab7f4d9f443c2db7a0dd16dd08b81752e6d99ac49fe619fb23e7d2cd145d30e8760988786ffa81647f9d666a73d0ac3a56e6c34ce921bd61a8339e
+  data.tar.gz: cbc90300023388689b9b5572f25dec507de1cc956f90c7406cfe8fbe685011200a248bbbb7574e564b3ff0b89a9726c60890c9175ac8e4832018b4766d1dabaf

data/README.md

@@ -24,15 +24,15 @@ Here is a quick example to get you started:
 require 'mistral_rb'
 
 # Initialize the API with your API key
-mistral = MistralAPI.new("your_api_key_here")
+mistral = MistralAPI.new(api_key: "your_api_key_here")
 
 # Create Chat Completion
-response = mistral.create_chat_completion("mistral-tiny", [{role: "user", content: "Who is Macron?"}])
+response = mistral.create_chat_completion(model: "mistral-tiny", messages: [{role: "user", content: "Who is Macron?"}])
 puts response.choices.to_s
 
 # Create Embeddings
-embedding_response = mistral.create_embeddings("mistral-embed", ["Hello", "world"])
-puts embedding_response.data.first.embedding
+embedding_response = mistral.create_embeddings(model: "mistral-embed", input: ["Hello", "world"])
+puts embedding_response.data.first.inspect
 
 # List Available Models
 model_list_response = mistral.list_available_models
@@ -41,6 +41,21 @@ model_list_response.data.each do |model|
 end
 ```
 
+Here is how to use streaming:
+
+```ruby
+api = MistralAPI.new("api_key")
+
+api.create_chat_completion(
+  model: 'mistral-tiny',
+  messages: [{ 'role' => 'user', 'content' => 'Who is Barack Obama ?' }],
+  stream: true
+) do |chunk|
+  puts chunk.inspect
+end
+
+```
+
 
 ## Development
 

data/lib/mistral_rb/response_models.rb

@@ -71,3 +71,34 @@ class Model
     @permissions = model_hash["permission"] # This could be further parsed into Permission objects if detailed parsing is required
   end
 end
+
+class StreamedCompletionResponse
+  attr_reader :id, :object, :created, :model, :choices
+
+  def initialize(response_hash)
+    @id = response_hash["id"]
+    @object = response_hash["object"]
+    @created = response_hash["created"]
+    @model = response_hash["model"]
+    @choices = response_hash["choices"].map { |choice| StreamedChoice.new(choice) }
+  end
+end
+
+class StreamedChoice
+  attr_reader :index, :delta, :finish_reason
+
+  def initialize(choice_hash)
+    @index = choice_hash["index"]
+    @delta = Delta.new(choice_hash["delta"]) if choice_hash["delta"]
+    @finish_reason = choice_hash["finish_reason"]
+  end
+end
+
+class Delta
+  attr_reader :role, :content
+
+  def initialize(delta_hash)
+    @role = delta_hash["role"]
+    @content = delta_hash["content"]
+  end
+end

data/lib/mistral_rb/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module MistralRb
-  VERSION = "0.1.1"
+  VERSION = "0.1.3"
 end

data/lib/mistral_rb.rb

@@ -8,7 +8,7 @@ require_relative "mistral_rb/response_models"
 class MistralAPI
   include HTTParty
 
-  def initialize(api_key, base_uri = "https://api.mistral.ai/v1")
+  def initialize(api_key:, base_uri: "https://api.mistral.ai/v1")
     @headers = {
       "Authorization" => "Bearer #{api_key}",
       "Content-Type" => "application/json"
@@ -16,7 +16,7 @@ class MistralAPI
     self.class.base_uri base_uri
   end
 
-  def create_chat_completion(model, messages, temperature = 0.7, top_p = 1, max_tokens = nil, stream = false, safe_mode = false, random_seed = nil)
+  def create_chat_completion(model:, messages:, temperature: 0.7, top_p: 1, max_tokens: nil, stream: false, safe_mode: false, random_seed: nil)
     body = {
       model: model,
       messages: messages,
@@ -26,14 +26,23 @@ class MistralAPI
       stream: stream,
       safe_mode: safe_mode,
       random_seed: random_seed
-    }.compact.to_json  # compact to remove nil values
+    }.compact.to_json
 
-    response = self.class.post("/chat/completions", body: body, headers: @headers)
-    parsed_response = handle_response(response)
-    CompletionResponse.new(parsed_response)
+    if stream
+      # Use on_data callback for streaming
+      self.class.post("/chat/completions", body: body, headers: @headers, stream_body: true) do |fragment, _, _|
+        processed_chunk = handle_stream_chunk(fragment)
+        yield(processed_chunk) if block_given? && processed_chunk
+      end
+    else
+      # Handle non-streaming response
+      response = self.class.post("/chat/completions", body: body, headers: @headers)
+      parsed_response = handle_response(response)
+      CompletionResponse.new(parsed_response)
+    end
   end
 
-  def create_embeddings(model, input, encoding_format = "float")
+  def create_embeddings(model:, input:, encoding_format: "float")
     body = {
       model: model,
       input: input,
@@ -60,4 +69,20 @@ class MistralAPI
       raise "API Error: #{response.code} - #{response.body}"
     end
   end
+
+  def handle_stream_chunk(chunk)
+    # Skip processing if the chunk indicates the end of the stream.
+    return nil if chunk.strip == "data: [DONE]"
+
+    if chunk.strip.start_with?("data:")
+      data_content = chunk.split("data:").last.strip
+      begin
+        # Only parse the JSON content if it's not the end-of-stream indicator
+        json_content = JSON.parse(data_content)
+        StreamedCompletionResponse.new(json_content)
+      rescue JSON::ParserError => e
+        puts "Error parsing JSON: #{e.message}"
+      end
+    end
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: mistral_rb
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.1.3
 platform: ruby
 authors:
 - Franck Stephane Ndzomga
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-12-22 00:00:00.000000000 Z
+date: 2023-12-24 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: httparty