feat: improve research

application/agents/research_agent.py

@@ -42,6 +42,7 @@ def _load_prompt(name: str) -> str:
         return f.read()
 
 
+CLARIFICATION_PROMPT = _load_prompt("clarification.txt")
 PLANNING_PROMPT = _load_prompt("planning.txt")
 STEP_PROMPT = _load_prompt("step.txt")
 SYNTHESIS_PROMPT = _load_prompt("synthesis.txt")
@@ -163,6 +164,19 @@ class ResearchAgent(BaseAgent):
         self._start_time = time.monotonic()
         tools_dict = self._setup_tools()
 
+        # Phase 0: Clarification (skip if user is responding to a prior clarification)
+        if not self._is_follow_up():
+            clarification = self._clarification_phase(query)
+            if clarification:
+                yield {"metadata": {"is_clarification": True}}
+                yield {"answer": clarification}
+                yield {"sources": []}
+                yield {"tool_calls": []}
+                log_context.stacks.append(
+                    {"component": "agent", "data": {"clarification": True}}
+                )
+                return
+
         # Phase 1: Planning (with adaptive depth)
         yield {"type": "research_progress", "data": {"status": "planning"}}
         plan, complexity = self._planning_phase(query)
@@ -276,6 +290,98 @@ class ResearchAgent(BaseAgent):
         self._prepare_tools(tools_dict)
         return tools_dict
 
+    # ------------------------------------------------------------------
+    # Phase 0: Clarification
+    # ------------------------------------------------------------------
+
+    def _is_follow_up(self) -> bool:
+        """Check if the user is responding to a prior clarification.
+
+        Uses the metadata flag stored in the conversation DB — no string matching.
+        Only skip clarification when the last query was explicitly flagged
+        as a clarification by this agent.
+        """
+        if not self.chat_history:
+            return False
+        last = self.chat_history[-1]
+        meta = last.get("metadata", {})
+        return bool(meta.get("is_clarification"))
+
+    def _clarification_phase(self, question: str) -> Optional[str]:
+        """Ask the LLM whether the question needs clarification.
+
+        Returns formatted clarification text if needed, or None to proceed.
+        Uses response_format to force valid JSON output.
+        """
+        messages = [
+            {"role": "system", "content": CLARIFICATION_PROMPT},
+            {"role": "user", "content": question},
+        ]
+
+        try:
+            response = self.llm.gen(
+                model=self.model_id,
+                messages=messages,
+                tools=None,
+                response_format={"type": "json_object"},
+            )
+            text = self._extract_text(response)
+            self._track_tokens(self._snapshot_llm_tokens())
+            logger.info(f"ResearchAgent clarification response: {text[:300]}")
+
+            data = self._parse_clarification_json(text)
+            if not data or not data.get("needs_clarification"):
+                return None
+
+            questions = data.get("questions", [])
+            if not questions:
+                return None
+
+            # Format as a friendly response
+            lines = [
+                "Before I begin researching, I'd like to clarify a few things:\n"
+            ]
+            for i, q in enumerate(questions[:3], 1):
+                lines.append(f"{i}. {q}")
+            lines.append(
+                "\nPlease provide these details and I'll start the research."
+            )
+            return "\n".join(lines)
+
+        except Exception as e:
+            logger.error(f"Clarification phase failed: {e}", exc_info=True)
+            return None  # proceed with research on failure
+
+    def _parse_clarification_json(self, text: str) -> Optional[Dict]:
+        """Parse clarification JSON from LLM response."""
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            pass
+
+        # Try extracting from code fences
+        for marker in ["```json", "```"]:
+            if marker in text:
+                start = text.index(marker) + len(marker)
+                end = text.index("```", start) if "```" in text[start:] else len(text)
+                try:
+                    return json.loads(text[start:end].strip())
+                except (json.JSONDecodeError, ValueError):
+                    pass
+
+        # Try finding JSON object
+        for i, ch in enumerate(text):
+            if ch == "{":
+                for j in range(len(text) - 1, i, -1):
+                    if text[j] == "}":
+                        try:
+                            return json.loads(text[i : j + 1])
+                        except json.JSONDecodeError:
+                            continue
+                break
+
+        return None
+
     # ------------------------------------------------------------------
     # Phase 1: Planning (with adaptive depth)
     # ------------------------------------------------------------------
@@ -292,7 +398,10 @@ class ResearchAgent(BaseAgent):
 
         try:
             response = self.llm.gen(
-                model=self.model_id, messages=messages, tools=None
+                model=self.model_id,
+                messages=messages,
+                tools=None,
+                response_format={"type": "json_object"},
             )
             text = self._extract_text(response)
             self._track_tokens(self._snapshot_llm_tokens())

application/api/answer/routes/base.py

@@ -206,9 +206,12 @@ class BaseAnswerResource:
             is_structured = False
             schema_info = None
             structured_chunks = []
+            query_metadata = {}
 
             for line in agent.gen(query=question):
-                if "answer" in line:
+                if "metadata" in line:
+                    query_metadata.update(line["metadata"])
+                elif "answer" in line:
                     response_full += str(line["answer"])
                     if line.get("structured"):
                         is_structured = True
@@ -295,6 +298,7 @@ class BaseAnswerResource:
                     is_shared_usage=is_shared_usage,
                     shared_token=shared_token,
                     attachment_ids=attachment_ids,
+                    metadata=query_metadata if query_metadata else None,
                 )
                 # Persist compression metadata/summary if it exists and wasn't saved mid-execution
                 compression_meta = getattr(agent, "compression_metadata", None)
@@ -384,6 +388,7 @@ class BaseAnswerResource:
                         is_shared_usage=is_shared_usage,
                         shared_token=shared_token,
                         attachment_ids=attachment_ids,
+                        metadata=query_metadata if query_metadata else None,
                     )
                     compression_meta = getattr(agent, "compression_metadata", None)
                     compression_saved = getattr(agent, "compression_saved", False)

application/api/answer/services/conversation_service.py

@@ -60,6 +60,7 @@ class ConversationService:
         is_shared_usage: bool = False,
         shared_token: Optional[str] = None,
         attachment_ids: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> str:
         """Save or update a conversation in the database"""
         if decoded_token is None:
@@ -93,6 +94,11 @@ class ConversationService:
                         f"queries.{index}.timestamp": current_time,
                         f"queries.{index}.attachments": attachment_ids,
                         f"queries.{index}.model_id": model_id,
+                        **(
+                            {f"queries.{index}.metadata": metadata}
+                            if metadata
+                            else {}
+                        ),
                     }
                 },
             )
@@ -124,6 +130,7 @@ class ConversationService:
                             "timestamp": current_time,
                             "attachments": attachment_ids,
                             "model_id": model_id,
+                            **({"metadata": metadata} if metadata else {}),
                         }
                     }
                 },
@@ -156,22 +163,24 @@ class ConversationService:
             if not completion or not completion.strip():
                 completion = question[:50] if question else "New Conversation"
 
+            query_doc = {
+                "prompt": question,
+                "response": response,
+                "thought": thought,
+                "sources": sources,
+                "tool_calls": tool_calls,
+                "timestamp": current_time,
+                "attachments": attachment_ids,
+                "model_id": model_id,
+            }
+            if metadata:
+                query_doc["metadata"] = metadata
+
             conversation_data = {
                 "user": user_id,
                 "date": current_time,
                 "name": completion,
-                "queries": [
-                    {
-                        "prompt": question,
-                        "response": response,
-                        "thought": thought,
-                        "sources": sources,
-                        "tool_calls": tool_calls,
-                        "timestamp": current_time,
-                        "attachments": attachment_ids,
-                        "model_id": model_id,
-                    }
-                ],
+                "queries": [query_doc],
             }
 
             if api_key:

application/api/answer/services/stream_processor.py

@@ -158,9 +158,17 @@ class StreamProcessor:
             if settings.ENABLE_CONVERSATION_COMPRESSION:
                 self._handle_compression(conversation)
             else:
-                # Original behavior - load all history
+                # Original behavior - load all history (include metadata if present)
                 self.history = [
-                    {"prompt": query["prompt"], "response": query["response"]}
+                    {
+                        "prompt": query["prompt"],
+                        "response": query["response"],
+                        **(
+                            {"metadata": query["metadata"]}
+                            if "metadata" in query
+                            else {}
+                        ),
+                    }
                     for query in conversation.get("queries", [])
                 ]
         else:
@@ -181,7 +189,11 @@ class StreamProcessor:
             if not result.success:
                 logger.error(f"Compression failed: {result.error}, using full history")
                 self.history = [
-                    {"prompt": query["prompt"], "response": query["response"]}
+                    {
+                        "prompt": query["prompt"],
+                        "response": query["response"],
+                        **({"metadata": query["metadata"]} if "metadata" in query else {}),
+                    }
                     for query in conversation.get("queries", [])
                 ]
                 return
@@ -197,6 +209,14 @@ class StreamProcessor:
                 )
 
             self.history = result.as_history()
+            # Preserve metadata from recent queries (as_history only has prompt/response)
+            recent = result.recent_queries if result.recent_queries else conversation.get("queries", [])
+            for i, entry in enumerate(self.history):
+                # Match by index from the end of recent queries
+                offset = len(recent) - len(self.history)
+                qi = offset + i
+                if 0 <= qi < len(recent) and "metadata" in recent[qi]:
+                    entry["metadata"] = recent[qi]["metadata"]
 
         except Exception as e:
             logger.error(
@@ -204,7 +224,11 @@ class StreamProcessor:
                 exc_info=True,
             )
             self.history = [
-                {"prompt": query["prompt"], "response": query["response"]}
+                {
+                    "prompt": query["prompt"],
+                    "response": query["response"],
+                    **({"metadata": query["metadata"]} if "metadata" in query else {}),
+                }
                 for query in conversation.get("queries", [])
             ]
 

application/prompts/research/clarification.txt

@@ -0,0 +1,23 @@
+You are a research assistant evaluating whether a user's question is clear enough to begin in-depth research.
+
+Decide whether the question can be researched as-is, or whether you need to ask clarifying questions first.
+
+Proceed WITHOUT clarification if:
+- The question has a clear topic and intent
+- There is enough context to form a research plan
+- The question is broad but can be broken into sub-topics
+- Documents/sources are available to search
+
+Ask for clarification ONLY if:
+- The question is critically ambiguous (e.g., "review the thing" with no context about what "thing" means)
+- Key information is missing that would make research impossible (not just imperfect)
+- The user seems to reference something specific but hasn't provided it
+
+Err on the side of proceeding. Most questions are clear enough to start researching.
+
+You MUST respond with ONLY a valid JSON object (no markdown, no code fences):
+{
+    "needs_clarification": true or false,
+    "questions": ["question 1", "question 2"] (only if needs_clarification is true, max 3 questions),
+    "reason": "brief explanation of why clarification is or isn't needed"
+}

application/prompts/research/planning.txt

@@ -1,5 +1,7 @@
 You are a research planner. Your job is to analyze the user's question and create a focused research plan.
 
+IMPORTANT: Every step must be a concrete research action — something you can search for and find information about. Never generate steps that ask the user for more information or request documents. Work with what you have.
+
 Assess the question's complexity:
 - "simple": Can be answered with 1-2 focused searches (e.g., "What is our refund policy?")
 - "moderate": Needs 3-4 searches across different topics (e.g., "Compare our pricing across product lines")

frontend/src/conversation/conversationSlice.ts

@@ -77,12 +77,16 @@ export const fetchAnswer = createAsyncThunk<
           if (currentConversationId === state.conversation.conversationId) {
             if (data.type === 'end') {
               dispatch(conversationSlice.actions.setStatus('idle'));
-              dispatch(
-                updateResearchProgress({
-                  index: targetIndex,
-                  progress: { status: 'complete' },
-                }),
-              );
+              // Only update research status if this query has research data
+              const currentState = getState() as RootState;
+              if (currentState.conversation.queries[targetIndex]?.research) {
+                dispatch(
+                  updateResearchProgress({
+                    index: targetIndex,
+                    progress: { status: 'complete' },
+                  }),
+                );
+              }
               getConversations(state.preference.token)
                 .then((fetchedConversations) => {
                   dispatch(setConversations(fetchedConversations));