From 556b0a1da570e94f0b9a9c58efa6046c38014f8e Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 15:16:18 +0000 Subject: [PATCH] feat: research init --- application/agents/agent_creator.py | 4 + application/agents/agentic_agent.py | 76 +++ application/agents/research_agent.py | 584 ++++++++++++++++++ application/agents/tools/internal_search.py | 155 +++++ application/agents/tools/think.py | 68 ++ .../api/answer/services/stream_processor.py | 60 +- application/api/user/agents/routes.py | 2 + application/api/user/base.py | 10 +- application/prompts/agentic/creative.txt | 16 + application/prompts/agentic/default.txt | 15 + application/prompts/agentic/strict.txt | 16 + application/prompts/research/planning.txt | 21 + application/prompts/research/step.txt | 13 + application/prompts/research/synthesis.txt | 22 + frontend/src/agents/NewAgent.tsx | 2 + frontend/src/assets/research.svg | 18 + .../src/conversation/ConversationBubble.tsx | 66 +- .../src/conversation/ConversationMessages.tsx | 3 +- .../src/conversation/ResearchProgress.tsx | 147 +++++ .../src/conversation/conversationModels.ts | 15 + .../src/conversation/conversationSlice.ts | 103 ++- 21 files changed, 1401 insertions(+), 15 deletions(-) create mode 100644 application/agents/agentic_agent.py create mode 100644 application/agents/research_agent.py create mode 100644 application/agents/tools/internal_search.py create mode 100644 application/agents/tools/think.py create mode 100644 application/prompts/agentic/creative.txt create mode 100644 application/prompts/agentic/default.txt create mode 100644 application/prompts/agentic/strict.txt create mode 100644 application/prompts/research/planning.txt create mode 100644 application/prompts/research/step.txt create mode 100644 application/prompts/research/synthesis.txt create mode 100644 frontend/src/assets/research.svg create mode 100644 frontend/src/conversation/ResearchProgress.tsx diff --git a/application/agents/agent_creator.py b/application/agents/agent_creator.py index 6773babd..d623269f 100644 --- a/application/agents/agent_creator.py +++ b/application/agents/agent_creator.py @@ -1,6 +1,8 @@ import logging +from application.agents.agentic_agent import AgenticAgent from application.agents.classic_agent import ClassicAgent +from application.agents.research_agent import ResearchAgent from application.agents.workflow_agent import WorkflowAgent logger = logging.getLogger(__name__) @@ -10,6 +12,8 @@ class AgentCreator: agents = { "classic": ClassicAgent, "react": ClassicAgent, # backwards compat: react falls back to classic + "agentic": AgenticAgent, + "research": ResearchAgent, "workflow": WorkflowAgent, } diff --git a/application/agents/agentic_agent.py b/application/agents/agentic_agent.py new file mode 100644 index 00000000..f642402a --- /dev/null +++ b/application/agents/agentic_agent.py @@ -0,0 +1,76 @@ +import logging +from typing import Dict, Generator, Optional + +from application.agents.base import BaseAgent +from application.agents.tools.internal_search import ( + INTERNAL_TOOL_ENTRY, + INTERNAL_TOOL_ID, + build_internal_tool_config, +) +from application.logging import LogContext + +logger = logging.getLogger(__name__) + + +class AgenticAgent(BaseAgent): + """Agent where the LLM controls retrieval via tools. + + Unlike ClassicAgent which pre-fetches docs into the prompt, + AgenticAgent gives the LLM an internal_search tool so it can + decide when, what, and whether to search. + """ + + def __init__( + self, + retriever_config: Optional[Dict] = None, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.retriever_config = retriever_config or {} + + def _gen_inner( + self, query: str, log_context: LogContext + ) -> Generator[Dict, None, None]: + # 1. Get user tools (same as ClassicAgent) + tools_dict = self.tool_executor.get_tools() + + # 2. Add internal search as a synthetic tool (only if sources are configured) + source = self.retriever_config.get("source", {}) + has_sources = bool(source.get("active_docs")) + if self.retriever_config and has_sources: + internal_entry = dict(INTERNAL_TOOL_ENTRY) + internal_entry["config"] = build_internal_tool_config( + **self.retriever_config + ) + tools_dict[INTERNAL_TOOL_ID] = internal_entry + + # 3. Prepare all tools for the LLM + self._prepare_tools(tools_dict) + + # 4. Build messages (prompt has NO pre-fetched docs) + messages = self._build_messages(self.prompt, query) + + # 5. Call LLM — the handler manages the tool loop + llm_response = self._llm_gen(messages, log_context) + + yield from self._handle_response( + llm_response, tools_dict, messages, log_context + ) + + # 6. Collect sources from internal search tool results + self._collect_internal_sources() + + yield {"sources": self.retrieved_docs} + yield {"tool_calls": self._get_truncated_tool_calls()} + + log_context.stacks.append( + {"component": "agent", "data": {"tool_calls": self.tool_calls.copy()}} + ) + + def _collect_internal_sources(self): + """Collect retrieved docs from the cached InternalSearchTool instance.""" + cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{self.user or ''}" + tool = self.tool_executor._loaded_tools.get(cache_key) + if tool and hasattr(tool, "retrieved_docs") and tool.retrieved_docs: + self.retrieved_docs = tool.retrieved_docs diff --git a/application/agents/research_agent.py b/application/agents/research_agent.py new file mode 100644 index 00000000..1f58b73f --- /dev/null +++ b/application/agents/research_agent.py @@ -0,0 +1,584 @@ +import json +import logging +import os +import time +from typing import Dict, Generator, List, Optional + +from application.agents.base import BaseAgent +from application.agents.tool_executor import ToolExecutor +from application.agents.tools.internal_search import ( + INTERNAL_TOOL_ENTRY, + INTERNAL_TOOL_ID, + build_internal_tool_config, +) +from application.agents.tools.think import THINK_TOOL_ENTRY, THINK_TOOL_ID +from application.logging import LogContext + +logger = logging.getLogger(__name__) + +# Defaults (can be overridden via constructor) +DEFAULT_MAX_STEPS = 6 +DEFAULT_MAX_SUB_ITERATIONS = 5 +DEFAULT_TIMEOUT_SECONDS = 300 # 5 minutes +DEFAULT_TOKEN_BUDGET = 100_000 +DEFAULT_PARALLEL_WORKERS = 3 + +# Adaptive depth caps per complexity level +COMPLEXITY_CAPS = { + "simple": 2, + "moderate": 4, + "complex": 6, +} + +_PROMPTS_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "prompts", + "research", +) + + +def _load_prompt(name: str) -> str: + with open(os.path.join(_PROMPTS_DIR, name), "r") as f: + return f.read() + + +PLANNING_PROMPT = _load_prompt("planning.txt") +STEP_PROMPT = _load_prompt("step.txt") +SYNTHESIS_PROMPT = _load_prompt("synthesis.txt") + + +# --------------------------------------------------------------------------- +# CitationManager +# --------------------------------------------------------------------------- + + +class CitationManager: + """Tracks and deduplicates citations across research steps.""" + + def __init__(self): + self.citations: Dict[int, Dict] = {} + self._counter = 0 + + def add(self, doc: Dict) -> int: + """Register a source, return its citation number. Deduplicates by source.""" + source = doc.get("source", "") + title = doc.get("title", "") + for num, existing in self.citations.items(): + if existing.get("source") == source and existing.get("title") == title: + return num + self._counter += 1 + self.citations[self._counter] = doc + return self._counter + + def add_docs(self, docs: List[Dict]) -> str: + """Register multiple docs, return formatted citation mapping text.""" + mapping_lines = [] + for doc in docs: + num = self.add(doc) + title = doc.get("title", "Untitled") + mapping_lines.append(f"[{num}] {title}") + return "\n".join(mapping_lines) + + def format_references(self) -> str: + """Generate [N] -> source mapping for report footer.""" + if not self.citations: + return "No sources found." + lines = [] + for num, doc in sorted(self.citations.items()): + title = doc.get("title", "Untitled") + source = doc.get("source", "Unknown") + filename = doc.get("filename", "") + display = filename or title + lines.append(f"[{num}] {display} — {source}") + return "\n".join(lines) + + def get_all_docs(self) -> List[Dict]: + return list(self.citations.values()) + + +# --------------------------------------------------------------------------- +# ResearchAgent +# --------------------------------------------------------------------------- + + +class ResearchAgent(BaseAgent): + """Multi-step research agent with parallel execution and budget controls. + + Orchestrates: Plan -> Research (per step, optionally parallel) -> Synthesize. + """ + + def __init__( + self, + retriever_config: Optional[Dict] = None, + max_steps: int = DEFAULT_MAX_STEPS, + max_sub_iterations: int = DEFAULT_MAX_SUB_ITERATIONS, + timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS, + token_budget: int = DEFAULT_TOKEN_BUDGET, + parallel_workers: int = DEFAULT_PARALLEL_WORKERS, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.retriever_config = retriever_config or {} + self.max_steps = max_steps + self.max_sub_iterations = max_sub_iterations + self.timeout_seconds = timeout_seconds + self.token_budget = token_budget + self.parallel_workers = parallel_workers + self.citations = CitationManager() + self._start_time: float = 0 + self._tokens_used: int = 0 + + # ------------------------------------------------------------------ + # Budget & timeout helpers + # ------------------------------------------------------------------ + + def _is_timed_out(self) -> bool: + return (time.monotonic() - self._start_time) >= self.timeout_seconds + + def _elapsed(self) -> float: + return round(time.monotonic() - self._start_time, 1) + + def _track_tokens(self, count: int): + self._tokens_used += count + + def _budget_remaining(self) -> int: + return max(self.token_budget - self._tokens_used, 0) + + def _is_over_budget(self) -> bool: + return self._tokens_used >= self.token_budget + + def _snapshot_llm_tokens(self) -> int: + """Read current token usage from LLM and return delta since last snapshot.""" + current = self.llm.token_usage.get("prompt_tokens", 0) + self.llm.token_usage.get("generated_tokens", 0) + return current + + # ------------------------------------------------------------------ + # Main orchestration + # ------------------------------------------------------------------ + + def _gen_inner( + self, query: str, log_context: LogContext + ) -> Generator[Dict, None, None]: + self._start_time = time.monotonic() + tools_dict = self._setup_tools() + + # Phase 1: Planning (with adaptive depth) + yield {"type": "research_progress", "data": {"status": "planning"}} + plan, complexity = self._planning_phase(query) + + if not plan: + logger.warning("ResearchAgent: Planning produced no steps, falling back") + plan = [{"query": query, "rationale": "Direct investigation"}] + complexity = "simple" + + yield { + "type": "research_plan", + "data": {"steps": plan, "complexity": complexity}, + } + + # Phase 2: Research each step (yields progress events in real-time) + intermediate_reports = [] + for i, step in enumerate(plan): + step_num = i + 1 + step_query = step.get("query", query) + + if self._is_timed_out(): + logger.warning( + f"ResearchAgent: Timeout at step {step_num}/{len(plan)} " + f"({self._elapsed()}s)" + ) + break + if self._is_over_budget(): + logger.warning( + f"ResearchAgent: Token budget exhausted at step {step_num}/{len(plan)}" + ) + break + + yield { + "type": "research_progress", + "data": { + "step": step_num, + "total": len(plan), + "query": step_query, + "status": "researching", + }, + } + + report = self._research_step(step_query, tools_dict) + intermediate_reports.append({"step": step, "content": report}) + + yield { + "type": "research_progress", + "data": { + "step": step_num, + "total": len(plan), + "query": step_query, + "status": "complete", + }, + } + + # Phase 3: Synthesis (streaming) + if self._is_timed_out(): + logger.warning( + f"ResearchAgent: Timeout ({self._elapsed()}s) before synthesis, " + f"synthesizing with {len(intermediate_reports)} reports" + ) + yield { + "type": "research_progress", + "data": { + "status": "synthesizing", + "elapsed_seconds": self._elapsed(), + "tokens_used": self._tokens_used, + }, + } + yield from self._synthesis_phase( + query, plan, intermediate_reports, tools_dict, log_context + ) + + # Sources and tool calls + self.retrieved_docs = self.citations.get_all_docs() + yield {"sources": self.retrieved_docs} + yield {"tool_calls": self._get_truncated_tool_calls()} + + logger.info( + f"ResearchAgent completed: {len(intermediate_reports)}/{len(plan)} steps, " + f"{self._elapsed()}s, ~{self._tokens_used} tokens" + ) + log_context.stacks.append( + {"component": "agent", "data": {"tool_calls": self.tool_calls.copy()}} + ) + + # ------------------------------------------------------------------ + # Tool setup + # ------------------------------------------------------------------ + + def _setup_tools(self) -> Dict: + """Build tools_dict with user tools + internal search + think.""" + tools_dict = self.tool_executor.get_tools() + + # Only add internal search if sources are configured + source = self.retriever_config.get("source", {}) + has_sources = bool(source.get("active_docs")) + if self.retriever_config and has_sources: + internal_entry = dict(INTERNAL_TOOL_ENTRY) + internal_entry["config"] = build_internal_tool_config( + **self.retriever_config + ) + tools_dict[INTERNAL_TOOL_ID] = internal_entry + elif self.retriever_config and not has_sources: + logger.info("ResearchAgent: No sources configured, skipping internal_search tool") + + think_entry = dict(THINK_TOOL_ENTRY) + think_entry["config"] = {} + tools_dict[THINK_TOOL_ID] = think_entry + + self._prepare_tools(tools_dict) + return tools_dict + + # ------------------------------------------------------------------ + # Phase 1: Planning (with adaptive depth) + # ------------------------------------------------------------------ + + def _planning_phase(self, question: str) -> tuple[List[Dict], str]: + """Decompose the question into research steps via LLM. + + Returns (steps, complexity) where complexity is simple/moderate/complex. + """ + messages = [ + {"role": "system", "content": PLANNING_PROMPT}, + {"role": "user", "content": question}, + ] + + try: + response = self.llm.gen( + model=self.model_id, messages=messages, tools=None + ) + text = self._extract_text(response) + self._track_tokens(self._snapshot_llm_tokens()) + logger.info(f"ResearchAgent planning LLM response: {text[:500]}") + + plan_data = self._parse_plan_json(text) + if isinstance(plan_data, dict): + complexity = plan_data.get("complexity", "moderate") + steps = plan_data.get("steps", []) + else: + complexity = "moderate" + steps = plan_data + + # Adaptive depth: cap steps based on assessed complexity + cap = COMPLEXITY_CAPS.get(complexity, self.max_steps) + cap = min(cap, self.max_steps) + steps = steps[:cap] + + logger.info( + f"ResearchAgent plan: complexity={complexity}, " + f"steps={len(steps)} (cap={cap})" + ) + return steps, complexity + + except Exception as e: + logger.error(f"Planning phase failed: {e}", exc_info=True) + return ( + [{"query": question, "rationale": "Direct investigation (planning failed)"}], + "simple", + ) + + def _parse_plan_json(self, text: str): + """Extract JSON plan from LLM response. Returns dict or list.""" + # Try direct parse + try: + data = json.loads(text) + if isinstance(data, dict) and "steps" in data: + return data + if isinstance(data, list): + return data + except json.JSONDecodeError: + pass + + # Try extracting from markdown 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: + data = json.loads(text[start:end].strip()) + if isinstance(data, dict) and "steps" in data: + return data + if isinstance(data, list): + return data + except (json.JSONDecodeError, ValueError): + pass + + # Try finding JSON object in text + for i, ch in enumerate(text): + if ch == "{": + for j in range(len(text) - 1, i, -1): + if text[j] == "}": + try: + data = json.loads(text[i : j + 1]) + if isinstance(data, dict) and "steps" in data: + return data + except json.JSONDecodeError: + continue + break + + logger.warning(f"Could not parse plan JSON from: {text[:200]}") + return [] + + # ------------------------------------------------------------------ + # Phase 2: Research step (core loop) + # ------------------------------------------------------------------ + + def _research_step(self, step_query: str, tools_dict: Dict) -> str: + """Run a focused research loop for one sub-question (sequential path).""" + report = self._research_step_with_executor( + step_query, tools_dict, self.tool_executor + ) + self._collect_step_sources() + return report + + def _research_step_with_executor( + self, step_query: str, tools_dict: Dict, executor: ToolExecutor + ) -> str: + """Core research loop. Works with any ToolExecutor instance.""" + system_prompt = STEP_PROMPT.replace("{step_query}", step_query) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": step_query}, + ] + + last_search_empty = False + + for iteration in range(self.max_sub_iterations): + # Check timeout and budget + if self._is_timed_out(): + logger.info( + f"Research step '{step_query[:50]}' timed out at iteration {iteration}" + ) + break + if self._is_over_budget(): + logger.info( + f"Research step '{step_query[:50]}' hit token budget at iteration {iteration}" + ) + break + + try: + response = self.llm.gen( + model=self.model_id, + messages=messages, + tools=self.tools if self.tools else None, + ) + self._track_tokens(self._snapshot_llm_tokens()) + except Exception as e: + logger.error( + f"Research step LLM call failed (iteration {iteration}): {e}", + exc_info=True, + ) + break + + parsed = self.llm_handler.parse_response(response) + + if not parsed.requires_tool_call: + return parsed.content or "No findings for this step." + + # Execute tool calls + messages, last_search_empty = self._execute_step_tools_with_refinement( + parsed.tool_calls, tools_dict, messages, executor, last_search_empty + ) + + # Max iterations / timeout / budget — ask for summary + messages.append( + { + "role": "user", + "content": "Please summarize your findings so far based on the information gathered.", + } + ) + try: + response = self.llm.gen( + model=self.model_id, messages=messages, tools=None + ) + self._track_tokens(self._snapshot_llm_tokens()) + text = self._extract_text(response) + return text or "Research step completed." + except Exception: + return "Research step completed." + + def _execute_step_tools_with_refinement( + self, + tool_calls, + tools_dict: Dict, + messages: List[Dict], + executor: ToolExecutor, + last_search_empty: bool, + ) -> tuple[List[Dict], bool]: + """Execute tool calls with query refinement on empty results. + + Returns (updated_messages, was_last_search_empty). + """ + search_returned_empty = False + + for call in tool_calls: + gen = executor.execute( + tools_dict, call, self.llm.__class__.__name__ + ) + result = None + call_id = None + while True: + try: + next(gen) + except StopIteration as e: + result, call_id = e.value + break + + # Detect empty search results for refinement + is_search = "search" in (call.name or "").lower() + result_str = str(result) if result else "" + if is_search and "No documents found" in result_str: + search_returned_empty = True + if last_search_empty: + # Two consecutive empty searches — inject refinement hint + result_str += ( + "\n\nHint: Previous search also returned no results. " + "Try a very different query with different keywords, " + "or broaden your search terms." + ) + result = result_str + + function_call_content = { + "function_call": { + "name": call.name, + "args": call.arguments, + "call_id": call_id, + } + } + messages.append( + {"role": "assistant", "content": [function_call_content]} + ) + tool_message = self.llm_handler.create_tool_message(call, result) + messages.append(tool_message) + + return messages, search_returned_empty + + def _collect_step_sources(self): + """Collect sources from InternalSearchTool and register with CitationManager.""" + cache_key = f"internal_search:{INTERNAL_TOOL_ID}:{self.user or ''}" + tool = self.tool_executor._loaded_tools.get(cache_key) + if tool and hasattr(tool, "retrieved_docs"): + for doc in tool.retrieved_docs: + self.citations.add(doc) + + # ------------------------------------------------------------------ + # Phase 3: Synthesis + # ------------------------------------------------------------------ + + def _synthesis_phase( + self, + question: str, + plan: List[Dict], + intermediate_reports: List[Dict], + tools_dict: Dict, + log_context: LogContext, + ) -> Generator[Dict, None, None]: + """Compile all findings into a final cited report (streaming).""" + plan_lines = [] + for i, step in enumerate(plan, 1): + plan_lines.append( + f"{i}. {step.get('query', 'Unknown')} — {step.get('rationale', '')}" + ) + plan_summary = "\n".join(plan_lines) + + findings_parts = [] + for i, report in enumerate(intermediate_reports, 1): + step_query = report["step"].get("query", "Unknown") + content = report["content"] + findings_parts.append( + f"--- Step {i}: {step_query} ---\n{content}" + ) + findings = "\n\n".join(findings_parts) + + references = self.citations.format_references() + + synthesis_prompt = SYNTHESIS_PROMPT.replace("{question}", question) + synthesis_prompt = synthesis_prompt.replace("{plan_summary}", plan_summary) + synthesis_prompt = synthesis_prompt.replace("{findings}", findings) + synthesis_prompt = synthesis_prompt.replace("{references}", references) + + messages = [ + {"role": "system", "content": synthesis_prompt}, + {"role": "user", "content": f"Please write the research report for: {question}"}, + ] + + llm_response = self.llm.gen_stream( + model=self.model_id, messages=messages, tools=None + ) + + if log_context: + from application.logging import build_stack_data + + log_context.stacks.append( + {"component": "synthesis_llm", "data": build_stack_data(self.llm)} + ) + + yield from self._handle_response( + llm_response, tools_dict, messages, log_context + ) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _extract_text(self, response) -> str: + """Extract text content from a non-streaming LLM response.""" + if isinstance(response, str): + return response + if hasattr(response, "message") and hasattr(response.message, "content"): + return response.message.content or "" + if hasattr(response, "choices") and response.choices: + choice = response.choices[0] + if hasattr(choice, "message") and hasattr(choice.message, "content"): + return choice.message.content or "" + if hasattr(response, "content") and isinstance(response.content, list): + if response.content and hasattr(response.content[0], "text"): + return response.content[0].text or "" + return str(response) if response else "" diff --git a/application/agents/tools/internal_search.py b/application/agents/tools/internal_search.py new file mode 100644 index 00000000..0b8c2c47 --- /dev/null +++ b/application/agents/tools/internal_search.py @@ -0,0 +1,155 @@ +import logging +from typing import Dict, List, Optional + +from application.agents.tools.base import Tool +from application.core.settings import settings +from application.retriever.retriever_creator import RetrieverCreator + +logger = logging.getLogger(__name__) + + +class InternalSearchTool(Tool): + """Wraps the ClassicRAG retriever as an LLM-callable tool. + + Instead of pre-fetching docs into the prompt, the LLM decides + when and what to search. Supports multiple searches per session. + """ + + def __init__(self, config: Dict): + self.config = config + self.retrieved_docs: List[Dict] = [] + self._retriever = None + + def _get_retriever(self): + if self._retriever is None: + self._retriever = RetrieverCreator.create_retriever( + self.config.get("retriever_name", "classic"), + source=self.config.get("source", {}), + chat_history=[], # no rephrasing — LLM controls the query + prompt="", + chunks=int(self.config.get("chunks", 2)), + doc_token_limit=int(self.config.get("doc_token_limit", 50000)), + model_id=self.config.get("model_id", "docsgpt-local"), + user_api_key=self.config.get("user_api_key"), + agent_id=self.config.get("agent_id"), + llm_name=self.config.get("llm_name", settings.LLM_PROVIDER), + api_key=self.config.get("api_key", settings.API_KEY), + decoded_token=self.config.get("decoded_token"), + ) + return self._retriever + + def execute_action(self, action_name: str, **kwargs): + if action_name != "search": + return f"Unknown action: {action_name}" + + query = kwargs.get("query", "") + if not query: + return "Error: 'query' parameter is required." + + try: + retriever = self._get_retriever() + docs = retriever.search(query) + except Exception as e: + logger.error(f"Internal search failed: {e}", exc_info=True) + return f"Search failed: {str(e)}" + + if not docs: + return "No documents found matching your query." + + # Accumulate for source tracking + for doc in docs: + if doc not in self.retrieved_docs: + self.retrieved_docs.append(doc) + + # Format results for the LLM + formatted = [] + for i, doc in enumerate(docs, 1): + title = doc.get("title", "Untitled") + text = doc.get("text", "") + source = doc.get("source", "Unknown") + filename = doc.get("filename", "") + header = filename or title + formatted.append(f"[{i}] {header} (source: {source})\n{text}") + + return "\n\n---\n\n".join(formatted) + + def get_actions_metadata(self): + return [ + { + "name": "search", + "description": ( + "Search the user's uploaded documents and knowledge base. " + "Use this to find relevant information before answering questions. " + "You can call this multiple times with different queries." + ), + "parameters": { + "properties": { + "query": { + "type": "string", + "description": "The search query. Be specific and focused.", + "filled_by_llm": True, + "required": True, + } + } + }, + } + ] + + def get_config_requirements(self): + return {} + + +# Constants for building synthetic tools_dict entries +INTERNAL_TOOL_ID = "internal" + +INTERNAL_TOOL_ENTRY = { + "name": "internal_search", + "actions": [ + { + "name": "search", + "description": ( + "Search the user's uploaded documents and knowledge base. " + "Use this to find relevant information before answering questions. " + "You can call this multiple times with different queries." + ), + "active": True, + "parameters": { + "properties": { + "query": { + "type": "string", + "description": "The search query. Be specific and focused.", + "filled_by_llm": True, + "required": True, + } + } + }, + } + ], +} + + +def build_internal_tool_config( + source: Dict, + retriever_name: str = "classic", + chunks: int = 2, + doc_token_limit: int = 50000, + model_id: str = "docsgpt-local", + user_api_key: Optional[str] = None, + agent_id: Optional[str] = None, + llm_name: str = None, + api_key: str = None, + decoded_token: Optional[Dict] = None, +) -> Dict: + """Build the config dict for InternalSearchTool.""" + return { + "source": source, + "retriever_name": retriever_name, + "chunks": chunks, + "doc_token_limit": doc_token_limit, + "model_id": model_id, + "user_api_key": user_api_key, + "agent_id": agent_id, + "llm_name": llm_name or settings.LLM_PROVIDER, + "api_key": api_key or settings.API_KEY, + "decoded_token": decoded_token, + } diff --git a/application/agents/tools/think.py b/application/agents/tools/think.py new file mode 100644 index 00000000..7c1fc2b6 --- /dev/null +++ b/application/agents/tools/think.py @@ -0,0 +1,68 @@ +from application.agents.tools.base import Tool + + +THINK_TOOL_ID = "think" + +THINK_TOOL_ENTRY = { + "name": "think", + "actions": [ + { + "name": "reason", + "description": ( + "Use this tool to think through your reasoning step by step " + "before deciding on your next action. Always reason before " + "searching or answering." + ), + "active": True, + "parameters": { + "properties": { + "reasoning": { + "type": "string", + "description": "Your step-by-step reasoning and analysis", + "filled_by_llm": True, + "required": True, + } + } + }, + } + ], +} + + +class ThinkTool(Tool): + """Pseudo-tool that captures chain-of-thought reasoning. + + Returns a short acknowledgment so the LLM can continue. + The reasoning content is captured in tool_call data for transparency. + """ + + def __init__(self, config=None): + pass + + def execute_action(self, action_name: str, **kwargs): + return "Continue." + + def get_actions_metadata(self): + return [ + { + "name": "reason", + "description": ( + "Use this tool to think through your reasoning step by step " + "before deciding on your next action. Always reason before " + "searching or answering." + ), + "parameters": { + "properties": { + "reasoning": { + "type": "string", + "description": "Your step-by-step reasoning and analysis", + "filled_by_llm": True, + "required": True, + } + } + }, + } + ] + + def get_config_requirements(self): + return {} diff --git a/application/api/answer/services/stream_processor.py b/application/api/answer/services/stream_processor.py index 61a4d4c1..904f6d3f 100644 --- a/application/api/answer/services/stream_processor.py +++ b/application/api/answer/services/stream_processor.py @@ -38,13 +38,23 @@ def get_prompt(prompt_id: str, prompts_collection=None) -> str: current_dir = Path(__file__).resolve().parents[3] prompts_dir = current_dir / "prompts" - preset_mapping = { + # Maps for classic agent types + CLASSIC_PRESETS = { "default": "chat_combine_default.txt", "creative": "chat_combine_creative.txt", "strict": "chat_combine_strict.txt", "reduce": "chat_reduce_prompt.txt", } + # Agentic counterparts — same styles, but with search tool instructions + AGENTIC_PRESETS = { + "default": "agentic/default.txt", + "creative": "agentic/creative.txt", + "strict": "agentic/strict.txt", + } + + preset_mapping = {**CLASSIC_PRESETS, **{f"agentic_{k}": v for k, v in AGENTIC_PRESETS.items()}} + if prompt_id in preset_mapping: file_path = os.path.join(prompts_dir, preset_mapping[prompt_id]) try: @@ -119,6 +129,14 @@ class StreamProcessor: create_agent() into a single convenience method. """ self.initialize() + + agent_type = self.agent_config.get("agent_type", "classic") + + # Agentic/research agents skip pre-fetch — the LLM searches on-demand via tools + if agent_type in ("agentic", "research"): + tools_data = self.pre_fetch_tools() + return self.create_agent(tools_data=tools_data) + docs_together, docs_list = self.pre_fetch_docs(question) tools_data = self.pre_fetch_tools() return self.create_agent( @@ -735,11 +753,21 @@ class StreamProcessor: tools_data: Optional[Dict[str, Any]] = None, ): """Create and return the configured agent with rendered prompt""" + agent_type = self.agent_config["agent_type"] + + # For agentic agents, swap standard presets for their agentic + # counterparts (which include search tool instructions instead of + # {summaries}). Custom / user-provided prompts pass through as-is. raw_prompt = self._get_prompt_content() if raw_prompt is None: - raw_prompt = get_prompt( - self.agent_config["prompt_id"], self.prompts_collection - ) + prompt_id = self.agent_config.get("prompt_id", "default") + agentic_presets = {"default", "creative", "strict"} + if agent_type in ("agentic", "research") and prompt_id in agentic_presets: + raw_prompt = get_prompt( + f"agentic_{prompt_id}", self.prompts_collection + ) + else: + raw_prompt = get_prompt(prompt_id, self.prompts_collection) self._prompt_content = raw_prompt rendered_prompt = self.prompt_renderer.render_prompt( @@ -759,8 +787,6 @@ class StreamProcessor: ) system_api_key = get_api_key_for_provider(provider or settings.LLM_PROVIDER) - agent_type = self.agent_config["agent_type"] - # Create LLM and handler (dependency injection) from application.llm.llm_creator import LLMCreator from application.llm.handlers.handler_creator import LLMHandlerCreator @@ -806,8 +832,26 @@ class StreamProcessor: "tool_executor": tool_executor, } - # Workflow-specific kwargs for workflow agents - if agent_type == "workflow": + # Type-specific kwargs + if agent_type in ("agentic", "research"): + agent_kwargs["retriever_config"] = { + "source": self.source, + "retriever_name": self.retriever_config.get( + "retriever_name", "classic" + ), + "chunks": self.retriever_config.get("chunks", 2), + "doc_token_limit": self.retriever_config.get( + "doc_token_limit", 50000 + ), + "model_id": self.model_id, + "user_api_key": self.agent_config["user_api_key"], + "agent_id": self.agent_id, + "llm_name": provider or settings.LLM_PROVIDER, + "api_key": system_api_key, + "decoded_token": self.decoded_token, + } + + elif agent_type == "workflow": workflow_config = self.agent_config.get("workflow") if isinstance(workflow_config, str): agent_kwargs["workflow_id"] = workflow_config diff --git a/application/api/user/agents/routes.py b/application/api/user/agents/routes.py index 02b185c2..c95f225d 100644 --- a/application/api/user/agents/routes.py +++ b/application/api/user/agents/routes.py @@ -104,6 +104,8 @@ AGENT_TYPE_SCHEMAS = { } AGENT_TYPE_SCHEMAS["react"] = AGENT_TYPE_SCHEMAS["classic"] +AGENT_TYPE_SCHEMAS["agentic"] = AGENT_TYPE_SCHEMAS["classic"] +AGENT_TYPE_SCHEMAS["research"] = AGENT_TYPE_SCHEMAS["classic"] AGENT_TYPE_SCHEMAS["openai"] = AGENT_TYPE_SCHEMAS["classic"] diff --git a/application/api/user/base.py b/application/api/user/base.py index 920af776..afc22268 100644 --- a/application/api/user/base.py +++ b/application/api/user/base.py @@ -145,9 +145,15 @@ def resolve_tool_details(tool_ids): Returns: List of tool details with id, name, and display_name """ + valid_ids = [] + for tid in tool_ids: + try: + valid_ids.append(ObjectId(tid)) + except Exception: + continue tools = user_tools_collection.find( - {"_id": {"$in": [ObjectId(tid) for tid in tool_ids]}} - ) + {"_id": {"$in": valid_ids}} + ) if valid_ids else [] return [ { "id": str(tool["_id"]), diff --git a/application/prompts/agentic/creative.txt b/application/prompts/agentic/creative.txt new file mode 100644 index 00000000..da2f8ff2 --- /dev/null +++ b/application/prompts/agentic/creative.txt @@ -0,0 +1,16 @@ +You are a helpful AI assistant, DocsGPT. You are proactive and helpful. Try to use tools, if they are available to you, +be proactive and fill in missing information. +Users can Upload documents for your context as attachments or sources via UI using the Conversation input box. +If appropriate, your answers can include code examples, formatted as follows: +```(language) +(code) +``` +Users are also able to see charts and diagrams if you use them with valid mermaid syntax in your responses. +Try to respond with mermaid charts if visualization helps with users queries. +You effectively utilize chat history, ensuring relevant and tailored responses. +You have access to a search tool that searches the user's uploaded documents and knowledge base. +Use the search_internal tool to find relevant information before answering questions. +You may search multiple times with different queries if needed. +Do not guess when documents are available — search first, then answer based on what you find. +If no relevant documents are found, use your general knowledge and tool capabilities. +Allow yourself to be very creative and use your imagination. \ No newline at end of file diff --git a/application/prompts/agentic/default.txt b/application/prompts/agentic/default.txt new file mode 100644 index 00000000..41339fe6 --- /dev/null +++ b/application/prompts/agentic/default.txt @@ -0,0 +1,15 @@ +You are a helpful AI assistant, DocsGPT. You are proactive and helpful. Try to use tools, if they are available to you, +be proactive and fill in missing information. +Users can Upload documents for your context as attachments or sources via UI using the Conversation input box. +If appropriate, your answers can include code examples, formatted as follows: +```(language) +(code) +``` +Users are also able to see charts and diagrams if you use them with valid mermaid syntax in your responses. +Try to respond with mermaid charts if visualization helps with users queries. +You effectively utilize chat history, ensuring relevant and tailored responses. +You have access to a search tool that searches the user's uploaded documents and knowledge base. +Use the search_internal tool to find relevant information before answering questions. +You may search multiple times with different queries if needed. +Do not guess when documents are available — search first, then answer based on what you find. +If no relevant documents are found, use your general knowledge and tool capabilities. \ No newline at end of file diff --git a/application/prompts/agentic/strict.txt b/application/prompts/agentic/strict.txt new file mode 100644 index 00000000..8878b323 --- /dev/null +++ b/application/prompts/agentic/strict.txt @@ -0,0 +1,16 @@ +You are a helpful AI assistant, DocsGPT. You are proactive and helpful. Try to use tools, if they are available to you, +be proactive and fill in missing information. +Users can Upload documents for your context as attachments or sources via UI using the Conversation input box. +If appropriate, your answers can include code examples, formatted as follows: +```(language) +(code) +``` +Users are also able to see charts and diagrams if you use them with valid mermaid syntax in your responses. +Try to respond with mermaid charts if visualization helps with users queries. +You effectively utilize chat history, ensuring relevant and tailored responses. +You have access to a search tool that searches the user's uploaded documents and knowledge base. +Use the search_internal tool to find relevant information before answering questions. +You may search multiple times with different queries if needed. +You MUST search before answering any factual question. Do not guess or use general knowledge when documents are available. +If you dont have enough information from the search results or tools, answer "I don't know" or "I don't have enough information". +Never make up information or provide false information! \ No newline at end of file diff --git a/application/prompts/research/planning.txt b/application/prompts/research/planning.txt new file mode 100644 index 00000000..65f9b0cb --- /dev/null +++ b/application/prompts/research/planning.txt @@ -0,0 +1,21 @@ +You are a research planner. Your job is to analyze the user's question and create a focused research plan. + +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") +- "complex": Requires 5-6 searches with synthesis (e.g., "Audit our compliance docs against regulation X") + +Break the question into sub-questions appropriate to its complexity. Don't over-decompose simple questions. + +Consider: +- What distinct aspects of the question need separate investigation? +- What order should they be investigated in? +- Are there any dependencies between sub-questions? + +You MUST respond with ONLY a valid JSON object in this exact format (no markdown, no code fences): +{ + "complexity": "simple|moderate|complex", + "steps": [ + {"query": "specific sub-question to investigate", "rationale": "why this step is needed"} + ] +} \ No newline at end of file diff --git a/application/prompts/research/step.txt b/application/prompts/research/step.txt new file mode 100644 index 00000000..aecaf2b8 --- /dev/null +++ b/application/prompts/research/step.txt @@ -0,0 +1,13 @@ +You are a research agent investigating a specific topic. Your goal is to find comprehensive, accurate information. + +Your current research task: {step_query} + +Instructions: +1. Use the available tools to search for and gather relevant information. +2. If you have a search_internal tool, use it to search the knowledge base. You may search multiple times with different queries. +3. If you have a reason_think tool, use it to analyze what you've found before drawing conclusions. +4. After gathering sufficient information, provide a detailed summary of your findings. +5. Cite specific documents and passages you found. Reference sources by their titles or filenames. +6. If you cannot find relevant information through tools, use your general knowledge but clearly indicate this. + +Be thorough — prefer completeness over brevity. Include all relevant details you find. \ No newline at end of file diff --git a/application/prompts/research/synthesis.txt b/application/prompts/research/synthesis.txt new file mode 100644 index 00000000..66bb1cad --- /dev/null +++ b/application/prompts/research/synthesis.txt @@ -0,0 +1,22 @@ +You are compiling a comprehensive research report based on multiple research steps. + +Original question: {question} + +Research plan: +{plan_summary} + +Intermediate findings: +{findings} + +Write a well-structured, thorough report that: +1. Directly addresses the user's original question +2. Synthesizes findings from all research steps into a coherent narrative +3. Uses inline citations [N] for every factual claim, where N maps to the source list below +4. Highlights key insights and connections across different research steps +5. Notes any gaps or areas where information was insufficient +6. Includes a "References" section at the end listing all cited sources + +Available sources for citation: +{references} + +Format the report with clear headings and sections. Be comprehensive but well-organized. \ No newline at end of file diff --git a/frontend/src/agents/NewAgent.tsx b/frontend/src/agents/NewAgent.tsx index d451d9ee..c36f51e0 100644 --- a/frontend/src/agents/NewAgent.tsx +++ b/frontend/src/agents/NewAgent.tsx @@ -138,6 +138,8 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) { const chunks = ['0', '2', '4', '6', '8', '10']; const agentTypes = [ { label: t('agents.form.agentTypes.classic'), value: 'classic' }, + { label: 'Agentic', value: 'agentic' }, + { label: 'Research', value: 'research' }, ]; const isPublishable = () => { diff --git a/frontend/src/assets/research.svg b/frontend/src/assets/research.svg new file mode 100644 index 00000000..1e3f9120 --- /dev/null +++ b/frontend/src/assets/research.svg @@ -0,0 +1,18 @@ + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/frontend/src/conversation/ConversationBubble.tsx b/frontend/src/conversation/ConversationBubble.tsx index bf058581..a894d3ba 100644 --- a/frontend/src/conversation/ConversationBubble.tsx +++ b/frontend/src/conversation/ConversationBubble.tsx @@ -37,7 +37,8 @@ import { selectSelectedDocs, } from '../preferences/preferenceSlice'; import classes from './ConversationBubble.module.css'; -import { FEEDBACK, MESSAGE_TYPE } from './conversationModels'; +import { FEEDBACK, MESSAGE_TYPE, ResearchState } from './conversationModels'; +import ResearchProgress from './ResearchProgress'; import { ToolCallsType } from './types'; const DisableSourceFE = import.meta.env.VITE_DISABLE_SOURCE_FE || false; @@ -53,6 +54,7 @@ const ConversationBubble = forwardRef< thought?: string; sources?: { title: string; text: string; link: string }[]; toolCalls?: ToolCallsType[]; + research?: ResearchState; retryBtn?: React.ReactElement; questionNumber?: number; isStreaming?: boolean; @@ -74,6 +76,7 @@ const ConversationBubble = forwardRef< thought, sources, toolCalls, + research, retryBtn, questionNumber, isStreaming, @@ -261,7 +264,15 @@ const ConversationBubble = forwardRef< return inlineProcessedContent; }; const processMarkdownContent = (content: string) => { - const processedContent = preprocessLaTeX(content); + let processedContent = preprocessLaTeX(content); + + // Convert citation references [N] into markdown links [N](#cite-N) + // so ReactMarkdown renders them as tags we can style. + // Avoid matching inside code blocks or existing links. + processedContent = processedContent.replace( + /(? `[${num}](#cite-${num})`, + ); const contentSegments: Array<{ type: 'text' | 'mermaid'; @@ -290,6 +301,7 @@ const ConversationBubble = forwardRef< return contentSegments; }; + bubble = (
{sources?.slice(0, 3)?.map((source, index) => ( -
+
setActiveTooltip(index)} @@ -393,6 +405,7 @@ const ConversationBubble = forwardRef<
)} + {research && } {toolCalls && toolCalls.length > 0 && ( )} @@ -474,6 +487,30 @@ const ConversationBubble = forwardRef< remarkPlugins={[remarkGfm, remarkMath]} rehypePlugins={[rehypeKatex]} components={{ + a({ href, children }) { + if (href?.startsWith('#cite-')) { + const num = href.replace('#cite-', ''); + const sourceIdx = parseInt(num, 10) - 1; + return ( + + ); + } + return {children}; + }, code(props) { const { children, @@ -649,6 +686,29 @@ const ConversationBubble = forwardRef<
+ {research && message && ( +
+ +
+ )}
diff --git a/frontend/src/conversation/ConversationMessages.tsx b/frontend/src/conversation/ConversationMessages.tsx index 1a666667..1d9fd6b6 100644 --- a/frontend/src/conversation/ConversationMessages.tsx +++ b/frontend/src/conversation/ConversationMessages.tsx @@ -140,7 +140,7 @@ export default function ConversationMessages({ ? LAST_BUBBLE_MARGIN : DEFAULT_BUBBLE_MARGIN; - if (query.thought || query.response || query.tool_calls) { + if (query.thought || query.response || query.tool_calls || query.research) { const isCurrentlyStreaming = status === 'loading' && index === queries.length - 1; return ( @@ -152,6 +152,7 @@ export default function ConversationMessages({ thought={query.thought} sources={query.sources} toolCalls={query.tool_calls} + research={query.research} onOpenArtifact={onOpenArtifact} feedback={query.feedback} isStreaming={isCurrentlyStreaming} diff --git a/frontend/src/conversation/ResearchProgress.tsx b/frontend/src/conversation/ResearchProgress.tsx new file mode 100644 index 00000000..cce6ca0b --- /dev/null +++ b/frontend/src/conversation/ResearchProgress.tsx @@ -0,0 +1,147 @@ +import { useEffect, useState } from 'react'; +import ResearchIcon from '../assets/research.svg'; +import Avatar from '../components/Avatar'; +import { ResearchState } from './conversationModels'; + +const SmallCheck = () => ( + + + +); + +const SmallSpinner = () => ( + + + + +); + +const SmallPending = () => ( +
+); + +function StatusText({ status, elapsed }: { status: string; elapsed?: number }) { + const labels: Record = { + planning: 'Planning research...', + researching: 'Researching...', + synthesizing: 'Writing report...', + complete: 'Complete', + }; + const elapsed_str = elapsed ? ` \u00B7 ${Math.round(elapsed)}s` : ''; + return ( + + {status === 'complete' ? ( + <> + {labels.complete} + {elapsed_str} + + ) : ( + <> + {labels[status] || status} + {elapsed_str} + + )} + + ); +} + +export default function ResearchProgress({ + research, +}: { + research: ResearchState; +}) { + const { plan, status, elapsed_seconds } = research; + const [isExpanded, setIsExpanded] = useState(true); + + // Auto-collapse when synthesis starts (report is streaming) + useEffect(() => { + if (status === 'synthesizing' || status === 'complete') { + setIsExpanded(false); + } + }, [status]); + + if (!plan && !status) return null; + + const completedSteps = plan?.filter((s) => s.status === 'complete').length ?? 0; + const totalSteps = plan?.length ?? 0; + + // Collapsed: single-line summary + const summaryText = totalSteps > 0 + ? `Researched ${completedSteps} topic${completedSteps !== 1 ? 's' : ''}` + : 'Research'; + + return ( +
+ {/* Header row — matches Reasoning / Sources pattern */} +
+ + } + /> + + {status && } +
+ + {/* Expanded: vertical timeline of steps */} + {isExpanded && plan && plan.length > 0 && ( +
+
+ {plan.map((step, i) => { + const isLast = i === plan.length - 1; + return ( +
+ {/* Timeline: dot + vertical line */} +
+
+ {step.status === 'complete' ? ( + + ) : step.status === 'researching' ? ( + + ) : ( + + )} +
+ {!isLast && ( +
+ )} +
+ {/* Step content */} +
+

+ {step.query} +

+
+
+ ); + })} +
+
+ )} +
+ ); +} diff --git a/frontend/src/conversation/conversationModels.ts b/frontend/src/conversation/conversationModels.ts index 904683a7..cfa4ec8e 100644 --- a/frontend/src/conversation/conversationModels.ts +++ b/frontend/src/conversation/conversationModels.ts @@ -18,6 +18,20 @@ export interface Attachment { token_count?: number; } +export interface ResearchStep { + query: string; + rationale?: string; + status: 'pending' | 'researching' | 'complete'; +} + +export interface ResearchState { + plan?: ResearchStep[]; + complexity?: string; + status?: string; + elapsed_seconds?: number; + tokens_used?: number; +} + export interface ConversationState { queries: Query[]; status: Status; @@ -50,6 +64,7 @@ export interface Query { attachments?: { id: string; fileName: string }[]; structured?: boolean; schema?: object; + research?: ResearchState; } export interface RetrievalPayload { diff --git a/frontend/src/conversation/conversationSlice.ts b/frontend/src/conversation/conversationSlice.ts index 936d2de8..fed5e6cf 100644 --- a/frontend/src/conversation/conversationSlice.ts +++ b/frontend/src/conversation/conversationSlice.ts @@ -11,7 +11,13 @@ import { handleFetchAnswer, handleFetchAnswerSteaming, } from './conversationHandlers'; -import { Answer, ConversationState, Query, Status } from './conversationModels'; +import { + Answer, + ConversationState, + Query, + ResearchStep, + Status, +} from './conversationModels'; import { ToolCallsType } from './types'; const initialState: ConversationState = { @@ -71,6 +77,12 @@ 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' }, + }), + ); getConversations(state.preference.token) .then((fetchedConversations) => { dispatch(setConversations(fetchedConversations)); @@ -122,6 +134,21 @@ export const fetchAnswer = createAsyncThunk< tool_call: data.data as ToolCallsType, }), ); + } else if (data.type === 'research_plan') { + dispatch( + updateResearchPlan({ + index: targetIndex, + plan: data.data.steps, + complexity: data.data.complexity, + }), + ); + } else if (data.type === 'research_progress') { + dispatch( + updateResearchProgress({ + index: targetIndex, + progress: data.data, + }), + ); } else if (data.type === 'error') { // set status to 'failed' dispatch(conversationSlice.actions.setStatus('failed')); @@ -256,6 +283,7 @@ export const conversationSlice = createSlice({ delete state.queries[index].structured; delete state.queries[index].schema; delete state.queries[index].feedback; + delete state.queries[index].research; }, updateStreamingQuery( state, @@ -337,6 +365,77 @@ export const conversationSlice = createSlice({ }; } else state.queries[index].tool_calls.push(tool_call); }, + updateResearchPlan( + state, + action: PayloadAction<{ + index: number; + plan: ResearchStep[]; + complexity?: string; + }>, + ) { + const { index, plan, complexity } = action.payload; + if (!state.queries[index].research) { + state.queries[index].research = {}; + } + state.queries[index].research!.plan = plan.map((step) => ({ + ...step, + status: 'pending', + })); + if (complexity) { + state.queries[index].research!.complexity = complexity; + } + }, + updateResearchProgress( + state, + action: PayloadAction<{ + index: number; + progress: { + status?: string; + step?: number; + total?: number; + query?: string; + elapsed_seconds?: number; + tokens_used?: number; + }; + }>, + ) { + const { index, progress } = action.payload; + if (!state.queries[index].research) { + state.queries[index].research = {}; + } + const research = state.queries[index].research!; + if (progress.elapsed_seconds !== undefined) { + research.elapsed_seconds = progress.elapsed_seconds; + } + if (progress.tokens_used !== undefined) { + research.tokens_used = progress.tokens_used; + } + // Update individual step status when step number is present + if (progress.step !== undefined) { + if (!research.plan) { + research.plan = []; + } + const stepIndex = progress.step - 1; + // Dynamically add step if it doesn't exist yet + while (research.plan.length <= stepIndex) { + research.plan.push({ + query: progress.query || `Step ${research.plan.length + 1}`, + status: 'pending', + }); + } + if (progress.status === 'researching' || progress.status === 'complete') { + research.plan[stepIndex].status = progress.status; + } + if (progress.query) { + research.plan[stepIndex].query = progress.query; + } + // Keep top-level status as "researching" while steps are running + research.status = 'researching'; + } else if (progress.status) { + // Top-level status updates (planning, synthesizing) + research.status = progress.status; + } + }, updateQuery( state, action: PayloadAction<{ index: number; query: Partial }>, @@ -405,6 +504,8 @@ export const { updateThought, updateStreamingSource, updateToolCall, + updateResearchPlan, + updateResearchProgress, setConversation, setStatus, raiseError,