raimei-nim 0.1.1 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 34cd5b478c9e975b72d433101f62c1cf5949097b844dfaa3bbbd8b1512803263
-  data.tar.gz: bd9081697d14c34807ece3d9e50beebcb9da994a7660d9617d165fa60655cabb
+  metadata.gz: abc44189bf4b4c7c97e31627d738bcadfb4a74b216f434e86e03a3e0b7b59b62
+  data.tar.gz: 7a042805a56694d8d2a2a917c9094cd2c79ac5986b0afc0e694265f8b5e376d5
 SHA512:
-  metadata.gz: d3c9e767389ab1856d05285067ed38f459a61d6ffa736e8727b69b36651909e3f6bbda27eeacaf6daa2c8f8794e3d7a6c803d6dc9bff6a2cb6e5a44b05af795f
-  data.tar.gz: d4e21bcce8052155ff694636547c9ab79cdc418da4635d1f6b66935e44a8dc9ca6855da04f29ceee727c788decdea7ff2a26eb6b59d32213837b67f2d231dcc5
+  metadata.gz: 2319cd31659f2412eecf42b7875891de1d6511dda190ec25b3bba269c1d5560926fc5d4b939a6bc458a45a6b5ee2814b69b5f2064e9f3091ca276cf8af7c0928
+  data.tar.gz: 32d33a87c2263b84713805cb866d1ac99cbb1f2f121f169a333caf9abd6345268da6abc83b770117a05b61ad06d94fcb266930872350683a130c5e9336450430

data/Gemfile

@@ -0,0 +1,13 @@
+# 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"

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