From 34e929c7056240b1fc64ec006634f897f4f0215e Mon Sep 17 00:00:00 2001
From: Pranav <pranavrajs@gmail.com>
Date: Sun, 11 May 2025 00:37:59 -0700
Subject: [PATCH] Move copilot to websockets

---
 app/channels/copilot_channel.rb               |  35 ++++++
 .../dashboard/api/inbox/conversation.js       |   4 -
 .../components-next/copilot/Copilot.vue       |  22 ++--
 .../components/copilot/CopilotContainer.vue   |  91 ++++++++++------
 .../copilot/CopilotThinkingBlock.vue          |  24 +++++
 .../helpers/CopilotActionCableConnector.js    |  49 +++++++++
 config/routes.rb                              |   1 -
 .../v1/accounts/captain/copilot_controller.rb |  36 -------
 enterprise/app/helpers/captain/chat_helper.rb | 101 ++++++++----------
 .../captain/process_copilot_message_job.rb    |  50 +++++++++
 .../services/captain/copilot/chat_service.rb  |  36 ++++---
 .../captain/llm/assistant_chat_service.rb     |   4 +-
 .../captain/llm/system_prompts_service.rb     |   5 +-
 .../services/captain/tool_registry_service.rb |  32 ++++++
 .../services/captain/tools/base_service.rb    |  34 ++++++
 .../tools/search_documentation_service.rb     |  49 +++++++++
 16 files changed, 422 insertions(+), 151 deletions(-)
 create mode 100644 app/channels/copilot_channel.rb
 create mode 100644 app/javascript/dashboard/components/copilot/CopilotThinkingBlock.vue
 create mode 100644 app/javascript/dashboard/helpers/CopilotActionCableConnector.js
 delete mode 100644 enterprise/app/controllers/api/v1/accounts/captain/copilot_controller.rb
 create mode 100644 enterprise/app/jobs/captain/process_copilot_message_job.rb
 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

diff --git a/app/channels/copilot_channel.rb b/app/channels/copilot_channel.rb
new file mode 100644
index 000000000..2fde1c32f
--- /dev/null
+++ b/app/channels/copilot_channel.rb
@@ -0,0 +1,35 @@
+class CopilotChannel < ApplicationCable::Channel
+  def subscribed
+    current_user
+    current_account
+    ensure_stream
+  end
+
+  def message(data)
+    return if @current_account.blank?
+
+    Captain::ProcessCopilotMessageJob.perform_later(
+      assistant_id: data['assistant_id'],
+      message: data['message'],
+      options: {
+        user_id: @current_user.id,
+        conversation_id: data['conversation_id'],
+        previous_history: data['previous_history']
+      }
+    )
+  end
+
+  def ensure_stream
+    stream_from "copilot_#{@current_account.id}_#{@current_user.id}"
+  end
+
+  def current_user
+    @current_user ||= User.find(params[:user_id])
+  end
+
+  def current_account
+    return if current_user.blank?
+
+    @current_account ||= @current_user.accounts.find(params[:account_id])
+  end
+end
diff --git a/app/javascript/dashboard/api/inbox/conversation.js b/app/javascript/dashboard/api/inbox/conversation.js
index 39546096f..1a87e6d96 100644
--- a/app/javascript/dashboard/api/inbox/conversation.js
+++ b/app/javascript/dashboard/api/inbox/conversation.js
@@ -134,10 +134,6 @@ class ConversationApi extends ApiClient {
     return axios.get(`${this.url}/${conversationId}/attachments`);
   }
 
-  requestCopilot(conversationId, body) {
-    return axios.post(`${this.url}/${conversationId}/copilot`, body);
-  }
-
   getInboxAssistant(conversationId) {
     return axios.get(`${this.url}/${conversationId}/inbox_assistant`);
   }
diff --git a/app/javascript/dashboard/components-next/copilot/Copilot.vue b/app/javascript/dashboard/components-next/copilot/Copilot.vue
index cc445257a..903450bdd 100644
--- a/app/javascript/dashboard/components-next/copilot/Copilot.vue
+++ b/app/javascript/dashboard/components-next/copilot/Copilot.vue
@@ -8,6 +8,7 @@ import CopilotInput from './CopilotInput.vue';
 import CopilotLoader from './CopilotLoader.vue';
 import CopilotAgentMessage from './CopilotAgentMessage.vue';
 import CopilotAssistantMessage from './CopilotAssistantMessage.vue';
+import CopilotThinkingBlock from 'dashboard/components/copilot/CopilotThinkingBlock.vue';
 import ToggleCopilotAssistant from './ToggleCopilotAssistant.vue';
 import Icon from '../icon/Icon.vue';
 import Button from '../button/Button.vue';
@@ -43,8 +44,6 @@ const emit = defineEmits(['sendMessage', 'reset', 'setAssistant', 'close']);
 
 const { t } = useI18n();
 
-const COPILOT_USER_ROLES = ['assistant', 'system'];
-
 const sendMessage = message => {
   emit('sendMessage', message);
   useTrack(COPILOT_EVENTS.SEND_MESSAGE);
@@ -113,11 +112,11 @@ watch(
         </div>
       </div>
     </div>
-    <div ref="chatContainer" class="flex-1 flex px-4 py-4 overflow-y-auto">
-      <div
-        v-if="messages.length"
-        class="space-y-6 flex-1 flex items-start justify-center"
-      >
+    <div
+      ref="chatContainer"
+      class="flex-1 flex px-4 py-4 overflow-y-auto items-start"
+    >
+      <div v-if="messages.length" class="space-y-6 flex-1 flex flex-col">
         <template v-for="message in messages" :key="message.id">
           <CopilotAgentMessage
             v-if="message.role === 'user'"
@@ -125,10 +124,17 @@ watch(
             :message="message"
           />
           <CopilotAssistantMessage
-            v-else-if="COPILOT_USER_ROLES.includes(message.role)"
+            v-else-if="
+              message.role === 'assistant' || message.role === 'system'
+            "
             :message="message"
             :conversation-inbox-type="conversationInboxType"
           />
+          <CopilotThinkingBlock
+            v-else-if="message.role === 'assistant_thinking'"
+            :content="message.content"
+            :reasoning="message.reasoning"
+          />
         </template>
 
         <CopilotLoader v-if="isCaptainTyping" />
diff --git a/app/javascript/dashboard/components/copilot/CopilotContainer.vue b/app/javascript/dashboard/components/copilot/CopilotContainer.vue
index 1fed0e861..fd10489dc 100644
--- a/app/javascript/dashboard/components/copilot/CopilotContainer.vue
+++ b/app/javascript/dashboard/components/copilot/CopilotContainer.vue
@@ -1,8 +1,8 @@
 <script setup>
-import { ref, computed, onMounted, watchEffect } from 'vue';
+import { ref, computed, onMounted, onBeforeUnmount } from 'vue';
 import { useStore } from 'dashboard/composables/store';
 import Copilot from 'dashboard/components-next/copilot/Copilot.vue';
-import ConversationAPI from 'dashboard/api/inbox/conversation';
+import CopilotActionCableConnector from 'dashboard/helpers/CopilotActionCableConnector';
 import { useMapGetter } from 'dashboard/composables/store';
 import { useUISettings } from 'dashboard/composables/useUISettings';
 
@@ -29,6 +29,7 @@ const isSidebarOpen = computed(() => getUIState.value('isCopilotSidebarOpen'));
 const messages = ref([]);
 const isCaptainTyping = ref(false);
 const selectedAssistantId = ref(null);
+const copilotConnector = ref(null);
 
 const activeAssistant = computed(() => {
   const preferredId = uiSettings.value.preferred_captain_assistant_id;
@@ -62,7 +63,12 @@ const handleReset = () => {
   messages.value = [];
 };
 
-const sendMessage = async message => {
+const sendMessage = message => {
+  // Ensure WebSocket is connected before sending
+  if (!copilotConnector.value || !isSidebarOpen.value) {
+    return;
+  }
+
   // Add user message
   messages.value.push({
     id: messages.value.length + 1,
@@ -71,47 +77,66 @@ const sendMessage = async message => {
   });
   isCaptainTyping.value = true;
 
-  try {
-    const { data } = await ConversationAPI.requestCopilot(
-      props.conversationId,
-      {
-        previous_history: messages.value
-          .map(m => ({
-            role: m.role,
-            content: m.content,
-          }))
-          .slice(0, -1),
-        message,
-        assistant_id: selectedAssistantId.value,
+  copilotConnector.value.sendMessage({
+    assistantId: activeAssistant.value.id,
+    conversationId: props.conversationId,
+    previousHistory: messages.value
+      .filter(m => m.role !== 'assistant_thinking')
+      .map(m => ({
+        role: m.role,
+        content: m.content,
+      }))
+      .slice(0, -1),
+    message,
+  });
+};
+
+const initializeWebSocket = () => {
+  copilotConnector.value = new CopilotActionCableConnector({
+    accountId: currentUser.value.account_id,
+    userId: currentUser.value.id,
+    onDisconnect: () => {
+      // copilotConnector.value = null;
+    },
+    onCopilotResponse: data => {
+      if (data.type === 'final_response') {
+        messages.value.push({
+          id: new Date().getTime(),
+          role: 'assistant',
+          content: data.response.response,
+        });
+        isCaptainTyping.value = false;
+      } else {
+        messages.value.push({
+          id: new Date().getTime(),
+          role: 'assistant_thinking',
+          content: data.response.response,
+          reasoning: data.response.reasoning,
+        });
       }
-    );
-    messages.value.push({
-      id: new Date().getTime(),
-      role: 'assistant',
-      content: data.message,
-    });
-  } catch (error) {
-    // eslint-disable-next-line
-    console.log(error);
-  } finally {
-    isCaptainTyping.value = false;
+    },
+  });
+};
+
+const disconnectWebSocket = () => {
+  if (copilotConnector.value) {
+    copilotConnector.value.disconnect();
+    copilotConnector.value = null;
   }
 };
 
 onMounted(() => {
+  initializeWebSocket();
   store.dispatch('captainAssistants/get');
 });
 
+onBeforeUnmount(() => {
+  disconnectWebSocket();
+});
+
 const handleClose = () => {
   store.dispatch('uiState/set', { isCopilotSidebarOpen: false });
 };
-
-watchEffect(() => {
-  if (props.conversationId) {
-    store.dispatch('getInboxCaptainAssistantById', props.conversationId);
-    selectedAssistantId.value = activeAssistant.value?.id;
-  }
-});
 </script>
 
 <template>
diff --git a/app/javascript/dashboard/components/copilot/CopilotThinkingBlock.vue b/app/javascript/dashboard/components/copilot/CopilotThinkingBlock.vue
new file mode 100644
index 000000000..9f2308d7d
--- /dev/null
+++ b/app/javascript/dashboard/components/copilot/CopilotThinkingBlock.vue
@@ -0,0 +1,24 @@
+<script setup>
+defineProps({
+  content: {
+    type: String,
+    required: true,
+  },
+  reasoning: {
+    type: String,
+    required: true,
+  },
+});
+</script>
+
+<template>
+  <div class="flex flex-col gap-2 p-2 rounded bg-n-background">
+    <div
+      v-if="reasoning"
+      class="text-xs text-n-slate-11 border-t border-n-weak pt-2"
+    >
+      {{ reasoning }}
+    </div>
+    <p class="text-sm text-n-slate-12">{{ content }}</p>
+  </div>
+</template>
diff --git a/app/javascript/dashboard/helpers/CopilotActionCableConnector.js b/app/javascript/dashboard/helpers/CopilotActionCableConnector.js
new file mode 100644
index 000000000..6a64ee59f
--- /dev/null
+++ b/app/javascript/dashboard/helpers/CopilotActionCableConnector.js
@@ -0,0 +1,49 @@
+import { createConsumer } from '@rails/actioncable';
+
+class CopilotActionCableConnector {
+  constructor(
+    { accountId, userId, onDisconnect, onCopilotResponse },
+    websocketHost = ''
+  ) {
+    const websocketURL = websocketHost ? `${websocketHost}/cable` : undefined;
+
+    this.consumer = createConsumer(websocketURL);
+    this.subscription = this.consumer.subscriptions.create(
+      {
+        channel: 'CopilotChannel',
+        account_id: accountId,
+        user_id: userId,
+      },
+      {
+        received: this.onReceived,
+        disconnected: onDisconnect,
+      }
+    );
+
+    this.events = {
+      'copilot.response': onCopilotResponse,
+    };
+  }
+
+  onReceived = ({ event, data } = {}) => {
+    if (this.events[event] && typeof this.events[event] === 'function') {
+      this.events[event](data);
+    }
+  };
+
+  disconnect() {
+    this.subscription.unsubscribe();
+    this.consumer.disconnect();
+  }
+
+  sendMessage({ message, assistantId, conversationId, previousHistory }) {
+    this.subscription.perform('message', {
+      message,
+      assistant_id: assistantId,
+      conversation_id: conversationId,
+      previous_history: previousHistory,
+    });
+  }
+}
+
+export default CopilotActionCableConnector;
diff --git a/config/routes.rb b/config/routes.rb
index 044347f47..829eb9a7e 100644
--- a/config/routes.rb
+++ b/config/routes.rb
@@ -61,7 +61,6 @@ Rails.application.routes.draw do
             resources :documents, only: [:index, :show, :create, :destroy]
             resources :assistant_responses
             resources :bulk_actions, only: [:create]
-            resources :copilot, only: [:create]
           end
           resources :agent_bots, only: [:index, :create, :show, :update, :destroy] do
             delete :avatar, on: :member
diff --git a/enterprise/app/controllers/api/v1/accounts/captain/copilot_controller.rb b/enterprise/app/controllers/api/v1/accounts/captain/copilot_controller.rb
deleted file mode 100644
index a2759f122..000000000
--- a/enterprise/app/controllers/api/v1/accounts/captain/copilot_controller.rb
+++ /dev/null
@@ -1,36 +0,0 @@
-class Api::V1::Accounts::Captain::CopilotController < Api::V1::Accounts::BaseController
-  before_action :fetch_conversation, only: [:create]
-
-  def create
-    # First try to get the user's preferred assistant from UI settings or from the request
-    assistant_id = copilot_params[:assistant_id] || current_user.ui_settings&.dig('preferred_captain_assistant_id')
-
-    # Find the assistant either by ID or from inbox
-    assistant = if assistant_id.present?
-                  Captain::Assistant.find_by(id: assistant_id, account_id: Current.account.id)
-                else
-                  @conversation.inbox.captain_assistant
-                end
-
-    return render json: { message: I18n.t('captain.copilot_error') } unless assistant
-
-    response = Captain::Copilot::ChatService.new(
-      assistant,
-      previous_messages: copilot_params[:previous_messages],
-      conversation_history: @conversation.to_llm_text,
-      language: @conversation.account.locale_english_name
-    ).generate_response(copilot_params[:message])
-
-    render json: { message: response['response'] }
-  end
-
-  private
-
-  def copilot_params
-    params.permit(:message, :assistant_id, previous_messages: [])
-  end
-
-  def fetch_conversation
-    @conversation = Current.account.conversations.find(params[:conversation_id])
-  end
-end
diff --git a/enterprise/app/helpers/captain/chat_helper.rb b/enterprise/app/helpers/captain/chat_helper.rb
index 146e8f813..a2151f610 100644
--- a/enterprise/app/helpers/captain/chat_helper.rb
+++ b/enterprise/app/helpers/captain/chat_helper.rb
@@ -1,47 +1,30 @@
 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}" }
 
+    available_tools = @tool_registry&.registered_tools || []
     response = @client.chat(
       parameters: {
         model: @model,
         messages: @messages,
-        tools: [search_documentation_tool],
+        tools: available_tools,
         response_format: { type: 'json_object' }
       }
     )
 
     handle_response(response)
-    @response
+  rescue StandardError => e
+    Rails.logger.error { "[CAPTAIN][ChatCompletion] #{e}" }
+    raise e
   end
 
   def handle_response(response)
-    Rails.logger.debug { "[CAPTAIN][ChatCompletion] #{response}" }
+    Rails.logger.info { "[CAPTAIN][ChatCompletion] #{response}" }
     message = response.dig('choices', 0, 'message')
     if message['tool_calls']
       process_tool_calls(message['tool_calls'])
     else
-      @response = JSON.parse(message['content'].strip)
+      JSON.parse(message['content'].strip)
     end
   end
 
@@ -55,39 +38,16 @@ module Captain::ChatHelper
 
   def process_tool_call(tool_call)
     tool_call_id = tool_call['id']
+    function_name = tool_call['function']['name']
+    arguments = JSON.parse(tool_call['function']['arguments'])
 
-    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_call(tool_call_id, function_name, arguments)
     else
-      append_tool_response('', tool_call_id)
+      process_invalid_tool_call(tool_call_id, function_name)
     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
-  end
-
   def append_tool_calls(tool_calls)
     @messages << {
       role: 'assistant',
@@ -95,11 +55,44 @@ module Captain::ChatHelper
     }
   end
 
-  def append_tool_response(sections, tool_call_id)
+  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
+
+  def publish_to_stream(response)
+    @stream_writer&.call(response)
+  end
+
+  def execute_tool_call(tool_call_id, function_name, arguments)
+    publish_to_stream(
+      {
+        response: { response: "Processing tool call #{function_name}" },
+        type: 'tool_calls_start'
+      }
+    )
+    result = @tool_registry.send(function_name, arguments)
+    append_tool_response(result, tool_call_id)
+    publish_to_stream(
+      {
+        response: { response: "Received tool response #{function_name}" },
+        type: 'tool_response',
+        tool: function_name
+      }
+    )
+  end
+
+  def process_invalid_tool_call(tool_call_id, function_name)
+    append_tool_response('Tool not implemented', tool_call_id)
+    publish_to_stream(
+      {
+        response: { response: 'Tool not implemented' },
+        type: 'tool_error',
+        tool: function_name
+      }
+    )
+  end
 end
diff --git a/enterprise/app/jobs/captain/process_copilot_message_job.rb b/enterprise/app/jobs/captain/process_copilot_message_job.rb
new file mode 100644
index 000000000..da8064dbc
--- /dev/null
+++ b/enterprise/app/jobs/captain/process_copilot_message_job.rb
@@ -0,0 +1,50 @@
+class Captain::ProcessCopilotMessageJob < ApplicationJob
+  queue_as :default
+
+  def perform(assistant_id:, message:, options: {})
+    ensure_assistant(assistant_id)
+    ensure_user(options[:user_id])
+    process_message(message, options)
+  end
+
+  private
+
+  def ensure_assistant(assistant_id)
+    @assistant = Captain::Assistant.find(assistant_id)
+    @account = @assistant.account
+  end
+
+  def ensure_user(user_id)
+    @user = @account.users.find(user_id)
+  end
+
+  def process_message(message, options)
+    return unless @assistant
+
+    conversation = find_conversation(options[:conversation_id])
+    generate_chat_response(message, conversation, options[:previous_history])
+  end
+
+  def find_conversation(conversation_id)
+    return unless conversation_id
+
+    @account.conversations.find_by(display_id: conversation_id)
+  end
+
+  def generate_chat_response(message, conversation, previous_history)
+    Captain::Copilot::ChatService.new(
+      @assistant,
+      previous_messages: previous_history || [],
+      conversation_history: conversation&.to_llm_text,
+      language: @account.locale_english_name,
+      stream_writer: ->(data) { broadcast_response(data) }
+    ).generate_response(message)
+  end
+
+  def broadcast_response(data)
+    ActionCable.server.broadcast(
+      "copilot_#{@account.id}_#{@user.id}",
+      { event: 'copilot.response', data: data }
+    )
+  end
+end
diff --git a/enterprise/app/services/captain/copilot/chat_service.rb b/enterprise/app/services/captain/copilot/chat_service.rb
index 6fd3c4e68..469a91f6e 100644
--- a/enterprise/app/services/captain/copilot/chat_service.rb
+++ b/enterprise/app/services/captain/copilot/chat_service.rb
@@ -1,17 +1,21 @@
-require 'openai'
-
 class Captain::Copilot::ChatService < Llm::BaseOpenAiService
   include Captain::ChatHelper
 
-  def initialize(assistant, config)
-    super()
+  attr_reader :assistant, :language
 
+  def initialize(assistant, config = {})
+    super()
     @assistant = assistant
-    @conversation_history = config[:conversation_history]
-    @previous_messages = config[:previous_messages] || []
+    @tool_registry = Captain::ToolRegistryService.new(@assistant)
     @language = config[:language] || 'english'
-    @messages = [system_message, conversation_history_context] + @previous_messages
-    @response = ''
+    @messages = build_initial_messages(config)
+    @stream_writer = config[:stream_writer]
+  end
+
+  def build_initial_messages(config)
+    messages = [system_message]
+    messages << conversation_history_context if config[:conversation_history].present?
+    messages + (config[:previous_messages] || [])
   end
 
   def generate_response(input)
@@ -20,6 +24,12 @@ class Captain::Copilot::ChatService < Llm::BaseOpenAiService
     Rails.logger.info("[CAPTAIN][CopilotChatService] Incrementing response usage for #{@assistant.account.id}")
     @assistant.account.increment_response_usage
 
+    publish_to_stream(
+      {
+        response: response,
+        type: 'final_response'
+      }
+    )
     response
   end
 
@@ -33,12 +43,14 @@ class Captain::Copilot::ChatService < Llm::BaseOpenAiService
   end
 
   def conversation_history_context
+    return if @conversation_history.blank?
+
     {
       role: 'system',
-      content: "
-      Message History with the user is below:
-      #{@conversation_history}
-      "
+      content: <<~HISTORY.strip
+        Message History with the user is below:
+        #{@conversation_history}
+      HISTORY
     }
   end
 end
diff --git a/enterprise/app/services/captain/llm/assistant_chat_service.rb b/enterprise/app/services/captain/llm/assistant_chat_service.rb
index 1e45cb1d8..61ae3f719 100644
--- a/enterprise/app/services/captain/llm/assistant_chat_service.rb
+++ b/enterprise/app/services/captain/llm/assistant_chat_service.rb
@@ -3,10 +3,10 @@ require 'openai'
 class Captain::Llm::AssistantChatService < Llm::BaseOpenAiService
   include Captain::ChatHelper
 
-  def initialize(assistant: nil)
+  def initialize(assistant)
     super()
-
     @assistant = assistant
+    @tool_registry = Captain::ToolRegistryService.new(assistant)
     @messages = [system_message]
     @response = ''
   end
diff --git a/enterprise/app/services/captain/llm/system_prompts_service.rb b/enterprise/app/services/captain/llm/system_prompts_service.rb
index b1e627275..71bb87414 100644
--- a/enterprise/app/services/captain/llm/system_prompts_service.rb
+++ b/enterprise/app/services/captain/llm/system_prompts_service.rb
@@ -63,7 +63,8 @@ class Captain::Llm::SystemPromptsService
         You should only provide information related to #{product_name} and must not address queries about other products or external events.
 
         [Context]
-        You will be provided with the message history between the support agent and the customer. Use this context to understand the conversation flow, identify unresolved queries, and ensure responses are relevant and consistent with previous interactions. Always maintain a coherent and professional tone throughout the conversation.
+        Identify unresolved queries, and ensure responses are relevant and consistent with previous interactions. Always maintain a coherent and professional tone throughout the conversation.
+
 
         [Response Guidelines]
         - Use natural, polite, and conversational language that is clear and easy to follow. Keep sentences short and use simple words.
@@ -90,6 +91,8 @@ class Captain::Llm::SystemPromptsService
         7. Write the response in multiple paragraphs and in markdown format.
         8. DO NOT use headings in Markdown
         9. Cite the sources if you used a tool to find the response.
+        10. Do not use your own training data or assumptions to answer queries. Base responses strictly on the provided information.
+        11. Always provide a reasoning for the response.
 
         ```json
         {
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..6063242b4
--- /dev/null
+++ b/enterprise/app/services/captain/tool_registry_service.rb
@@ -0,0 +1,32 @@
+class Captain::ToolRegistryService
+  attr_reader :registered_tools, :tools
+
+  def initialize(assistant)
+    @assistant = assistant
+    @registered_tools = []
+    @tools = {}
+    register_default_tools
+  end
+
+  def register_tool(tool_class)
+    tool = tool_class.new(@assistant)
+    @tools[tool.name] = tool
+    @registered_tools << tool.to_registry_format
+  end
+
+  def register_default_tools
+    register_tool(Captain::Tools::SearchDocumentationService)
+  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..72a85fdab
--- /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.debug { "[CAPTAIN][DocumentationSearch] #{query}" }
+    assistant
+      .responses
+      .approved
+      .search(query)
+      .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