From 1976bc421c52e669316395decd92b7874b949e8b Mon Sep 17 00:00:00 2001
From: Obie Fernandez <obiefernandez@gmail.com>
Date: Thu, 28 Mar 2024 16:55:14 -0600
Subject: [PATCH] Raises `ServerError` on errors. Accepts multiple models for
 fallback

---
 README.md                 | 26 +++++++++++++++++++++++++-
 lib/open_router/client.rb | 24 +++++++++++++++++++-----
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index f0d83c1..983678f 100644
--- a/README.md
+++ b/README.md
@@ -91,7 +91,27 @@ puts response["choices"][0]["message"]["content"]
 
 ### Models
 
-Fetch the list of available models from the OpenRouter API:
+Pass an array to the `model` parameter to enable [explicit model routing](https://openrouter.ai/docs#model-routing).
+
+```ruby
+OpenRouter::Client.new.chat_completion(
+  [
+    { role: "system", content: SYSTEM_PROMPT },
+    { role: "user", content: "Provide analysis of the data formatted as JSON:" }
+  ],
+  model: [
+    "mistralai/mixtral-8x7b-instruct:nitro",
+    "mistralai/mixtral-8x7b-instruct"
+  ],
+  extras: {
+    response_format: {
+      type: "json_object"
+    }
+  }
+)
+```
+
+[Browse full list of models available](https://openrouter.ai/models) or fetch from the OpenRouter API:
 
 ```ruby
 models = client.models
@@ -110,6 +130,10 @@ puts stats
 # => {"id"=>"generation-abcdefg", "object"=>"generation", "created"=>1684195200, "model"=>"openrouter/auto", "usage"=>{"prompt_tokens"=>10, "completion_tokens"=>50, "total_tokens"=>60}, "cost"=>0.0006}
 ```
 
+## Errors
+
+The client will raise an `OpenRouter::ServerError` in the case of an error returned from a completion (or empty response).
+
 ## Contributing
 
 Bug reports and pull requests are welcome on GitHub at <https://github.com/OlympiaAI/open_router>. This project is intended to be a safe, welcoming space for collaboration, and contributors are expected to adhere to the [code of conduct](https://github.com/OlympiaAI/open_router/blob/main/CODE_OF_CONDUCT.md).
diff --git a/lib/open_router/client.rb b/lib/open_router/client.rb
index 5ba4101..ea0eb38 100644
--- a/lib/open_router/client.rb
+++ b/lib/open_router/client.rb
@@ -1,8 +1,10 @@
 # frozen_string_literal: true
 
-require_relative "http"
+require_relative 'http'
 
 module OpenRouter
+  class ServerError < StandardError; end
+
   class Client
     include OpenRouter::HTTP
 
@@ -13,19 +15,31 @@ def initialize
 
     # Performs a chat completion request to the OpenRouter API.
     # @param messages [Array<Hash>] Array of message hashes with role and content, like [{role: "user", content: "What is the meaning of life?"}]
-    # @param model [String] Model identifier, defaults to 'openrouter/auto'
+    # @param model [String|Array] Model identifier, or array of model identifiers if you want to fallback to the next model in case of failure
     # @param providers [Array<String>] Optional array of provider identifiers, ordered by priority
     # @param transforms [Array<String>] Optional array of strings that tell OpenRouter to apply a series of transformations to the prompt before sending it to the model. Transformations are applied in-order
     # @param extras [Hash] Optional hash of model-specific parameters to send to the OpenRouter API
     # @param stream [Proc, nil] Optional callable object for streaming
     # @return [Hash] The completion response.
-    def complete(messages, model: "openrouter/auto", providers: [], transforms: [], extras: {}, stream: nil) # rubocop:disable Metrics/ParameterLists
-      parameters = { model:, messages: }
+    def chat_completion(messages, model: 'openrouter/auto', providers: [], transforms: [], extras: {}, stream: nil)
+      parameters = { messages: }
+      if model.is_a?(String)
+        parameters[:model] = model
+      elsif model.is_a?(Array)
+        parameters[:models] = model
+        parameters[:route] = "fallback"
+      end
       parameters[:provider] = { provider: { order: providers } } if providers.any?
       parameters[:transforms] = transforms if transforms.any?
       parameters[:stream] = stream if stream
       parameters.merge!(extras)
-      json_post(path: "/chat/completions", parameters:)
+
+      puts parameters if Rails.env.local?
+
+      json_post(path: "/chat/completions", parameters:).tap do |response|
+        raise ServerError, "Empty response from OpenRouter. Might be worth retrying once or twice." if response.blank?
+        raise ServerError, response.dig("error", "message") if response.dig("error", "message").present?
+      end
     end
 
     # Fetches the list of available models from the OpenRouter API.