diff --git a/lib/captain/rewrite_service.rb b/lib/captain/rewrite_service.rb index 653b7465f..3a217d3c6 100644 --- a/lib/captain/rewrite_service.rb +++ b/lib/captain/rewrite_service.rb @@ -1,8 +1,20 @@ class Captain::RewriteService < Captain::BaseTaskService pattr_initialize [:account!, :content!, :operation!, { conversation_display_id: nil }] + TONE_OPERATIONS = %i[casual professional friendly confident straightforward].freeze + ALLOWED_OPERATIONS = (%i[fix_spelling_grammar improve] + TONE_OPERATIONS).freeze + def perform - send(operation) + operation_sym = operation.to_sym + raise ArgumentError, "Invalid operation: #{operation}" unless ALLOWED_OPERATIONS.include?(operation_sym) + + send(operation_sym) + end + + TONE_OPERATIONS.each do |tone| + define_method(tone) do + call_llm_with_prompt(tone_rewrite_prompt(tone.to_s)) + end end private @@ -11,26 +23,6 @@ class Captain::RewriteService < Captain::BaseTaskService call_llm_with_prompt(prompt_from_file('fix_spelling_grammar')) end - def casual - call_llm_with_prompt(tone_rewrite_prompt('casual')) - end - - def professional - call_llm_with_prompt(tone_rewrite_prompt('professional')) - end - - def friendly - call_llm_with_prompt(tone_rewrite_prompt('friendly')) - end - - def confident - call_llm_with_prompt(tone_rewrite_prompt('confident')) - end - - def straightforward - call_llm_with_prompt(tone_rewrite_prompt('straightforward')) - end - def improve template = prompt_from_file('improve') diff --git a/spec/lib/captain/rewrite_service_spec.rb b/spec/lib/captain/rewrite_service_spec.rb index cf53fc8f1..d7900d3cb 100644 --- a/spec/lib/captain/rewrite_service_spec.rb +++ b/spec/lib/captain/rewrite_service_spec.rb @@ -139,4 +139,28 @@ RSpec.describe Captain::RewriteService do expect(result[:message]).to eq('Rewritten text') end end + + describe '#perform with invalid operation' do + it 'raises ArgumentError for unknown operation' do + invalid_service = described_class.new( + account: account, + content: content, + operation: 'invalid_operation', + conversation_display_id: conversation.display_id + ) + + expect { invalid_service.perform }.to raise_error(ArgumentError, /Invalid operation/) + end + + it 'prevents method injection attacks' do + dangerous_service = described_class.new( + account: account, + content: content, + operation: 'perform', + conversation_display_id: conversation.display_id + ) + + expect { dangerous_service.perform }.to raise_error(ArgumentError, /Invalid operation/) + end + end end