From d657f35a7680c24f8e9ab3329b9ddc85899f1fc3 Mon Sep 17 00:00:00 2001
From: Pranav <pranav@chatwoot.com>
Date: Mon, 19 May 2025 15:26:38 -0700
Subject: [PATCH] feat: Introduce the concept of tool registry within Captain
 (#11516)

This PR introduces the concept of a tool registry. The implementation is
straightforward: you can define a tool by creating a class with a
function name. The function name gets registered in the registry and can
be referenced during LLM calls. When the LLM invokes a tool using the
registered name, the registry locates and executes the appropriate tool.
If the LLM calls an unregistered tool, the registry returns an error
indicating that the tool is not defined.
---
 enterprise/app/helpers/captain/chat_helper.rb | 67 ++++----------
 .../services/captain/copilot/chat_service.rb  |  7 ++
 .../captain/llm/assistant_chat_service.rb     |  6 ++
 .../services/captain/tool_registry_service.rb | 27 ++++++
 .../services/captain/tools/base_service.rb    | 34 ++++++++
 .../tools/search_documentation_service.rb     | 49 +++++++++++
 .../captain/tool_registry_service_spec.rb     | 87 +++++++++++++++++++
 .../search_documentation_service_spec.rb      | 77 ++++++++++++++++
 8 files changed, 303 insertions(+), 51 deletions(-)
 create mode 100644 enterprise/app/services/captain/tool_registry_service.rb
 create mode 100644 enterprise/app/services/captain/tools/base_service.rb
 create mode 100644 enterprise/app/services/captain/tools/search_documentation_service.rb
 create mode 100644 spec/enterprise/services/captain/tool_registry_service_spec.rb
 create mode 100644 spec/enterprise/services/captain/tools/search_documentation_service_spec.rb

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