Merge branch 'ai-assistant-ui-context' into 'develop'

AI Assistant: update UI context

See merge request baserow/baserow!3739

Author: silvestrid

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

@@ -17,13 +17,76 @@ class AssistantChatsRequestSerializer(serializers.Serializer):
     limit = serializers.IntegerField(default=100, max_value=100, min_value=1)
 
 
+class UIContextApplicationSerializer(serializers.Serializer):
+    id = serializers.CharField(help_text="The unique ID of the application.")
+    name = serializers.CharField(help_text="The name of the application.")
+
+
+class UIContextTableSerializer(serializers.Serializer):
+    id = serializers.IntegerField(help_text="The ID of the table.")
+    name = serializers.CharField(help_text="The name of the table.")
+
+
+class UIContextViewSerializer(serializers.Serializer):
+    id = serializers.IntegerField(help_text="The ID of the view.")
+    name = serializers.CharField(help_text="The name of the view.")
+    type = serializers.CharField(help_text="The type of the view.")
+
+
 class UIContextWorkspaceSerializer(serializers.Serializer):
     id = serializers.IntegerField(help_text="The ID of the workspace.")
     name = serializers.CharField(help_text="The name of the workspace.")
 
 
+class UIContextPageSerializer(serializers.Serializer):
+    id = serializers.CharField(help_text="The unique ID of the page.")
+    name = serializers.CharField(help_text="The name of the page.")
+
+
+class UIContextWorkflowSerializer(serializers.Serializer):
+    id = serializers.CharField(help_text="The unique ID of the workflow.")
+    name = serializers.CharField(help_text="The name of the workflow.")
+
+
 class UIContextSerializer(serializers.Serializer):
     workspace = UIContextWorkspaceSerializer()
+    # database builder context
+    database = UIContextApplicationSerializer(
+        required=False,
+        help_text="The application the user is currently in, e.g. 'database'.",
+    )
+    table = UIContextTableSerializer(
+        required=False,
+        help_text="The table the user is currently viewing, if any.",
+    )
+    view = UIContextViewSerializer(
+        required=False,
+        help_text="The view the user is currently viewing, if any.",
+    )
+    # application builder context
+    application = UIContextApplicationSerializer(
+        required=False,
+        help_text="The application the user is currently in, e.g. 'application'.",
+    )
+    page = UIContextPageSerializer(
+        required=False,
+        help_text="The page the user is currently viewing, if any.",
+    )
+    # automation builder context
+    automation = UIContextApplicationSerializer(
+        required=False,
+        help_text="The application the user is currently in, e.g. 'automation'.",
+    )
+    workflow = UIContextWorkflowSerializer(
+        required=False,
+        help_text="The workflow the user is currently viewing, if any.",
+    )
+    # dashboard builder context
+    dashboard = UIContextApplicationSerializer(
+        required=False,
+        help_text="The application the user is currently in, e.g. 'dashboard'.",
+    )
+    # user context
     timezone = serializers.CharField(
         required=False,
         help_text="The timezone of the user, e.g. 'Europe/Amsterdam'.",

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

@@ -31,7 +31,11 @@
 )
 from baserow_enterprise.assistant.handler import AssistantHandler
 from baserow_enterprise.assistant.operations import ChatAssistantChatOperationType
-from baserow_enterprise.assistant.types import AssistantMessageUnion, UIContext
+from baserow_enterprise.assistant.types import (
+    AssistantMessageUnion,
+    HumanMessage,
+    UIContext,
+)
 from baserow_enterprise.features import ASSISTANT
 
 from .serializers import (
@@ -142,7 +146,7 @@ class AssistantChatView(APIView):
     def post(self, request: Request, chat_uuid: str, data) -> StreamingHttpResponse:
         feature_flag_is_enabled(FF_ASSISTANT, raise_if_disabled=True)
 
-        ui_context = UIContext(**data["ui_context"])
+        ui_context = UIContext.from_validate_request(request, data["ui_context"])
         workspace_id = ui_context.workspace.id
         workspace = CoreHandler().get_workspace(workspace_id)
         LicenseHandler.raise_if_user_doesnt_have_feature(
@@ -158,12 +162,10 @@ def post(self, request: Request, chat_uuid: str, data) -> StreamingHttpResponse:
         handler = AssistantHandler()
         chat, _ = handler.get_or_create_chat(request.user, workspace, chat_uuid)
         assistant = handler.get_assistant(chat)
-        human_message = data["content"]
+        human_message = HumanMessage(content=data["content"], ui_context=ui_context)
 
         async def stream_assistant_messages():
-            async for msg in assistant.astream_messages(
-                human_message, ui_context=ui_context
-            ):
+            async for msg in assistant.astream_messages(human_message):
                 yield self._stream_assistant_message(msg)
 
         response = StreamingHttpResponse(

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

@@ -26,6 +26,13 @@
 class ChatSignature(dspy.Signature):
     question: str = dspy.InputField()
     history: dspy.History = dspy.InputField()
+    ui_context: UIContext = dspy.InputField(
+        default=None,
+        description=(
+            "The frontend UI content the user is currently in. "
+            "Whenever make sense, use it to ground your answer."
+        ),
+    )
     answer: str = dspy.OutputField()
 
 
@@ -131,12 +138,13 @@ def _get_chat_signature(self) -> dspy.Signature:
             history, and answer fields.
         """
 
+        chat_signature_instructions = "## INSTRUCTIONS\n\nGiven the fields `question`, `history`, and `ui_context`, produce the fields `answer`"
         if self._chat.title:  # only inject our base system prompt
             return ChatSignature.with_instructions(
                 "\n".join(
                     [
                         ASSISTANT_SYSTEM_PROMPT,
-                        "Given the fields `question`, `history`, produce the fields `answer`.",
+                        f"{chat_signature_instructions}.",
                     ]
                 )
             )
@@ -155,7 +163,7 @@ def _get_chat_signature(self) -> dspy.Signature:
                 instructions="\n".join(
                     [
                         ASSISTANT_SYSTEM_PROMPT,
-                        "Given the fields `question`, `history`, produce the fields `answer`, `chat_title`.",
+                        f"{chat_signature_instructions}, `chat_title`.",
                     ]
                 ),
             )
@@ -250,13 +258,12 @@ async def aload_chat_history(self, limit=20):
         self.history = dspy.History(messages=messages)
 
     async def astream_messages(
-        self, user_message: str, ui_context: UIContext
+        self, human_message: HumanMessage
     ) -> AsyncGenerator[AssistantMessageUnion, None]:
         """
         Streams the response to a user message.
 
-        :param user_message: The message from the user.
-        :param ui_context: The UI context where the message was sent.
+        :param human_message: The message from the user.
         :return: An async generator that yields the response messages.
         """
 
@@ -281,10 +288,16 @@ async def astream_messages(
                 self._assistant,
                 stream_listeners=stream_listeners,
             )
-            output_stream = stream_predict(question=user_message, history=self.history)
+            output_stream = stream_predict(
+                history=self.history,
+                question=human_message.content,
+                ui_context=human_message.ui_context.model_dump_json(
+                    exclude_none=True, indent=2
+                ),
+            )
 
             await self.acreate_chat_message(
-                AssistantChatMessage.Role.HUMAN, user_message
+                AssistantChatMessage.Role.HUMAN, human_message.content
             )
 
             chat_title, answer = "", ""

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

@@ -16,7 +16,7 @@
 **Key concepts**:
 • **Fields**: Define schema (30+ types including link_row for relationships); one primary field per table
 • **Views**: Present data with filters/sorts/grouping/colors; can be shared, personal, or public
-• **Snapshots**: Database backups; **Data sync**: Table replication; **Webhooks**: Row/field event triggers
+• **Snapshots**: Database backups; **Data sync**: Table replication; **Webhooks**: Row/field/view event triggers
 • **Permissions**: RBAC at workspace/database/table/field levels; database tokens for API
 """
 
@@ -40,11 +40,12 @@
 
 You know:
 1. **Core concepts** (below) - answer directly
-2. **Detailed docs** - use knowledge retrieval tool to search when needed
+2. **Detailed docs** - use search_docs tool to search when needed
 3. **API specs** - guide users to https://api.baserow.io/api/schema.json
 
 ## HOW TO HELP
 
+• Use American English spelling and grammar
 • Be clear, concise, and actionable
 • For troubleshooting: ask for error messages or describe expected vs actual results
 • If uncertain: acknowledge it, then suggest how to find the answer (search docs, check API, etc.)
@@ -63,7 +64,5 @@
     + APPLICATION_BUILDER_CONCEPTS
     + "\n"
     """
-
-## INSTRUCTIONS
 """
 )

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

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import StrEnum
 from typing import Literal, Optional
 
@@ -10,12 +10,76 @@ class WorkspaceUIContext(BaseModel):
     name: str
 
 
+class ApplicationUIContext(BaseModel):
+    id: str
+    name: str
+
+
+class TableUIContext(BaseModel):
+    id: int
+    name: str
+
+
+class ViewUIContext(BaseModel):
+    id: int
+    name: str
+    type: str
+
+
+class UserUIContext(BaseModel):
+    id: int
+    name: str
+    email: str
+
+    @classmethod
+    def from_user(cls, user) -> "UserUIContext":
+        return cls(id=user.id, name=user.first_name, email=user.email)
+
+
+class PageUIContext(BaseModel):
+    id: str
+    name: str
+
+
+class WorkflowUIContext(BaseModel):
+    id: str
+    name: str
+
+
+class DashboardUIContext(BaseModel):
+    id: str
+    name: str
+
+
 class UIContext(BaseModel):
     workspace: WorkspaceUIContext
-    timezone: Optional[str] = Field(
+    # database builder context
+    database: Optional[ApplicationUIContext] = None
+    table: Optional[TableUIContext] = None
+    view: Optional[ViewUIContext] = None
+    # application builder context
+    application: Optional[ApplicationUIContext] = None
+    page: Optional[PageUIContext] = None
+    # automation context
+    automation: Optional[ApplicationUIContext] = None
+    workflow: Optional[WorkflowUIContext] = None
+    # dashboard context
+    dashboard: Optional[DashboardUIContext] = None
+    # user and time context
+    user: UserUIContext
+    timestamp: datetime = Field(
+        default_factory=lambda: datetime.now(tz=timezone.utc),
+        description="The UTC timestamp when the message was sent",
+    )
+    timezone: str = Field(
         default="UTC", description="The timezone of the user, e.g. 'Europe/Amsterdam'"
     )
 
+    @classmethod
+    def from_validate_request(cls, request, ui_context_data) -> "UIContext":
+        user_context = UserUIContext.from_user(request.user)
+        return cls(user=user_context, **ui_context_data)
+
 
 class AssistantMessageType(StrEnum):
     HUMAN = "human"