From 2349f1ecc56b1b4d25b08d9dccd9080dc742bd06 Mon Sep 17 00:00:00 2001 From: aakashb95 Date: Fri, 26 Jun 2026 12:37:48 +0530 Subject: [PATCH] feat(captain): route llm calls by provider --- .../app/services/llm/base_ai_service.rb | 19 +++- lib/captain/base_task_service.rb | 99 +++++++++---------- lib/captain/tool_instrumentation.rb | 4 +- lib/integrations/llm_base_service.rb | 21 ++-- .../llm_instrumentation_helpers.rb | 2 +- lib/llm/credential_resolver.rb | 26 +++++ lib/llm/provider_chat.rb | 37 +++++++ spec/lib/captain/base_task_service_spec.rb | 35 +++++-- 8 files changed, 172 insertions(+), 71 deletions(-) create mode 100644 lib/llm/credential_resolver.rb create mode 100644 lib/llm/provider_chat.rb diff --git a/enterprise/app/services/llm/base_ai_service.rb b/enterprise/app/services/llm/base_ai_service.rb index bec3b5cb9..a77f16aef 100644 --- a/enterprise/app/services/llm/base_ai_service.rb +++ b/enterprise/app/services/llm/base_ai_service.rb @@ -6,7 +6,7 @@ class Llm::BaseAiService DEFAULT_MODEL = Llm::Config::DEFAULT_MODEL DEFAULT_TEMPERATURE = 1.0 - attr_reader :model, :temperature + attr_reader :model, :provider, :temperature def initialize(feature: nil, account: nil, fallback_model: nil) @llm_feature = feature @@ -19,7 +19,8 @@ class Llm::BaseAiService end def chat(model: @model, temperature: @temperature) - RubyLLM.chat(model: model).with_temperature(temperature) + chat = RubyLLM.chat(model: model, provider: provider_for_model(model), assume_model_exists: true).with_temperature(temperature) + Llm::ProviderChat.new(chat, provider: provider_for_model(model)) end private @@ -34,9 +35,13 @@ class Llm::BaseAiService def setup_model route = feature_route - return @model = route[:model] if account_override_route?(route) + if account_override_route?(route) + @model = route[:model] + return setup_provider(route) + end @model = @fallback_model.presence || installation_model.presence || route&.dig(:model) || DEFAULT_MODEL + setup_provider(route) end def feature_route @@ -53,6 +58,14 @@ class Llm::BaseAiService InstallationConfig.find_by(name: 'CAPTAIN_OPEN_AI_MODEL')&.value end + def setup_provider(route) + @provider = provider_for_model(@model, route&.dig(:provider)) + end + + def provider_for_model(model, fallback_provider = Llm::Config::DEFAULT_PROVIDER) + Llm::Models.provider_for(model) || fallback_provider || Llm::Config::DEFAULT_PROVIDER + end + def setup_temperature @temperature = DEFAULT_TEMPERATURE end diff --git a/lib/captain/base_task_service.rb b/lib/captain/base_task_service.rb index cfeb4e427..7992f9ef7 100644 --- a/lib/captain/base_task_service.rb +++ b/lib/captain/base_task_service.rb @@ -31,24 +31,24 @@ class Captain::BaseTaskService @conversation ||= account.conversations.find_by(display_id: conversation_display_id) end - def api_base - endpoint = InstallationConfig.find_by(name: 'CAPTAIN_OPEN_AI_ENDPOINT')&.value.presence || 'https://api.openai.com/' - endpoint = endpoint.chomp('/') - "#{endpoint}/v1" - end + def api_base = Llm::Config.api_base_for(llm_provider) def make_api_call(messages:, model: nil, feature: nil, schema: nil, tools: []) + llm_route = resolved_llm_route(model: model, feature: feature) + # Community edition prerequisite checks # Enterprise module handles these with more specific error messages (cloud vs self-hosted) return { error: I18n.t('captain.disabled'), error_code: 403 } unless captain_tasks_enabled? - return { error: I18n.t('captain.api_key_missing'), error_code: 401 } unless api_key_configured? + return { error: I18n.t('captain.api_key_missing'), error_code: 401 } unless api_key_configured?(llm_route[:provider]) - model = resolved_model(model: model, feature: feature) - instrumentation_params = build_instrumentation_params(model, messages) - instrumentation_method = tools.any? ? :instrument_tool_session : :instrument_llm_call + @llm_provider = llm_route[:provider] + model = llm_route[:model] + request_tools = Llm::Config.supports_tools_and_schema?(llm_route[:provider]) ? tools : [] + instrumentation_params = build_instrumentation_params(model, messages, llm_route[:provider]) + instrumentation_method = request_tools.any? ? :instrument_tool_session : :instrument_llm_call response = send(instrumentation_method, instrumentation_params) do - execute_ruby_llm_request(model: model, messages: messages, schema: schema, tools: tools) + execute_ruby_llm_request(llm_route: llm_route, messages: messages, schema: schema, tools: request_tools) end return response unless build_follow_up_context? && response[:message].present? @@ -56,20 +56,28 @@ class Captain::BaseTaskService response.merge(follow_up_context: build_follow_up_context(messages, response)) end - def resolved_model(model:, feature:) - return model if feature.blank? + def resolved_llm_route(model:, feature:) + return explicit_model_route(model) if feature.blank? route = Llm::FeatureRouter.resolve(feature: feature, account: account) - return model if model.present? && route[:source] == :default + resolved_model = model.present? && route[:source] == :default ? model : route[:model] - route[:model] + route.merge(model: resolved_model, provider: provider_for_model(resolved_model, route[:provider])) end - def execute_ruby_llm_request(model:, messages:, schema: nil, tools: []) - credential = llm_credential + def explicit_model_route(model) + resolved_model = model.presence || GPT_MODEL + { model: resolved_model, provider: provider_for_model(resolved_model), source: :explicit } + end - Llm::Config.with_api_key(credential[:api_key], api_base: api_base) do |context| - chat = build_chat(context, model: model, messages: messages, schema: schema, tools: tools) + def provider_for_model(model, fallback_provider = Llm::Config::DEFAULT_PROVIDER) = Llm::Models.provider_for(model) || fallback_provider + + def execute_ruby_llm_request(llm_route:, messages:, schema: nil, tools: []) + provider = llm_route[:provider] + credential = llm_credential(provider) + + Llm::Config.with_api_key(credential[:api_key], provider: provider, api_base: api_base) do |context| + chat = build_chat(context, llm_route: llm_route, messages: messages, schema: schema, tools: tools) conversation_messages = messages.reject { |m| m[:role] == 'system' } return { error: 'No conversation messages provided', error_code: 400, request_messages: messages } if conversation_messages.empty? @@ -82,15 +90,17 @@ class Captain::BaseTaskService { error: e.message, request_messages: messages } end - def build_chat(context, model:, messages:, schema: nil, tools: []) - chat = context.chat(model: model) + def build_chat(context, llm_route:, messages:, schema: nil, tools: []) + model = llm_route[:model] + provider = llm_route[:provider] + chat = Llm::ProviderChat.new(context.chat(model: model, provider: provider, assume_model_exists: true), provider: provider) system_msg = messages.find { |m| m[:role] == 'system' } chat.with_instructions(system_msg[:content]) if system_msg chat.with_schema(schema) if schema if tools.any? tools.each { |tool| chat = chat.with_tool(tool) } - chat.on_end_message { |message| record_generation(chat, message, model) } + chat.on_end_message { |message| record_generation(chat, message, model, provider) } end chat @@ -116,13 +126,14 @@ class Captain::BaseTaskService } end - def build_instrumentation_params(model, messages) + def build_instrumentation_params(model, messages, provider) { span_name: "llm.#{event_name}", account_id: account.id, conversation_id: conversation&.display_id, feature_name: event_name, model: model, + provider: provider, messages: messages, temperature: nil, metadata: instrumentation_metadata @@ -155,9 +166,7 @@ class Captain::BaseTaskService messages end - def captain_tasks_enabled? - account.feature_enabled?('captain_tasks') - end + def captain_tasks_enabled? = account.feature_enabled?('captain_tasks') # Extension point consulted by the Enterprise quota wrapper. Subclasses # whose calls should not consume captain_responses should override this to @@ -168,43 +177,27 @@ class Captain::BaseTaskService llm_credential&.dig(:source) != :hook end - def api_key_configured? - llm_credential.present? + def api_key_configured?(provider = llm_provider) = llm_credential(provider).present? + + def api_key = llm_credential&.dig(:api_key) + + def llm_provider = @llm_provider || Llm::Config::DEFAULT_PROVIDER + + def llm_credential(provider = llm_provider) + @llm_credentials ||= {} + @llm_credentials[provider.to_s] ||= Llm::CredentialResolver.new(provider: provider, openai_hook: resolved_openai_hook(provider)).resolve end - def api_key - llm_credential&.dig(:api_key) - end + def resolved_openai_hook(provider) = use_account_openai_hook? && Llm::Config.openai_provider?(provider) ? openai_hook : nil - def llm_credential - @llm_credential ||= if use_account_openai_hook? - hook_llm_credential || system_llm_credential - else - system_llm_credential - end - end + def use_account_openai_hook? = false - def use_account_openai_hook? - false - end - - def hook_llm_credential - key = openai_hook&.settings&.dig('api_key').presence - { api_key: key, source: :hook } if key - end - - def system_llm_credential - { api_key: system_api_key, source: :system } if system_api_key.present? - end + def system_llm_credential(provider = llm_provider) = Llm::CredentialResolver.new(provider: provider).resolve def openai_hook @openai_hook ||= account.hooks.find_by(app_id: 'openai', status: 'enabled') end - def system_api_key - @system_api_key ||= InstallationConfig.find_by(name: 'CAPTAIN_OPEN_AI_API_KEY')&.value - end - def exception_tracking_account account end diff --git a/lib/captain/tool_instrumentation.rb b/lib/captain/tool_instrumentation.rb index 3c4a1c76e..f244d99c2 100644 --- a/lib/captain/tool_instrumentation.rb +++ b/lib/captain/tool_instrumentation.rb @@ -38,13 +38,13 @@ module Captain::ToolInstrumentation span.status = OpenTelemetry::Trace::Status.error(error.to_s.truncate(1000)) end - def record_generation(chat, message, model) + def record_generation(chat, message, model, provider = Llm::Config::DEFAULT_PROVIDER) return unless ChatwootApp.otel_enabled? return unless message.respond_to?(:role) && message.role.to_s == 'assistant' tracer.in_span("llm.#{event_name}.generation") do |span| apply_current_langfuse_attributes(span) - span.set_attribute(ATTR_GEN_AI_PROVIDER, 'openai') + span.set_attribute(ATTR_GEN_AI_PROVIDER, provider) span.set_attribute(ATTR_GEN_AI_REQUEST_MODEL, model) span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, message.input_tokens) span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, message.output_tokens) if message.respond_to?(:output_tokens) diff --git a/lib/integrations/llm_base_service.rb b/lib/integrations/llm_base_service.rb index 8410130ee..fbc92cf14 100644 --- a/lib/integrations/llm_base_service.rb +++ b/lib/integrations/llm_base_service.rb @@ -84,13 +84,12 @@ class Integrations::LlmBaseService end def api_base - endpoint = InstallationConfig.find_by(name: 'CAPTAIN_OPEN_AI_ENDPOINT')&.value.presence || 'https://api.openai.com/' - endpoint = endpoint.chomp('/') - "#{endpoint}/v1" + Llm::Config.api_base_for(llm_provider) end def make_api_call(body) parsed_body = JSON.parse(body) + @llm_provider = provider_for_model(parsed_body['model']) instrumentation_params = build_instrumentation_params(parsed_body) instrument_llm_call(instrumentation_params) do @@ -102,9 +101,10 @@ class Integrations::LlmBaseService messages = parsed_body['messages'] model = parsed_body['model'] credential = llm_credential + return { error: I18n.t('captain.api_key_missing'), error_code: 401, request_messages: messages } if credential.blank? - Llm::Config.with_api_key(credential[:api_key], api_base: api_base) do |context| - chat = context.chat(model: model) + Llm::Config.with_api_key(credential[:api_key], provider: llm_provider, api_base: api_base) do |context| + chat = Llm::ProviderChat.new(context.chat(model: model, provider: llm_provider, assume_model_exists: true), provider: llm_provider) setup_chat_with_messages(chat, messages) end rescue StandardError => e @@ -161,13 +161,22 @@ class Integrations::LlmBaseService conversation_id: conversation&.display_id, feature_name: event_name, model: parsed_body['model'], + provider: llm_provider, messages: parsed_body['messages'], temperature: parsed_body['temperature'] } end def llm_credential - @llm_credential ||= { api_key: hook.settings['api_key'], source: :hook } + @llm_credential ||= Llm::CredentialResolver.new(provider: llm_provider, openai_hook: hook).resolve + end + + def llm_provider + @llm_provider || Llm::Config::DEFAULT_PROVIDER + end + + def provider_for_model(model) + Llm::Models.provider_for(model) || Llm::Config::DEFAULT_PROVIDER end def exception_tracking_account diff --git a/lib/integrations/llm_instrumentation_helpers.rb b/lib/integrations/llm_instrumentation_helpers.rb index debbfaeda..b2be8dc5a 100644 --- a/lib/integrations/llm_instrumentation_helpers.rb +++ b/lib/integrations/llm_instrumentation_helpers.rb @@ -35,7 +35,7 @@ module Integrations::LlmInstrumentationHelpers end def set_request_attributes(span, params) - provider = determine_provider(params[:model]) + provider = params[:provider] || determine_provider(params[:model]) span.set_attribute(ATTR_GEN_AI_PROVIDER, provider) span.set_attribute(ATTR_GEN_AI_REQUEST_MODEL, params[:model]) span.set_attribute(ATTR_GEN_AI_REQUEST_TEMPERATURE, params[:temperature]) if params[:temperature] diff --git a/lib/llm/credential_resolver.rb b/lib/llm/credential_resolver.rb new file mode 100644 index 000000000..411408890 --- /dev/null +++ b/lib/llm/credential_resolver.rb @@ -0,0 +1,26 @@ +class Llm::CredentialResolver + def initialize(provider:, openai_hook: nil) + @provider = provider.to_s + @openai_hook = openai_hook + end + + def resolve + hook_llm_credential || system_llm_credential + end + + private + + attr_reader :provider, :openai_hook + + def hook_llm_credential + return unless Llm::Config.openai_provider?(provider) + + key = openai_hook&.settings&.dig('api_key').presence + { api_key: key, provider: provider, source: :hook } if key + end + + def system_llm_credential + key = Llm::Config.api_key_for(provider).presence + { api_key: key, provider: provider, source: :system } if key + end +end diff --git a/lib/llm/provider_chat.rb b/lib/llm/provider_chat.rb new file mode 100644 index 000000000..0ca6c6a72 --- /dev/null +++ b/lib/llm/provider_chat.rb @@ -0,0 +1,37 @@ +require 'delegate' + +class Llm::ProviderChat < SimpleDelegator + def initialize(chat, provider:) + @provider = provider.to_s + super(chat) + end + + def with_schema(schema) + return self unless supports_tools_and_schema? + + __setobj__(__getobj__.with_schema(schema)) + self + end + + def with_tool(tool) + return self unless supports_tools_and_schema? + + __setobj__(__getobj__.with_tool(tool)) + self + end + + def with_params(**params) + filtered_params = params.dup + filtered_params.delete(:response_format) unless supports_tools_and_schema? + return self if filtered_params.blank? + + __setobj__(__getobj__.with_params(**filtered_params)) + self + end + + private + + def supports_tools_and_schema? + Llm::Config.supports_tools_and_schema?(@provider) + end +end diff --git a/spec/lib/captain/base_task_service_spec.rb b/spec/lib/captain/base_task_service_spec.rb index b24a5c49c..5f023caea 100644 --- a/spec/lib/captain/base_task_service_spec.rb +++ b/spec/lib/captain/base_task_service_spec.rb @@ -171,22 +171,22 @@ RSpec.describe Captain::BaseTaskService do it 'uses the resolved feature model for the request and instrumentation' do account.update!(captain_models: { 'editor' => 'gpt-4.1' }) - expect(mock_context).to receive(:chat).with(model: 'gpt-4.1').and_return(mock_chat) + expect(mock_context).to receive(:chat).with(model: 'gpt-4.1', provider: 'openai', assume_model_exists: true).and_return(mock_chat) expect(service).to receive(:instrument_llm_call).with( - hash_including(model: 'gpt-4.1', feature_name: 'test_event') + hash_including(model: 'gpt-4.1', provider: 'openai', feature_name: 'test_event') ).and_call_original service.send(:make_api_call, feature: 'editor', messages: messages) end it 'uses the supplied model as a feature fallback when there is no account override' do - expect(mock_context).to receive(:chat).with(model: 'gpt-5.2').and_return(mock_chat) + expect(mock_context).to receive(:chat).with(model: 'gpt-5.2', provider: 'openai', assume_model_exists: true).and_return(mock_chat) service.send(:make_api_call, feature: 'document_faq_generation', model: 'gpt-5.2', messages: messages) end it 'uses the help center article generation feature default' do - expect(mock_context).to receive(:chat).with(model: 'gpt-5.2').and_return(mock_chat) + expect(mock_context).to receive(:chat).with(model: 'gpt-5.2', provider: 'openai', assume_model_exists: true).and_return(mock_chat) service.send(:make_api_call, feature: 'help_center_article_generation', messages: messages) end @@ -194,11 +194,34 @@ RSpec.describe Captain::BaseTaskService do it 'prefers account overrides over supplied feature fallback models' do account.update!(captain_models: { 'help_center_article_generation' => 'gpt-4.1' }) - expect(mock_context).to receive(:chat).with(model: 'gpt-4.1').and_return(mock_chat) + expect(mock_context).to receive(:chat).with(model: 'gpt-4.1', provider: 'openai', assume_model_exists: true).and_return(mock_chat) service.send(:make_api_call, feature: 'help_center_article_generation', model: 'gpt-5.2', messages: messages) end + it 'uses the model provider for account overrides' do + create(:installation_config, name: 'CAPTAIN_ANTHROPIC_API_KEY', value: 'anthropic-key') + account.update!(captain_models: { 'assistant' => 'claude-haiku-4.5' }) + + expect(Llm::Config).to receive(:with_api_key).with('anthropic-key', provider: 'anthropic', api_base: nil).and_yield(mock_context) + expect(mock_context).to receive(:chat).with(model: 'claude-haiku-4.5', provider: 'anthropic', assume_model_exists: true).and_return(mock_chat) + + service.send(:make_api_call, feature: 'assistant', messages: messages) + end + + it 'does not attach schemas or tools for non-OpenAI providers' do + create(:installation_config, name: 'CAPTAIN_ANTHROPIC_API_KEY', value: 'anthropic-key') + account.update!(captain_models: { 'assistant' => 'claude-haiku-4.5' }) + + expect(mock_context).to receive(:chat).with(model: 'claude-haiku-4.5', provider: 'anthropic', assume_model_exists: true).and_return(mock_chat) + expect(mock_chat).not_to receive(:with_schema) + expect(mock_chat).not_to receive(:with_tool) + expect(service).not_to receive(:instrument_tool_session) + expect(service).to receive(:instrument_llm_call).and_call_original + + service.send(:make_api_call, feature: 'assistant', messages: messages, schema: Class.new, tools: [Class.new]) + end + it 'returns formatted response with tokens' do result = service.send(:make_api_call, model: model, messages: messages) @@ -295,7 +318,7 @@ RSpec.describe Captain::BaseTaskService do it 'tracks exceptions against the system key when an account hook exists' do create(:integrations_hook, :openai, account: account, settings: { 'api_key' => 'hook-key' }) - expect(Llm::Config).to receive(:with_api_key).with('test-key', api_base: anything).and_raise(error) + expect(Llm::Config).to receive(:with_api_key).with('test-key', provider: 'openai', api_base: nil).and_raise(error) expect(ChatwootExceptionTracker).to receive(:new).with(error, account: account).and_return(exception_tracker) expect(exception_tracker).to receive(:capture_exception)