mistral_rb 0.1.1 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 69d426e25d1c19668cbde210ee7269a98e5769cb852c5184220c99460c7c2657
-  data.tar.gz: d0b6828a2313d6bdfd18e0879c8ba8b05a1559a3e4c045d0d26970a2e40fa1c9
+  metadata.gz: 8826573e4021a69c7f6c560357c79a75fc80a63ff97b6610143fe3505f4d9138
+  data.tar.gz: 04bf840637ab24c388b28b175b987083f9786579d58dec9f8e5366f0f7a1d541
 SHA512:
-  metadata.gz: 647b05e72b0699d8e6c3bf74408cdb3b8cbc633b37221f0b53e074d1a28e7bf21ff74023c921f83b6a3df8882693d726716a328644e9b9bec53fc4d8dbebc333
-  data.tar.gz: ae4b88b5773aa88981220b8cf3c1c301ca262275be1f67c6d37b26e195896df099512632d405afd0ec85a4e985873e48670b9a282614b9e74a20289c7b0e70ef
+  metadata.gz: b3aee45181ce35cd5b34df3b10f091ef084bc00bba078452f0a5485e47eb0b19374d07219f6a8a258884410c9cc4557d4e8ac07f0bf0cfcccbe25def5f5baea4
+  data.tar.gz: 886665f3b0e2db01b155421c6680b6523cbae26c61ef260de303cb9ad90a6e69c1a60b433d2fcfc78167d450e842016629f40a02d51d714cf4d89bfb58d98bd4

data/README.md

@@ -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.2"
 end

data/lib/mistral_rb.rb

@@ -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,11 +26,20 @@ 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")
@@ -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.2
 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