Class: Train::Model

Inherits:
Object
  • Object
show all
Defined in:
lib/train.rb

Overview

Train::Model represents a connection to an AI model endpoint

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(api_key:, url:, model_name:, model_id_changed: nil) ⇒ Model

Returns a new instance of Model.



56
57
58
59
60
61
62
63
 
# File 'lib/train.rb', line 56

def initialize(api_key:, url:, model_name:, model_id_changed: nil)
  @api_key = api_key
  @url = URI.parse(url)

  # 🔁 override if custom model ID provided
  default_id = Train::MODELS.get(model_name)
  @model_name = model_id_changed || default_id || model_name
end

Instance Attribute Details

#api_key ⇒ Object

Returns the value of attribute api_key.



54
55
56
 
# File 'lib/train.rb', line 54

def api_key
  @api_key
end

#model_name ⇒ Object

Returns the value of attribute model_name.



54
55
56
 
# File 'lib/train.rb', line 54

def model_name
  @model_name
end

#url ⇒ Object

Returns the value of attribute url.



54
55
56
 
# File 'lib/train.rb', line 54

def url
  @url
end

Instance Method Details

#chat(prompt, system: "You are a helpful AI.") ⇒ String

Sends a chat message to the AI model

Parameters:

  • prompt (String or Symbol) —

    User message or Train::Chat::INPUT for interactive mode

  • system (String) (defaults to: "You are a helpful AI.") —

    Optional system instruction

Returns:

  • (String) —

    AI-generated response or fallback message



72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
 
# File 'lib/train.rb', line 72

def chat(prompt, system: "You are a helpful AI.")
  return Chat.start(self) if prompt == Chat::INPUT

  http = Net::HTTP.new(url.host, url.port)
  http.use_ssl = (url.scheme == "https")

  headers = {
    "Content-Type" => "application/json",
    "Authorization" => "Bearer #{api_key}"
  }

  body = {
    model: model_name,
    messages: [
      { role: "system", content: system },
      { role: "user", content: prompt }
    ]
  }.to_json

  response = http.post(url.path, body, headers)
  json = JSON.parse(response.body)

  json["choices"]&.first&.dig("message", "content") || "⚠️ No response"
rescue => e
  "❌ Train error: #{e.message}"
end