chore(AI Assistant): improve context awereness (baserow#4261)

Author: silvestrid

enterprise/backend/src/baserow_enterprise/assistant/assistant.py

@@ -122,9 +122,12 @@ class ChatSignature(udspy.Signature):
     __doc__ = f"{ASSISTANT_SYSTEM_PROMPT}\n TASK INSTRUCTIONS: \n"
 
     question: str = udspy.InputField()
+    context: str = udspy.InputField(
+        description="Context and facts extracted from the history to help answer the question."
+    )
     ui_context: dict[str, Any] | None = udspy.InputField(
         default=None,
-        desc=(
+        description=(
             "The context the user is currently in. "
             "It contains information about the user, the workspace, open table, view, etc."
             "Whenever make sense, use it to ground your answer."
@@ -133,6 +136,29 @@ class ChatSignature(udspy.Signature):
     answer: str = udspy.OutputField()
 
 
+class QuestionContextSummarizationSignature(udspy.Signature):
+    """
+    Extract relevant facts from conversation history that provide context for answering
+    the current question. Do not answer the question or modify it - only extract and
+    summarize the relevant historical facts that will help in decision-making.
+    """
+
+    question: str = udspy.InputField(
+        description="The current user question that needs context from history."
+    )
+    previous_messages: list[str] = udspy.InputField(
+        description="Conversation history as alternating user/assistant messages."
+    )
+    facts: str = udspy.OutputField(
+        description=(
+            "Relevant facts extracted from the conversation history as a concise "
+            "paragraph. Include only information that provides necessary context for "
+            "answering the question. Do not answer the question itself, do not modify "
+            "the question, and do not include irrelevant details."
+        )
+    )
+
+
 def get_assistant_cancellation_key(chat_uuid: str) -> str:
     """
     Get the Redis cache key for cancellation tracking.
@@ -176,7 +202,7 @@ def _init_assistant(self):
         )
         self.callbacks = AssistantCallbacks(self.tool_helpers)
         self._assistant = udspy.ReAct(ChatSignature, tools=tools, max_iters=20)
-        self.history = None
+        self.history: list[str] = []
 
     async def acreate_chat_message(
         self,
@@ -265,7 +291,6 @@ async def aload_chat_history(self, limit=30):
             msg async for msg in self._chat.messages.order_by("-created_on")[:limit]
         ]
 
-        self.history = udspy.History()
         while len(last_saved_messages) >= 2:
             # Pop the oldest message pair to respect chronological order.
             first_message = last_saved_messages.pop()
@@ -276,9 +301,9 @@ async def aload_chat_history(self, limit=30):
             ):
                 continue
 
-            self.history.add_user_message(first_message.content)
+            self.history.append(f"Human: {first_message.content}")
             ai_answer = last_saved_messages.pop()
-            self.history.add_assistant_message(ai_answer.content)
+            self.history.append(f"AI: {ai_answer.content}")
 
     @lru_cache(maxsize=1)
     def check_llm_ready_or_raise(self):
@@ -376,17 +401,24 @@ def _check_cancellation(self, cache_key: str, message_id: str) -> None:
             cache.delete(cache_key)
             raise AssistantMessageCancelled(message_id=message_id)
 
-    async def _enhance_question_with_history(self, question: str) -> str:
-        """Enhance the user question with chat history context if available."""
+    async def _summarize_context_from_history(self, question: str) -> str:
+        """
+        Extract relevant facts from chat history to provide context for the question or
+        return an empty string if there is no history.
+
+        :param question: The current user question that needs context from history.
+        :return: A string containing relevant facts from the conversation history.
+        """
 
-        if not self.history.messages:
-            return question
+        if not self.history:
+            return ""
 
-        predictor = udspy.Predict("question, context -> enhanced_question")
+        predictor = udspy.Predict(QuestionContextSummarizationSignature)
         result = await predictor.aforward(
-            question=question, context=self.history.messages
+            question=question,
+            previous_messages=self.history,
         )
-        return result.enhanced_question
+        return result.facts
 
     async def _process_stream_event(
         self,
@@ -458,12 +490,13 @@ async def astream_messages(
             if self.history is None:
                 await self.aload_chat_history()
 
-            user_question = await self._enhance_question_with_history(
+            context_from_history = await self._summarize_context_from_history(
                 human_message.content
             )
 
             output_stream = self._assistant.astream(
-                question=user_question,
+                question=human_message.content,
+                context=context_from_history,
                 ui_context=human_message.ui_context.model_dump_json(exclude_none=True),
             )
 

enterprise/backend/src/baserow_enterprise/assistant/tools/search_docs/tools.py

@@ -43,7 +43,7 @@ def __init__(self):
         self.rag = udspy.ChainOfThought(SearchDocsSignature)
 
     def forward(self, question: str, *args, **kwargs):
-        context = KnowledgeBaseHandler().search(question, num_results=10)
+        context = KnowledgeBaseHandler().search(question, num_results=7)
         return self.rag(context=context, question=question)
 
 

enterprise/backend/tests/baserow_enterprise_tests/assistant/test_assistant.py

@@ -207,25 +207,18 @@ def test_aload_chat_history_formats_as_question_answer_pairs(
 
         # History should contain user/assistant message pairs
         assert assistant.history is not None
-        assert len(assistant.history.messages) == 4
+        assert len(assistant.history) == 4
 
         # First pair
-        assert assistant.history.messages[0]["content"] == "What is Baserow?"
-        assert assistant.history.messages[0]["role"] == "user"
-        assert (
-            assistant.history.messages[1]["content"]
-            == "Baserow is a no-code database platform."
-        )
-        assert assistant.history.messages[1]["role"] == "assistant"
+        assert assistant.history[0] == "Human: What is Baserow?"
+        assert assistant.history[1] == "AI: Baserow is a no-code database platform."
 
         # Second pair
-        assert assistant.history.messages[2]["content"] == "How do I create a table?"
-        assert assistant.history.messages[2]["role"] == "user"
+        assert assistant.history[2] == "Human: How do I create a table?"
         assert (
-            assistant.history.messages[3]["content"]
-            == "You can create a table by clicking the + button."
+            assistant.history[3]
+            == "AI: You can create a table by clicking the + button."
         )
-        assert assistant.history.messages[3]["role"] == "assistant"
 
     def test_aload_chat_history_respects_limit(self, enterprise_data_fixture):
         """Test that history loading respects the limit parameter"""
@@ -251,7 +244,7 @@ def test_aload_chat_history_respects_limit(self, enterprise_data_fixture):
         async_to_sync(assistant.aload_chat_history)(limit=6)  # Last 6 messages
 
         # Should only load the most recent 6 messages (3 pairs)
-        assert len(assistant.history.messages) == 6
+        assert len(assistant.history) == 6
 
     def test_aload_chat_history_handles_incomplete_pairs(self, enterprise_data_fixture):
         """
@@ -281,11 +274,9 @@ def test_aload_chat_history_handles_incomplete_pairs(self, enterprise_data_fixtu
         async_to_sync(assistant.aload_chat_history)()
 
         # Should only include the complete pair (2 messages: user + assistant)
-        assert len(assistant.history.messages) == 2
-        assert assistant.history.messages[0]["content"] == "Question 1"
-        assert assistant.history.messages[0]["role"] == "user"
-        assert assistant.history.messages[1]["content"] == "Answer 1"
-        assert assistant.history.messages[1]["role"] == "assistant"
+        assert len(assistant.history) == 2
+        assert assistant.history[0] == "Human: Question 1"
+        assert assistant.history[1] == "AI: Answer 1"
 
 
 @pytest.mark.django_db

enterprise/web-frontend/modules/baserow_enterprise/components/assistant/AssistantMessageActions.vue

@@ -39,10 +39,6 @@
       </template>
     </div>
 
-    <div v-if="message.can_submit_feedback" class="assistant__disclaimer">
-      {{ $t('assistantMessageActions.disclaimer') }}
-    </div>
-
     <!-- Additional user feedback context for the thumb down button -->
     <Context
       ref="feedbackContext"

enterprise/web-frontend/modules/baserow_enterprise/components/assistant/AssistantMessageList.vue

@@ -1,7 +1,7 @@
 <template>
   <div class="assistant__messages-list">
     <div
-      v-for="message in messages"
+      v-for="(message, index) in messages"
       :key="message.id"
       class="assistant__message"
       :class="{
@@ -33,7 +33,6 @@
               ></div>
             </div>
 
-            <!-- Sources section - only show for AI messages with sources -->
             <AssistantMessageSources
               v-if="message.role === 'ai'"
               :sources="message.sources"
@@ -44,6 +43,12 @@
         </div>
 
         <AssistantMessageActions :message="message" />
+        <div
+          v-if="message.can_submit_feedback && isLastMessage(index)"
+          class="assistant__disclaimer"
+        >
+          {{ $t('assistantMessageList.disclaimer') }}
+        </div>
       </div>
     </div>
   </div>
@@ -129,6 +134,10 @@ export default {
         !this.expandedSources[messageId]
       )
     },
+
+    isLastMessage(index) {
+      return index === this.messages.length - 1
+    },
   },
 }
 </script>

enterprise/web-frontend/modules/baserow_enterprise/locales/en.json

@@ -368,13 +368,15 @@
   "assistantMessageSources": {
     "sources": "{count} source | {count} sources"
   },
+  "assistantMessageList": {
+    "disclaimer": "Kuma can make mistakes, please double-check responses"
+  },
   "assistantMessageActions": {
     "feedbackContextTitle": "Help us improve",
     "feedbackContextPlaceholder": "What could we improve? (optional)",
     "copiedToClipboard": "Copied to clipboard",
     "copiedContentToast": "The Assistant's response content has been copied to your clipboard",
-    "copyFailed": "Failed to copy to clipboard",
-    "disclaimer": "Kuma can make mistakes, please double-check responses"
+    "copyFailed": "Failed to copy to clipboard"
   },
   "chatwootSupportSidebarWorkspace": {
     "directSupport": "Direct support"
@@ -683,7 +685,7 @@
     "fieldInvalidTitle": "Date dependency field error"
   },
   "dateDependency": {
-   "invalidChildRow": "Successor row is invalid",
+    "invalidChildRow": "Successor row is invalid",
     "invalidParentRow": "Predecessor row is invalid",
     "invalidParentEndDateAfterChildStartDate": "Predecessor row end date is after successor start date",
     "invalidStartDateEmpty": "Start date is empty",