diff --git a/application/agents/research_agent.py b/application/agents/research_agent.py index 1f58b73f..1e4f9a73 100644 --- a/application/agents/research_agent.py +++ b/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()) diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 2f9c9950..29d3d66c 100644 --- a/application/api/answer/routes/base.py +++ b/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) diff --git a/application/api/answer/services/conversation_service.py b/application/api/answer/services/conversation_service.py index 09baf3d1..00fc52cc 100644 --- a/application/api/answer/services/conversation_service.py +++ b/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: diff --git a/application/api/answer/services/stream_processor.py b/application/api/answer/services/stream_processor.py index 904f6d3f..f7ee03d6 100644 --- a/application/api/answer/services/stream_processor.py +++ b/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", []) ] diff --git a/application/prompts/research/clarification.txt b/application/prompts/research/clarification.txt new file mode 100644 index 00000000..000c73c2 --- /dev/null +++ b/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" +} \ No newline at end of file diff --git a/application/prompts/research/planning.txt b/application/prompts/research/planning.txt index 65f9b0cb..8263bcaf 100644 --- a/application/prompts/research/planning.txt +++ b/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") diff --git a/frontend/src/conversation/conversationSlice.ts b/frontend/src/conversation/conversationSlice.ts index fed5e6cf..ca2e2cbc 100644 --- a/frontend/src/conversation/conversationSlice.ts +++ b/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));