llm_gateway 0.2.0 → 0.4.0

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

@@ -28,6 +28,29 @@ module LlmGateway
     class UnknownError < ClientError; end
     class PromptTooLong < BadRequestError; end
     class UnsupportedModel < ClientError; end
+    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
 

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.2.0"
+  VERSION = "0.4.0"
 end

data/lib/llm_gateway.rb

@@ -8,17 +8,176 @@ require_relative "llm_gateway/client"
 require_relative "llm_gateway/prompt"
 require_relative "llm_gateway/tool"
 
-# Load adapters - order matters for inheritance
-require_relative "llm_gateway/adapters/claude/client"
-require_relative "llm_gateway/adapters/claude/input_mapper"
-require_relative "llm_gateway/adapters/claude/output_mapper"
-require_relative "llm_gateway/adapters/groq/client"
-require_relative "llm_gateway/adapters/groq/input_mapper"
-require_relative "llm_gateway/adapters/groq/output_mapper"
-require_relative "llm_gateway/adapters/open_ai/client"
-require_relative "llm_gateway/adapters/open_ai/input_mapper"
-require_relative "llm_gateway/adapters/open_ai/output_mapper"
+# Load clients - order matters for inheritance
+require_relative "llm_gateway/clients/anthropic"
+require_relative "llm_gateway/clients/claude_code/oauth_flow"
+require_relative "llm_gateway/clients/claude_code/token_manager"
+require_relative "llm_gateway/clients/openai"
+require_relative "llm_gateway/clients/openai_codex/oauth_flow"
+require_relative "llm_gateway/clients/openai_codex/token_manager"
+require_relative "llm_gateway/clients/groq"
+
+# Load adapters
+require_relative "llm_gateway/adapters/option_mapper"
+require_relative "llm_gateway/adapters/anthropic_option_mapper"
+require_relative "llm_gateway/adapters/structs"
+
+require_relative "llm_gateway/adapters/anthropic/input_mapper"
+require_relative "llm_gateway/adapters/anthropic/output_mapper"
+require_relative "llm_gateway/adapters/openai/file_output_mapper"
+require_relative "llm_gateway/adapters/openai/prompt_cache_option_mapper"
+require_relative "llm_gateway/adapters/openai/chat_completions/input_mapper"
+require_relative "llm_gateway/adapters/openai/chat_completions/output_mapper"
+require_relative "llm_gateway/adapters/openai/chat_completions/option_mapper"
+require_relative "llm_gateway/adapters/openai/file_output_mapper"
+require_relative "llm_gateway/adapters/openai/responses/input_mapper"
+require_relative "llm_gateway/adapters/openai/responses/output_mapper"
+require_relative "llm_gateway/adapters/openai/responses/option_mapper"
+
+# Load adapter classes
+require_relative "llm_gateway/adapters/adapter"
+require_relative "llm_gateway/adapters/anthropic/messages_adapter"
+require_relative "llm_gateway/adapters/openai/chat_completions_adapter"
+require_relative "llm_gateway/adapters/openai/responses_adapter"
+require_relative "llm_gateway/adapters/openai_codex/responses_adapter"
+require_relative "llm_gateway/adapters/groq/chat_completions_adapter"
+
+# Load provider registry
+require_relative "llm_gateway/provider_registry"
 
 module LlmGateway
   class Error < StandardError; end
+
+  # Direction constants for message mappers
+  DIRECTION_IN = :in
+  DIRECTION_OUT = :out
+
+  # Backward-compatible aliases for renamed clients/adapters
+  module Clients
+    Claude = Anthropic
+    OpenAi = OpenAI
+  end
+
+  module Adapters
+    module Claude
+      Client = LlmGateway::Clients::Anthropic
+      MessagesAdapter = LlmGateway::Adapters::Anthropic::MessagesAdapter
+      InputMapper = LlmGateway::Adapters::Anthropic::InputMapper
+      OutputMapper = LlmGateway::Adapters::Anthropic::OutputMapper
+      StreamMapper = LlmGateway::Adapters::Anthropic::StreamMapper
+      BidirectionalMessageMapper = LlmGateway::Adapters::Anthropic::BidirectionalMessageMapper
+      FileOutputMapper = LlmGateway::Adapters::Anthropic::FileOutputMapper
+    end
+
+    module Anthropic
+      Client = LlmGateway::Clients::Anthropic
+    end
+
+    module OpenAI
+      Client = LlmGateway::Clients::OpenAI
+    end
+
+    module OpenAi
+      Client = LlmGateway::Clients::OpenAI
+      ChatCompletionsAdapter = LlmGateway::Adapters::OpenAI::ChatCompletionsAdapter
+      ResponsesAdapter = LlmGateway::Adapters::OpenAI::ResponsesAdapter
+      PromptCacheOptionMapper = LlmGateway::Adapters::OpenAI::PromptCacheOptionMapper
+      FileOutputMapper = LlmGateway::Adapters::OpenAI::FileOutputMapper
+      ChatCompletions = LlmGateway::Adapters::OpenAI::ChatCompletions
+      Responses = LlmGateway::Adapters::OpenAI::Responses
+    end
+
+    module OpenAICodex
+      Client = LlmGateway::Clients::OpenAI
+    end
+
+    module OpenAiCodex
+      Client = LlmGateway::Clients::OpenAI
+      ResponsesAdapter = LlmGateway::Adapters::OpenAICodex::ResponsesAdapter
+      InputMapper = LlmGateway::Adapters::OpenAICodex::InputMapper
+      OptionMapper = LlmGateway::Adapters::OpenAICodex::OptionMapper
+    end
+
+    module Groq
+      Client = LlmGateway::Clients::Groq
+    end
+  end
+
+  def self.build_provider(config)
+    config = config.transform_keys(&:to_sym)
+    provider_name = config.delete(:provider)
+    entry = ProviderRegistry.resolve(provider_name)
+
+    client = entry[:client].new(**config)
+    entry[:adapter].new(client)
+  end
+
+  def self.configure(configs)
+    @configured_clients ||= {}
+
+    configs.each do |entry|
+      name = entry[:name] || entry["name"]
+      config = entry[:config] || entry["config"]
+
+      raise ArgumentError, "Each config entry must have a :name" unless name
+
+      client = build_provider(config)
+      @configured_clients[name.to_sym] = client
+
+      define_singleton_method(name.to_sym) { @configured_clients[name.to_sym] }
+    end
+  end
+
+  def self.configured_clients
+    @configured_clients ||= {}
+  end
+
+  def self.reset_configuration!
+    @configured_clients&.each_key do |name|
+      singleton_class.remove_method(name) if respond_to?(name)
+    end
+    @configured_clients = {}
+  end
+
+  # Register built-in providers (canonical keys)
+  ProviderRegistry.register("anthropic_messages",
+    client: Clients::Anthropic,
+    adapter: Adapters::Anthropic::MessagesAdapter)
+
+  ProviderRegistry.register("openai_completions",
+    client: Clients::OpenAI,
+    adapter: Adapters::OpenAI::ChatCompletionsAdapter)
+
+  ProviderRegistry.register("openai_responses",
+    client: Clients::OpenAI,
+    adapter: Adapters::OpenAI::ResponsesAdapter)
+
+  ProviderRegistry.register("groq_completions",
+    client: Clients::Groq,
+    adapter: Adapters::Groq::ChatCompletionsAdapter)
+
+  ProviderRegistry.register("openai_codex",
+    client: Clients::OpenAI,
+    adapter: Adapters::OpenAICodex::ResponsesAdapter)
+
+  # Backward-compatible aliases (deprecated)
+  ProviderRegistry.register("anthropic_apikey_messages",
+    client: Clients::Anthropic,
+    adapter: Adapters::Anthropic::MessagesAdapter)
+
+  ProviderRegistry.register("openai_apikey_completions",
+    client: Clients::OpenAI,
+    adapter: Adapters::OpenAI::ChatCompletionsAdapter)
+
+  ProviderRegistry.register("openai_apikey_responses",
+    client: Clients::OpenAI,
+    adapter: Adapters::OpenAI::ResponsesAdapter)
+
+  ProviderRegistry.register("groq_apikey_completions",
+    client: Clients::Groq,
+    adapter: Adapters::Groq::ChatCompletionsAdapter)
+
+  ProviderRegistry.register("openai_oauth_codex",
+    client: Clients::OpenAI,
+    adapter: Adapters::OpenAICodex::ResponsesAdapter)
 end