raimei-nim 0.1.1 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 34cd5b478c9e975b72d433101f62c1cf5949097b844dfaa3bbbd8b1512803263
-  data.tar.gz: bd9081697d14c34807ece3d9e50beebcb9da994a7660d9617d165fa60655cabb
+  metadata.gz: 676783e23d825f284282d6a366566b0a4029e7314e87c23524841e4a3dbcbd6c
+  data.tar.gz: b86cd974c8ab09e1c46ad49fc0c6f2acaa8c6f7f704de5aeb014415ecf18ec2f
 SHA512:
-  metadata.gz: d3c9e767389ab1856d05285067ed38f459a61d6ffa736e8727b69b36651909e3f6bbda27eeacaf6daa2c8f8794e3d7a6c803d6dc9bff6a2cb6e5a44b05af795f
-  data.tar.gz: d4e21bcce8052155ff694636547c9ab79cdc418da4635d1f6b66935e44a8dc9ca6855da04f29ceee727c788decdea7ff2a26eb6b59d32213837b67f2d231dcc5
+  metadata.gz: e39c9f8dcd33f14ef651eda0228145c1136b5ae23e099f2338ee3697f9905094b45917fab77476a674519f3a078605406a6253c768a9c5514f4844577970dbcb
+  data.tar.gz: 1e7bc601c8dd19f8d53245841812d29cfafe8034c098c820219bf2417a4f5ef46e82f3ff0e3f034e834518ef1db6b2a499977cc3c33c6d1f541897aaff841f2d

data/Gemfile

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+source "https://rubygems.org"
+
+# Specify your gem's dependencies in raimei-nim.gemspec
+gemspec
+
+gem "irb"
+gem "rake", "~> 13.0"
+
+gem "minitest", "~> 5.16"
+
+gem "rubocop", "~> 1.21"
+
+gem "webrick", "~> 1.8"

data/exe/raimei-chat

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+
+# frozen_string_literal: true
+
+require "optparse"
+require "raimei/nim"
+
+opts = { url: ENV["URL"], model: ENV["MODEL"] || "llama3", api_key: ENV["API_KEY"], stream: true }
+OptionParser.new do |o|
+  o.banner = "Usage: raimei-chat [--url URL] [--model MODEL] [--api-key KEY] [--no-stream] [prompt...]\n" \
+             "Reads from ARGV or stdin."
+  o.on("--url URL")       { |v| opts[:url] = v }
+  o.on("--model MODEL")   { |v| opts[:model] = v }
+  o.on("--api-key KEY")   { |v| opts[:api_key] = v }
+  o.on("--no-stream")     { opts[:stream] = false }
+end.parse!
+
+abort "Set --url or URL env" unless opts[:url]
+prompt = if !STDIN.tty? && ARGV.empty? then STDIN.read else ARGV.join(" ") end
+abort "Provide a prompt (args or stdin)" if prompt.to_s.strip.empty?
+
+client = Raimei::NIM::Client.new(url: opts[:url], api_key: opts[:api_key])
+if opts[:stream]
+  client.chat(model: opts[:model], messages: [{ role: "user", content: prompt }], stream: true).each { |d| print d }
+  puts
+else
+  puts client.chat(model: opts[:model], messages: [{ role: "user", content: prompt }], stream: false)
+end
+

data/lib/raimei/nim/version.rb

@@ -2,6 +2,6 @@
 
 module Raimei
   module Nim
-    VERSION = "0.1.1"
+    VERSION = "0.1.3"
   end
 end

data/lib/raimei/nim.rb

@@ -1,31 +1,75 @@
 # frozen_string_literal: true
 
-require "json"
-require "net/http"
-require "uri"
+require "json"; require "net/http"; require "uri"
 
 module Raimei
   module NIM
     class Client
-      def initialize(url:, api_key: nil, timeout: 60)
-        @uri = URI(url) # e.g., http://localhost:11434/v1/chat/completions
-        @api_key = api_key
-        @timeout = timeout
+      def initialize(url:, api_key: nil, timeout: 60, retries: 2)
+        @uri = URI(url); @api_key = api_key; @timeout = timeout; @retries = retries
       end
 
       # messages: [{role:"user", content:"Hi"}]
-      def chat(model:, messages:, **opts)
-        body = { model: model, messages: messages, stream: false }.merge(opts)
+      # stream: true -> Enumerator (or yield chunks if block given)
+      # stream: false -> String
+      def chat(model:, messages:, stream: true, **opts, &blk)
+        body = { model: model, messages: messages, stream: stream }.merge(opts)
         req  = Net::HTTP::Post.new(@uri)
         req["Content-Type"] = "application/json"
+        req["Accept"]       = stream ? "text/event-stream" : "application/json"
         req["Authorization"] = "Bearer #{@api_key}" if @api_key
         req.body = JSON.dump(body)
 
-        Net::HTTP.start(@uri.host, @uri.port, use_ssl: @uri.scheme == "https", read_timeout: @timeout) do |http|
+        with_retries do
+          if stream
+            if block_given?
+              _sse(req) { |delta| blk.call(delta) }
+            else
+              Enumerator.new { |y| _sse(req) { |delta| y << delta } }
+            end
+          else
+            _json(req)
+          end
+        end
+      end
+
+      private
+
+      def with_retries
+        attempts = 0
+        begin
+          return yield
+        rescue => e
+          attempts += 1
+          raise e if attempts > @retries
+          sleep(0.5 * attempts) # tiny backoff
+          retry
+        end
+      end
+
+      def _json(req)
+        Net::HTTP.start(@uri.host, @uri.port, use_ssl: @uri.scheme=="https", read_timeout:@timeout) do |http|
           res = http.request(req)
           raise "HTTP #{res.code}: #{res.body}" unless res.is_a?(Net::HTTPSuccess)
-          json = JSON.parse(res.body)
-          json.dig("choices", 0, "message", "content").to_s
+          j = JSON.parse(res.body)
+          j.dig("choices",0,"message","content").to_s
+        end
+      end
+
+      def _sse(req)
+        Net::HTTP.start(@uri.host, @uri.port, use_ssl:@uri.scheme=="https", read_timeout:@timeout) do |http|
+          http.request(req) do |res|
+            res.read_body do |chunk|
+              chunk.each_line do |line|
+                next unless line.start_with?("data:")
+                data = line.sub("data:","").strip
+                break if data == "[DONE]"
+                json = JSON.parse(data) rescue nil
+                next unless json
+                yield(json.dig("choices",0,"delta","content") || "")
+              end
+            end
+          end
         end
       end
     end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: raimei-nim
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.1.3
 platform: ruby
 authors:
 - Hal Fulton
@@ -13,14 +13,17 @@ description: Streaming SSE client with timeouts/retries; works with NIM/TRT-LLM,
   Ollama, and other OpenAI-style servers.
 email:
 - rubyhacker@gmail.com
-executables: []
+executables:
+- raimei-chat
 extensions: []
 extra_rdoc_files: []
 files:
 - CHANGELOG.md
+- Gemfile
 - LICENSE.txt
 - README.md
 - Rakefile
+- exe/raimei-chat
 - lib/raimei/nim.rb
 - lib/raimei/nim/version.rb
 - sig/raimei/nim.rbs