fix (AI Assistant): improve docs search accuracy (baserow#4538)

Author: silvestrid

changelog/entries/unreleased/bug/improve_docs_search_accuracy_for_the_ai_assistant.json

@@ -0,0 +1,9 @@
+{
+    "type": "bug",
+    "message": "Improve docs search accuracy for the AI Assistant",
+    "issue_origin": "github",
+    "issue_number": null,
+    "domain": "core",
+    "bullet_points": [],
+    "created_at": "2026-01-14"
+}

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

@@ -5,7 +5,6 @@
 from django.conf import settings
 from django.core.cache import cache
 from django.utils import translation
-from django.utils.translation import gettext as _
 
 import udspy
 from udspy.callback import BaseCallback
@@ -21,7 +20,7 @@
 from baserow_enterprise.assistant.tools.registries import assistant_tool_registry
 
 from .models import AssistantChat, AssistantChatMessage, AssistantChatPrediction
-from .signatures import ChatSignature, RequestRouter
+from .signatures import ChatSignature
 from .types import (
     AiMessage,
     AiMessageChunk,
@@ -200,33 +199,10 @@ def _init_assistant(self):
             "temperature": settings.BASEROW_ENTERPRISE_ASSISTANT_LLM_TEMPERATURE,
             "response_format": {"type": "json_object"},
         }
-        self.search_user_docs_tool = self._get_search_user_docs_tool(tools)
-        self.agent_tools = tools
-        self._request_router = udspy.ChainOfThought(RequestRouter, **module_kwargs)
         self._assistant = udspy.ReAct(
-            ChatSignature, tools=self.agent_tools, max_iters=20, **module_kwargs
+            ChatSignature, tools=tools, max_iters=20, **module_kwargs
         )
 
-    def _get_search_user_docs_tool(
-        self, tools: list[udspy.Tool | Callable]
-    ) -> udspy.Tool | None:
-        """
-        Retrieves the search_user_docs tool from the list of tools if available.
-
-        :param tools: The list of tools to search through.
-        :return: The search_user_docs as udspy.Tool or None if not found.
-        """
-
-        search_user_docs_tool = next(
-            (tool for tool in tools if tool.name == "search_user_docs"), None
-        )
-        if search_user_docs_tool is None or isinstance(
-            search_user_docs_tool, udspy.Tool
-        ):
-            return search_user_docs_tool
-
-        return udspy.Tool(search_user_docs_tool)
-
     async def acreate_chat_message(
         self,
         role: AssistantChatMessage.Role,
@@ -300,14 +276,14 @@ def list_chat_messages(
                 )
         return list(reversed(messages))
 
-    async def afetch_chat_history(self, limit=30):
+    async def afetch_chat_history(self, limit: int = 50) -> udspy.History:
         """
         Loads the chat history into a udspy.History object. It only loads complete
         message pairs (human + AI). The history will be in chronological order and must
         respect the module signature (question, answer).
 
         :param limit: The maximum number of message pairs to load.
-        :return: None
+        :return: A udspy.History instance containing the chat history.
         """
 
         history = udspy.History()
@@ -425,82 +401,6 @@ def _check_cancellation(self, cache_key: str, message_id: str) -> None:
             cache.delete(cache_key)
             raise AssistantMessageCancelled(message_id=message_id)
 
-    async def get_router_stream(
-        self, message: HumanMessage
-    ) -> AsyncGenerator[Any, None]:
-        """
-        Returns an async generator that streams the router's response to a user
-
-        :param message: The current user message that needs context from history.
-        :return: An async generator that yields stream events.
-        """
-
-        self.history = await self.afetch_chat_history()
-
-        return self._request_router.astream(
-            question=message.content,
-            conversation_history=RequestRouter.format_conversation_history(
-                self.history
-            ),
-        )
-
-    async def _process_router_stream(
-        self,
-        event: Any,
-        human_msg: AssistantChatMessage,
-    ) -> Tuple[list[AssistantMessageUnion], bool, udspy.Prediction | None]:
-        """
-        Process a single event from the smart router output stream.
-
-        :param event: The event to process.
-        :param human_msg: The human message instance.
-        :return: a tuple of (messages_to_yield, prediction).
-        """
-
-        messages = []
-        prediction = None
-
-        if isinstance(event, (AiThinkingMessage, AiNavigationMessage)):
-            messages.append(event)
-            return messages, prediction
-
-        # Stream the final answer
-        if isinstance(event, udspy.OutputStreamChunk):
-            if event.field_name == "answer" and event.content.strip():
-                messages.append(
-                    AiMessageChunk(
-                        content=event.content,
-                        sources=self._assistant_callbacks.sources,
-                    )
-                )
-
-        elif isinstance(event, udspy.Prediction):
-            if hasattr(event, "routing_decision"):
-                prediction = event
-
-            if getattr(event, "routing_decision", None) == "delegate_to_agent":
-                messages.append(AiThinkingMessage(content=_("Thinking...")))
-            elif getattr(event, "routing_decision", None) == "search_user_docs":
-                if self.search_user_docs_tool is not None:
-                    await self.search_user_docs_tool(question=event.search_query)
-                else:
-                    messages.append(
-                        AiMessage(
-                            content=_(
-                                "I wanted to search the documentation for you, "
-                                "but the search tool isn't currently available.\n\n"
-                                "To enable documentation search, you'll need to set up "
-                                "the local knowledge base. \n\n"
-                                "You can find setup instructions at: https://baserow.io/user-docs"
-                            ),
-                        )
-                    )
-            elif getattr(event, "answer", None):
-                ai_msg = await self._acreate_ai_message_response(human_msg, event)
-                messages.append(ai_msg)
-
-        return messages, prediction
-
     async def _process_agent_stream(
         self,
         event: Any,
@@ -547,7 +447,7 @@ async def _process_agent_stream(
         return messages, prediction
 
     def get_agent_stream(
-        self, message: HumanMessage, extracted_context: str
+        self, message: HumanMessage, conversation_history: udspy.History | None = None
     ) -> AsyncGenerator[Any, None]:
         """
         Returns an async generator that streams the ReAct agent's response to a user
@@ -557,12 +457,19 @@ def get_agent_stream(
         :return: An async generator that yields stream events.
         """
 
-        ui_context = message.ui_context.format() if message.ui_context else None
+        formatted_history = (
+            ChatSignature.format_conversation_history(conversation_history)
+            if conversation_history
+            else []
+        )
+        formatted_ui_context = (
+            message.ui_context.format() if message.ui_context else None
+        )
 
         return self._assistant.astream(
             question=message.content,
-            context=extracted_context,
-            ui_context=ui_context,
+            conversation_history=formatted_history,
+            ui_context=formatted_ui_context,
         )
 
     async def _process_stream(
@@ -618,31 +525,18 @@ async def astream_messages(
             message_id = str(human_msg.id)
             yield AiStartedMessage(message_id=message_id)
 
-            router_stream = await self.get_router_stream(message)
-            routing_decision, extracted_context = None, ""
+            history = await self.afetch_chat_history(limit=30)
 
-            async for msg, prediction in self._process_stream(
-                human_msg, router_stream, self._process_router_stream
+            agent_stream = self.get_agent_stream(message, history)
+
+            async for msg, __ in self._process_stream(
+                human_msg, agent_stream, self._process_agent_stream
             ):
-                if prediction is not None:
-                    routing_decision = prediction.routing_decision
-                    extracted_context = prediction.extracted_context
                 yield msg
 
-            if routing_decision == "delegate_to_agent":
-                agent_stream = self.get_agent_stream(
-                    message,
-                    extracted_context=extracted_context,
-                )
-
-                async for msg, __ in self._process_stream(
-                    human_msg, agent_stream, self._process_agent_stream
-                ):
-                    yield msg
-
-                # Generate chat title if needed
-                if not self._chat.title:
-                    chat_title = await self._generate_chat_title(human_msg.content)
-                    self._chat.title = chat_title
-                    await self._chat.asave(update_fields=["title", "updated_on"])
-                    yield ChatTitleMessage(content=chat_title)
+            # Generate chat title if needed
+            if not self._chat.title:
+                chat_title = await self._generate_chat_title(human_msg.content)
+                self._chat.title = chat_title
+                await self._chat.asave(update_fields=["title", "updated_on"])
+                yield ChatTitleMessage(content=chat_title)

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

@@ -110,103 +110,58 @@
 AGENT_SYSTEM_PROMPT = (
     ASSISTANT_SYSTEM_PROMPT_BASE
     + """
-**CRITICAL:** You MUST use your action tools to fulfill the request, loading additional tools if needed.
-
-### YOUR TOOLS:
-- **Action tools**: Navigate, list databases, tables, fields, views, filters, workflows, rows, etc.
-- **Tool loaders**: Load additional specialized tools (e.g., load_rows_tools, load_views_tools). Use them to access capabilities not currently available.
-
-**IMPORTANT - HOW TO UNDERSTAND YOUR TOOLS:**
-- Read each tool's NAME, DESCRIPTION, and ARGUMENTS carefully
-- Tool names and descriptions tell you what they do (e.g., "list_tables", "create_rows_in_table_X")
-- Arguments show what inputs they need
-- **NEVER use search_user_docs to learn about tools** - it contains end-user documentation, NOT information about which tools to use or how to call them
-- Inspect available tools directly to decide what to use
-
-### HOW TO WORK:
-1. **Use action tools** to accomplish the user's goal
-2. **If a needed tool isn't available**, call a tool loader to load it (e.g., if you need to create a field but don't have the tool, load field creation tools)
-3. **Keep using tools** until the goal is reached or you confirm NO tool can help and NO tool loader can provide the needed tool
-
-### EXAMPLE - CORRECT USE OF TOOL LOADERS:
-**User request:** "Change all 'Done' tasks to 'Todo'"
-
-**CORRECT approach:**
-✓ Step 1: Identify that Tasks is a table in the open database, and status is the field to update
-✓ Step 2: Notice you need to update rows but don't have the tool
-✓ Step 3: Call the row tool loader (e.g., `load_rows_tools` for table X, requesting update capabilities)
-✓ Step 4: Use the newly loaded `update_rows` tool to update the rows
-✓ Step 5: Complete the task
-
-**CRITICAL:** Before giving up, ALWAYS check if a tool loader can provide the necessary tools to complete the task.
-
-### IF YOU CANNOT COMPLETE THE REQUEST:
-If you've exhausted all available tools and loaders and cannot complete the task, offer: "I wasn't able to complete this using my available tools. Would you like me to search the documentation for instructions on how to do this manually?"
-
-### YOUR PRIORITY:
-1. **First**: Use action tools to complete the request
-2. **If tool missing**: Try loading it with a tool loader (scan all available loaders)
-3. **If truly unable**: Explain the issue and offer to search documentation (never provide instructions from memory)
-
-The router determined this requires action. You were chosen because the user wants you to DO something, not provide information.
-
-Be aware of your limitations. If users ask for something outside your capabilities, finish immediately, explain what you can and cannot do based on the limitations below, and offer to search the documentation for further help.
+## YOUR TOOLS
+
+**CRITICAL - Understanding your tools:**
+- Learn what each tool does ONLY from its **name** and **description**
+- **NEVER use `search_user_docs` to learn about your tools** - it contains end-user documentation, NOT information about your available tools or how to call them
+- `search_user_docs` is ONLY for answering user questions about Baserow features and providing manual instructions
+
+## REQUEST HANDLING
+
+### ACTION REQUESTS - CHECK FIRST
+
+**CRITICAL: Before treating a request as a question, determine if it's an action you can perform.**
+
+Recognize action requests by:
+- Imperative verbs: "Show...", "Filter...", "Create...", "Add...", "Delete...", "Update...", "Sort...", "Hide..."
+- Desired states: "I want only...", "I need a field that...", "Make it show..."
+- Example: "Show only rows where the primary field is empty" → This is an ACTION (create a filter), not a question about filtering
+
+**DO vs EXPLAIN:**
+- If you have tools to do it → **DO IT**
+- If you lack tools → **THEN explain** how to do it manually
+- **NEVER explain how to do something you can do yourself**
+
+**Workflow:**
+1. Check your tools - can you fulfill this?
+2. **YES**: Execute (ask for clarification only if request is ambiguous)
+3. **NO** (see LIMITATIONS): Explain you can't, then provide manual instructions from docs
+
+### QUESTIONS (only after ruling out action requests)
+
+**FACTUAL QUESTIONS** - asking what Baserow IS or HAS:
+- Examples: "Does Baserow have X feature?", "How does Y work?", "What options exist for Z?"
+- These have objectively correct/incorrect answers that must come from documentation
+- **ALWAYS search documentation first** using `search_user_docs`
+- Check the `reliability_note` in the response:
+  - **HIGH CONFIDENCE**: Present the answer confidently with sources
+  - **PARTIAL MATCH**: Provide the answer but note some details may be incomplete
+  - **LOW CONFIDENCE / NOTHING FOUND**: Tell the user you couldn't find this in the documentation. **DO NOT guess or assume features exist** - if docs don't mention it (e.g., a "barcode field"), it likely doesn't exist. Suggest checking the community forum or contacting support.
+- **NEVER fabricate Baserow features or capabilities**
+
+**ADVISORY QUESTIONS** - asking how to USE or APPLY Baserow:
+- Examples: "How should I structure X?", "What's a good approach for Y?", "Help me build Z", "Which field type works best for W?"
+- These ask for your expertise in applying Baserow to solve problems - there's no single correct answer
+- **Use your knowledge** of Baserow's real capabilities (field types, views, formulas, automations, linking, etc.) to provide helpful recommendations
+- You may search docs for reference, but can also directly advise based on your understanding of Baserow
+- Focus on practical solutions using actual Baserow functionality
+
+**Key principle**: Never fabricate what Baserow CAN do. Freely advise on HOW to use what Baserow actually offers.
 """
     + AGENT_LIMITATIONS
     + """
-### TASK INSTRUCTIONS:
-"""
-)
 
-
-REQUEST_ROUTER_PROMPT = (
-    ASSISTANT_SYSTEM_PROMPT_BASE
-    + """
-Route based on what the user wants YOU to do:
-
-**delegate_to_agent** (DEFAULT) - User wants YOU to perform an action
-- Commands/requests for YOU: "Create...", "Delete...", "Update...", "Add...", "Show me...", "List...", "Find..."
-- Vague/unclear requests
-- Anything not explicitly asking for instructions
-
-**search_user_docs** - User wants to learn HOW TO do something themselves
-- ONLY when explicitly asking for instructions: "How do I...", "How can I...", "What are the steps to..."
-- ONLY when asking for explanations: "What is...", "What does... mean", "Explain..."
-- NOT for action requests even if phrased as questions
-
-## Critical Rules
-- "Create X" → delegate_to_agent (action request for YOU)
-- "How do I create X?" → search_user_docs (asking for instructions)
-- When uncertain → delegate_to_agent
-
-## Output Requirements
-**delegate_to_agent:**
-- extracted_context: Comprehensive details from conversation history (IDs, names, actions, specs)
-- search_query: empty
-
-**search_user_docs:**
-- search_query: Clear question using Baserow terminology and the answer language if not English
-- extracted_context: empty
-
-## Examples
-
-**Example 1 - delegate_to_agent (action):**
-question: "Create a calendar view"
-→ routing_decision: "delegate_to_agent"
-→ search_query: ""
-→ extracted_context: "User wants to create a calendar view."
-
-**Example 2 - search_user_docs (instructions):**
-question: "How do I create a calendar view?"
-→ routing_decision: "search_user_docs"
-→ search_query: "How to create a calendar view in Baserow"
-→ extracted_context: ""
-
-**Example 3 - delegate_to_agent (with history):**
-question: "Assign them to Bob"
-conversation_history: ["[0] (user): Show urgent tasks", "[1] (assistant): Found 5 tasks in table 'Tasks' (ID: 123)"]
-→ routing_decision: "delegate_to_agent"
-→ search_query: ""
-→ extracted_context: "User wants to assign urgent tasks to Bob. Tasks in table 'Tasks' (ID: 123). Found 5 urgent tasks."
+## TASK INSTRUCTIONS:
 """
 )