diff --git a/enterprise/app/helpers/captain/chat_helper.rb b/enterprise/app/helpers/captain/chat_helper.rb index 146e8f813..bbe3851f0 100644 --- a/enterprise/app/helpers/captain/chat_helper.rb +++ b/enterprise/app/helpers/captain/chat_helper.rb @@ -1,24 +1,4 @@ module Captain::ChatHelper - def search_documentation_tool - { - type: 'function', - function: { - name: 'search_documentation', - description: "Use this function to get documentation on functionalities you don't know about.", - parameters: { - type: 'object', - properties: { - search_query: { - type: 'string', - description: 'The search query to look up in the documentation.' - } - }, - required: ['search_query'] - } - } - } - end - def request_chat_completion Rails.logger.debug { "[CAPTAIN][ChatCompletion] #{@messages}" } @@ -26,13 +6,12 @@ module Captain::ChatHelper parameters: { model: @model, messages: @messages, - tools: [search_documentation_tool], + tools: @tool_registry&.registered_tools || [], response_format: { type: 'json_object' } } ) handle_response(response) - @response end def handle_response(response) @@ -41,7 +20,7 @@ module Captain::ChatHelper if message['tool_calls'] process_tool_calls(message['tool_calls']) else - @response = JSON.parse(message['content'].strip) + JSON.parse(message['content'].strip) end end @@ -54,38 +33,20 @@ module Captain::ChatHelper end def process_tool_call(tool_call) + arguments = JSON.parse(tool_call['function']['arguments']) + function_name = tool_call['function']['name'] tool_call_id = tool_call['id'] - if tool_call['function']['name'] == 'search_documentation' - query = JSON.parse(tool_call['function']['arguments'])['search_query'] - sections = fetch_documentation(query) - append_tool_response(sections, tool_call_id) + if @tool_registry.respond_to?(function_name) + execute_tool(function_name, arguments, tool_call_id) else - append_tool_response('', tool_call_id) + process_invalid_tool_call(tool_call_id) end end - def fetch_documentation(query) - Rails.logger.debug { "[CAPTAIN][DocumentationSearch] #{query}" } - @assistant - .responses - .approved - .search(query) - .map { |response| format_response(response) }.join - end - - def format_response(response) - formatted_response = " - Question: #{response.question} - Answer: #{response.answer} - " - if response.documentable.present? && response.documentable.try(:external_link) - formatted_response += " - Source: #{response.documentable.external_link} - " - end - - formatted_response + def execute_tool(function_name, arguments, tool_call_id) + result = @tool_registry.send(function_name, arguments) + append_tool_response(result, tool_call_id) end def append_tool_calls(tool_calls) @@ -95,11 +56,15 @@ module Captain::ChatHelper } end - def append_tool_response(sections, tool_call_id) + def process_invalid_tool_call(tool_call_id) + append_tool_response('Tool not available', tool_call_id) + end + + def append_tool_response(content, tool_call_id) @messages << { role: 'tool', tool_call_id: tool_call_id, - content: "Found the following FAQs in the documentation:\n #{sections}" + content: content } end end diff --git a/enterprise/app/services/captain/copilot/chat_service.rb b/enterprise/app/services/captain/copilot/chat_service.rb index 6fd3c4e68..202348825 100644 --- a/enterprise/app/services/captain/copilot/chat_service.rb +++ b/enterprise/app/services/captain/copilot/chat_service.rb @@ -10,6 +10,8 @@ class Captain::Copilot::ChatService < Llm::BaseOpenAiService @conversation_history = config[:conversation_history] @previous_messages = config[:previous_messages] || [] @language = config[:language] || 'english' + + register_tools @messages = [system_message, conversation_history_context] + @previous_messages @response = '' end @@ -25,6 +27,11 @@ class Captain::Copilot::ChatService < Llm::BaseOpenAiService private + def register_tools + @tool_registry = Captain::ToolRegistryService.new(@assistant) + @tool_registry.register_tool(Captain::Tools::SearchDocumentationService) + end + def system_message { role: 'system', diff --git a/enterprise/app/services/captain/llm/assistant_chat_service.rb b/enterprise/app/services/captain/llm/assistant_chat_service.rb index 50688f18b..8077b22c9 100644 --- a/enterprise/app/services/captain/llm/assistant_chat_service.rb +++ b/enterprise/app/services/captain/llm/assistant_chat_service.rb @@ -9,6 +9,7 @@ class Captain::Llm::AssistantChatService < Llm::BaseOpenAiService @assistant = assistant @messages = [system_message] @response = '' + register_tools end def generate_response(input, previous_messages = [], role = 'user') @@ -19,6 +20,11 @@ class Captain::Llm::AssistantChatService < Llm::BaseOpenAiService private + def register_tools + @tool_registry = Captain::ToolRegistryService.new(@assistant) + @tool_registry.register_tool(Captain::Tools::SearchDocumentationService) + end + def system_message { role: 'system', diff --git a/enterprise/app/services/captain/tool_registry_service.rb b/enterprise/app/services/captain/tool_registry_service.rb new file mode 100644 index 000000000..088d89483 --- /dev/null +++ b/enterprise/app/services/captain/tool_registry_service.rb @@ -0,0 +1,27 @@ +class Captain::ToolRegistryService + attr_reader :registered_tools, :tools + + def initialize(assistant) + @assistant = assistant + @registered_tools = [] + @tools = {} + end + + def register_tool(tool_class) + tool = tool_class.new(@assistant) + @tools[tool.name] = tool + @registered_tools << tool.to_registry_format + end + + def method_missing(method_name, *arguments) + if @tools.key?(method_name.to_s) + @tools[method_name.to_s].execute(*arguments) + else + super + end + end + + def respond_to_missing?(method_name, include_private = false) + @tools.key?(method_name.to_s) || super + end +end diff --git a/enterprise/app/services/captain/tools/base_service.rb b/enterprise/app/services/captain/tools/base_service.rb new file mode 100644 index 000000000..10e44d2f4 --- /dev/null +++ b/enterprise/app/services/captain/tools/base_service.rb @@ -0,0 +1,34 @@ +class Captain::Tools::BaseService + attr_accessor :assistant + + def initialize(assistant) + @assistant = assistant + end + + def name + raise NotImplementedError, "#{self.class} must implement name" + end + + def description + raise NotImplementedError, "#{self.class} must implement description" + end + + def parameters + raise NotImplementedError, "#{self.class} must implement parameters" + end + + def execute(arguments) + raise NotImplementedError, "#{self.class} must implement execute" + end + + def to_registry_format + { + type: 'function', + function: { + name: name, + description: description, + parameters: parameters + } + } + end +end diff --git a/enterprise/app/services/captain/tools/search_documentation_service.rb b/enterprise/app/services/captain/tools/search_documentation_service.rb new file mode 100644 index 000000000..672baf24a --- /dev/null +++ b/enterprise/app/services/captain/tools/search_documentation_service.rb @@ -0,0 +1,49 @@ +class Captain::Tools::SearchDocumentationService < Captain::Tools::BaseService + def name + 'search_documentation' + end + + def description + 'Search and retrieve documentation from knowledge base' + end + + def parameters + { + type: 'object', + properties: { + search_query: { + type: 'string', + description: 'The search query to look up in the documentation.' + } + }, + required: ['search_query'] + } + end + + def execute(arguments) + query = arguments['search_query'] + Rails.logger.info { "#{self.class.name}: #{query}" } + + responses = assistant.responses.approved.search(query) + + return 'No FAQs found for the given query' if responses.empty? + + responses.map { |response| format_response(response) }.join + end + + private + + def format_response(response) + formatted_response = " + Question: #{response.question} + Answer: #{response.answer} + " + if response.documentable.present? && response.documentable.try(:external_link) + formatted_response += " + Source: #{response.documentable.external_link} + " + end + + formatted_response + end +end diff --git a/spec/enterprise/services/captain/tool_registry_service_spec.rb b/spec/enterprise/services/captain/tool_registry_service_spec.rb new file mode 100644 index 000000000..e355dea0a --- /dev/null +++ b/spec/enterprise/services/captain/tool_registry_service_spec.rb @@ -0,0 +1,87 @@ +require 'rails_helper' + +# Test tool implementation +class TestTool < Captain::Tools::BaseService + def name + 'test_tool' + end + + def description + 'A test tool for specs' + end + + def parameters + { + type: 'object', + properties: { + test_param: { + type: 'string' + } + } + } + end + + def execute(*args) + args + end +end + +RSpec.describe Captain::ToolRegistryService do + let(:assistant) { create(:captain_assistant) } + let(:service) { described_class.new(assistant) } + + describe '#initialize' do + it 'initializes with empty tools and registered_tools' do + expect(service.tools).to be_empty + expect(service.registered_tools).to be_empty + end + end + + describe '#register_tool' do + let(:tool_class) { TestTool } + + it 'registers a new tool' do + service.register_tool(tool_class) + + expect(service.tools['test_tool']).to be_a(TestTool) + expect(service.registered_tools).to include( + { + type: 'function', + function: { + name: 'test_tool', + description: 'A test tool for specs', + parameters: { + type: 'object', + properties: { + test_param: { + type: 'string' + } + } + } + } + } + ) + end + end + + describe 'method_missing' do + let(:tool_class) { TestTool } + + before do + service.register_tool(tool_class) + end + + context 'when method corresponds to a registered tool' do + it 'executes the tool with given arguments' do + result = service.test_tool(test_param: 'arg1') + expect(result).to eq([{ test_param: 'arg1' }]) + end + end + + context 'when method does not correspond to a registered tool' do + it 'raises NoMethodError' do + expect { service.unknown_tool }.to raise_error(NoMethodError) + end + end + end +end diff --git a/spec/enterprise/services/captain/tools/search_documentation_service_spec.rb b/spec/enterprise/services/captain/tools/search_documentation_service_spec.rb new file mode 100644 index 000000000..9f5586e6b --- /dev/null +++ b/spec/enterprise/services/captain/tools/search_documentation_service_spec.rb @@ -0,0 +1,77 @@ +require 'rails_helper' + +RSpec.describe Captain::Tools::SearchDocumentationService do + let(:assistant) { create(:captain_assistant) } + let(:service) { described_class.new(assistant) } + let(:question) { 'How to create a new account?' } + let(:answer) { 'You can create a new account by clicking on the Sign Up button.' } + let(:external_link) { 'https://example.com/docs/create-account' } + + describe '#name' do + it 'returns the correct service name' do + expect(service.name).to eq('search_documentation') + end + end + + describe '#description' do + it 'returns the service description' do + expect(service.description).to eq('Search and retrieve documentation from knowledge base') + end + end + + describe '#parameters' do + it 'returns the required parameters schema' do + expected_schema = { + type: 'object', + properties: { + search_query: { + type: 'string', + description: 'The search query to look up in the documentation.' + } + }, + required: ['search_query'] + } + + expect(service.parameters).to eq(expected_schema) + end + end + + describe '#execute' do + let!(:response) do + create( + :captain_assistant_response, + assistant: assistant, + question: question, + answer: answer, + status: 'approved' + ) + end + + let(:documentable) { create(:captain_document, external_link: external_link) } + + context 'when matching responses exist' do + before do + response.update(documentable: documentable) + allow(Captain::AssistantResponse).to receive(:search).with(question).and_return([response]) + end + + it 'returns formatted responses for the search query' do + result = service.execute({ 'search_query' => question }) + + expect(result).to include(question) + expect(result).to include(answer) + expect(result).to include(external_link) + end + end + + context 'when no matching responses exist' do + before do + allow(Captain::AssistantResponse).to receive(:search).with(question).and_return([]) + end + + it 'returns an empty string' do + expect(service.execute({ 'search_query' => question })).to eq('No FAQs found for the given query') + end + end + end +end