From 94b81b49f96f7a37a7d63aa627b0c8b0e721c0f5 Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Fri, 4 Apr 2025 22:49:03 +0530
Subject: [PATCH] feat: scope assistant

---
 .../captain/llm/assistant_chat_service.rb     |  3 +-
 .../captain/tools/documentation_search.rb     | 35 ++++++++++++-------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/enterprise/app/services/captain/llm/assistant_chat_service.rb b/enterprise/app/services/captain/llm/assistant_chat_service.rb
index 14eccb5e6..8423b1687 100644
--- a/enterprise/app/services/captain/llm/assistant_chat_service.rb
+++ b/enterprise/app/services/captain/llm/assistant_chat_service.rb
@@ -2,7 +2,8 @@ class Captain::Llm::AssistantChatService < Llm::BaseService
   def initialize(assistant: nil)
     super()
     @assistant = assistant
-    @chat = ::RubyLLM.chat(model: @model).with_tool(Captain::Tools::DocumentationSearch)
+    search_tool = Captain::Tools::DocumentationSearch.new(assistant)
+    @chat = ::RubyLLM.chat(model: @model).with_tool(search_tool)
     @chat.with_instructions(system_message)
   end
 
diff --git a/enterprise/app/services/captain/tools/documentation_search.rb b/enterprise/app/services/captain/tools/documentation_search.rb
index 65fdf3e71..dc8fef7dd 100644
--- a/enterprise/app/services/captain/tools/documentation_search.rb
+++ b/enterprise/app/services/captain/tools/documentation_search.rb
@@ -5,23 +5,34 @@ class Captain::Tools::DocumentationSearch < RubyLLM::Tool
         type: :string,
         desc: 'The search query to find relevant documentation'
 
-  def execute(search_query:)
-    responses = ::Captain::AssistantResponse.approved.search(search_query)
+  def initialize(assistant)
+    super()
+    @assistant = assistant
+  end
+
+  def execute(query:)
+    @assistant
+      .responses
+      .approved
+      .search(query)
+      .map { |response| format_response(response) }.join
+
     format_responses(responses)
   end
 
   private
 
-  def format_responses(responses)
-    responses.map do |response|
-      formatted_response = {
-        question: response.question,
-        answer: response.answer
-      }
-
-      formatted_response[:source] = response.documentable.external_link if response.documentable.present? && response.documentable.try(:external_link)
-
-      formatted_response
+  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