llm_gateway 0.3.0 → 0.4.0

data/lib/llm_gateway/clients/groq.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require_relative "../base_client"
+
+module LlmGateway
+  module Clients
+    class Groq < BaseClient
+      def initialize(model_key: "openai/gpt-oss-20b", api_key: ENV["GROQ_API_KEY"])
+        @base_endpoint = "https://api.groq.com/openai/v1"
+        super(model_key: model_key, api_key: api_key)
+      end
+
+      def chat(messages, tools: nil, system: [], **options)
+        body = {
+          model: model_key,
+          messages: system + messages,
+          tools: tools
+        }
+        body.merge!(options)
+
+        post("chat/completions", body)
+      end
+
+      private
+
+      def build_headers
+        {
+          "content-type" => "application/json",
+          "Authorization" => "Bearer #{api_key}"
+        }
+      end
+
+      def handle_client_specific_errors(response, error)
+        # Groq likely uses 'code' like OpenAI since it's OpenAI-compatible
+        error_code = error["code"]
+        error_message = error["message"]
+
+        if Errors.context_overflow_message?(error_message)
+          raise Errors::PromptTooLong.new(error_message, error["type"])
+        end
+
+        case response.code.to_i
+        when 429
+          raise Errors::RateLimitError.new(error["type"], error_code) if error_code == "rate_limit_exceeded"
+
+          raise Errors::OverloadError.new(error_message, error_code)
+        end
+
+        # If we get here, we didn't handle it specifically
+        raise Errors::APIStatusError.new(error_message, error_code)
+      end
+    end
+  end
+end

data/lib/llm_gateway/clients/openai.rb

@@ -0,0 +1,208 @@
+# frozen_string_literal: true
+
+require_relative "../base_client"
+
+module LlmGateway
+  module Clients
+    class OpenAI < BaseClient
+      CODEX_BASE_ENDPOINT = "https://chatgpt.com/backend-api/codex"
+
+      attr_reader :account_id
+
+      def initialize(model_key: "gpt-4o", api_key: ENV["OPENAI_API_KEY"], account_id: nil)
+        @base_endpoint = "https://api.openai.com/v1"
+        @account_id = account_id
+        super(model_key: model_key, api_key: api_key)
+      end
+
+      def chat(messages, tools: nil, system: [], **options)
+        body = {
+          model: model_key,
+          messages: system + messages
+        }
+        body[:tools] = tools if tools
+        body.merge!(options)
+
+        post("chat/completions", body)
+      end
+
+      def stream(messages, tools: nil, system: [], **options, &block)
+        body = {
+          model: model_key,
+          messages: system + messages
+        }
+        body[:tools] = tools if tools
+        body.merge!(options)
+        body[:stream_options] = (body[:stream_options] || {}).merge(include_usage: true)
+
+        post_stream("chat/completions", body, &block)
+      end
+
+      def responses(messages, tools: nil, system: [], **options)
+        body = {
+          model: model_key,
+          input: messages.flatten
+        }
+        body[:instructions] = system[0][:content] if system.any?
+        body[:tools] = tools if tools
+        body.merge!(options)
+
+        post("responses", body)
+      end
+
+      def stream_responses(messages, tools: nil, system: [], **options, &block)
+        body = {
+          model: model_key,
+          input: messages.flatten
+        }
+        body[:instructions] = system[0][:content] if system.any?
+        body[:tools] = tools if tools
+        body.merge!(options)
+
+        post_stream("responses", body, &block)
+      end
+
+      def get_oauth_access_token(access_token:, refresh_token:, expires_at:, account_id: nil, &block)
+        token_manager = LlmGateway::Clients::OpenAI::TokenManager.new(
+          access_token: access_token,
+          refresh_token: refresh_token,
+          expires_at: expires_at,
+          account_id: account_id
+        )
+        token_manager.on_token_refresh = block if block_given?
+        token_manager.ensure_valid_token
+        token_manager.access_token
+      end
+
+      def chat_codex(messages, tools: nil, system: [], account_id: nil, **options)
+        body = build_codex_body(messages, system, tools, **options)
+
+        completed_response = nil
+        post_codex_stream("responses", body, account_id: account_id) do |raw_sse|
+          if raw_sse[:event] == "response.completed"
+            completed_response = raw_sse.dig(:data, :response)
+          end
+        end
+
+        completed_response
+      end
+
+      def stream_codex(messages, tools: nil, system: [], account_id: nil, **options, &block)
+        body = build_codex_body(messages, system, tools, **options)
+        post_codex_stream("responses", body, account_id: account_id, &block)
+      end
+
+      def download_file(file_id)
+        get("files/#{file_id}/content")
+      end
+
+      def generate_embeddings(input)
+        body = {
+          input:,
+          model: model_key
+        }
+        post("embeddings", body)
+      end
+
+      def upload_file(filename, content, mime_type = "application/octet-stream", purpose: "user_data")
+        post_file("files", content, filename, purpose: purpose, mime_type: mime_type)
+      end
+
+      private
+
+      def build_codex_body(messages, system, tools, **options)
+        instructions = Array(system).filter_map { |s| s.is_a?(Hash) ? s[:content] : s }.join("\n")
+        instructions = "You are a helpful assistant." if instructions.empty?
+
+        body = {
+          model: model_key,
+          instructions: instructions,
+          input: messages,
+          store: false,
+          include: [ "reasoning.encrypted_content" ],
+          stream: true
+        }
+
+        body[:tools] = tools if tools
+        body.merge!(options)
+
+        body
+      end
+
+      def codex_headers(account_id: nil, **options)
+        effective_account_id = account_id || @account_id
+
+        headers = {
+          "content-type" => "application/json",
+          "Authorization" => "Bearer #{api_key}",
+          "OpenAI-Beta" => "responses=experimental"
+        }
+        headers["chatgpt-account-id"] = effective_account_id if effective_account_id
+        headers
+      end
+
+      def post_codex_stream(url_part, body = nil, account_id: nil, &block)
+        endpoint = "#{CODEX_BASE_ENDPOINT}/#{url_part.sub(%r{^/}, "")}"
+        uri = URI(endpoint)
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = true
+        http.read_timeout = 480
+        http.open_timeout = 10
+
+        body.merge!(stream: true)
+        request = Net::HTTP::Post.new(uri)
+        codex_headers(account_id: account_id).each { |key, value| request[key] = value }
+        prompt_cache_key = body.delete(:prompt_cache_key)
+        request[:session_id] = prompt_cache_key if prompt_cache_key
+
+        request.body = body.to_json if body
+
+        http.request(request) do |response|
+          unless response.code.to_i == 200
+            full_body = +""
+            response.read_body { |chunk| full_body << chunk }
+            response.instance_variable_set(:@body, full_body)
+            response.instance_variable_set(:@read, true)
+            handle_error(response)
+          end
+
+          parse_sse_stream(response, &block)
+        end
+      end
+
+      def build_headers
+        {
+          "content-type" => "application/json",
+          "Authorization" => "Bearer #{api_key}"
+        }
+      end
+
+      def handle_client_specific_errors(response, error)
+        # OpenAI uses 'code' instead of 'type' for error codes
+        error_code = error["code"]
+        error_message = error["message"]
+
+        if Errors.context_overflow_message?(error_message)
+          raise Errors::PromptTooLong.new(error_message, error_code)
+        end
+
+        case response.code.to_i
+        when 429
+          raise Errors::RateLimitError.new(error_message, error_code)
+        when 503
+          raise Errors::OverloadError.new(error_message, error_code)
+        end
+        # If we get here, we didn't handle it specifically
+        fallback_body = response.body.to_s.strip
+        fallback_message = if fallback_body.empty?
+          "OpenAI request failed with status #{response.code}"
+        else
+          "OpenAI request failed with status #{response.code}: #{fallback_body}"
+        end
+
+        message = error["message"] || fallback_message
+        raise Errors::APIStatusError.new(message, error_code)
+      end
+    end
+  end
+end

data/lib/llm_gateway/clients/openai_codex/oauth_flow.rb

@@ -0,0 +1,258 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "json"
+require "securerandom"
+require "digest"
+require "base64"
+require "uri"
+require "time"
+
+module LlmGateway
+  module Clients
+    class OpenAI
+      class OAuthFlow
+        CLIENT_ID    = "app_EMoamEEZ73f0CkXaXp7hrann"
+        AUTHORIZE_URL = "https://auth.openai.com/oauth/authorize"
+        TOKEN_URL     = "https://auth.openai.com/oauth/token"
+        REDIRECT_URI  = "http://localhost:1455/auth/callback"
+        SCOPE         = "openid profile email offline_access"
+        JWT_CLAIM_PATH = "https://api.openai.com/auth"
+
+        attr_reader :client_id, :redirect_uri, :scope
+
+        def initialize(
+          client_id: CLIENT_ID,
+          redirect_uri: REDIRECT_URI,
+          scope: SCOPE
+        )
+          @client_id   = client_id
+          @redirect_uri = redirect_uri
+          @scope        = scope
+        end
+
+        # Step 1: Generate the authorization URL and PKCE values.
+        # Returns { authorization_url:, code_verifier:, state: }
+        def start(state: SecureRandom.hex(16))
+          code_verifier, code_challenge = generate_pkce
+
+          {
+            authorization_url: build_authorization_url(code_challenge, state),
+            code_verifier: code_verifier,
+            state: state
+          }
+        end
+
+        # Step 2: Exchange the authorization code for tokens.
+        # Accepts a raw code string, a full redirect URL, or code#state format.
+        # Returns { access_token:, refresh_token:, expires_at:, account_id: }
+        def exchange_code(input, code_verifier, expected_state: nil)
+          code = parse_authorization_input(input, expected_state)
+          raise ArgumentError, "Missing authorization code" unless code
+
+          uri     = URI(TOKEN_URL)
+          http    = Net::HTTP.new(uri.host, uri.port)
+          http.use_ssl      = true
+          http.read_timeout = 30
+          http.open_timeout = 10
+
+          request = Net::HTTP::Post.new(uri)
+          request["Content-Type"] = "application/x-www-form-urlencoded"
+          request.body = URI.encode_www_form(
+            grant_type: "authorization_code",
+            client_id: @client_id,
+            code: code,
+            code_verifier: code_verifier,
+            redirect_uri: @redirect_uri
+          )
+
+          response = http.request(request)
+
+          if response.code.to_i == 200
+            data = JSON.parse(response.body)
+
+            unless data["access_token"] && data["refresh_token"] && data["expires_in"]
+              raise "Token response missing required fields: #{data.keys.join(", ")}"
+            end
+
+            expires_at = Time.now + data["expires_in"].to_i
+            account_id = self.class.extract_account_id_from_token(data["access_token"])
+            raise "Failed to extract account_id from access token" unless account_id
+
+            {
+              access_token: data["access_token"],
+              refresh_token: data["refresh_token"],
+              expires_at: expires_at,
+              account_id: account_id
+            }
+          else
+            error_body = parse_error_body(response.body)
+            raise "OAuth token exchange failed (#{response.code}): #{error_body["error_description"] || error_body["error"] || response.body}"
+          end
+        end
+
+        # Parse a callback URL (or query string) into { code:, state: }
+        def parse_callback(callback_url)
+          uri    = URI.parse(callback_url)
+          params = URI.decode_www_form(uri.query.to_s).to_h
+          code   = params["code"]
+          raise ArgumentError, "Callback URL is missing code parameter" if code.nil? || code.empty?
+
+          { code: code, state: params["state"] }
+        rescue URI::InvalidURIError => e
+          raise ArgumentError, "Invalid callback URL: #{e.message}"
+        end
+
+        # Interactive OAuth flow: print URL, wait for paste, return tokens.
+        # Returns { access_token:, refresh_token:, expires_at:, account_id: }
+        def login
+          flow = start
+
+          puts "Open this URL to authorize with OpenAI:"
+          puts flow[:authorization_url]
+          puts
+          puts "After logging in your browser will redirect to localhost (the page won't load)."
+          puts "Copy the full URL from your browser's address bar and paste it below."
+          puts
+          print "Paste the redirect URL (or authorization code): "
+
+          tty   = File.open("/dev/tty", "r")
+          input = tty.gets&.strip
+          tty.close
+
+          raise "No authorization code provided" if input.nil? || input.empty?
+
+          exchange_code(input, flow[:code_verifier], expected_state: flow[:state])
+        end
+
+        # Refresh an existing access token (class method).
+        # Returns { access_token:, refresh_token:, expires_at:, account_id: }
+        def self.refresh_access_token(refresh_token, client_id: CLIENT_ID)
+          uri  = URI(TOKEN_URL)
+          http = Net::HTTP.new(uri.host, uri.port)
+          http.use_ssl      = true
+          http.read_timeout = 30
+          http.open_timeout = 10
+
+          request = Net::HTTP::Post.new(uri)
+          request["Content-Type"] = "application/x-www-form-urlencoded"
+          request.body = URI.encode_www_form(
+            grant_type: "refresh_token",
+            refresh_token: refresh_token,
+            client_id: client_id
+          )
+
+          response = http.request(request)
+
+          if response.code.to_i == 200
+            data = JSON.parse(response.body)
+
+            unless data["access_token"] && data["refresh_token"] && data["expires_in"]
+              raise "Token refresh response missing required fields"
+            end
+
+            expires_at = Time.now + data["expires_in"].to_i
+            account_id = extract_account_id_from_token(data["access_token"])
+
+            {
+              access_token: data["access_token"],
+              refresh_token: data["refresh_token"],
+              expires_at: expires_at,
+              account_id: account_id
+            }
+          else
+            error_body = begin
+              JSON.parse(response.body)
+            rescue StandardError
+              {}
+            end
+            raise "Token refresh failed (#{response.code}): #{error_body["error_description"] || error_body["error"] || response.body}"
+          end
+        end
+
+        # Extract the ChatGPT account_id from a JWT access token.
+        def self.extract_account_id_from_token(token)
+          parts = token.to_s.split(".")
+          return nil unless parts.length == 3
+
+          payload_b64 = parts[1]
+          # Re-pad to a multiple of 4 for base64 decoding
+          payload_b64 += "=" * ((4 - payload_b64.length % 4) % 4)
+          payload = JSON.parse(Base64.urlsafe_decode64(payload_b64))
+
+          auth       = payload[JWT_CLAIM_PATH]
+          account_id = auth&.dig("chatgpt_account_id")
+
+          account_id.is_a?(String) && !account_id.empty? ? account_id : nil
+        rescue StandardError
+          nil
+        end
+
+        private
+
+        def generate_pkce
+          code_verifier  = SecureRandom.urlsafe_base64(32).tr("=", "")
+          digest         = Digest::SHA256.digest(code_verifier)
+          code_challenge = Base64.urlsafe_encode64(digest).tr("=", "")
+          [ code_verifier, code_challenge ]
+        end
+
+        def build_authorization_url(code_challenge, state)
+          params = {
+            response_type: "code",
+            client_id: @client_id,
+            redirect_uri: @redirect_uri,
+            scope: @scope,
+            code_challenge: code_challenge,
+            code_challenge_method: "S256",
+            state: state,
+            id_token_add_organizations: "true",
+            codex_cli_simplified_flow: "true",
+            originator: "llm_gateway"
+          }
+          "#{AUTHORIZE_URL}?#{URI.encode_www_form(params)}"
+        end
+
+        def parse_authorization_input(input, expected_state = nil)
+          return nil if input.nil? || input.empty?
+
+          value = input.to_s.strip
+
+          # Full URL
+          if value.start_with?("http://", "https://")
+            parsed = parse_callback(value)
+            if expected_state && parsed[:state] && parsed[:state] != expected_state
+              raise "State mismatch"
+            end
+            return parsed[:code]
+          end
+
+          # code#state shorthand
+          if value.include?("#")
+            code, state = value.split("#", 2)
+            raise "State mismatch" if expected_state && state && state != expected_state
+            return code
+          end
+
+          # Query-string fragment
+          if value.include?("code=")
+            params = URI.decode_www_form(value).to_h
+            if expected_state && params["state"] && params["state"] != expected_state
+              raise "State mismatch"
+            end
+            return params["code"]
+          end
+
+          # Raw code
+          value
+        end
+
+        def parse_error_body(body)
+          JSON.parse(body)
+        rescue StandardError
+          {}
+        end
+      end
+    end
+  end
+end

data/lib/llm_gateway/clients/openai_codex/token_manager.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "json"
+require "time"
+
+module LlmGateway
+  module Clients
+    class OpenAI
+      class TokenManager
+        attr_reader :access_token, :refresh_token, :expires_at, :account_id, :client_id
+        attr_accessor :on_token_refresh
+
+        def initialize(
+          access_token: nil,
+          refresh_token:,
+          expires_at: nil,
+          account_id: nil,
+          client_id: OAuthFlow::CLIENT_ID
+        )
+          @access_token  = access_token
+          @refresh_token = refresh_token
+          @expires_at    = parse_expires_at(expires_at)
+          @account_id    = account_id
+          @client_id     = client_id
+          @on_token_refresh = nil
+        end
+
+        def token_expired?
+          return true if @expires_at.nil?
+
+          Time.now >= @expires_at
+        end
+
+        def ensure_valid_token
+          refresh_access_token! if token_expired?
+        end
+
+        def refresh_access_token!
+          raise ArgumentError, "Cannot refresh token: refresh_token not provided" unless @refresh_token
+
+          result = OAuthFlow.refresh_access_token(@refresh_token, client_id: @client_id)
+
+          @access_token  = result[:access_token]
+          @refresh_token = result[:refresh_token]
+          @expires_at    = result[:expires_at]
+          @account_id    = result[:account_id] if result[:account_id]
+
+          @on_token_refresh&.call(@access_token, @refresh_token, @expires_at)
+
+          @access_token
+        end
+
+        private
+
+        def parse_expires_at(expires)
+          case expires
+          when Time
+            expires
+          when String
+            Time.parse(expires)
+          when Integer, Float
+            Time.at(expires.to_i)
+          else
+            nil
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_gateway/errors.rb

@@ -31,6 +31,27 @@ module LlmGateway
     class UnsupportedProvider < ClientError; end
     class MissingMapperForProvider < ClientError; end
 
+    OVERFLOW_PATTERNS = [
+      /prompt is too long/i, # Anthropic
+      /exceeds the context window/i, # OpenAI
+      /reduce the length of the messages/i, # Groq
+      /maximum context length is \d+ tokens/i,
+      /context[_ ]length[_ ]exceeded/i,
+      /too many tokens/i,
+      /token limit exceeded/i,
+      /request too large.*tokens per min/i, # OpenAI TPM wording
+      /input tokens per minute/i, # Anthropic TPM wording
+      /reduce the prompt length/i,
+      /input or output tokens must be reduced/i
+    ].freeze
+
+    def self.context_overflow_message?(message)
+      text = message.to_s
+      return false if text.empty?
+
+      OVERFLOW_PATTERNS.any? { |pattern| pattern.match?(text) }
+    end
+
     class PromptError < BaseError; end
 
     class HallucinationError < PromptError; end

data/lib/llm_gateway/prompt.rb

@@ -30,6 +30,13 @@ module LlmGateway
 
     def initialize(model)
       @model = model
+      @connection = if model.is_a?(String)
+        LlmGateway.configured_clients.values.find do |client|
+          client.client.model_key == model
+        end
+      else
+        model
+      end
     end
 
     def run
@@ -55,7 +62,11 @@ module LlmGateway
     end
 
     def post
-      LlmGateway::Client.chat(model, prompt, tools: tools, system: system_prompt)
+      if @connection
+        @connection.chat(prompt, tools: tools, system: system_prompt)
+      else
+        LlmGateway::Client.chat(model, prompt, tools: tools, system: system_prompt)
+      end
     end
 
     def tools

data/lib/llm_gateway/provider_registry.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module LlmGateway
+  class ProviderRegistry
+    class << self
+      def register(name, client:, adapter:)
+        registry[name.to_s] = { client: client, adapter: adapter }
+      end
+
+      def resolve(name)
+        name = name.to_s
+        entry = registry[name]
+        raise Errors::UnsupportedProvider, "Unknown provider: #{name}" unless entry
+
+        entry
+      end
+
+      def registered?(name)
+        registry.key?(name.to_s)
+      end
+
+      def providers
+        registry.keys
+      end
+
+      def reset!
+        @registry = {}
+      end
+
+      private
+
+      def registry
+        @registry ||= {}
+      end
+    end
+  end
+end

data/lib/llm_gateway/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LlmGateway
-  VERSION = "0.3.0"
+  VERSION = "0.4.0"
 end