feat(max): deep research template system

ee/api/conversation.py

@@ -49,6 +49,7 @@ class MessageSerializer(serializers.Serializer):
     trace_id = serializers.UUIDField(required=True)
     session_id = serializers.CharField(required=False)
     deep_research_mode = serializers.BooleanField(required=False, default=False)
+    deep_research_template = serializers.JSONField(required=False)
 
     def validate(self, data):
         if data["content"] is not None:
@@ -146,8 +147,10 @@ def create(self, request: Request, *args, **kwargs):
                     {"error": "Cannot access other users' conversations"}, status=status.HTTP_400_BAD_REQUEST
                 )
         except Conversation.DoesNotExist:
-            # Conversation doesn't exist, create it if we have a message
-            if not has_message:
+            # Allow creation with either a message or a deep research template
+            if not has_message and not (
+                is_deep_research and serializer.validated_data.get("deep_research_template") is not None
+            ):
                 return Response(
                     {"error": "Cannot stream from non-existent conversation"}, status=status.HTTP_400_BAD_REQUEST
                 )
@@ -164,6 +167,10 @@ def create(self, request: Request, *args, **kwargs):
         if has_message and not is_idle:
             raise Conflict("Cannot resume streaming with a new message")
 
+        deep_research_template = None
+        if is_deep_research:
+            deep_research_template = serializer.validated_data.get("deep_research_template")
+
         workflow_inputs = AssistantConversationRunnerWorkflowInputs(
             team_id=self.team_id,
             user_id=cast(User, request.user).pk,  # Use pk instead of id for User model
@@ -175,6 +182,7 @@ def create(self, request: Request, *args, **kwargs):
             session_id=request.headers.get("X-POSTHOG-SESSION-ID"),  # Relies on posthog-js __add_tracing_headers
             billing_context=serializer.validated_data.get("billing_context"),
             mode=mode,
+            deep_research_template=deep_research_template,
         )
 
         async def async_stream(

ee/hogai/assistant/assistant.py

@@ -9,6 +9,8 @@
 from ee.hogai.assistant.deep_research_assistant import DeepResearchAssistant
 from ee.hogai.assistant.insights_assistant import InsightsAssistant
 from ee.hogai.assistant.main_assistant import MainAssistant
+from ee.hogai.graph.deep_research.types import DeepResearchState, PartialDeepResearchState
+from ee.hogai.utils.types import AssistantState, PartialAssistantState
 from ee.hogai.utils.types.base import AssistantMode
 from ee.hogai.utils.types.composed import AssistantMaxGraphState, AssistantMaxPartialGraphState
 from ee.models import Conversation
@@ -30,23 +32,59 @@ def create(
         trace_id: Optional[str | UUID] = None,
         initial_state: Optional[AssistantMaxGraphState | AssistantMaxPartialGraphState] = None,
         billing_context: Optional[MaxBillingContext] = None,
+        deep_research_template: Optional[dict[str, Any]] = None,
     ) -> BaseAssistant:
-        assistant_class: type[BaseAssistant]
         if mode == AssistantMode.ASSISTANT:
-            assistant_class = MainAssistant
+            assistant_initial_state: Optional[AssistantState | PartialAssistantState] = None
+            if initial_state is not None:
+                if isinstance(initial_state, (AssistantState | PartialAssistantState)):
+                    assistant_initial_state = initial_state
+            return MainAssistant(
+                team,
+                conversation,
+                new_message=new_message,
+                user=user,
+                session_id=session_id,
+                contextual_tools=contextual_tools,
+                is_new_conversation=is_new_conversation,
+                trace_id=trace_id,
+                billing_context=billing_context,
+                initial_state=assistant_initial_state,
+            )
         elif mode == AssistantMode.INSIGHTS_TOOL:
-            assistant_class = InsightsAssistant
+            assistant_initial_state = None
+            if initial_state is not None:
+                if isinstance(initial_state, (AssistantState | PartialAssistantState)):
+                    assistant_initial_state = initial_state
+            return InsightsAssistant(
+                team,
+                conversation,
+                new_message=new_message,
+                user=user,
+                session_id=session_id,
+                contextual_tools=contextual_tools,
+                is_new_conversation=is_new_conversation,
+                trace_id=trace_id,
+                billing_context=billing_context,
+                initial_state=assistant_initial_state,
+            )
         elif mode == AssistantMode.DEEP_RESEARCH:
-            assistant_class = DeepResearchAssistant
-        return assistant_class(
-            team,
-            conversation,
-            new_message=new_message,
-            user=user,
-            session_id=session_id,
-            contextual_tools=contextual_tools,
-            is_new_conversation=is_new_conversation,
-            trace_id=trace_id,
-            billing_context=billing_context,
-            initial_state=initial_state,  # type: ignore
-        )
+            deep_research_initial_state: Optional[DeepResearchState | PartialDeepResearchState] = None
+            if initial_state is not None:
+                if isinstance(initial_state, (DeepResearchState | PartialDeepResearchState)):
+                    deep_research_initial_state = initial_state
+            return DeepResearchAssistant(
+                team,
+                conversation,
+                new_message=new_message,
+                user=user,
+                session_id=session_id,
+                contextual_tools=contextual_tools,
+                is_new_conversation=is_new_conversation,
+                trace_id=trace_id,
+                billing_context=billing_context,
+                initial_state=deep_research_initial_state,
+                deep_research_template=deep_research_template,
+            )
+        else:
+            raise ValueError(f"Unknown assistant mode: {mode}")

ee/hogai/assistant/base.py

@@ -101,6 +101,7 @@ def __init__(
         billing_context: Optional[MaxBillingContext] = None,
         initial_state: Optional[AssistantMaxGraphState | AssistantMaxPartialGraphState] = None,
         callback_handler: Optional[BaseCallbackHandler] = None,
+        deep_research_template: Optional[dict[str, Any]] = None,
     ):
         self._team = team
         self._contextual_tools = contextual_tools or {}
@@ -138,6 +139,7 @@ def __init__(
         self._mode = mode
         self._initial_state = initial_state
         self._commentary_chunk = None
+        self._deep_research_template = deep_research_template
 
     @property
     @abstractmethod

ee/hogai/assistant/deep_research_assistant.py

@@ -41,6 +41,7 @@ def __init__(
         trace_id: Optional[str | UUID] = None,
         billing_context: Optional[MaxBillingContext] = None,
         initial_state: Optional[DeepResearchState | PartialDeepResearchState] = None,
+        deep_research_template: Optional[dict[str, Any]] = None,
     ):
         super().__init__(
             team,
@@ -57,6 +58,7 @@ def __init__(
             trace_id=trace_id,
             billing_context=billing_context,
             initial_state=initial_state,
+            deep_research_template=deep_research_template,
         )
 
     @property
@@ -110,15 +112,35 @@ def _should_persist_commentary_message(self, node_name: MaxNodeName) -> bool:
         return False
 
     def get_initial_state(self) -> DeepResearchState:
-        if self._latest_message:
-            return DeepResearchState(
-                messages=[self._latest_message],
-                start_id=self._latest_message.id,
-                graph_status=None,
-                notebook_short_id=None,
-            )
-        else:
-            return DeepResearchState(messages=[])
+        # Inject a default human message when a template is selected without user input,
+        # and stream it immediately by setting _latest_message.
+        message_for_state = self._latest_message
+        if not self._latest_message and self._deep_research_template:
+            from uuid import uuid4
+
+            title = None
+            if isinstance(self._deep_research_template, dict):
+                title = self._deep_research_template.get("notebook_title")
+
+            content = f"Load template: {title}" if title else "Load template"
+            message_for_state = HumanMessage(content=content, id=str(uuid4()))
+            self._latest_message = message_for_state
+
+        base_state = DeepResearchState(
+            messages=[message_for_state] if message_for_state else [],
+            start_id=message_for_state.id if message_for_state else None,
+            graph_status=None,
+            notebook_short_id=None,
+        )
+
+        if self._deep_research_template:
+            if isinstance(self._deep_research_template, dict):
+                notebook_short_id = self._deep_research_template.get("notebook_short_id")
+                if notebook_short_id:
+                    base_state.template_notebook_short_id = notebook_short_id
+                    base_state.skip_onboarding = True
+
+        return base_state
 
     def get_resumed_state(self) -> PartialDeepResearchState:
         if not self._latest_message:

ee/hogai/graph/deep_research/notebook/nodes.py

@@ -1,7 +1,19 @@
+from uuid import uuid4
+
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import RunnableConfig
+from posthoganalytics import capture_exception
+
+from posthog.schema import (
+    DeepResearchNotebook,
+    DeepResearchType,
+    HumanMessage,
+    NotebookUpdateMessage,
+    ProsemirrorJSONContent,
+)
 
-from posthog.schema import DeepResearchNotebook, DeepResearchType, HumanMessage
+from posthog.models.notebook.notebook import Notebook
+from posthog.sync import database_sync_to_async
 
 from ee.hogai.graph.deep_research.base.nodes import DeepResearchNode
 from ee.hogai.graph.deep_research.notebook.prompts import DEEP_RESEARCH_NOTEBOOK_PLANNING_PROMPT
@@ -15,18 +27,47 @@ def node_name(self) -> MaxNodeName:
         return DeepResearchNodeName.NOTEBOOK_PLANNING
 
     async def arun(self, state: DeepResearchState, config: RunnableConfig) -> PartialDeepResearchState:
+        # Load template
+        template_markdown = await self._retrieve_template_markdown(state)
         # We use instructions with the OpenAI Responses API
         instructions = DEEP_RESEARCH_NOTEBOOK_PLANNING_PROMPT.format(
             core_memory=await self._aget_core_memory(),
         )
 
+        # Get last message if available, otherwise use empty string for template-only mode
+        if not state.messages:
+            raise IndexError("No messages in state")
+
         last_message = state.messages[-1]
         if not isinstance(last_message, HumanMessage):
             raise ValueError("Last message is not a human message.")
 
+        human_content = last_message.content if last_message else ""
+        # If a template was provided, emit a synthetic "loaded notebook" message once
+        pre_messages: list = []
+        if template_markdown and not state.has_emitted_template_loaded:
+            serializer = self._get_notebook_serializer()
+            json_content = serializer.from_markdown_to_json(template_markdown)
+            loaded_message = NotebookUpdateMessage(
+                id=str(uuid4()),
+                notebook_id=str(state.template_notebook_short_id or ""),
+                content=ProsemirrorJSONContent.model_validate(json_content.model_dump(exclude_none=True)),
+                notebook_type="deep_research",
+                event="loaded",
+            )
+            await self._write_message(loaded_message)
+            pre_messages.append(loaded_message)
+
+        # If template exists, use it (with or without additional human content)
+        # If no template, use human content (which should exist in this case)
+        if template_markdown:
+            human_message = f"{template_markdown}\n\n{human_content}" if human_content else template_markdown
+        else:
+            human_message = human_content
+
         prompt = ChatPromptTemplate.from_messages(
             [
-                ("human", last_message.content),
+                ("human", human_message),
             ]
         )
 
@@ -49,8 +90,46 @@ async def arun(self, state: DeepResearchState, config: RunnableConfig) -> Partia
         notebook_update_message.current_run_notebooks = current_run_notebooks
 
         return PartialDeepResearchState(
-            messages=[notebook_update_message],
+            messages=[*pre_messages, notebook_update_message],
             previous_response_id=None,  # we reset the previous response id because we're starting a new conversation after the onboarding
             conversation_notebooks=[notebook_info],
             current_run_notebooks=current_run_notebooks,
+            has_emitted_template_loaded=True if pre_messages else state.has_emitted_template_loaded,
         )
+
+    @database_sync_to_async
+    def get_notebook(self, state: DeepResearchState) -> Notebook:
+        return Notebook.objects.filter(
+            team=self._team, short_id=str(state.template_notebook_short_id), deleted=False
+        ).first()
+
+    async def _retrieve_template_markdown(self, state: DeepResearchState) -> str | None:
+        if not (
+            state.template_notebook_short_id and not state.template_markdown and not state.has_emitted_template_loaded
+        ):
+            return state.template_markdown
+
+        try:
+            notebook = await self.get_notebook(state)
+
+            if not notebook:
+                return state.template_markdown
+
+            text_content = getattr(notebook, "text_content", None)
+            if text_content:
+                return text_content
+
+            content = getattr(notebook, "content", None)
+            if content:
+                try:
+                    nb_json = ProsemirrorJSONContent.model_validate(notebook.content)
+                    from ee.hogai.notebook.notebook_serializer import NotebookSerializer
+
+                    return NotebookSerializer().from_json_to_markdown(nb_json)
+                except Exception:
+                    return state.template_markdown
+
+            return state.template_markdown
+        except Exception as e:
+            capture_exception(e)
+            return state.template_markdown

ee/hogai/graph/deep_research/onboarding/nodes.py

@@ -20,6 +20,12 @@ def node_name(self) -> MaxNodeName:
         return DeepResearchNodeName.NOTEBOOK_PLANNING
 
     def should_run_onboarding_at_start(self, state: DeepResearchState) -> Literal["onboarding", "planning", "continue"]:
+        # Skipping onboarding when provided a template
+        if state.skip_onboarding:
+            if state.current_run_notebooks:
+                return "continue"
+            return "planning"
+
         if not state.messages:
             return "onboarding"
 

ee/hogai/graph/deep_research/types.py

@@ -63,6 +63,22 @@ class _SharedDeepResearchState(BaseStateWithMessages, BaseStateWithTasks):
     """
     Notebooks created in the current deep research run (reset on new run).
     """
+    skip_onboarding: Annotated[Optional[bool], replace] = Field(default=None)
+    """
+    If true, skip the onboarding node routing and go straight to planning.
+    """
+    template_markdown: Annotated[Optional[str], replace] = Field(default=None)
+    """
+    Template markdown content when deep research starts from a template.
+    """
+    template_notebook_short_id: Annotated[Optional[str], replace] = Field(default=None)
+    """
+    Template notebook ID when deep research starts from a template.
+    """
+    has_emitted_template_loaded: Annotated[bool, replace] = Field(default=False)
+    """
+    Whether the template loaded message has been emitted.
+    """
 
 
 class DeepResearchState(_SharedDeepResearchState):