Support gpustack (#142)

Author: graysonchen

bin/console

@@ -19,6 +19,8 @@ RubyLLM.configure do |config|
   config.bedrock_secret_key = ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
   config.bedrock_region = ENV.fetch('AWS_REGION', nil)
   config.bedrock_session_token = ENV.fetch('AWS_SESSION_TOKEN', nil)
+  config.gpustack_api_base = ENV.fetch('GPUSTACK_API_BASE', nil)
+  config.gpustack_api_key = ENV.fetch('GPUSTACK_API_KEY', nil)
 end
 
 IRB.start(__FILE__)

lib/ruby_llm.rb

@@ -18,6 +18,7 @@
   'deepseek' => 'DeepSeek',
   'bedrock' => 'Bedrock',
   'openrouter' => 'OpenRouter',
+  'gpustack' => 'GPUStack',
   'pdf' => 'PDF'
 )
 loader.ignore("#{__dir__}/tasks")
@@ -83,6 +84,7 @@ def logger
 RubyLLM::Provider.register :bedrock, RubyLLM::Providers::Bedrock
 RubyLLM::Provider.register :openrouter, RubyLLM::Providers::OpenRouter
 RubyLLM::Provider.register :ollama, RubyLLM::Providers::Ollama
+RubyLLM::Provider.register :gpustack, RubyLLM::Providers::GPUStack
 
 if defined?(Rails::Railtie)
   require 'ruby_llm/railtie'

lib/ruby_llm/configuration.rb

@@ -24,6 +24,8 @@ class Configuration
                   :bedrock_session_token,
                   :openrouter_api_key,
                   :ollama_api_base,
+                  :gpustack_api_base,
+                  :gpustack_api_key,
                   # Default models
                   :default_model,
                   :default_embedding_model,

lib/ruby_llm/providers/gpustack.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # GPUStack API integration based on Ollama.
+    module GPUStack
+      extend OpenAI
+      extend GPUStack::Chat
+      extend GPUStack::Models
+
+      module_function
+
+      def api_base(config)
+        config.gpustack_api_base
+      end
+
+      def headers(config)
+        {
+          'Authorization' => "Bearer #{config.gpustack_api_key}"
+        }
+      end
+
+      def slug
+        'gpustack'
+      end
+
+      def local?
+        true
+      end
+
+      def configuration_requirements
+        %i[gpustack_api_base gpustack_api_key]
+      end
+    end
+  end
+end

lib/ruby_llm/providers/gpustack/chat.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module GPUStack
+      # Chat methods of the GPUStack API integration
+      module Chat
+        module_function
+
+        def format_role(role)
+          # GPUStack doesn't use the new OpenAI convention for system prompts
+          role.to_s
+        end
+      end
+    end
+  end
+end

lib/ruby_llm/providers/gpustack/models.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module GPUStack
+      # Models methods of the GPUStack API integration
+      module Models
+        module_function
+
+        def models_url
+          'models'
+        end
+
+        def parse_list_models_response(response, slug, _capabilities)
+          items = response.body['items'] || []
+          items.map do |model|
+            Model::Info.new(
+              id: model['name'],
+              created_at: model['created_at'] ? Time.parse(model['created_at']) : nil,
+              display_name: "#{model['source']}/#{model['name']}",
+              provider: slug,
+              type: determine_model_type(model),
+              metadata: {
+                description: model['description'],
+                source: model['source'],
+                huggingface_repo_id: model['huggingface_repo_id'],
+                ollama_library_model_name: model['ollama_library_model_name'],
+                backend: model['backend'],
+                meta: model['meta'],
+                categories: model['categories']
+              },
+              context_window: model.dig('meta', 'n_ctx'),
+              # Using context window as max tokens since it's not explicitly provided
+              max_tokens: model.dig('meta', 'n_ctx'),
+              supports_vision: model.dig('meta', 'support_vision') || false,
+              supports_functions: model.dig('meta', 'support_tool_calls') || false,
+              supports_json_mode: true, # Assuming all models support JSON mode
+              input_price_per_million: 0.0,  # Price information not available in new format
+              output_price_per_million: 0.0  # Price information not available in new format
+            )
+          end
+        end
+
+        private
+
+        def determine_model_type(model)
+          return 'embedding' if model['categories']&.include?('embedding')
+          return 'chat' if model['categories']&.include?('llm')
+
+          'other'
+        end
+      end
+    end
+  end
+end