raimei-nim 0.1.0 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d78531725801b2a89ac2a2e74343c8e07050f1ff68e0afcf3477a91d2696dba2
-  data.tar.gz: 3e54451e5a36baeaab8af0aa43b4eba36615fddc164cbdec09922d49a9596fda
+  metadata.gz: abc44189bf4b4c7c97e31627d738bcadfb4a74b216f434e86e03a3e0b7b59b62
+  data.tar.gz: 7a042805a56694d8d2a2a917c9094cd2c79ac5986b0afc0e694265f8b5e376d5
 SHA512:
-  metadata.gz: 232bc484fadcec0d57ace1baf1b12917549cc79f9452e2475dcc3d6ac04091c90c5196f81acc9a33edaa2aa7d3493d3468f8a1e9223f7b7f01838c6598466113
-  data.tar.gz: d52ab3f98d35710d1ffd6a98a01f0637e1bd0f106257d087126b9ecfb39c6691c5f7a9aa0936f88eb07b7fd1df2ba685d917c3c84e7d2cfa0856c49ccec40daf
+  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/README.md

@@ -1,5 +1,8 @@
 # Raimei::Nim
 
+> ⚠️ Experimental: APIs may change.
+
+
 TODO: Delete this and the text below, and describe your gem
 
 Welcome to your new gem! In this directory, you'll find the files you need to be able to package up your Ruby library into a gem. Put your Ruby code in the file `lib/raimei/nim`. To experiment with that code, run `bin/console` for an interactive prompt.

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.0"
+    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.0
+  version: 0.1.2
 platform: ruby
 authors:
 - Hal Fulton
@@ -9,17 +9,21 @@ bindir: exe
 cert_chain: []
 date: 1980-01-02 00:00:00.000000000 Z
 dependencies: []
-description: NVIDIA NIM / OpenAI-compatible client for Ruby (Raimei).
+description: Streaming SSE client with timeouts/retries; works with NIM/TRT-LLM, vLLM,
+  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
@@ -36,7 +40,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: 3.2.0
+      version: '3.2'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
@@ -45,5 +49,5 @@ required_rubygems_version: !ruby/object:Gem::Requirement
 requirements: []
 rubygems_version: 3.7.1
 specification_version: 4
-summary: NVIDIA NIM / OpenAI-compatible client for Ruby (Raimei).
+summary: OpenAI/NIM-compatible client for Ruby.
 test_files: []