llm_gateway 0.2.0 → 0.4.0

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

@@ -0,0 +1,162 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "json"
+require "securerandom"
+require "digest"
+require "base64"
+require "uri"
+require "time"
+
+module LlmGateway
+  module Clients
+    module ClaudeCode
+      class OAuthFlow
+        CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
+        TOKEN_URL = "https://api.anthropic.com/v1/oauth/token"
+        AUTH_URL = "https://claude.ai/oauth/authorize"
+        REDIRECT_URI = "https://console.anthropic.com/oauth/code/callback"
+        DEFAULT_SCOPES = "org:create_api_key user:profile user:inference"
+
+        attr_reader :client_id, :redirect_uri, :scopes
+
+        def initialize(
+          client_id: CLIENT_ID,
+          redirect_uri: REDIRECT_URI,
+          scopes: DEFAULT_SCOPES
+        )
+          @client_id = client_id
+          @redirect_uri = redirect_uri
+          @scopes = scopes
+        end
+
+        # Step 1: Generate the authorization URL for the user to visit.
+        # Returns a hash with everything needed to complete the flow later.
+        def start(state: SecureRandom.hex(16))
+          code_verifier, code_challenge = generate_pkce
+
+          auth_url = build_authorization_url(code_challenge, state)
+
+          {
+            authorization_url: auth_url,
+            code_verifier: code_verifier,
+            state: state
+          }
+        end
+
+        # Step 2: Exchange the authorization code for tokens.
+        # Accepts one of:
+        # - "code#state" (legacy format)
+        # - a raw authorization code, with state passed separately
+        # - a full callback URL containing ?code=...&state=...
+        # Returns { access_token:, refresh_token:, expires_at: }
+        def exchange_code(auth_code_or_callback, code_verifier, state: nil)
+          code, resolved_state = extract_code_and_state(auth_code_or_callback, state)
+
+          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/json"
+
+          request.body = {
+            grant_type: "authorization_code",
+            client_id: @client_id,
+            code: code,
+            state: resolved_state || "",
+            redirect_uri: @redirect_uri,
+            code_verifier: code_verifier
+          }.to_json
+
+          response = http.request(request)
+
+          if response.code.to_i == 200
+            data = JSON.parse(response.body)
+
+            expires_at = if data["expires_in"]
+                           Time.now + data["expires_in"].to_i
+            elsif data["expires_at"]
+                           Time.parse(data["expires_at"])
+            end
+
+            {
+              access_token: data["access_token"],
+              refresh_token: data["refresh_token"],
+              expires_at: expires_at
+            }
+          else
+            error_body = begin
+              JSON.parse(response.body)
+            rescue StandardError
+              {}
+            end
+            raise Errors::AuthenticationError.new(
+              "OAuth token exchange failed: #{error_body["error_description"] || error_body["error"] || response.body}",
+              error_body["error"]
+            )
+          end
+        end
+
+        def parse_callback(callback_url)
+          uri = URI(callback_url)
+          code = uri.query && URI.decode_www_form(uri.query).to_h["code"]
+          state = uri.query && URI.decode_www_form(uri.query).to_h["state"]
+
+          raise ArgumentError, "Callback URL is missing code parameter" if code.nil? || code.empty?
+
+          { code: code, state: state }
+        rescue URI::InvalidURIError => e
+          raise ArgumentError, "Invalid callback URL: #{e.message}"
+        end
+
+        private
+
+        def extract_code_and_state(auth_code_or_callback, state)
+          value = auth_code_or_callback.to_s.strip
+          raise ArgumentError, "Authorization code is required" if value.empty?
+
+          if looks_like_url?(value)
+            callback = parse_callback(value)
+            [ callback[:code], callback[:state] || state ]
+          elsif value.include?("#")
+            code, parsed_state = value.split("#", 2)
+            [ code, parsed_state || state ]
+          else
+            [ value, state ]
+          end
+        end
+
+        def looks_like_url?(value)
+          value.start_with?("http://", "https://")
+        end
+
+        def generate_pkce
+          code_verifier = [ SecureRandom.random_bytes(32) ].pack("m0").tr("+/", "-_").tr("=", "")
+
+          digest = Digest::SHA256.digest(code_verifier)
+          code_challenge = [ digest ].pack("m0").tr("+/", "-_").tr("=", "")
+
+          [ code_verifier, code_challenge ]
+        end
+
+        def build_authorization_url(code_challenge, state)
+          params = {
+            code: "true",
+            client_id: @client_id,
+            response_type: "code",
+            redirect_uri: @redirect_uri,
+            scope: @scopes,
+            code_challenge: code_challenge,
+            code_challenge_method: "S256",
+            state: state
+          }
+
+          "#{AUTH_URL}?#{URI.encode_www_form(params)}"
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "json"
+require "time"
+
+module LlmGateway
+  module Clients
+    module ClaudeCode
+      class TokenManager
+        TOKEN_URL = "https://api.anthropic.com/v1/oauth/token"
+        CLIENT_ID = OAuthFlow::CLIENT_ID
+
+        attr_reader :refresh_token, :expires_at, :client_id, :client_secret, :access_token
+        attr_accessor :on_token_refresh
+
+        def initialize(
+          access_token: nil,
+          refresh_token:,
+          expires_at: nil,
+          client_id: CLIENT_ID,
+          client_secret: nil
+        )
+          @access_token = access_token
+          @refresh_token = refresh_token
+          @expires_at = parse_expires_at(expires_at)
+          @client_id = client_id
+          @client_secret = client_secret
+          @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
+          raise ArgumentError, "Cannot refresh token: client_id not provided" unless @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/json"
+
+          request_body = {
+            grant_type: "refresh_token",
+            client_id: @client_id,
+            refresh_token: @refresh_token
+          }
+          request_body[:client_secret] = @client_secret if @client_secret
+
+          request.body = request_body.to_json
+
+          response = http.request(request)
+
+          if response.code.to_i == 200
+            data = JSON.parse(response.body)
+            @access_token = data["access_token"]
+
+            if data["refresh_token"]
+              @refresh_token = data["refresh_token"]
+            end
+
+            if data["expires_in"]
+              @expires_at = Time.now + data["expires_in"].to_i
+            elsif data["expires_at"]
+              @expires_at = Time.parse(data["expires_at"])
+            end
+
+            @on_token_refresh&.call(@access_token, @refresh_token, @expires_at)
+
+            @access_token
+          else
+            error_body = begin
+              JSON.parse(response.body)
+            rescue StandardError
+              {}
+            end
+            raise Errors::AuthenticationError.new(
+              "Failed to refresh token: #{error_body['error'] || response.body}",
+              error_body["error_code"]
+            )
+          end
+        end
+
+        private
+
+        def parse_expires_at(expires)
+          case expires
+          when Time
+            expires
+          when String
+            Time.parse(expires)
+          when Integer
+            Time.at(expires)
+          else
+            nil
+          end
+        end
+      end
+    end
+  end
+end

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