From 976700bd578a885a49e8c217670be6230d4ff64d Mon Sep 17 00:00:00 2001 From: Justin Bowen Date: Mon, 27 Oct 2025 11:38:42 -0700 Subject: [PATCH] Adding xAI provider for grok models --- .../generation_provider/x_ai_provider.rb | 144 +++++++++++ .../generation_provider/x_ai_provider_test.rb | 227 ++++++++++++++++++ 2 files changed, 371 insertions(+) create mode 100644 lib/active_agent/generation_provider/x_ai_provider.rb create mode 100644 test/generation_provider/x_ai_provider_test.rb diff --git a/lib/active_agent/generation_provider/x_ai_provider.rb b/lib/active_agent/generation_provider/x_ai_provider.rb new file mode 100644 index 00000000..c1a47fb8 --- /dev/null +++ b/lib/active_agent/generation_provider/x_ai_provider.rb @@ -0,0 +1,144 @@ +# frozen_string_literal: true + +begin + gem "ruby-openai", ">= 8.1.0" + require "openai" +rescue LoadError + raise LoadError, "The 'ruby-openai >= 8.1.0' gem is required for XAIProvider. Please add it to your Gemfile and run `bundle install`." +end + +require "active_agent/action_prompt/action" +require_relative "base" +require_relative "response" +require_relative "stream_processing" +require_relative "message_formatting" +require_relative "tool_management" + +module ActiveAgent + module GenerationProvider + # XAI (Grok) Generation Provider + # Uses OpenAI-compatible API format with xAI's endpoint + class XAIProvider < Base + include StreamProcessing + include MessageFormatting + include ToolManagement + + XAI_API_HOST = "https://api.x.ai" + + def initialize(config) + super + # Support both api_key and access_token for backwards compatibility + @access_token = config["api_key"] || config["access_token"] || ENV["XAI_API_KEY"] || ENV["GROK_API_KEY"] + + unless @access_token + raise ArgumentError, "XAI API key is required. Set it in config as 'api_key', 'access_token', or via XAI_API_KEY/GROK_API_KEY environment variable." + end + + # xAI uses OpenAI-compatible client with custom endpoint + @client = OpenAI::Client.new( + access_token: @access_token, + uri_base: config["host"] || XAI_API_HOST, + log_errors: Rails.env.development? + ) + + # Default to grok-2-latest but allow configuration + @model_name = config["model"] || "grok-2-latest" + end + + def generate(prompt) + @prompt = prompt + + with_error_handling do + chat_prompt(parameters: prompt_parameters) + end + end + + def embed(prompt) + # xAI doesn't currently provide embedding models + raise NotImplementedError, "xAI does not currently support embeddings. Use a different provider for embedding tasks." + end + + protected + + # Override from StreamProcessing module - uses OpenAI format + def process_stream_chunk(chunk, message, agent_stream) + new_content = chunk.dig("choices", 0, "delta", "content") + if new_content && !new_content.blank? + message.generation_id = chunk.dig("id") + message.content += new_content + agent_stream&.call(message, new_content, false, prompt.action_name) + elsif chunk.dig("choices", 0, "delta", "tool_calls") && chunk.dig("choices", 0, "delta", "role") + message = handle_message(chunk.dig("choices", 0, "delta")) + prompt.messages << message + @response = ActiveAgent::GenerationProvider::Response.new( + prompt:, + message:, + raw_response: chunk, + raw_request: @streaming_request_params + ) + end + + if chunk.dig("choices", 0, "finish_reason") + finalize_stream(message, agent_stream) + end + end + + # Override from MessageFormatting module to handle image format (if xAI adds vision support) + def format_image_content(message) + [ { + type: "image_url", + image_url: { url: message.content } + } ] + end + + private + + # Override from ParameterBuilder to add xAI-specific parameters if needed + def build_provider_parameters + params = {} + + # Add any xAI-specific parameters here + # For now, xAI follows OpenAI's format closely + + params + end + + def chat_response(response, request_params = nil) + return @response if prompt.options[:stream] + + message_json = response.dig("choices", 0, "message") + message_json["id"] = response.dig("id") if message_json["id"].blank? + message = handle_message(message_json) + + update_context(prompt: prompt, message: message, response: response) + + @response = ActiveAgent::GenerationProvider::Response.new( + prompt: prompt, + message: message, + raw_response: response, + raw_request: request_params + ) + end + + def handle_message(message_json) + ActiveAgent::ActionPrompt::Message.new( + generation_id: message_json["id"], + content: message_json["content"], + role: message_json["role"].intern, + action_requested: message_json["finish_reason"] == "tool_calls", + raw_actions: message_json["tool_calls"] || [], + requested_actions: handle_actions(message_json["tool_calls"]), + content_type: prompt.output_schema.present? ? "application/json" : "text/plain" + ) + end + + def chat_prompt(parameters: prompt_parameters) + if prompt.options[:stream] || config["stream"] + parameters[:stream] = provider_stream + @streaming_request_params = parameters + end + chat_response(@client.chat(parameters: parameters), parameters) + end + end + end +end diff --git a/test/generation_provider/x_ai_provider_test.rb b/test/generation_provider/x_ai_provider_test.rb new file mode 100644 index 00000000..01f03170 --- /dev/null +++ b/test/generation_provider/x_ai_provider_test.rb @@ -0,0 +1,227 @@ +require "test_helper" +require "active_agent/generation_provider/x_ai_provider" + +# Test for xAI Provider gem loading and configuration +class XAIProviderTest < ActiveAgentTestCase + # Test the gem load rescue block + test "gem load rescue block provides correct error message" do + # Since we can't easily simulate the gem not being available without complex mocking, + # we'll test that the error message is correct by creating a minimal reproduction + expected_message = "The 'ruby-openai >= 8.1.0' gem is required for XAIProvider. Please add it to your Gemfile and run `bundle install`." + + # Verify the rescue block pattern exists in the source code + provider_file_path = File.join(Rails.root, "../../lib/active_agent/generation_provider/x_ai_provider.rb") + provider_source = File.read(provider_file_path) + + assert_includes provider_source, "begin" + assert_includes provider_source, 'gem "ruby-openai"' + assert_includes provider_source, 'require "openai"' + assert_includes provider_source, "rescue LoadError" + assert_includes provider_source, expected_message + + # Test the actual error by creating a minimal scenario + test_code = <<~RUBY + begin + gem "nonexistent-openai-gem" + require "nonexistent-openai-gem" + rescue LoadError + raise LoadError, "#{expected_message}" + end + RUBY + + error = assert_raises(LoadError) do + eval(test_code) + end + + assert_equal expected_message, error.message + end + + test "loads successfully when ruby-openai gem is available" do + # This test ensures the provider loads correctly when the gem is present + # Since the gem is already loaded in our test environment, this should work + + # Verify the class exists and can be instantiated with valid config + assert defined?(ActiveAgent::GenerationProvider::XAIProvider) + + config = { + "service" => "XAI", + "api_key" => "test-key", + "model" => "grok-2-latest" + } + + assert_nothing_raised do + ActiveAgent::GenerationProvider::XAIProvider.new(config) + end + end + + # Test configuration loading and presence + test "raises error when xAI API key is missing" do + config = { + "service" => "XAI", + "model" => "grok-2-latest" + # Missing api_key + } + + error = assert_raises(ArgumentError) do + ActiveAgent::GenerationProvider::XAIProvider.new(config) + end + + assert_includes error.message, "XAI API key is required" + end + + test "loads configuration from active_agent.yml when present" do + # Mock a configuration + mock_config = { + "test" => { + "xai" => { + "service" => "XAI", + "api_key" => "test-api-key", + "model" => "grok-2-latest", + "temperature" => 0.7 + } + } + } + + ActiveAgent.instance_variable_set(:@config, mock_config) + + # Set Rails environment for testing + rails_env = ENV["RAILS_ENV"] + ENV["RAILS_ENV"] = "test" + + config = ApplicationAgent.configuration(:xai) + + assert_equal "XAI", config.config["service"] + assert_equal "test-api-key", config.config["api_key"] + assert_equal "grok-2-latest", config.config["model"] + assert_equal 0.7, config.config["temperature"] + + # Restore original environment + ENV["RAILS_ENV"] = rails_env + end + + test "loads configuration from environment-specific section" do + mock_config = { + "development" => { + "xai" => { + "service" => "XAI", + "api_key" => "dev-api-key", + "model" => "grok-2-latest" + } + }, + "test" => { + "xai" => { + "service" => "XAI", + "api_key" => "test-api-key", + "model" => "grok-2-latest" + } + } + } + + ActiveAgent.instance_variable_set(:@config, mock_config) + + # Test development configuration + original_env = ENV["RAILS_ENV"] + ENV["RAILS_ENV"] = "development" + + config = ApplicationAgent.configuration(:xai) + assert_equal "dev-api-key", config.config["api_key"] + + # Test test configuration + ENV["RAILS_ENV"] = "test" + config = ApplicationAgent.configuration(:xai) + assert_equal "test-api-key", config.config["api_key"] + + ENV["RAILS_ENV"] = original_env + end + + test "xAI provider initialization with API key from environment variable" do + # Test with XAI_API_KEY env var + original_xai_key = ENV["XAI_API_KEY"] + original_grok_key = ENV["GROK_API_KEY"] + + ENV["XAI_API_KEY"] = "env-xai-key" + ENV["GROK_API_KEY"] = nil + + config = { + "service" => "XAI", + "model" => "grok-2-latest" + } + + provider = ActiveAgent::GenerationProvider::XAIProvider.new(config) + assert_equal "env-xai-key", provider.instance_variable_get(:@access_token) + + # Test with GROK_API_KEY env var + ENV["XAI_API_KEY"] = nil + ENV["GROK_API_KEY"] = "env-grok-key" + + provider = ActiveAgent::GenerationProvider::XAIProvider.new(config) + assert_equal "env-grok-key", provider.instance_variable_get(:@access_token) + + # Restore original environment + ENV["XAI_API_KEY"] = original_xai_key + ENV["GROK_API_KEY"] = original_grok_key + end + + test "xAI provider initialization with custom host" do + config = { + "service" => "XAI", + "api_key" => "test-key", + "model" => "grok-2-latest", + "host" => "https://custom-xai-host.com" + } + + provider = ActiveAgent::GenerationProvider::XAIProvider.new(config) + client = provider.instance_variable_get(:@client) + + # The OpenAI client should be configured with the custom host + assert_not_nil client + end + + test "xAI provider defaults to grok-2-latest model" do + config = { + "service" => "XAI", + "api_key" => "test-key" + # Model not specified + } + + provider = ActiveAgent::GenerationProvider::XAIProvider.new(config) + assert_equal "grok-2-latest", provider.instance_variable_get(:@model_name) + end + + test "xAI provider uses configured model" do + config = { + "service" => "XAI", + "api_key" => "test-key", + "model" => "grok-1" + } + + provider = ActiveAgent::GenerationProvider::XAIProvider.new(config) + assert_equal "grok-1", provider.instance_variable_get(:@model_name) + end + + test "xAI provider defaults to correct API host" do + config = { + "service" => "XAI", + "api_key" => "test-key" + } + + provider = ActiveAgent::GenerationProvider::XAIProvider.new(config) + assert_equal "https://api.x.ai", ActiveAgent::GenerationProvider::XAIProvider::XAI_API_HOST + end + + test "embed method raises NotImplementedError for xAI" do + config = { + "service" => "XAI", + "api_key" => "test-key" + } + + provider = ActiveAgent::GenerationProvider::XAIProvider.new(config) + mock_prompt = Minitest::Mock.new + + error = assert_raises(NotImplementedError) do + provider.embed(mock_prompt) + end + + assert_includes error.message, "xAI does not currently support embeddings" + end +end