diff --git a/ee/lib/api/chat.rb b/ee/lib/api/chat.rb index d93c8990ccb598bf7fd572cf2539ee403ceb0b8a..bf42c781e53a3f223a50d91a981ba6bb62c21903 100644 --- a/ee/lib/api/chat.rb +++ b/ee/lib/api/chat.rb @@ -40,6 +40,28 @@ def find_resource(parameters) optional :client_subscription_id, type: String, limit: 500, desc: 'Client Subscription ID' optional :with_clean_history, type: Boolean, desc: 'Indicates if we need to reset the history before and after the request' + optional :current_file, type: Hash do + optional :file_name, type: String, limit: 1000, desc: 'The name of the current file' + optional :content_above_cursor, type: String, + limit: ::API::CodeSuggestions::MAX_CONTENT_SIZE, desc: 'The content above cursor' + optional :content_below_cursor, type: String, + limit: ::API::CodeSuggestions::MAX_CONTENT_SIZE, desc: 'The content below cursor' + optional :selected_text, type: String, + limit: ::API::CodeSuggestions::MAX_CONTENT_SIZE, + desc: 'The content currently selected by the user' + end + optional :additional_context, type: Array, allow_blank: true, + desc: 'List of additional context to be passed for the chat' do + requires :type, type: String, + values: ::CodeSuggestions::Prompts::CodeGeneration::AnthropicMessages::CONTENT_TYPES.values, + desc: 'Type of the additional context.' + requires :name, type: String, + limit: ::API::CodeSuggestions::MAX_CONTEXT_NAME_SIZE, allow_blank: false, + desc: 'Name of the additional context.' + requires :content, type: String, + limit: ::API::CodeSuggestions::MAX_BODY_SIZE, allow_blank: false, + desc: 'Content of the additional context.' + end end post do safe_params = declared_params(include_missing: false) diff --git a/ee/lib/gitlab/duo/chat/completions.rb b/ee/lib/gitlab/duo/chat/completions.rb index a8297a7fc297609f61cb37199be0b8d3f20f6c9b..ebe903a0cd19fae23a5cf9815e3dee8818f64f3d 100644 --- a/ee/lib/gitlab/duo/chat/completions.rb +++ b/ee/lib/gitlab/duo/chat/completions.rb @@ -12,7 +12,7 @@ def initialize(user, resource: nil) def execute(safe_params: {}) action_name = 'chat' - options = safe_params.slice(:referer_url) + options = safe_params.slice(:referer_url, :current_file, :additional_context).compact_blank message_attributes = { request_id: SecureRandom.uuid, role: ::Gitlab::Llm::AiMessage::ROLE_USER, diff --git a/ee/spec/requests/api/chat_spec.rb b/ee/spec/requests/api/chat_spec.rb index a0172f60c4d620245d6ae99367fd437c17ebd82c..349c16e85be77a67a7992ce68dc01d52950e5eea 100644 --- a/ee/spec/requests/api/chat_spec.rb +++ b/ee/spec/requests/api/chat_spec.rb @@ -22,7 +22,17 @@ let(:headers) { {} } let(:resource) { issue } let(:content) { 'what is this issue about' } - let(:params) { { content: content, resource_type: resource.class.to_s.downcase, resource_id: resource.id } } + let(:additional_context) { nil } + let(:current_file) { nil } + let(:params) do + { + content: content, + resource_type: resource.class.to_s.downcase, + resource_id: resource.id, + additional_context: additional_context, + current_file: current_file + } + end before do group.add_member(authorized_user, :developer) @@ -279,6 +289,34 @@ post_api end end + + context 'with additional context' do + let(:additional_context) { [{ type: "file", name: "test.py", content: "print('hello world')" }] } + let(:options) { { additional_context: additional_context } } + + it 'sends additional context to the chat' do + expect(chat_message).to receive(:save!) + expect(Gitlab::Llm::ChatMessage).to receive(:new).with(chat_message_params).and_return(chat_message) + expect(Llm::Internal::CompletionService).to receive(:new).with(chat_message, options).and_return(chat) + expect(chat).to receive(:execute) + + post_api + end + end + + context 'with current file' do + let(:current_file) { { file_name: "test.py", selected_text: "print('hello world')" } } + let(:options) { { current_file: current_file } } + + it 'sends current file to the chat' do + expect(chat_message).to receive(:save!) + expect(Gitlab::Llm::ChatMessage).to receive(:new).with(chat_message_params).and_return(chat_message) + expect(Llm::Internal::CompletionService).to receive(:new).with(chat_message, options).and_return(chat) + expect(chat).to receive(:execute) + + post_api + end + end end end end