From 32c268a21e49f4057bd939fd2ed529443fe2f10e Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 12:47:17 +0000 Subject: [PATCH 1/7] refactor: simplify agent architecture and remove ReActAgent --- application/agents/agent_creator.py | 3 +- application/agents/base.py | 395 +++----------- application/agents/classic_agent.py | 6 +- application/agents/react_agent.py | 238 -------- application/agents/tool_executor.py | 307 +++++++++++ application/agents/workflows/node_agent.py | 24 +- application/api/answer/routes/answer.py | 13 +- application/api/answer/routes/stream.py | 9 +- .../api/answer/services/stream_processor.py | 134 ++--- application/app.py | 5 + application/prompts/react_final_prompt.txt | 3 - application/prompts/react_planning_prompt.txt | 13 - application/seed/config/premade_agents.yaml | 2 +- frontend/src/agents/NewAgent.tsx | 1 - frontend/src/agents/WorkflowBuilder.tsx | 3 +- .../src/agents/workflow/WorkflowBuilder.tsx | 5 +- results.txt | 12 + tests/agents/test_agent_creator.py | 8 +- tests/agents/test_react_agent.py | 514 ------------------ tests/conftest.py | 2 +- 20 files changed, 467 insertions(+), 1230 deletions(-) delete mode 100644 application/agents/react_agent.py create mode 100644 application/agents/tool_executor.py delete mode 100644 application/prompts/react_final_prompt.txt delete mode 100644 application/prompts/react_planning_prompt.txt create mode 100644 results.txt delete mode 100644 tests/agents/test_react_agent.py diff --git a/application/agents/agent_creator.py b/application/agents/agent_creator.py index e15165bf..6773babd 100644 --- a/application/agents/agent_creator.py +++ b/application/agents/agent_creator.py @@ -1,7 +1,6 @@ import logging from application.agents.classic_agent import ClassicAgent -from application.agents.react_agent import ReActAgent from application.agents.workflow_agent import WorkflowAgent logger = logging.getLogger(__name__) @@ -10,7 +9,7 @@ logger = logging.getLogger(__name__) class AgentCreator: agents = { "classic": ClassicAgent, - "react": ReActAgent, + "react": ClassicAgent, # backwards compat: react falls back to classic "workflow": WorkflowAgent, } diff --git a/application/agents/base.py b/application/agents/base.py index ee55a449..8e6363fc 100644 --- a/application/agents/base.py +++ b/application/agents/base.py @@ -3,20 +3,15 @@ import uuid from abc import ABC, abstractmethod from typing import Dict, Generator, List, Optional -from bson.objectid import ObjectId - -from application.agents.tools.tool_action_parser import ToolActionParser -from application.agents.tools.tool_manager import ToolManager +from application.agents.tool_executor import ToolExecutor from application.core.json_schema_utils import ( JsonSchemaValidationError, normalize_json_schema_payload, ) -from application.core.mongo_db import MongoDB from application.core.settings import settings from application.llm.handlers.handler_creator import LLMHandlerCreator from application.llm.llm_creator import LLMCreator from application.logging import build_stack_data, log_activity, LogContext -from application.security.encryption import decrypt_credentials logger = logging.getLogger(__name__) @@ -41,6 +36,9 @@ class BaseAgent(ABC): limited_request_mode: Optional[bool] = False, request_limit: Optional[int] = settings.DEFAULT_AGENT_LIMITS["request_limit"], compressed_summary: Optional[str] = None, + llm=None, + llm_handler=None, + tool_executor: Optional[ToolExecutor] = None, ): self.endpoint = endpoint self.llm_name = llm_name @@ -51,22 +49,41 @@ class BaseAgent(ABC): self.prompt = prompt self.decoded_token = decoded_token or {} self.user: str = self.decoded_token.get("sub") - self.tool_config: Dict = {} self.tools: List[Dict] = [] - self.tool_calls: List[Dict] = [] self.chat_history: List[Dict] = chat_history if chat_history is not None else [] - self.llm = LLMCreator.create_llm( - llm_name, - api_key=api_key, - user_api_key=user_api_key, - decoded_token=decoded_token, - model_id=model_id, - agent_id=agent_id, - ) + + # Dependency injection for LLM — fall back to creating if not provided + if llm is not None: + self.llm = llm + else: + self.llm = LLMCreator.create_llm( + llm_name, + api_key=api_key, + user_api_key=user_api_key, + decoded_token=decoded_token, + model_id=model_id, + agent_id=agent_id, + ) + self.retrieved_docs = retrieved_docs or [] - self.llm_handler = LLMHandlerCreator.create_handler( - llm_name if llm_name else "default" - ) + + if llm_handler is not None: + self.llm_handler = llm_handler + else: + self.llm_handler = LLMHandlerCreator.create_handler( + llm_name if llm_name else "default" + ) + + # Tool executor — injected or created + if tool_executor is not None: + self.tool_executor = tool_executor + else: + self.tool_executor = ToolExecutor( + user_api_key=user_api_key, + user=self.user, + decoded_token=decoded_token, + ) + self.attachments = attachments or [] self.json_schema = None if json_schema is not None: @@ -94,333 +111,72 @@ class BaseAgent(ABC): ) -> Generator[Dict, None, None]: pass + # ---- Tool delegation (thin wrappers around ToolExecutor) ---- + + @property + def tool_calls(self) -> List[Dict]: + return self.tool_executor.tool_calls + + @tool_calls.setter + def tool_calls(self, value: List[Dict]): + self.tool_executor.tool_calls = value + def _get_tools(self, api_key: str = None) -> Dict[str, Dict]: - mongo = MongoDB.get_client() - db = mongo[settings.MONGO_DB_NAME] - agents_collection = db["agents"] - tools_collection = db["user_tools"] - - agent_data = agents_collection.find_one({"key": api_key or self.user_api_key}) - tool_ids = agent_data.get("tools", []) if agent_data else [] - - tools = ( - tools_collection.find( - {"_id": {"$in": [ObjectId(tool_id) for tool_id in tool_ids]}} - ) - if tool_ids - else [] - ) - tools = list(tools) - tools_by_id = {str(tool["_id"]): tool for tool in tools} if tools else {} - - return tools_by_id + return self.tool_executor._get_tools_by_api_key(api_key or self.user_api_key) def _get_user_tools(self, user="local"): - mongo = MongoDB.get_client() - db = mongo[settings.MONGO_DB_NAME] - user_tools_collection = db["user_tools"] - user_tools = user_tools_collection.find({"user": user, "status": True}) - user_tools = list(user_tools) - - return {str(i): tool for i, tool in enumerate(user_tools)} + return self.tool_executor._get_user_tools(user) def _build_tool_parameters(self, action): - params = {"type": "object", "properties": {}, "required": []} - for param_type in ["query_params", "headers", "body", "parameters"]: - if param_type in action and action[param_type].get("properties"): - for k, v in action[param_type]["properties"].items(): - if v.get("filled_by_llm", True): - params["properties"][k] = { - key: value - for key, value in v.items() - if key not in ("filled_by_llm", "value", "required") - } - if v.get("required", False): - params["required"].append(k) - return params + return self.tool_executor._build_tool_parameters(action) def _prepare_tools(self, tools_dict): - self.tools = [ - { - "type": "function", - "function": { - "name": f"{action['name']}_{tool_id}", - "description": action["description"], - "parameters": self._build_tool_parameters(action), - }, - } - for tool_id, tool in tools_dict.items() - if ( - (tool["name"] == "api_tool" and "actions" in tool.get("config", {})) - or (tool["name"] != "api_tool" and "actions" in tool) - ) - for action in ( - tool["config"]["actions"].values() - if tool["name"] == "api_tool" - else tool["actions"] - ) - if action.get("active", True) - ] + self.tools = self.tool_executor.prepare_tools_for_llm(tools_dict) def _execute_tool_action(self, tools_dict, call): - parser = ToolActionParser(self.llm.__class__.__name__) - tool_id, action_name, call_args = parser.parse_args(call) - - call_id = getattr(call, "id", None) or str(uuid.uuid4()) - - # Check if parsing failed - - if tool_id is None or action_name is None: - error_message = f"Error: Failed to parse LLM tool call. Tool name: {getattr(call, 'name', 'unknown')}" - logger.error(error_message) - - tool_call_data = { - "tool_name": "unknown", - "call_id": call_id, - "action_name": getattr(call, "name", "unknown"), - "arguments": call_args or {}, - "result": f"Failed to parse tool call. Invalid tool name format: {getattr(call, 'name', 'unknown')}", - } - yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} - self.tool_calls.append(tool_call_data) - return "Failed to parse tool call.", call_id - # Check if tool_id exists in available tools - - if tool_id not in tools_dict: - error_message = f"Error: Tool ID '{tool_id}' extracted from LLM call not found in available tools_dict. Available IDs: {list(tools_dict.keys())}" - logger.error(error_message) - - # Return error result - - tool_call_data = { - "tool_name": "unknown", - "call_id": call_id, - "action_name": f"{action_name}_{tool_id}", - "arguments": call_args, - "result": f"Tool with ID {tool_id} not found. Available tools: {list(tools_dict.keys())}", - } - yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} - self.tool_calls.append(tool_call_data) - return f"Tool with ID {tool_id} not found.", call_id - tool_call_data = { - "tool_name": tools_dict[tool_id]["name"], - "call_id": call_id, - "action_name": f"{action_name}_{tool_id}", - "arguments": call_args, - } - yield {"type": "tool_call", "data": {**tool_call_data, "status": "pending"}} - - tool_data = tools_dict[tool_id] - action_data = ( - tool_data["config"]["actions"][action_name] - if tool_data["name"] == "api_tool" - else next( - action - for action in tool_data["actions"] - if action["name"] == action_name - ) + return self.tool_executor.execute( + tools_dict, call, self.llm.__class__.__name__ ) - query_params, headers, body, parameters = {}, {}, {}, {} - param_types = { - "query_params": query_params, - "headers": headers, - "body": body, - "parameters": parameters, - } - - for param_type, target_dict in param_types.items(): - if param_type in action_data and action_data[param_type].get("properties"): - for param, details in action_data[param_type]["properties"].items(): - if ( - param not in call_args - and "value" in details - and details["value"] - ): - target_dict[param] = details["value"] - for param, value in call_args.items(): - for param_type, target_dict in param_types.items(): - if param_type in action_data and param in action_data[param_type].get( - "properties", {} - ): - target_dict[param] = value - tm = ToolManager(config={}) - - # Prepare tool_config and add tool_id for memory tools - - if tool_data["name"] == "api_tool": - action_config = tool_data["config"]["actions"][action_name] - tool_config = { - "url": action_config["url"], - "method": action_config["method"], - "headers": headers, - "query_params": query_params, - } - if "body_content_type" in action_config: - tool_config["body_content_type"] = action_config.get( - "body_content_type", "application/json" - ) - tool_config["body_encoding_rules"] = action_config.get( - "body_encoding_rules", {} - ) - else: - tool_config = tool_data["config"].copy() if tool_data["config"] else {} - if tool_config.get("encrypted_credentials") and self.user: - decrypted = decrypt_credentials( - tool_config["encrypted_credentials"], self.user - ) - tool_config.update(decrypted) - tool_config["auth_credentials"] = decrypted - tool_config.pop("encrypted_credentials", None) - tool_config["tool_id"] = str(tool_data.get("_id", tool_id)) - if hasattr(self, "conversation_id") and self.conversation_id: - tool_config["conversation_id"] = self.conversation_id - if tool_data["name"] == "mcp_tool": - tool_config["query_mode"] = True - tool = tm.load_tool( - tool_data["name"], - tool_config=tool_config, - user_id=self.user, - ) - resolved_arguments = ( - {"query_params": query_params, "headers": headers, "body": body} - if tool_data["name"] == "api_tool" - else parameters - ) - if tool_data["name"] == "api_tool": - logger.debug( - f"Executing api: {action_name} with query_params: {query_params}, headers: {headers}, body: {body}" - ) - result = tool.execute_action(action_name, **body) - else: - logger.debug(f"Executing tool: {action_name} with args: {call_args}") - result = tool.execute_action(action_name, **parameters) - - get_artifact_id = ( - getattr(tool, "get_artifact_id", None) - if tool_data["name"] != "api_tool" - else None - ) - - artifact_id = None - if callable(get_artifact_id): - try: - artifact_id = get_artifact_id(action_name, **parameters) - except Exception: - logger.exception( - "Failed to extract artifact_id from tool %s for action %s", - tool_data["name"], - action_name, - ) - - artifact_id = str(artifact_id).strip() if artifact_id is not None else "" - if artifact_id: - tool_call_data["artifact_id"] = artifact_id - result_full = str(result) - tool_call_data["resolved_arguments"] = resolved_arguments - tool_call_data["result_full"] = result_full - tool_call_data["result"] = ( - f"{result_full[:50]}..." if len(result_full) > 50 else result_full - ) - - stream_tool_call_data = { - key: value - for key, value in tool_call_data.items() - if key not in {"result_full", "resolved_arguments"} - } - yield {"type": "tool_call", "data": {**stream_tool_call_data, "status": "completed"}} - self.tool_calls.append(tool_call_data) - - return result, call_id - def _get_truncated_tool_calls(self): - return [ - { - "tool_name": tool_call.get("tool_name"), - "call_id": tool_call.get("call_id"), - "action_name": tool_call.get("action_name"), - "arguments": tool_call.get("arguments"), - "artifact_id": tool_call.get("artifact_id"), - "result": ( - f"{str(tool_call['result'])[:50]}..." - if len(str(tool_call["result"])) > 50 - else tool_call["result"] - ), - "status": "completed", - } - for tool_call in self.tool_calls - ] + return self.tool_executor.get_truncated_tool_calls() + + # ---- Context / token management ---- def _calculate_current_context_tokens(self, messages: List[Dict]) -> int: - """ - Calculate total tokens in current context (messages). - - Args: - messages: List of message dicts - - Returns: - Total token count - """ from application.api.answer.services.compression.token_counter import ( TokenCounter, ) - return TokenCounter.count_message_tokens(messages) def _check_context_limit(self, messages: List[Dict]) -> bool: - """ - Check if we're approaching context limit (80%). - - Args: - messages: Current message list - - Returns: - True if at or above 80% of context limit - """ from application.core.model_utils import get_token_limit - from application.core.settings import settings try: - # Calculate current tokens current_tokens = self._calculate_current_context_tokens(messages) self.current_token_count = current_tokens - - # Get context limit for model context_limit = get_token_limit(self.model_id) - - # Calculate threshold (80%) threshold = int(context_limit * settings.COMPRESSION_THRESHOLD_PERCENTAGE) - # Check if we've reached the limit if current_tokens >= threshold: logger.warning( f"Context limit approaching: {current_tokens}/{context_limit} tokens " f"({(current_tokens/context_limit)*100:.1f}%)" ) return True - return False - except Exception as e: logger.error(f"Error checking context limit: {str(e)}", exc_info=True) return False def _validate_context_size(self, messages: List[Dict]) -> None: - """ - Pre-flight validation before calling LLM. Logs warnings but never raises errors. - - Args: - messages: Messages to be sent to LLM - """ from application.core.model_utils import get_token_limit current_tokens = self._calculate_current_context_tokens(messages) self.current_token_count = current_tokens context_limit = get_token_limit(self.model_id) - percentage = (current_tokens / context_limit) * 100 - # Log based on usage level if current_tokens >= context_limit: logger.warning( f"Context at limit: {current_tokens:,}/{context_limit:,} tokens " @@ -435,44 +191,32 @@ class BaseAgent(ABC): ) def _truncate_text_middle(self, text: str, max_tokens: int) -> str: - """ - Truncate text by removing content from the middle, preserving start and end. - - Args: - text: Text to truncate - max_tokens: Maximum tokens allowed - - Returns: - Truncated text with middle removed if needed - """ from application.utils import num_tokens_from_string current_tokens = num_tokens_from_string(text) if current_tokens <= max_tokens: return text - # Estimate chars per token (roughly 4 chars per token for English) chars_per_token = len(text) / current_tokens if current_tokens > 0 else 4 - target_chars = int(max_tokens * chars_per_token * 0.95) # 5% safety margin + target_chars = int(max_tokens * chars_per_token * 0.95) if target_chars <= 0: return "" - # Split: keep 40% from start, 40% from end, remove middle start_chars = int(target_chars * 0.4) end_chars = int(target_chars * 0.4) truncation_marker = "\n\n[... content truncated to fit context limit ...]\n\n" - truncated = text[:start_chars] + truncation_marker + text[-end_chars:] logger.info( f"Truncated text from {current_tokens:,} to ~{max_tokens:,} tokens " f"(removed middle section)" ) - return truncated + # ---- Message building ---- + def _build_messages( self, system_prompt: str, @@ -482,7 +226,6 @@ class BaseAgent(ABC): from application.core.model_utils import get_token_limit from application.utils import num_tokens_from_string - # Append compression summary to system prompt if present if self.compressed_summary: compression_context = ( "\n\n---\n\n" @@ -496,23 +239,18 @@ class BaseAgent(ABC): context_limit = get_token_limit(self.model_id) system_tokens = num_tokens_from_string(system_prompt) - # Reserve 10% for response/tools safety_buffer = int(context_limit * 0.1) available_after_system = context_limit - system_tokens - safety_buffer - # Max tokens for query: 80% of available space (leave room for history) max_query_tokens = int(available_after_system * 0.8) query_tokens = num_tokens_from_string(query) - # Truncate query from middle if it exceeds 80% of available context if query_tokens > max_query_tokens: query = self._truncate_text_middle(query, max_query_tokens) query_tokens = num_tokens_from_string(query) - # Calculate remaining budget for chat history available_for_history = max(available_after_system - query_tokens, 0) - # Truncate chat history to fit within available budget working_history = self._truncate_history_to_fit( self.chat_history, available_for_history, @@ -557,16 +295,6 @@ class BaseAgent(ABC): history: List[Dict], max_tokens: int, ) -> List[Dict]: - """ - Truncate chat history to fit within token budget, keeping most recent messages. - - Args: - history: Full chat history - max_tokens: Maximum tokens allowed for history - - Returns: - Truncated history (most recent messages that fit) - """ from application.utils import num_tokens_from_string if not history or max_tokens <= 0: @@ -575,7 +303,6 @@ class BaseAgent(ABC): truncated = [] current_tokens = 0 - # Iterate from newest to oldest for message in reversed(history): message_tokens = 0 @@ -595,7 +322,7 @@ class BaseAgent(ABC): if current_tokens + message_tokens <= max_tokens: current_tokens += message_tokens - truncated.insert(0, message) # Maintain chronological order + truncated.insert(0, message) else: break @@ -607,13 +334,13 @@ class BaseAgent(ABC): return truncated + # ---- LLM generation ---- + def _llm_gen(self, messages: List[Dict], log_context: Optional[LogContext] = None): - # Pre-flight context validation - fail fast if over limit self._validate_context_size(messages) gen_kwargs = {"model": self.model_id, "messages": messages} if self.attachments: - # Usage accounting only; stripped before provider invocation. gen_kwargs["_usage_attachments"] = self.attachments if ( diff --git a/application/agents/classic_agent.py b/application/agents/classic_agent.py index af8e585f..d6dced9e 100644 --- a/application/agents/classic_agent.py +++ b/application/agents/classic_agent.py @@ -15,11 +15,7 @@ class ClassicAgent(BaseAgent): ) -> Generator[Dict, None, None]: """Core generator function for ClassicAgent execution flow""" - tools_dict = ( - self._get_user_tools(self.user) - if not self.user_api_key - else self._get_tools(self.user_api_key) - ) + tools_dict = self.tool_executor.get_tools() self._prepare_tools(tools_dict) messages = self._build_messages(self.prompt, query) diff --git a/application/agents/react_agent.py b/application/agents/react_agent.py deleted file mode 100644 index 92be75f6..00000000 --- a/application/agents/react_agent.py +++ /dev/null @@ -1,238 +0,0 @@ -import logging -import os -from typing import Any, Dict, Generator, List - -from application.agents.base import BaseAgent -from application.logging import build_stack_data, LogContext - -logger = logging.getLogger(__name__) - -MAX_ITERATIONS_REASONING = 10 - -current_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -) -with open( - os.path.join(current_dir, "application/prompts", "react_planning_prompt.txt"), "r" -) as f: - PLANNING_PROMPT_TEMPLATE = f.read() -with open( - os.path.join(current_dir, "application/prompts", "react_final_prompt.txt"), "r" -) as f: - FINAL_PROMPT_TEMPLATE = f.read() - - -class ReActAgent(BaseAgent): - """ - Research and Action (ReAct) Agent - Advanced reasoning agent with iterative planning. - - Implements a think-act-observe loop for complex problem-solving: - 1. Creates a strategic plan based on the query - 2. Executes tools and gathers observations - 3. Iteratively refines approach until satisfied - 4. Synthesizes final answer from all observations - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.plan: str = "" - self.observations: List[str] = [] - - def _gen_inner( - self, query: str, log_context: LogContext - ) -> Generator[Dict, None, None]: - """Execute ReAct reasoning loop with planning, action, and observation cycles""" - - self._reset_state() - - tools_dict = ( - self._get_tools(self.user_api_key) - if self.user_api_key - else self._get_user_tools(self.user) - ) - self._prepare_tools(tools_dict) - - for iteration in range(1, MAX_ITERATIONS_REASONING + 1): - yield {"thought": f"Reasoning... (iteration {iteration})\n\n"} - - yield from self._planning_phase(query, log_context) - - if not self.plan: - logger.warning( - f"ReActAgent: No plan generated in iteration {iteration}" - ) - break - self.observations.append(f"Plan (iteration {iteration}): {self.plan}") - - satisfied = yield from self._execution_phase(query, tools_dict, log_context) - - if satisfied: - logger.info("ReActAgent: Goal satisfied, stopping reasoning loop") - break - yield from self._synthesis_phase(query, log_context) - - def _reset_state(self): - """Reset agent state for new query""" - self.plan = "" - self.observations = [] - - def _planning_phase( - self, query: str, log_context: LogContext - ) -> Generator[Dict, None, None]: - """Generate strategic plan for query""" - logger.info("ReActAgent: Creating plan...") - - plan_prompt = self._build_planning_prompt(query) - messages = [{"role": "user", "content": plan_prompt}] - - plan_stream = self.llm.gen_stream( - model=self.model_id, - messages=messages, - tools=self.tools if self.tools else None, - ) - - if log_context: - log_context.stacks.append( - {"component": "planning_llm", "data": build_stack_data(self.llm)} - ) - plan_parts = [] - for chunk in plan_stream: - content = self._extract_content(chunk) - if content: - plan_parts.append(content) - yield {"thought": content} - self.plan = "".join(plan_parts) - - def _execution_phase( - self, query: str, tools_dict: Dict, log_context: LogContext - ) -> Generator[bool, None, None]: - """Execute plan with tool calls and observations""" - execution_prompt = self._build_execution_prompt(query) - messages = self._build_messages(execution_prompt, query) - - llm_response = self._llm_gen(messages, log_context) - initial_content = self._extract_content(llm_response) - - if initial_content: - self.observations.append(f"Initial response: {initial_content}") - processed_response = self._llm_handler( - llm_response, tools_dict, messages, log_context - ) - - for tool_call in self.tool_calls: - observation = ( - f"Executed: {tool_call.get('tool_name', 'Unknown')} " - f"with args {tool_call.get('arguments', {})}. " - f"Result: {str(tool_call.get('result', ''))[:200]}" - ) - self.observations.append(observation) - final_content = self._extract_content(processed_response) - if final_content: - self.observations.append(f"Response after tools: {final_content}") - if log_context: - log_context.stacks.append( - { - "component": "agent_tool_calls", - "data": {"tool_calls": self.tool_calls.copy()}, - } - ) - yield {"sources": self.retrieved_docs} - yield {"tool_calls": self._get_truncated_tool_calls()} - - return "SATISFIED" in (final_content or "") - - def _synthesis_phase( - self, query: str, log_context: LogContext - ) -> Generator[Dict, None, None]: - """Synthesize final answer from all observations""" - logger.info("ReActAgent: Generating final answer...") - - final_prompt = self._build_final_answer_prompt(query) - messages = [{"role": "user", "content": final_prompt}] - - final_stream = self.llm.gen_stream( - model=self.model_id, messages=messages, tools=None - ) - - if log_context: - log_context.stacks.append( - {"component": "final_answer_llm", "data": build_stack_data(self.llm)} - ) - for chunk in final_stream: - content = self._extract_content(chunk) - if content: - yield {"answer": content} - - def _build_planning_prompt(self, query: str) -> str: - """Build planning phase prompt""" - prompt = PLANNING_PROMPT_TEMPLATE.replace("{query}", query) - prompt = prompt.replace("{prompt}", self.prompt or "") - prompt = prompt.replace("{summaries}", "") - prompt = prompt.replace("{observations}", "\n".join(self.observations)) - return prompt - - def _build_execution_prompt(self, query: str) -> str: - """Build execution phase prompt with plan and observations""" - observations_str = "\n".join(self.observations) - - if len(observations_str) > 20000: - observations_str = observations_str[:20000] + "\n...[truncated]" - return ( - f"{self.prompt or ''}\n\n" - f"Follow this plan:\n{self.plan}\n\n" - f"Observations:\n{observations_str}\n\n" - f"If sufficient data exists to answer '{query}', respond with 'SATISFIED'. " - f"Otherwise, continue executing the plan." - ) - - def _build_final_answer_prompt(self, query: str) -> str: - """Build final synthesis prompt""" - observations_str = "\n".join(self.observations) - - if len(observations_str) > 10000: - observations_str = observations_str[:10000] + "\n...[truncated]" - logger.warning("ReActAgent: Observations truncated for final answer") - return FINAL_PROMPT_TEMPLATE.format(query=query, observations=observations_str) - - def _extract_content(self, response: Any) -> str: - """Extract text content from various LLM response formats""" - if not response: - return "" - collected = [] - - if isinstance(response, str): - return response - if hasattr(response, "message") and hasattr(response.message, "content"): - if response.message.content: - return response.message.content - if hasattr(response, "choices") and response.choices: - if hasattr(response.choices[0], "message"): - content = response.choices[0].message.content - if content: - return content - if hasattr(response, "content") and isinstance(response.content, list): - if response.content and hasattr(response.content[0], "text"): - return response.content[0].text - try: - for chunk in response: - content_piece = "" - - if hasattr(chunk, "choices") and chunk.choices: - if hasattr(chunk.choices[0], "delta"): - delta_content = chunk.choices[0].delta.content - if delta_content: - content_piece = delta_content - elif hasattr(chunk, "type") and chunk.type == "content_block_delta": - if hasattr(chunk, "delta") and hasattr(chunk.delta, "text"): - content_piece = chunk.delta.text - elif isinstance(chunk, str): - content_piece = chunk - if content_piece: - collected.append(content_piece) - except (TypeError, AttributeError): - logger.debug( - f"Response not iterable or unexpected format: {type(response)}" - ) - except Exception as e: - logger.error(f"Error extracting content: {e}") - return "".join(collected) \ No newline at end of file diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py new file mode 100644 index 00000000..0fde060d --- /dev/null +++ b/application/agents/tool_executor.py @@ -0,0 +1,307 @@ +import logging +import uuid +from typing import Dict, List, Optional + +from bson.objectid import ObjectId + +from application.agents.tools.tool_action_parser import ToolActionParser +from application.agents.tools.tool_manager import ToolManager +from application.core.mongo_db import MongoDB +from application.core.settings import settings +from application.security.encryption import decrypt_credentials + +logger = logging.getLogger(__name__) + + +class ToolExecutor: + """Handles tool discovery, preparation, and execution. + + Extracted from BaseAgent to separate concerns and enable tool caching. + """ + + def __init__( + self, + user_api_key: Optional[str] = None, + user: Optional[str] = None, + decoded_token: Optional[Dict] = None, + ): + self.user_api_key = user_api_key + self.user = user + self.decoded_token = decoded_token + self.tool_calls: List[Dict] = [] + self._loaded_tools: Dict[str, object] = {} + self.conversation_id: Optional[str] = None + + def get_tools(self) -> Dict[str, Dict]: + """Load tool configs from DB based on user context.""" + if self.user_api_key: + return self._get_tools_by_api_key(self.user_api_key) + return self._get_user_tools(self.user or "local") + + def _get_tools_by_api_key(self, api_key: str) -> Dict[str, Dict]: + mongo = MongoDB.get_client() + db = mongo[settings.MONGO_DB_NAME] + agents_collection = db["agents"] + tools_collection = db["user_tools"] + + agent_data = agents_collection.find_one({"key": api_key}) + tool_ids = agent_data.get("tools", []) if agent_data else [] + + tools = ( + tools_collection.find( + {"_id": {"$in": [ObjectId(tool_id) for tool_id in tool_ids]}} + ) + if tool_ids + else [] + ) + tools = list(tools) + return {str(tool["_id"]): tool for tool in tools} if tools else {} + + def _get_user_tools(self, user: str = "local") -> Dict[str, Dict]: + mongo = MongoDB.get_client() + db = mongo[settings.MONGO_DB_NAME] + user_tools_collection = db["user_tools"] + user_tools = user_tools_collection.find({"user": user, "status": True}) + user_tools = list(user_tools) + return {str(i): tool for i, tool in enumerate(user_tools)} + + def prepare_tools_for_llm(self, tools_dict: Dict) -> List[Dict]: + """Convert tool configs to LLM function schemas.""" + return [ + { + "type": "function", + "function": { + "name": f"{action['name']}_{tool_id}", + "description": action["description"], + "parameters": self._build_tool_parameters(action), + }, + } + for tool_id, tool in tools_dict.items() + if ( + (tool["name"] == "api_tool" and "actions" in tool.get("config", {})) + or (tool["name"] != "api_tool" and "actions" in tool) + ) + for action in ( + tool["config"]["actions"].values() + if tool["name"] == "api_tool" + else tool["actions"] + ) + if action.get("active", True) + ] + + def _build_tool_parameters(self, action: Dict) -> Dict: + params = {"type": "object", "properties": {}, "required": []} + for param_type in ["query_params", "headers", "body", "parameters"]: + if param_type in action and action[param_type].get("properties"): + for k, v in action[param_type]["properties"].items(): + if v.get("filled_by_llm", True): + params["properties"][k] = { + key: value + for key, value in v.items() + if key not in ("filled_by_llm", "value", "required") + } + if v.get("required", False): + params["required"].append(k) + return params + + def execute(self, tools_dict: Dict, call, llm_class_name: str): + """Execute a tool call. Yields status events, returns (result, call_id).""" + parser = ToolActionParser(llm_class_name) + tool_id, action_name, call_args = parser.parse_args(call) + + call_id = getattr(call, "id", None) or str(uuid.uuid4()) + + if tool_id is None or action_name is None: + error_message = f"Error: Failed to parse LLM tool call. Tool name: {getattr(call, 'name', 'unknown')}" + logger.error(error_message) + + tool_call_data = { + "tool_name": "unknown", + "call_id": call_id, + "action_name": getattr(call, "name", "unknown"), + "arguments": call_args or {}, + "result": f"Failed to parse tool call. Invalid tool name format: {getattr(call, 'name', 'unknown')}", + } + yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} + self.tool_calls.append(tool_call_data) + return "Failed to parse tool call.", call_id + + if tool_id not in tools_dict: + error_message = f"Error: Tool ID '{tool_id}' extracted from LLM call not found in available tools_dict. Available IDs: {list(tools_dict.keys())}" + logger.error(error_message) + + tool_call_data = { + "tool_name": "unknown", + "call_id": call_id, + "action_name": f"{action_name}_{tool_id}", + "arguments": call_args, + "result": f"Tool with ID {tool_id} not found. Available tools: {list(tools_dict.keys())}", + } + yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}} + self.tool_calls.append(tool_call_data) + return f"Tool with ID {tool_id} not found.", call_id + + tool_call_data = { + "tool_name": tools_dict[tool_id]["name"], + "call_id": call_id, + "action_name": f"{action_name}_{tool_id}", + "arguments": call_args, + } + yield {"type": "tool_call", "data": {**tool_call_data, "status": "pending"}} + + tool_data = tools_dict[tool_id] + action_data = ( + tool_data["config"]["actions"][action_name] + if tool_data["name"] == "api_tool" + else next( + action + for action in tool_data["actions"] + if action["name"] == action_name + ) + ) + + query_params, headers, body, parameters = {}, {}, {}, {} + param_types = { + "query_params": query_params, + "headers": headers, + "body": body, + "parameters": parameters, + } + + for param_type, target_dict in param_types.items(): + if param_type in action_data and action_data[param_type].get("properties"): + for param, details in action_data[param_type]["properties"].items(): + if ( + param not in call_args + and "value" in details + and details["value"] + ): + target_dict[param] = details["value"] + for param, value in call_args.items(): + for param_type, target_dict in param_types.items(): + if param_type in action_data and param in action_data[param_type].get( + "properties", {} + ): + target_dict[param] = value + + # Load tool (with caching) + tool = self._get_or_load_tool(tool_data, tool_id, action_name) + + resolved_arguments = ( + {"query_params": query_params, "headers": headers, "body": body} + if tool_data["name"] == "api_tool" + else parameters + ) + if tool_data["name"] == "api_tool": + logger.debug( + f"Executing api: {action_name} with query_params: {query_params}, headers: {headers}, body: {body}" + ) + result = tool.execute_action(action_name, **body) + else: + logger.debug(f"Executing tool: {action_name} with args: {call_args}") + result = tool.execute_action(action_name, **parameters) + + get_artifact_id = ( + getattr(tool, "get_artifact_id", None) + if tool_data["name"] != "api_tool" + else None + ) + + artifact_id = None + if callable(get_artifact_id): + try: + artifact_id = get_artifact_id(action_name, **parameters) + except Exception: + logger.exception( + "Failed to extract artifact_id from tool %s for action %s", + tool_data["name"], + action_name, + ) + + artifact_id = str(artifact_id).strip() if artifact_id is not None else "" + if artifact_id: + tool_call_data["artifact_id"] = artifact_id + result_full = str(result) + tool_call_data["resolved_arguments"] = resolved_arguments + tool_call_data["result_full"] = result_full + tool_call_data["result"] = ( + f"{result_full[:50]}..." if len(result_full) > 50 else result_full + ) + + stream_tool_call_data = { + key: value + for key, value in tool_call_data.items() + if key not in {"result_full", "resolved_arguments"} + } + yield {"type": "tool_call", "data": {**stream_tool_call_data, "status": "completed"}} + self.tool_calls.append(tool_call_data) + + return result, call_id + + def _get_or_load_tool(self, tool_data: Dict, tool_id: str, action_name: str): + """Load a tool, using cache when possible.""" + cache_key = f"{tool_data['name']}:{tool_id}:{self.user or ''}" + if cache_key in self._loaded_tools: + return self._loaded_tools[cache_key] + + tm = ToolManager(config={}) + + if tool_data["name"] == "api_tool": + action_config = tool_data["config"]["actions"][action_name] + tool_config = { + "url": action_config["url"], + "method": action_config["method"], + "headers": {}, + "query_params": {}, + } + if "body_content_type" in action_config: + tool_config["body_content_type"] = action_config.get( + "body_content_type", "application/json" + ) + tool_config["body_encoding_rules"] = action_config.get( + "body_encoding_rules", {} + ) + else: + tool_config = tool_data["config"].copy() if tool_data["config"] else {} + if tool_config.get("encrypted_credentials") and self.user: + decrypted = decrypt_credentials( + tool_config["encrypted_credentials"], self.user + ) + tool_config.update(decrypted) + tool_config["auth_credentials"] = decrypted + tool_config.pop("encrypted_credentials", None) + tool_config["tool_id"] = str(tool_data.get("_id", tool_id)) + if self.conversation_id: + tool_config["conversation_id"] = self.conversation_id + if tool_data["name"] == "mcp_tool": + tool_config["query_mode"] = True + + tool = tm.load_tool( + tool_data["name"], + tool_config=tool_config, + user_id=self.user, + ) + + # Don't cache api_tool since config varies by action + if tool_data["name"] != "api_tool": + self._loaded_tools[cache_key] = tool + + return tool + + def get_truncated_tool_calls(self) -> List[Dict]: + return [ + { + "tool_name": tool_call.get("tool_name"), + "call_id": tool_call.get("call_id"), + "action_name": tool_call.get("action_name"), + "arguments": tool_call.get("arguments"), + "artifact_id": tool_call.get("artifact_id"), + "result": ( + f"{str(tool_call['result'])[:50]}..." + if len(str(tool_call["result"])) > 50 + else tool_call["result"] + ), + "status": "completed", + } + for tool_call in self.tool_calls + ] diff --git a/application/agents/workflows/node_agent.py b/application/agents/workflows/node_agent.py index d3c9d607..c26f6464 100644 --- a/application/agents/workflows/node_agent.py +++ b/application/agents/workflows/node_agent.py @@ -4,7 +4,6 @@ from typing import Any, Dict, List, Optional, Type from application.agents.base import BaseAgent from application.agents.classic_agent import ClassicAgent -from application.agents.react_agent import ReActAgent from application.agents.workflows.schemas import AgentType @@ -57,32 +56,11 @@ class WorkflowNodeClassicAgent(ToolFilterMixin, ClassicAgent): self._allowed_tool_ids = tool_ids or [] -class WorkflowNodeReActAgent(ToolFilterMixin, ReActAgent): - - def __init__( - self, - endpoint: str, - llm_name: str, - model_id: str, - api_key: str, - tool_ids: Optional[List[str]] = None, - **kwargs, - ): - super().__init__( - endpoint=endpoint, - llm_name=llm_name, - model_id=model_id, - api_key=api_key, - **kwargs, - ) - self._allowed_tool_ids = tool_ids or [] - - class WorkflowNodeAgentFactory: _agents: Dict[AgentType, Type[BaseAgent]] = { AgentType.CLASSIC: WorkflowNodeClassicAgent, - AgentType.REACT: WorkflowNodeReActAgent, + AgentType.REACT: WorkflowNodeClassicAgent, # backwards compat } @classmethod diff --git a/application/api/answer/routes/answer.py b/application/api/answer/routes/answer.py index b90ffa15..f3111605 100644 --- a/application/api/answer/routes/answer.py +++ b/application/api/answer/routes/answer.py @@ -74,21 +74,10 @@ class AnswerResource(Resource, BaseAnswerResource): decoded_token = getattr(request, "decoded_token", None) processor = StreamProcessor(data, decoded_token) try: - processor.initialize() + agent = processor.build_agent(data.get("question", "")) if not processor.decoded_token: return make_response({"error": "Unauthorized"}, 401) - docs_together, docs_list = processor.pre_fetch_docs( - data.get("question", "") - ) - tools_data = processor.pre_fetch_tools() - - agent = processor.create_agent( - docs_together=docs_together, - docs=docs_list, - tools_data=tools_data, - ) - if error := self.check_usage(processor.agent_config): return error diff --git a/application/api/answer/routes/stream.py b/application/api/answer/routes/stream.py index d1a71b25..c2cc0ec6 100644 --- a/application/api/answer/routes/stream.py +++ b/application/api/answer/routes/stream.py @@ -80,7 +80,7 @@ class StreamResource(Resource, BaseAnswerResource): decoded_token = getattr(request, "decoded_token", None) processor = StreamProcessor(data, decoded_token) try: - processor.initialize() + agent = processor.build_agent(data["question"]) if not processor.decoded_token: return Response( self.error_stream_generate("Unauthorized"), @@ -88,13 +88,6 @@ class StreamResource(Resource, BaseAnswerResource): mimetype="text/event-stream", ) - docs_together, docs_list = processor.pre_fetch_docs(data["question"]) - tools_data = processor.pre_fetch_tools() - - agent = processor.create_agent( - docs_together=docs_together, docs=docs_list, tools_data=tools_data - ) - if error := self.check_usage(processor.agent_config): return error return Response( diff --git a/application/api/answer/services/stream_processor.py b/application/api/answer/services/stream_processor.py index 0cf47094..61a4d4c1 100644 --- a/application/api/answer/services/stream_processor.py +++ b/application/api/answer/services/stream_processor.py @@ -91,6 +91,7 @@ class StreamProcessor: self.is_shared_usage = False self.shared_token = None self.agent_id = self.data.get("agent_id") + self.agent_key = None self.model_id: Optional[str] = None self.conversation_service = ConversationService() self.compression_orchestrator = CompressionOrchestrator( @@ -111,6 +112,21 @@ class StreamProcessor: self._load_conversation_history() self._process_attachments() + def build_agent(self, question: str): + """One call to go from request data to a ready-to-run agent. + + Combines initialize(), pre_fetch_docs(), pre_fetch_tools(), and + create_agent() into a single convenience method. + """ + self.initialize() + docs_together, docs_list = self.pre_fetch_docs(question) + tools_data = self.pre_fetch_tools() + return self.create_agent( + docs_together=docs_together, + docs=docs_list, + tools_data=tools_data, + ) + def _load_conversation_history(self): """Load conversation history either from DB or request""" if self.conversation_id and self.initial_user_id: @@ -135,14 +151,8 @@ class StreamProcessor: ) def _handle_compression(self, conversation: Dict[str, Any]): - """ - Handle conversation compression logic using orchestrator. - - Args: - conversation: Full conversation document - """ + """Handle conversation compression logic using orchestrator.""" try: - # Use orchestrator to handle all compression logic result = self.compression_orchestrator.compress_if_needed( conversation_id=self.conversation_id, user_id=self.initial_user_id, @@ -158,7 +168,6 @@ class StreamProcessor: ] return - # Set compressed summary if compression was performed if result.compression_performed and result.compressed_summary: self.compressed_summary = result.compressed_summary self.compressed_summary_tokens = TokenCounter.count_message_tokens( @@ -169,7 +178,6 @@ class StreamProcessor: f"+ {len(result.recent_queries)} recent messages" ) - # Build history from recent queries self.history = result.as_history() except Exception as e: @@ -177,7 +185,6 @@ class StreamProcessor: f"Error handling compression, falling back to standard history: {str(e)}", exc_info=True, ) - # Fallback to original behavior self.history = [ {"prompt": query["prompt"], "response": query["response"]} for query in conversation.get("queries", []) @@ -191,9 +198,6 @@ class StreamProcessor: ) def _get_attachments_content(self, attachment_ids, user_id): - """ - Retrieve content from attachment documents based on their IDs. - """ if not attachment_ids: return [] attachments = [] @@ -202,7 +206,6 @@ class StreamProcessor: attachment_doc = self.attachments_collection.find_one( {"_id": ObjectId(attachment_id), "user": user_id} ) - if attachment_doc: attachments.append(attachment_doc) except Exception as e: @@ -232,7 +235,6 @@ class StreamProcessor: ) self.model_id = requested_model else: - # Check if agent has a default model configured agent_default_model = self.agent_config.get("default_model_id", "") if agent_default_model and validate_model_id(agent_default_model): self.model_id = agent_default_model @@ -285,7 +287,6 @@ class StreamProcessor: data["source"] = "default" else: data["source"] = None - # Handle multiple sources sources = data.get("sources", []) if sources and isinstance(sources, list): @@ -311,7 +312,6 @@ class StreamProcessor: else: data["sources"] = [] - # Preserve model configuration from agent data["default_model_id"] = data.get("default_model_id", "") return data @@ -376,7 +376,10 @@ class StreamProcessor: return None def _configure_agent(self): - """Configure the agent based on request data""" + """Configure the agent based on request data. + + Unified flow: resolve the effective API key, then extract config once. + """ agent_id = self._resolve_agent_id() self.agent_key, self.is_shared_usage, self.shared_token = self._get_agent_key( @@ -384,55 +387,36 @@ class StreamProcessor: ) self.agent_id = str(agent_id) if agent_id else None - api_key = self.data.get("api_key") - if api_key: - data_key = self._get_data_from_api_key(api_key) + # Determine the effective API key (explicit > agent-derived) + effective_key = self.data.get("api_key") or self.agent_key + + if effective_key: + data_key = self._get_data_from_api_key(effective_key) if data_key.get("_id"): self.agent_id = str(data_key.get("_id")) + self.agent_config.update( { "prompt_id": data_key.get("prompt_id", "default"), "agent_type": data_key.get("agent_type", settings.AGENT_NAME), - "user_api_key": api_key, + "user_api_key": effective_key, "json_schema": data_key.get("json_schema"), "default_model_id": data_key.get("default_model_id", ""), } ) - self.initial_user_id = data_key.get("user") - self.decoded_token = {"sub": data_key.get("user")} - if data_key.get("source"): - self.source = {"active_docs": data_key["source"]} - if data_key.get("workflow"): - self.agent_config["workflow"] = data_key["workflow"] - self.agent_config["workflow_owner"] = data_key.get("user") - if data_key.get("retriever"): - self.retriever_config["retriever_name"] = data_key["retriever"] - if data_key.get("chunks") is not None: - try: - self.retriever_config["chunks"] = int(data_key["chunks"]) - except (ValueError, TypeError): - logger.warning( - f"Invalid chunks value: {data_key['chunks']}, using default value 2" - ) - self.retriever_config["chunks"] = 2 - elif self.agent_key: - data_key = self._get_data_from_api_key(self.agent_key) - if data_key.get("_id"): - self.agent_id = str(data_key.get("_id")) - self.agent_config.update( - { - "prompt_id": data_key.get("prompt_id", "default"), - "agent_type": data_key.get("agent_type", settings.AGENT_NAME), - "user_api_key": self.agent_key, - "json_schema": data_key.get("json_schema"), - "default_model_id": data_key.get("default_model_id", ""), - } - ) - self.decoded_token = ( - self.decoded_token - if self.is_shared_usage - else {"sub": data_key.get("user")} - ) + + # Set identity context + if self.data.get("api_key"): + # External API key: use the key owner's identity + self.initial_user_id = data_key.get("user") + self.decoded_token = {"sub": data_key.get("user")} + elif self.is_shared_usage: + # Shared agent: keep the caller's identity + pass + else: + # Owner using their own agent + self.decoded_token = {"sub": data_key.get("user")} + if data_key.get("source"): self.source = {"active_docs": data_key["source"]} if data_key.get("workflow"): @@ -449,6 +433,7 @@ class StreamProcessor: ) self.retriever_config["chunks"] = 2 else: + # No API key — default/workflow configuration agent_type = settings.AGENT_NAME if self.data.get("workflow") and isinstance( self.data.get("workflow"), dict @@ -531,12 +516,7 @@ class StreamProcessor: return None, None def pre_fetch_tools(self) -> Optional[Dict[str, Any]]: - """Pre-fetch tool data for template rendering before agent creation - - Can be controlled via: - 1. Global setting: ENABLE_TOOL_PREFETCH in .env - 2. Per-request: disable_tool_prefetch in request data - """ + """Pre-fetch tool data for template rendering before agent creation""" if not settings.ENABLE_TOOL_PREFETCH: logger.info( "Tool pre-fetching disabled globally via ENABLE_TOOL_PREFETCH setting" @@ -781,6 +761,31 @@ class StreamProcessor: 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 + from application.agents.tool_executor import ToolExecutor + + llm = LLMCreator.create_llm( + provider or settings.LLM_PROVIDER, + api_key=system_api_key, + user_api_key=self.agent_config["user_api_key"], + decoded_token=self.decoded_token, + model_id=self.model_id, + agent_id=self.agent_id, + ) + llm_handler = LLMHandlerCreator.create_handler( + provider if provider else "default" + ) + + user = self.decoded_token.get("sub") if self.decoded_token else None + tool_executor = ToolExecutor( + user_api_key=self.agent_config["user_api_key"], + user=user, + decoded_token=self.decoded_token, + ) + tool_executor.conversation_id = self.conversation_id + # Base agent kwargs agent_kwargs = { "endpoint": "stream", @@ -796,6 +801,9 @@ class StreamProcessor: "attachments": self.attachments, "json_schema": self.agent_config.get("json_schema"), "compressed_summary": self.compressed_summary, + "llm": llm, + "llm_handler": llm_handler, + "tool_executor": tool_executor, } # Workflow-specific kwargs for workflow agents diff --git a/application/app.py b/application/app.py index bb9b3b1b..aed069e0 100644 --- a/application/app.py +++ b/application/app.py @@ -72,6 +72,11 @@ def home(): return "Welcome to DocsGPT Backend!" +@app.route("/api/health") +def health(): + return jsonify({"status": "ok"}) + + @app.route("/api/config") def get_config(): response = { diff --git a/application/prompts/react_final_prompt.txt b/application/prompts/react_final_prompt.txt deleted file mode 100644 index 50916127..00000000 --- a/application/prompts/react_final_prompt.txt +++ /dev/null @@ -1,3 +0,0 @@ -Query: {query} -Observations: {observations} -Now, using the insights from the observations, formulate a well-structured and precise final answer. \ No newline at end of file diff --git a/application/prompts/react_planning_prompt.txt b/application/prompts/react_planning_prompt.txt deleted file mode 100644 index c3571845..00000000 --- a/application/prompts/react_planning_prompt.txt +++ /dev/null @@ -1,13 +0,0 @@ -You are an AI assistant and talk like you're thinking out loud. Given the following query, outline a concise thought process that includes key steps and considerations necessary for effective analysis and response. Avoid pointwise formatting. The goal is to break down the query into manageable components without excessive detail, focusing on clarity and logical progression. - -Include the following elements in your thought and execution process: -1. Identify the main objective of the query. -2. Determine any relevant context or background information needed to understand the query. -3. List potential approaches or methods to address the query. -4. Highlight any critical factors or constraints that may influence the outcome. -5. Plan with available tools to help you with the analysis but dont execute them. Tools will be executed by another AI. - -Query: {query} -Summaries: {summaries} -Prompt: {prompt} -Observations(potentially previous tool calls): {observations} diff --git a/application/seed/config/premade_agents.yaml b/application/seed/config/premade_agents.yaml index cd895bec..df404156 100644 --- a/application/seed/config/premade_agents.yaml +++ b/application/seed/config/premade_agents.yaml @@ -18,7 +18,7 @@ agents: - name: "Researcher" description: "A specialized research agent that performs deep dives into subjects." image: "https://d3dg1063dc54p9.cloudfront.net/imgs/agents/agent-researcher.svg" - agent_type: "react" + agent_type: "classic" prompt: name: "Researcher-Agent" content: | diff --git a/frontend/src/agents/NewAgent.tsx b/frontend/src/agents/NewAgent.tsx index c9a5e55d..d451d9ee 100644 --- a/frontend/src/agents/NewAgent.tsx +++ b/frontend/src/agents/NewAgent.tsx @@ -138,7 +138,6 @@ 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: t('agents.form.agentTypes.react'), value: 'react' }, ]; const isPublishable = () => { diff --git a/frontend/src/agents/WorkflowBuilder.tsx b/frontend/src/agents/WorkflowBuilder.tsx index 96a142f9..38a2ac67 100644 --- a/frontend/src/agents/WorkflowBuilder.tsx +++ b/frontend/src/agents/WorkflowBuilder.tsx @@ -55,7 +55,7 @@ import WorkflowPreview from './workflow/WorkflowPreview'; import type { Model } from '../models/types'; interface AgentNodeConfig { - agent_type: 'classic' | 'react'; + agent_type: 'classic'; llm_name?: string; model_id?: string; system_prompt: string; @@ -884,7 +884,6 @@ function WorkflowBuilderInner() { Classic - ReAct diff --git a/frontend/src/agents/workflow/WorkflowBuilder.tsx b/frontend/src/agents/workflow/WorkflowBuilder.tsx index 12c9a070..845145be 100644 --- a/frontend/src/agents/workflow/WorkflowBuilder.tsx +++ b/frontend/src/agents/workflow/WorkflowBuilder.tsx @@ -75,7 +75,7 @@ import type { Model } from '../../models/types'; const PRIMARY_ACTION_SPINNER_DELAY_MS = 180; interface AgentNodeConfig { - agent_type: 'classic' | 'react'; + agent_type: 'classic'; llm_name?: string; model_id?: string; system_prompt: string; @@ -1748,9 +1748,6 @@ function WorkflowBuilderInner() { Classic - - ReAct - diff --git a/results.txt b/results.txt new file mode 100644 index 00000000..beff8056 --- /dev/null +++ b/results.txt @@ -0,0 +1,12 @@ +Base URL:http://petstore.swagger.io,https://api.example.com +Path1: /pets +description: None +parameters: [] +methods: +get=A paged array of pets +post=Null response +Path2: /pets/{petId} +description: None +parameters: [] +methods: +get=Expected response to a valid request diff --git a/tests/agents/test_agent_creator.py b/tests/agents/test_agent_creator.py index 33eda73e..97f060b7 100644 --- a/tests/agents/test_agent_creator.py +++ b/tests/agents/test_agent_creator.py @@ -1,7 +1,6 @@ import pytest from application.agents.agent_creator import AgentCreator from application.agents.classic_agent import ClassicAgent -from application.agents.react_agent import ReActAgent @pytest.mark.unit @@ -14,11 +13,9 @@ class TestAgentCreator: assert agent.llm_name == agent_base_params["llm_name"] assert agent.model_id == agent_base_params["model_id"] - def test_create_react_agent(self, agent_base_params): + def test_create_react_falls_back_to_classic(self, agent_base_params): agent = AgentCreator.create_agent("react", **agent_base_params) - assert isinstance(agent, ReActAgent) - assert agent.endpoint == agent_base_params["endpoint"] - assert agent.llm_name == agent_base_params["llm_name"] + assert isinstance(agent, ClassicAgent) def test_create_agent_case_insensitive(self, agent_base_params): agent_upper = AgentCreator.create_agent("CLASSIC", **agent_base_params) @@ -35,7 +32,6 @@ class TestAgentCreator: assert "classic" in AgentCreator.agents assert "react" in AgentCreator.agents assert AgentCreator.agents["classic"] == ClassicAgent - assert AgentCreator.agents["react"] == ReActAgent def test_create_agent_with_optional_params(self, agent_base_params): agent_base_params["user_api_key"] = "user_key_123" diff --git a/tests/agents/test_react_agent.py b/tests/agents/test_react_agent.py deleted file mode 100644 index 0df38c3b..00000000 --- a/tests/agents/test_react_agent.py +++ /dev/null @@ -1,514 +0,0 @@ -from unittest.mock import Mock, mock_open, patch - -import pytest -from application.agents.react_agent import ReActAgent - - -@pytest.mark.unit -class TestReActAgent: - - def test_react_agent_initialization( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - assert isinstance(agent, ReActAgent) - assert agent.plan == "" - assert agent.observations == [] - - def test_react_agent_inherits_base_properties( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - assert agent.endpoint == agent_base_params["endpoint"] - assert agent.llm_name == agent_base_params["llm_name"] - assert agent.model_id == agent_base_params["model_id"] - - -@pytest.mark.unit -class TestReActAgentContentExtraction: - - def test_extract_content_from_string( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - response = "Simple string response" - content = agent._extract_content(response) - - assert content == "Simple string response" - - def test_extract_content_from_message_object( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - response = Mock() - response.message = Mock() - response.message.content = "Message content" - - content = agent._extract_content(response) - - assert content == "Message content" - - def test_extract_content_from_openai_response( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - response = Mock() - response.choices = [Mock()] - response.choices[0].message = Mock() - response.choices[0].message.content = "OpenAI content" - response.message = None - response.content = None - - content = agent._extract_content(response) - - assert content == "OpenAI content" - - def test_extract_content_from_anthropic_response( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - text_block = Mock() - text_block.text = "Anthropic content" - - response = Mock() - response.content = [text_block] - response.message = None - response.choices = None - - content = agent._extract_content(response) - - assert content == "Anthropic content" - - def test_extract_content_from_openai_stream( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - chunk1 = Mock() - chunk1.choices = [Mock()] - chunk1.choices[0].delta = Mock() - chunk1.choices[0].delta.content = "Part 1 " - - chunk2 = Mock() - chunk2.choices = [Mock()] - chunk2.choices[0].delta = Mock() - chunk2.choices[0].delta.content = "Part 2" - - response = iter([chunk1, chunk2]) - content = agent._extract_content(response) - - assert content == "Part 1 Part 2" - - def test_extract_content_from_anthropic_stream( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - chunk1 = Mock() - chunk1.type = "content_block_delta" - chunk1.delta = Mock() - chunk1.delta.text = "Stream 1 " - chunk1.choices = [] - - chunk2 = Mock() - chunk2.type = "content_block_delta" - chunk2.delta = Mock() - chunk2.delta.text = "Stream 2" - chunk2.choices = [] - - response = iter([chunk1, chunk2]) - content = agent._extract_content(response) - - assert content == "Stream 1 Stream 2" - - def test_extract_content_from_string_stream( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - response = iter(["chunk1", "chunk2", "chunk3"]) - content = agent._extract_content(response) - - assert content == "chunk1chunk2chunk3" - - def test_extract_content_handles_none_content( - self, agent_base_params, mock_llm_creator, mock_llm_handler_creator - ): - agent = ReActAgent(**agent_base_params) - - response = Mock() - response.message = Mock() - response.message.content = None - response.choices = None - response.content = None - - content = agent._extract_content(response) - - assert content == "" - - -@pytest.mark.unit -class TestReActAgentPlanning: - - @patch( - "builtins.open", - new_callable=mock_open, - read_data="Test planning prompt: {query} {summaries} {prompt} {observations}", - ) - def test_planning_phase( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_creator, - mock_llm_handler_creator, - log_context, - ): - def mock_gen_stream(*args, **kwargs): - # Return simple strings - _extract_content handles strings directly - - yield "Plan " - yield "content" - - mock_llm.gen_stream = Mock(return_value=mock_gen_stream()) - - agent = ReActAgent(**agent_base_params) - agent.observations = ["Observation 1"] - - plan_chunks = list(agent._planning_phase("Test query", log_context)) - - # Should yield thought dicts - - assert any("thought" in chunk for chunk in plan_chunks) - assert agent.plan == "Plan content" - - mock_llm.gen_stream.assert_called_once() - - @patch("builtins.open", new_callable=mock_open, read_data="Test: {query}") - def test_planning_phase_fills_template( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_creator, - mock_llm_handler_creator, - log_context, - ): - mock_llm.gen_stream = Mock(return_value=iter([])) - - agent = ReActAgent(**agent_base_params) - list(agent._planning_phase("My query", log_context)) - - call_args = mock_llm.gen_stream.call_args[1] - messages = call_args["messages"] - - assert "My query" in messages[0]["content"] - - -@pytest.mark.unit -class TestReActAgentFinalAnswer: - - @patch( - "builtins.open", - new_callable=mock_open, - read_data="Final answer for: {query} with {observations}", - ) - def test_synthesis_phase( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_creator, - mock_llm_handler_creator, - log_context, - ): - def mock_gen_stream(*args, **kwargs): - yield Mock(choices=[Mock(delta=Mock(content="Final "))]) - yield Mock(choices=[Mock(delta=Mock(content="answer"))]) - - mock_llm.gen_stream = Mock(return_value=mock_gen_stream()) - - agent = ReActAgent(**agent_base_params) - agent.observations = ["Obs 1", "Obs 2"] - - answer_chunks = list(agent._synthesis_phase("Test query", log_context)) - - # Should yield answer dicts - - assert any("answer" in chunk for chunk in answer_chunks) - - @patch("builtins.open", new_callable=mock_open, read_data="Answer: {observations}") - def test_synthesis_phase_truncates_long_observations( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_creator, - mock_llm_handler_creator, - log_context, - ): - mock_llm.gen_stream = Mock(return_value=iter([])) - - agent = ReActAgent(**agent_base_params) - agent.observations = ["A" * 15000] - - list(agent._synthesis_phase("Query", log_context)) - - call_args = mock_llm.gen_stream.call_args[1] - messages = call_args["messages"] - - assert "truncated" in messages[0]["content"] - - @patch("builtins.open", new_callable=mock_open, read_data="Test: {query}") - def test_synthesis_phase_no_tools( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_creator, - mock_llm_handler_creator, - log_context, - ): - mock_llm.gen_stream = Mock(return_value=iter([])) - - agent = ReActAgent(**agent_base_params) - agent.observations = ["Obs"] - list(agent._synthesis_phase("Query", log_context)) - - call_args = mock_llm.gen_stream.call_args[1] - - assert call_args["tools"] is None - - -@pytest.mark.unit -class TestReActAgentGenInner: - - @patch( - "builtins.open", new_callable=mock_open, read_data="Prompt template: {query}" - ) - def test_gen_inner_resets_state( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_handler, - mock_llm_creator, - mock_llm_handler_creator, - mock_mongo_db, - log_context, - ): - mock_llm.gen_stream = Mock(return_value=iter(["SATISFIED"])) - - def mock_handler(*args, **kwargs): - yield "SATISFIED" - - mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) - - agent = ReActAgent(**agent_base_params) - agent.plan = "Old plan" - agent.observations = ["Old obs"] - - list(agent._gen_inner("New query", log_context)) - - assert agent.plan != "Old plan" - assert len(agent.observations) > 0 - - @patch("builtins.open", new_callable=mock_open, read_data="Prompt: {query}") - def test_gen_inner_stops_on_satisfied( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_handler, - mock_llm_creator, - mock_llm_handler_creator, - mock_mongo_db, - log_context, - ): - iteration_count = 0 - - def mock_gen_stream(*args, **kwargs): - nonlocal iteration_count - iteration_count += 1 - if iteration_count == 1: - yield "Plan" - else: - yield "SATISFIED - done" - - mock_llm.gen_stream = Mock( - side_effect=lambda *args, **kwargs: mock_gen_stream(*args, **kwargs) - ) - - def mock_handler(*args, **kwargs): - yield "SATISFIED - done" - - mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) - - agent = ReActAgent(**agent_base_params) - results = list(agent._gen_inner("Test query", log_context)) - - assert any("answer" in r for r in results) - - @patch("builtins.open", new_callable=mock_open, read_data="Prompt: {query}") - def test_gen_inner_max_iterations( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_handler, - mock_llm_creator, - mock_llm_handler_creator, - mock_mongo_db, - log_context, - ): - call_count = 0 - - def mock_gen_stream(*args, **kwargs): - nonlocal call_count - call_count += 1 - yield f"Iteration {call_count}" - - mock_llm.gen_stream = Mock( - side_effect=lambda *args, **kwargs: mock_gen_stream(*args, **kwargs) - ) - - def mock_handler(*args, **kwargs): - yield "Continue..." - - mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) - - agent = ReActAgent(**agent_base_params) - - results = list(agent._gen_inner("Test query", log_context)) - - thought_results = [r for r in results if "thought" in r] - assert len(thought_results) > 0 - - @patch("builtins.open", new_callable=mock_open, read_data="Prompt: {query}") - def test_gen_inner_yields_sources( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_handler, - mock_llm_creator, - mock_llm_handler_creator, - mock_mongo_db, - log_context, - ): - mock_llm.gen_stream = Mock(return_value=iter(["SATISFIED"])) - - def mock_handler(*args, **kwargs): - yield "SATISFIED" - - mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) - - agent = ReActAgent(**agent_base_params) - results = list(agent._gen_inner("Test query", log_context)) - - sources = [r for r in results if "sources" in r] - assert len(sources) >= 1 - - @patch("builtins.open", new_callable=mock_open, read_data="Prompt: {query}") - def test_gen_inner_yields_tool_calls( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_handler, - mock_llm_creator, - mock_llm_handler_creator, - mock_mongo_db, - log_context, - ): - mock_llm.gen_stream = Mock(return_value=iter(["SATISFIED"])) - - def mock_handler(*args, **kwargs): - yield "SATISFIED" - - mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) - - agent = ReActAgent(**agent_base_params) - agent.tool_calls = [{"tool": "test", "result": "A" * 100}] - - results = list(agent._gen_inner("Test query", log_context)) - - tool_call_results = [r for r in results if "tool_calls" in r] - if tool_call_results: - assert len(tool_call_results[0]["tool_calls"][0]["result"]) <= 53 - - @patch("builtins.open", new_callable=mock_open, read_data="Prompt: {query}") - def test_gen_inner_logs_observations( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_handler, - mock_llm_creator, - mock_llm_handler_creator, - mock_mongo_db, - log_context, - ): - mock_llm.gen_stream = Mock(return_value=iter(["SATISFIED"])) - - def mock_handler(*args, **kwargs): - yield "SATISFIED" - - mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) - - agent = ReActAgent(**agent_base_params) - list(agent._gen_inner("Test query", log_context)) - - assert len(agent.observations) > 0 - - -@pytest.mark.integration -class TestReActAgentIntegration: - - @patch( - "builtins.open", - new_callable=mock_open, - read_data="Prompt: {query} {summaries} {prompt} {observations}", - ) - def test_full_react_workflow( - self, - mock_file, - agent_base_params, - mock_llm, - mock_llm_handler, - mock_llm_creator, - mock_llm_handler_creator, - mock_mongo_db, - log_context, - ): - call_sequence = [] - - def mock_gen_stream(*args, **kwargs): - call_sequence.append("gen_stream") - if len(call_sequence) <= 2: - yield "Planning..." - else: - yield "SATISFIED final answer" - - mock_llm.gen_stream = Mock( - side_effect=lambda *args, **kwargs: mock_gen_stream(*args, **kwargs) - ) - - def mock_handler(*args, **kwargs): - call_sequence.append("handler") - yield "SATISFIED final answer" - - mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) - - agent = ReActAgent(**agent_base_params) - results = list(agent._gen_inner("Complex query", log_context)) - - assert len(results) > 0 - assert any("thought" in r for r in results) - assert any("answer" in r for r in results) diff --git a/tests/conftest.py b/tests/conftest.py index 2355808e..229f455e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -188,7 +188,7 @@ def mock_tool_manager(mock_tool, monkeypatch): manager = Mock() manager.load_tool = Mock(return_value=mock_tool) monkeypatch.setattr( - "application.agents.base.ToolManager", Mock(return_value=manager) + "application.agents.tool_executor.ToolManager", Mock(return_value=manager) ) return manager From 556b0a1da570e94f0b9a9c58efa6046c38014f8e Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 15:16:18 +0000 Subject: [PATCH 2/7] 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, From 72393dc369a64af96979aea5ec034830f2f59177 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 17:42:24 +0000 Subject: [PATCH 3/7] feat: improve research --- application/agents/research_agent.py | 111 +++++++++++++++++- application/api/answer/routes/base.py | 7 +- .../answer/services/conversation_service.py | 33 ++++-- .../api/answer/services/stream_processor.py | 32 ++++- .../prompts/research/clarification.txt | 23 ++++ application/prompts/research/planning.txt | 2 + .../src/conversation/conversationSlice.ts | 16 ++- 7 files changed, 200 insertions(+), 24 deletions(-) create mode 100644 application/prompts/research/clarification.txt 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)); From e8d11fdfa6038f5f3ebdef08b7ddffcde33ee98f Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 19:21:46 +0000 Subject: [PATCH 4/7] feat: list files internal tool --- application/agents/agentic_agent.py | 47 +++- application/agents/research_agent.py | 11 +- application/agents/tools/internal_search.py | 270 ++++++++++++++++++-- 3 files changed, 301 insertions(+), 27 deletions(-) diff --git a/application/agents/agentic_agent.py b/application/agents/agentic_agent.py index f642402a..c8af7a2a 100644 --- a/application/agents/agentic_agent.py +++ b/application/agents/agentic_agent.py @@ -3,9 +3,9 @@ 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, + build_internal_tool_entry, ) from application.logging import LogContext @@ -39,9 +39,13 @@ class AgenticAgent(BaseAgent): 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) + has_dir = _sources_have_directory_structure(source) + internal_entry = build_internal_tool_entry( + has_directory_structure=has_dir + ) internal_entry["config"] = build_internal_tool_config( - **self.retriever_config + **self.retriever_config, + has_directory_structure=has_dir, ) tools_dict[INTERNAL_TOOL_ID] = internal_entry @@ -74,3 +78,40 @@ class AgenticAgent(BaseAgent): 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 + + +def _sources_have_directory_structure(source: Dict) -> bool: + """Check if any of the active sources have directory_structure in MongoDB.""" + active_docs = source.get("active_docs", []) + if not active_docs: + return False + + try: + from bson.objectid import ObjectId + from application.core.mongo_db import MongoDB + + mongo = MongoDB.get_client() + db = mongo[settings.MONGO_DB_NAME] + sources_collection = db["sources"] + + if isinstance(active_docs, str): + active_docs = [active_docs] + + for doc_id in active_docs: + try: + source_doc = sources_collection.find_one( + {"_id": ObjectId(doc_id)}, + {"directory_structure": 1}, + ) + if source_doc and source_doc.get("directory_structure"): + return True + except Exception: + continue + except Exception as e: + logger.debug(f"Could not check directory structure: {e}") + + return False + + +# Import settings at module level for _sources_have_directory_structure +from application.core.settings import settings # noqa: E402 diff --git a/application/agents/research_agent.py b/application/agents/research_agent.py index 1e4f9a73..b743cb16 100644 --- a/application/agents/research_agent.py +++ b/application/agents/research_agent.py @@ -6,10 +6,11 @@ from typing import Dict, Generator, List, Optional from application.agents.base import BaseAgent from application.agents.tool_executor import ToolExecutor +from application.agents.agentic_agent import _sources_have_directory_structure from application.agents.tools.internal_search import ( - INTERNAL_TOOL_ENTRY, INTERNAL_TOOL_ID, build_internal_tool_config, + build_internal_tool_entry, ) from application.agents.tools.think import THINK_TOOL_ENTRY, THINK_TOOL_ID from application.logging import LogContext @@ -275,9 +276,13 @@ class ResearchAgent(BaseAgent): 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) + has_dir = _sources_have_directory_structure(source) + internal_entry = build_internal_tool_entry( + has_directory_structure=has_dir + ) internal_entry["config"] = build_internal_tool_config( - **self.retriever_config + **self.retriever_config, + has_directory_structure=has_dir, ) tools_dict[INTERNAL_TOOL_ID] = internal_entry elif self.retriever_config and not has_sources: diff --git a/application/agents/tools/internal_search.py b/application/agents/tools/internal_search.py index 0b8c2c47..18b9c514 100644 --- a/application/agents/tools/internal_search.py +++ b/application/agents/tools/internal_search.py @@ -1,3 +1,4 @@ +import json import logging from typing import Dict, List, Optional @@ -13,19 +14,25 @@ class InternalSearchTool(Tool): Instead of pre-fetching docs into the prompt, the LLM decides when and what to search. Supports multiple searches per session. + + Optional capabilities (enabled when sources have directory_structure): + - path_filter on search: restrict results to a specific file/folder + - list_files action: browse the file/folder structure """ def __init__(self, config: Dict): self.config = config self.retrieved_docs: List[Dict] = [] self._retriever = None + self._directory_structure: Optional[Dict] = None + self._dir_structure_loaded = False 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 + chat_history=[], prompt="", chunks=int(self.config.get("chunks", 2)), doc_token_limit=int(self.config.get("doc_token_limit", 50000)), @@ -38,11 +45,65 @@ class InternalSearchTool(Tool): ) return self._retriever - def execute_action(self, action_name: str, **kwargs): - if action_name != "search": - return f"Unknown action: {action_name}" + def _get_directory_structure(self) -> Optional[Dict]: + """Load directory structure from MongoDB for the configured sources.""" + if self._dir_structure_loaded: + return self._directory_structure + self._dir_structure_loaded = True + source = self.config.get("source", {}) + active_docs = source.get("active_docs", []) + if not active_docs: + return None + + try: + from bson.objectid import ObjectId + from application.core.mongo_db import MongoDB + + mongo = MongoDB.get_client() + db = mongo[settings.MONGO_DB_NAME] + sources_collection = db["sources"] + + if isinstance(active_docs, str): + active_docs = [active_docs] + + merged_structure = {} + for doc_id in active_docs: + try: + source_doc = sources_collection.find_one( + {"_id": ObjectId(doc_id)} + ) + if not source_doc: + continue + dir_str = source_doc.get("directory_structure") + if dir_str: + if isinstance(dir_str, str): + dir_str = json.loads(dir_str) + source_name = source_doc.get("name", doc_id) + if len(active_docs) > 1: + merged_structure[source_name] = dir_str + else: + merged_structure = dir_str + except Exception as e: + logger.debug(f"Could not load dir structure for {doc_id}: {e}") + + self._directory_structure = merged_structure if merged_structure else None + except Exception as e: + logger.debug(f"Failed to load directory structures: {e}") + + return self._directory_structure + + def execute_action(self, action_name: str, **kwargs): + if action_name == "search": + return self._execute_search(**kwargs) + elif action_name == "list_files": + return self._execute_list_files(**kwargs) + return f"Unknown action: {action_name}" + + def _execute_search(self, **kwargs) -> str: query = kwargs.get("query", "") + path_filter = kwargs.get("path_filter", "") + if not query: return "Error: 'query' parameter is required." @@ -56,6 +117,19 @@ class InternalSearchTool(Tool): if not docs: return "No documents found matching your query." + # Apply path filter if specified + if path_filter: + path_lower = path_filter.lower() + docs = [ + d + for d in docs + if path_lower in d.get("source", "").lower() + or path_lower in d.get("filename", "").lower() + or path_lower in d.get("title", "").lower() + ] + if not docs: + return f"No documents found matching query '{query}' in path '{path_filter}'." + # Accumulate for source tracking for doc in docs: if doc not in self.retrieved_docs: @@ -73,8 +147,81 @@ class InternalSearchTool(Tool): return "\n\n---\n\n".join(formatted) + def _execute_list_files(self, **kwargs) -> str: + path = kwargs.get("path", "") + dir_structure = self._get_directory_structure() + + if not dir_structure: + return "No file structure available for the current sources." + + # Navigate to the requested path + current = dir_structure + if path: + for part in path.strip("/").split("/"): + if not part: + continue + if isinstance(current, dict) and part in current: + current = current[part] + else: + return f"Path '{path}' not found in the file structure." + + # Format the structure for the LLM + return self._format_structure(current, path or "/") + + def _format_structure(self, node: Dict, current_path: str) -> str: + if not isinstance(node, dict): + return f"'{current_path}' is a file, not a directory." + + lines = [f"File structure at '{current_path}':\n"] + folders = [] + files = [] + + for name, value in sorted(node.items()): + if isinstance(value, dict): + # Check if it's a file metadata dict or a folder + if "type" in value or "size_bytes" in value or "token_count" in value: + # It's a file with metadata + size = value.get("token_count", "") + ftype = value.get("type", "") + info_parts = [] + if ftype: + info_parts.append(ftype) + if size: + info_parts.append(f"{size} tokens") + info = f" ({', '.join(info_parts)})" if info_parts else "" + files.append(f" {name}{info}") + else: + # It's a folder + count = self._count_files(value) + folders.append(f" {name}/ ({count} items)") + else: + files.append(f" {name}") + + if folders: + lines.append("Folders:") + lines.extend(folders) + if files: + lines.append("Files:") + lines.extend(files) + if not folders and not files: + lines.append(" (empty)") + + return "\n".join(lines) + + def _count_files(self, node: Dict) -> int: + count = 0 + for value in node.values(): + if isinstance(value, dict): + if "type" in value or "size_bytes" in value or "token_count" in value: + count += 1 + else: + count += self._count_files(value) + else: + count += 1 + return count + def get_actions_metadata(self): - return [ + actions = [ { "name": "search", "description": ( @@ -89,12 +236,47 @@ class InternalSearchTool(Tool): "description": "The search query. Be specific and focused.", "filled_by_llm": True, "required": True, - } + }, } }, } ] + # Add path_filter and list_files only if directory structure exists + has_structure = self.config.get("has_directory_structure", False) + if has_structure: + actions[0]["parameters"]["properties"]["path_filter"] = { + "type": "string", + "description": ( + "Optional: filter results to a specific file or folder path. " + "Use list_files first to see available paths." + ), + "filled_by_llm": True, + "required": False, + } + actions.append( + { + "name": "list_files", + "description": ( + "Browse the file and folder structure of the knowledge base. " + "Use this to see what files are available before searching. " + "Optionally provide a path to browse a specific folder." + ), + "parameters": { + "properties": { + "path": { + "type": "string", + "description": "Optional: folder path to browse. Leave empty for root.", + "filled_by_llm": True, + "required": False, + } + } + }, + } + ) + + return actions + def get_config_requirements(self): return {} @@ -102,9 +284,25 @@ class InternalSearchTool(Tool): # Constants for building synthetic tools_dict entries INTERNAL_TOOL_ID = "internal" -INTERNAL_TOOL_ENTRY = { - "name": "internal_search", - "actions": [ + +def build_internal_tool_entry(has_directory_structure: bool = False) -> Dict: + """Build the tools_dict entry for InternalSearchTool. + + Dynamically includes list_files and path_filter based on + whether the sources have directory structure. + """ + search_params = { + "properties": { + "query": { + "type": "string", + "description": "The search query. Be specific and focused.", + "filled_by_llm": True, + "required": True, + } + } + } + + actions = [ { "name": "search", "description": ( @@ -113,19 +311,47 @@ INTERNAL_TOOL_ENTRY = { "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, - } - } - }, + "parameters": search_params, } - ], -} + ] + + if has_directory_structure: + search_params["properties"]["path_filter"] = { + "type": "string", + "description": ( + "Optional: filter results to a specific file or folder path. " + "Use list_files first to see available paths." + ), + "filled_by_llm": True, + "required": False, + } + actions.append( + { + "name": "list_files", + "description": ( + "Browse the file and folder structure of the knowledge base. " + "Use this to see what files are available before searching. " + "Optionally provide a path to browse a specific folder." + ), + "active": True, + "parameters": { + "properties": { + "path": { + "type": "string", + "description": "Optional: folder path to browse. Leave empty for root.", + "filled_by_llm": True, + "required": False, + } + } + }, + } + ) + + return {"name": "internal_search", "actions": actions} + + +# Keep backward compat +INTERNAL_TOOL_ENTRY = build_internal_tool_entry(has_directory_structure=False) def build_internal_tool_config( @@ -139,6 +365,7 @@ def build_internal_tool_config( llm_name: str = None, api_key: str = None, decoded_token: Optional[Dict] = None, + has_directory_structure: bool = False, ) -> Dict: """Build the config dict for InternalSearchTool.""" return { @@ -152,4 +379,5 @@ def build_internal_tool_config( "llm_name": llm_name or settings.LLM_PROVIDER, "api_key": api_key or settings.API_KEY, "decoded_token": decoded_token, + "has_directory_structure": has_directory_structure, } From c6ece177cd8d6ec0b4aad26262ac559bc507c7c3 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 20:04:03 +0000 Subject: [PATCH 5/7] fix: small issues --- .gitignore | 1 + application/agents/tools/internal_search.py | 2 +- results.txt | 12 ------------ 3 files changed, 2 insertions(+), 13 deletions(-) delete mode 100644 results.txt diff --git a/.gitignore b/.gitignore index 9b09303d..4d8f8c84 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ __pycache__/ *.py[cod] *$py.class +results.txt experiments/ experiments diff --git a/application/agents/tools/internal_search.py b/application/agents/tools/internal_search.py index 18b9c514..66f91e83 100644 --- a/application/agents/tools/internal_search.py +++ b/application/agents/tools/internal_search.py @@ -112,7 +112,7 @@ class InternalSearchTool(Tool): docs = retriever.search(query) except Exception as e: logger.error(f"Internal search failed: {e}", exc_info=True) - return f"Search failed: {str(e)}" + return "Search failed: an internal error occurred." if not docs: return "No documents found matching your query." diff --git a/results.txt b/results.txt deleted file mode 100644 index beff8056..00000000 --- a/results.txt +++ /dev/null @@ -1,12 +0,0 @@ -Base URL:http://petstore.swagger.io,https://api.example.com -Path1: /pets -description: None -parameters: [] -methods: -get=A paged array of pets -post=Null response -Path2: /pets/{petId} -description: None -parameters: [] -methods: -get=Expected response to a valid request From 462f2e9494f765957fe5d276dde65dbad8d2a26a Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 22:34:25 +0000 Subject: [PATCH 6/7] mini refactors --- application/agents/agentic_agent.py | 58 +- application/agents/research_agent.py | 34 +- application/agents/tools/internal_search.py | 53 ++ application/agents/workflows/node_agent.py | 19 +- application/agents/workflows/schemas.py | 2 + .../agents/workflows/workflow_engine.py | 39 +- application/prompts/agentic/creative.txt | 2 +- application/prompts/agentic/default.txt | 2 +- application/prompts/agentic/strict.txt | 2 +- .../prompts/research/clarification.txt | 2 +- application/prompts/research/planning.txt | 2 +- application/prompts/research/step.txt | 2 +- application/prompts/research/synthesis.txt | 2 +- frontend/src/agents/WorkflowBuilder.tsx | 8 +- .../src/agents/workflow/WorkflowBuilder.tsx | 8 +- tests/agents/test_agentic_agent.py | 139 +++++ tests/agents/test_internal_search_tool.py | 250 +++++++++ tests/agents/test_research_agent.py | 400 ++++++++++++++ tests/agents/test_think_tool.py | 52 ++ tests/agents/test_tool_executor.py | 279 ++++++++++ tests/agents/test_workflow_agent_types.py | 475 +++++++++++++++++ tests/integration/test_chat.py | 501 ++++++++++++++++++ 22 files changed, 2233 insertions(+), 98 deletions(-) create mode 100644 tests/agents/test_agentic_agent.py create mode 100644 tests/agents/test_internal_search_tool.py create mode 100644 tests/agents/test_research_agent.py create mode 100644 tests/agents/test_think_tool.py create mode 100644 tests/agents/test_tool_executor.py create mode 100644 tests/agents/test_workflow_agent_types.py diff --git a/application/agents/agentic_agent.py b/application/agents/agentic_agent.py index c8af7a2a..fe08277a 100644 --- a/application/agents/agentic_agent.py +++ b/application/agents/agentic_agent.py @@ -4,8 +4,7 @@ from typing import Dict, Generator, Optional from application.agents.base import BaseAgent from application.agents.tools.internal_search import ( INTERNAL_TOOL_ID, - build_internal_tool_config, - build_internal_tool_entry, + add_internal_search_tool, ) from application.logging import LogContext @@ -32,24 +31,8 @@ class AgenticAgent(BaseAgent): 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: - has_dir = _sources_have_directory_structure(source) - internal_entry = build_internal_tool_entry( - has_directory_structure=has_dir - ) - internal_entry["config"] = build_internal_tool_config( - **self.retriever_config, - has_directory_structure=has_dir, - ) - tools_dict[INTERNAL_TOOL_ID] = internal_entry - - # 3. Prepare all tools for the LLM + add_internal_search_tool(tools_dict, self.retriever_config) self._prepare_tools(tools_dict) # 4. Build messages (prompt has NO pre-fetched docs) @@ -78,40 +61,3 @@ class AgenticAgent(BaseAgent): 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 - - -def _sources_have_directory_structure(source: Dict) -> bool: - """Check if any of the active sources have directory_structure in MongoDB.""" - active_docs = source.get("active_docs", []) - if not active_docs: - return False - - try: - from bson.objectid import ObjectId - from application.core.mongo_db import MongoDB - - mongo = MongoDB.get_client() - db = mongo[settings.MONGO_DB_NAME] - sources_collection = db["sources"] - - if isinstance(active_docs, str): - active_docs = [active_docs] - - for doc_id in active_docs: - try: - source_doc = sources_collection.find_one( - {"_id": ObjectId(doc_id)}, - {"directory_structure": 1}, - ) - if source_doc and source_doc.get("directory_structure"): - return True - except Exception: - continue - except Exception as e: - logger.debug(f"Could not check directory structure: {e}") - - return False - - -# Import settings at module level for _sources_have_directory_structure -from application.core.settings import settings # noqa: E402 diff --git a/application/agents/research_agent.py b/application/agents/research_agent.py index b743cb16..280fa2cd 100644 --- a/application/agents/research_agent.py +++ b/application/agents/research_agent.py @@ -6,11 +6,9 @@ from typing import Dict, Generator, List, Optional from application.agents.base import BaseAgent from application.agents.tool_executor import ToolExecutor -from application.agents.agentic_agent import _sources_have_directory_structure from application.agents.tools.internal_search import ( INTERNAL_TOOL_ID, - build_internal_tool_config, - build_internal_tool_entry, + add_internal_search_tool, ) from application.agents.tools.think import THINK_TOOL_ENTRY, THINK_TOOL_ID from application.logging import LogContext @@ -130,6 +128,7 @@ class ResearchAgent(BaseAgent): self.citations = CitationManager() self._start_time: float = 0 self._tokens_used: int = 0 + self._last_token_snapshot: int = 0 # ------------------------------------------------------------------ # Budget & timeout helpers @@ -153,7 +152,9 @@ class ResearchAgent(BaseAgent): 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 + delta = current - self._last_token_snapshot + self._last_token_snapshot = current + return delta # ------------------------------------------------------------------ # Main orchestration @@ -272,21 +273,7 @@ class ResearchAgent(BaseAgent): """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: - has_dir = _sources_have_directory_structure(source) - internal_entry = build_internal_tool_entry( - has_directory_structure=has_dir - ) - internal_entry["config"] = build_internal_tool_config( - **self.retriever_config, - has_directory_structure=has_dir, - ) - 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") + add_internal_search_tool(tools_dict, self.retriever_config) think_entry = dict(THINK_TOOL_ENTRY) think_entry["config"] = {} @@ -580,7 +567,14 @@ class ResearchAgent(BaseAgent): call_id = None while True: try: - next(gen) + event = next(gen) + # Log tool_call status events instead of discarding them + if isinstance(event, dict) and event.get("type") == "tool_call": + logger.debug( + "Tool %s status: %s", + event.get("data", {}).get("action_name", ""), + event.get("data", {}).get("status", ""), + ) except StopIteration as e: result, call_id = e.value break diff --git a/application/agents/tools/internal_search.py b/application/agents/tools/internal_search.py index 66f91e83..2cd7915b 100644 --- a/application/agents/tools/internal_search.py +++ b/application/agents/tools/internal_search.py @@ -354,6 +354,59 @@ def build_internal_tool_entry(has_directory_structure: bool = False) -> Dict: INTERNAL_TOOL_ENTRY = build_internal_tool_entry(has_directory_structure=False) +def sources_have_directory_structure(source: Dict) -> bool: + """Check if any of the active sources have directory_structure in MongoDB.""" + active_docs = source.get("active_docs", []) + if not active_docs: + return False + + try: + from bson.objectid import ObjectId + from application.core.mongo_db import MongoDB + + mongo = MongoDB.get_client() + db = mongo[settings.MONGO_DB_NAME] + sources_collection = db["sources"] + + if isinstance(active_docs, str): + active_docs = [active_docs] + + for doc_id in active_docs: + try: + source_doc = sources_collection.find_one( + {"_id": ObjectId(doc_id)}, + {"directory_structure": 1}, + ) + if source_doc and source_doc.get("directory_structure"): + return True + except Exception: + continue + except Exception as e: + logger.debug(f"Could not check directory structure: {e}") + + return False + + +def add_internal_search_tool(tools_dict: Dict, retriever_config: Dict) -> None: + """Add the internal search tool to tools_dict if sources are configured. + + Shared by AgenticAgent and ResearchAgent to avoid duplicate setup logic. + Mutates tools_dict in place. + """ + source = retriever_config.get("source", {}) + has_sources = bool(source.get("active_docs")) + if not retriever_config or not has_sources: + return + + has_dir = sources_have_directory_structure(source) + internal_entry = build_internal_tool_entry(has_directory_structure=has_dir) + internal_entry["config"] = build_internal_tool_config( + **retriever_config, + has_directory_structure=has_dir, + ) + tools_dict[INTERNAL_TOOL_ID] = internal_entry + + def build_internal_tool_config( source: Dict, retriever_name: str = "classic", diff --git a/application/agents/workflows/node_agent.py b/application/agents/workflows/node_agent.py index c26f6464..67437e45 100644 --- a/application/agents/workflows/node_agent.py +++ b/application/agents/workflows/node_agent.py @@ -2,8 +2,10 @@ from typing import Any, Dict, List, Optional, Type +from application.agents.agentic_agent import AgenticAgent from application.agents.base import BaseAgent from application.agents.classic_agent import ClassicAgent +from application.agents.research_agent import ResearchAgent from application.agents.workflows.schemas import AgentType @@ -35,7 +37,8 @@ class ToolFilterMixin: return filtered_tools -class WorkflowNodeClassicAgent(ToolFilterMixin, ClassicAgent): +class _WorkflowNodeMixin: + """Common __init__ for all workflow node agents.""" def __init__( self, @@ -56,11 +59,25 @@ class WorkflowNodeClassicAgent(ToolFilterMixin, ClassicAgent): self._allowed_tool_ids = tool_ids or [] +class WorkflowNodeClassicAgent(ToolFilterMixin, _WorkflowNodeMixin, ClassicAgent): + pass + + +class WorkflowNodeAgenticAgent(ToolFilterMixin, _WorkflowNodeMixin, AgenticAgent): + pass + + +class WorkflowNodeResearchAgent(ToolFilterMixin, _WorkflowNodeMixin, ResearchAgent): + pass + + class WorkflowNodeAgentFactory: _agents: Dict[AgentType, Type[BaseAgent]] = { AgentType.CLASSIC: WorkflowNodeClassicAgent, AgentType.REACT: WorkflowNodeClassicAgent, # backwards compat + AgentType.AGENTIC: WorkflowNodeAgenticAgent, + AgentType.RESEARCH: WorkflowNodeResearchAgent, } @classmethod diff --git a/application/agents/workflows/schemas.py b/application/agents/workflows/schemas.py index 5355b88e..2a5bc79e 100644 --- a/application/agents/workflows/schemas.py +++ b/application/agents/workflows/schemas.py @@ -18,6 +18,8 @@ class NodeType(str, Enum): class AgentType(str, Enum): CLASSIC = "classic" REACT = "react" + AGENTIC = "agentic" + RESEARCH = "research" class ExecutionStatus(str, Enum): diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index 5444458a..00d471da 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -7,6 +7,7 @@ from application.agents.workflows.cel_evaluator import CelEvaluationError, evalu from application.agents.workflows.node_agent import WorkflowNodeAgentFactory from application.agents.workflows.schemas import ( AgentNodeConfig, + AgentType, ConditionNodeConfig, ExecutionStatus, NodeExecutionLog, @@ -223,18 +224,32 @@ class WorkflowEngine: f'Model "{node_model_id}" does not support structured output for node "{node.title}"' ) - node_agent = WorkflowNodeAgentFactory.create( - agent_type=node_config.agent_type, - endpoint=self.agent.endpoint, - llm_name=node_llm_name, - model_id=node_model_id, - api_key=node_api_key, - tool_ids=node_config.tools, - prompt=node_config.system_prompt, - chat_history=self.agent.chat_history, - decoded_token=self.agent.decoded_token, - json_schema=node_json_schema, - ) + factory_kwargs = { + "agent_type": node_config.agent_type, + "endpoint": self.agent.endpoint, + "llm_name": node_llm_name, + "model_id": node_model_id, + "api_key": node_api_key, + "tool_ids": node_config.tools, + "prompt": node_config.system_prompt, + "chat_history": self.agent.chat_history, + "decoded_token": self.agent.decoded_token, + "json_schema": node_json_schema, + } + + # Agentic/research agents need retriever_config for on-demand search + if node_config.agent_type in (AgentType.AGENTIC, AgentType.RESEARCH): + factory_kwargs["retriever_config"] = { + "source": {"active_docs": node_config.sources} if node_config.sources else {}, + "retriever_name": node_config.retriever or "classic", + "chunks": int(node_config.chunks) if node_config.chunks else 2, + "model_id": node_model_id, + "llm_name": node_llm_name, + "api_key": node_api_key, + "decoded_token": self.agent.decoded_token, + } + + node_agent = WorkflowNodeAgentFactory.create(**factory_kwargs) full_response_parts: List[str] = [] structured_response_parts: List[str] = [] diff --git a/application/prompts/agentic/creative.txt b/application/prompts/agentic/creative.txt index da2f8ff2..a1360b76 100644 --- a/application/prompts/agentic/creative.txt +++ b/application/prompts/agentic/creative.txt @@ -13,4 +13,4 @@ Use the search_internal tool to find relevant information before answering quest 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 +Allow yourself to be very creative and use your imagination. diff --git a/application/prompts/agentic/default.txt b/application/prompts/agentic/default.txt index 41339fe6..a2118414 100644 --- a/application/prompts/agentic/default.txt +++ b/application/prompts/agentic/default.txt @@ -12,4 +12,4 @@ You have access to a search tool that searches the user's uploaded documents and 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 +If no relevant documents are found, use your general knowledge and tool capabilities. diff --git a/application/prompts/agentic/strict.txt b/application/prompts/agentic/strict.txt index 8878b323..dd327918 100644 --- a/application/prompts/agentic/strict.txt +++ b/application/prompts/agentic/strict.txt @@ -13,4 +13,4 @@ Use the search_internal tool to find relevant information before answering quest 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 +Never make up information or provide false information! diff --git a/application/prompts/research/clarification.txt b/application/prompts/research/clarification.txt index 000c73c2..499d46e3 100644 --- a/application/prompts/research/clarification.txt +++ b/application/prompts/research/clarification.txt @@ -20,4 +20,4 @@ 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 8263bcaf..74c84a94 100644 --- a/application/prompts/research/planning.txt +++ b/application/prompts/research/planning.txt @@ -20,4 +20,4 @@ You MUST respond with ONLY a valid JSON object in this exact format (no markdown "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 index aecaf2b8..c0bf48a0 100644 --- a/application/prompts/research/step.txt +++ b/application/prompts/research/step.txt @@ -10,4 +10,4 @@ Instructions: 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 +Be thorough — prefer completeness over brevity. Include all relevant details you find. diff --git a/application/prompts/research/synthesis.txt b/application/prompts/research/synthesis.txt index 66bb1cad..30ea3958 100644 --- a/application/prompts/research/synthesis.txt +++ b/application/prompts/research/synthesis.txt @@ -19,4 +19,4 @@ Write a well-structured, thorough report that: Available sources for citation: {references} -Format the report with clear headings and sections. Be comprehensive but well-organized. \ No newline at end of file +Format the report with clear headings and sections. Be comprehensive but well-organized. diff --git a/frontend/src/agents/WorkflowBuilder.tsx b/frontend/src/agents/WorkflowBuilder.tsx index 38a2ac67..d3c7c8e1 100644 --- a/frontend/src/agents/WorkflowBuilder.tsx +++ b/frontend/src/agents/WorkflowBuilder.tsx @@ -55,7 +55,7 @@ import WorkflowPreview from './workflow/WorkflowPreview'; import type { Model } from '../models/types'; interface AgentNodeConfig { - agent_type: 'classic'; + agent_type: 'classic' | 'agentic' | 'research'; llm_name?: string; model_id?: string; system_prompt: string; @@ -884,6 +884,12 @@ function WorkflowBuilderInner() { Classic + + Agentic + + + Research +
diff --git a/frontend/src/agents/workflow/WorkflowBuilder.tsx b/frontend/src/agents/workflow/WorkflowBuilder.tsx index 845145be..5e765ea9 100644 --- a/frontend/src/agents/workflow/WorkflowBuilder.tsx +++ b/frontend/src/agents/workflow/WorkflowBuilder.tsx @@ -75,7 +75,7 @@ import type { Model } from '../../models/types'; const PRIMARY_ACTION_SPINNER_DELAY_MS = 180; interface AgentNodeConfig { - agent_type: 'classic'; + agent_type: 'classic' | 'agentic' | 'research'; llm_name?: string; model_id?: string; system_prompt: string; @@ -1748,6 +1748,12 @@ function WorkflowBuilderInner() { Classic + + Agentic + + + Research +
diff --git a/tests/agents/test_agentic_agent.py b/tests/agents/test_agentic_agent.py new file mode 100644 index 00000000..f742fffd --- /dev/null +++ b/tests/agents/test_agentic_agent.py @@ -0,0 +1,139 @@ +"""Tests for AgenticAgent — LLM-controlled retrieval agent.""" + +from unittest.mock import Mock + +import pytest +from application.agents.agentic_agent import AgenticAgent + + +@pytest.mark.unit +class TestAgenticAgentInit: + + def test_initialization( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = AgenticAgent(**agent_base_params) + assert isinstance(agent, AgenticAgent) + assert agent.retriever_config == {} + + def test_initialization_with_retriever_config( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + rc = {"source": {"active_docs": ["abc"]}, "retriever_name": "classic"} + agent = AgenticAgent(retriever_config=rc, **agent_base_params) + assert agent.retriever_config == rc + + def test_inherits_base_properties( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = AgenticAgent(**agent_base_params) + assert agent.endpoint == agent_base_params["endpoint"] + assert agent.llm_name == agent_base_params["llm_name"] + assert agent.model_id == agent_base_params["model_id"] + + +@pytest.mark.unit +class TestAgenticAgentGenInner: + + def test_basic_flow_yields_sources_and_tool_calls( + self, + agent_base_params, + mock_llm, + mock_llm_handler, + mock_llm_creator, + mock_llm_handler_creator, + mock_mongo_db, + log_context, + ): + mock_llm.gen_stream = Mock(return_value=iter(["Answer"])) + + def mock_handler(*args, **kwargs): + yield "Processed" + + mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) + + agent = AgenticAgent(**agent_base_params) + results = list(agent._gen_inner("Test query", log_context)) + + sources = [r for r in results if "sources" in r] + tool_calls = [r for r in results if "tool_calls" in r] + assert len(sources) == 1 + assert len(tool_calls) == 1 + + def test_logs_agent_component( + self, + agent_base_params, + mock_llm, + mock_llm_handler, + mock_llm_creator, + mock_llm_handler_creator, + mock_mongo_db, + log_context, + ): + mock_llm.gen_stream = Mock(return_value=iter(["Answer"])) + + def mock_handler(*args, **kwargs): + yield "Done" + + mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) + + agent = AgenticAgent(**agent_base_params) + list(agent._gen_inner("Query", log_context)) + + agent_logs = [s for s in log_context.stacks if s["component"] == "agent"] + assert len(agent_logs) == 1 + assert "tool_calls" in agent_logs[0]["data"] + + def test_no_pre_fetched_docs_in_messages( + self, + agent_base_params, + mock_llm, + mock_llm_handler, + mock_llm_creator, + mock_llm_handler_creator, + mock_mongo_db, + log_context, + ): + mock_llm.gen_stream = Mock(return_value=iter(["Answer"])) + + def mock_handler(*args, **kwargs): + yield "Done" + + mock_llm_handler.process_message_flow = Mock(side_effect=mock_handler) + + agent = AgenticAgent(**agent_base_params) + list(agent._gen_inner("Query", log_context)) + + call_kwargs = mock_llm.gen_stream.call_args[1] + messages = call_kwargs["messages"] + # System prompt should not contain {summaries} replacement + assert messages[0]["role"] == "system" + assert messages[-1]["role"] == "user" + assert messages[-1]["content"] == "Query" + + +@pytest.mark.unit +class TestAgenticAgentCollectSources: + + def test_collect_internal_sources_from_cache( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = AgenticAgent(**agent_base_params) + + mock_tool = Mock() + mock_tool.retrieved_docs = [ + {"text": "Found", "title": "Doc", "source": "test"}, + ] + cache_key = f"internal_search:internal:{agent.user or ''}" + agent.tool_executor._loaded_tools[cache_key] = mock_tool + + agent._collect_internal_sources() + assert len(agent.retrieved_docs) == 1 + assert agent.retrieved_docs[0]["title"] == "Doc" + + def test_collect_internal_sources_no_cache( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = AgenticAgent(**agent_base_params) + agent._collect_internal_sources() + assert agent.retrieved_docs == [] diff --git a/tests/agents/test_internal_search_tool.py b/tests/agents/test_internal_search_tool.py new file mode 100644 index 00000000..b9e4fce0 --- /dev/null +++ b/tests/agents/test_internal_search_tool.py @@ -0,0 +1,250 @@ +"""Tests for InternalSearchTool and its helper functions.""" + +from unittest.mock import Mock, patch + +import pytest +from application.agents.tools.internal_search import ( + INTERNAL_TOOL_ID, + InternalSearchTool, + add_internal_search_tool, + build_internal_tool_config, + build_internal_tool_entry, +) + + +@pytest.mark.unit +class TestInternalSearchToolSearch: + + def _make_tool(self, **config_overrides): + config = {"source": {}, "retriever_name": "classic", "chunks": 2} + config.update(config_overrides) + return InternalSearchTool(config) + + def test_search_no_query_returns_error(self): + tool = self._make_tool() + result = tool.execute_action("search", query="") + assert "required" in result.lower() + + def test_search_returns_formatted_docs(self): + tool = self._make_tool() + mock_retriever = Mock() + mock_retriever.search.return_value = [ + {"text": "Hello world", "title": "Doc1", "source": "test", "filename": "doc1.md"}, + ] + tool._retriever = mock_retriever + + result = tool.execute_action("search", query="hello") + assert "doc1.md" in result + assert "Hello world" in result + assert len(tool.retrieved_docs) == 1 + + def test_search_no_results(self): + tool = self._make_tool() + mock_retriever = Mock() + mock_retriever.search.return_value = [] + tool._retriever = mock_retriever + + result = tool.execute_action("search", query="nonexistent") + assert "No documents found" in result + + def test_search_accumulates_docs(self): + tool = self._make_tool() + mock_retriever = Mock() + tool._retriever = mock_retriever + + mock_retriever.search.return_value = [ + {"text": "A", "title": "D1", "source": "s1"}, + ] + tool.execute_action("search", query="first") + + mock_retriever.search.return_value = [ + {"text": "B", "title": "D2", "source": "s2"}, + ] + tool.execute_action("search", query="second") + + assert len(tool.retrieved_docs) == 2 + + def test_search_deduplicates_docs(self): + tool = self._make_tool() + doc = {"text": "Same", "title": "Same", "source": "same"} + mock_retriever = Mock() + mock_retriever.search.return_value = [doc] + tool._retriever = mock_retriever + + tool.execute_action("search", query="q1") + tool.execute_action("search", query="q2") + + assert len(tool.retrieved_docs) == 1 + + def test_search_with_path_filter(self): + tool = self._make_tool() + mock_retriever = Mock() + mock_retriever.search.return_value = [ + {"text": "A", "title": "T", "source": "src/main.py", "filename": "main.py"}, + {"text": "B", "title": "T", "source": "docs/readme.md", "filename": "readme.md"}, + ] + tool._retriever = mock_retriever + + result = tool.execute_action("search", query="code", path_filter="src/") + assert "main.py" in result + assert "readme.md" not in result + + def test_search_path_filter_no_match(self): + tool = self._make_tool() + mock_retriever = Mock() + mock_retriever.search.return_value = [ + {"text": "A", "title": "T", "source": "other/file.txt"}, + ] + tool._retriever = mock_retriever + + result = tool.execute_action("search", query="code", path_filter="src/") + assert "No documents found" in result + + def test_search_retriever_error(self): + tool = self._make_tool() + mock_retriever = Mock() + mock_retriever.search.side_effect = Exception("Connection error") + tool._retriever = mock_retriever + + result = tool.execute_action("search", query="test") + assert "failed" in result.lower() or "error" in result.lower() + + def test_unknown_action(self): + tool = self._make_tool() + result = tool.execute_action("nonexistent") + assert "Unknown action" in result + + +@pytest.mark.unit +class TestInternalSearchToolListFiles: + + def test_list_files_no_structure(self): + tool = InternalSearchTool({"source": {}}) + tool._dir_structure_loaded = True + tool._directory_structure = None + + result = tool.execute_action("list_files") + assert "No file structure" in result + + def test_list_files_root(self): + tool = InternalSearchTool({"source": {}}) + tool._dir_structure_loaded = True + tool._directory_structure = { + "src": {"main.py": {}}, + "README.md": {"type": "md", "token_count": 100}, + } + + result = tool.execute_action("list_files") + assert "src/" in result + assert "README.md" in result + + def test_list_files_nested_path(self): + tool = InternalSearchTool({"source": {}}) + tool._dir_structure_loaded = True + tool._directory_structure = { + "src": { + "utils": {"helper.py": {}}, + }, + } + + result = tool.execute_action("list_files", path="src") + assert "utils/" in result + + def test_list_files_invalid_path(self): + tool = InternalSearchTool({"source": {}}) + tool._dir_structure_loaded = True + tool._directory_structure = {"src": {}} + + result = tool.execute_action("list_files", path="nonexistent") + assert "not found" in result + + +@pytest.mark.unit +class TestInternalSearchToolMetadata: + + def test_actions_without_directory_structure(self): + tool = InternalSearchTool({"has_directory_structure": False}) + meta = tool.get_actions_metadata() + + action_names = [a["name"] for a in meta] + assert "search" in action_names + assert "list_files" not in action_names + + # search should not have path_filter + search = meta[0] + assert "path_filter" not in search["parameters"]["properties"] + + def test_actions_with_directory_structure(self): + tool = InternalSearchTool({"has_directory_structure": True}) + meta = tool.get_actions_metadata() + + action_names = [a["name"] for a in meta] + assert "search" in action_names + assert "list_files" in action_names + + # search should have path_filter + search = next(a for a in meta if a["name"] == "search") + assert "path_filter" in search["parameters"]["properties"] + + +@pytest.mark.unit +class TestBuildHelpers: + + def test_build_entry_without_directory_structure(self): + entry = build_internal_tool_entry(has_directory_structure=False) + assert entry["name"] == "internal_search" + action_names = [a["name"] for a in entry["actions"]] + assert "search" in action_names + assert "list_files" not in action_names + + def test_build_entry_with_directory_structure(self): + entry = build_internal_tool_entry(has_directory_structure=True) + action_names = [a["name"] for a in entry["actions"]] + assert "list_files" in action_names + + def test_build_config(self): + config = build_internal_tool_config( + source={"active_docs": ["abc"]}, + retriever_name="semantic", + chunks=4, + ) + assert config["source"] == {"active_docs": ["abc"]} + assert config["retriever_name"] == "semantic" + assert config["chunks"] == 4 + + def test_internal_tool_id(self): + assert INTERNAL_TOOL_ID == "internal" + + def test_add_internal_search_tool_with_sources(self): + tools_dict = {} + retriever_config = { + "source": {"active_docs": ["abc"]}, + "retriever_name": "classic", + "chunks": 2, + "model_id": "gpt-4", + "llm_name": "openai", + "api_key": "key", + } + + with patch( + "application.agents.tools.internal_search.sources_have_directory_structure", + return_value=False, + ): + add_internal_search_tool(tools_dict, retriever_config) + + assert INTERNAL_TOOL_ID in tools_dict + assert tools_dict[INTERNAL_TOOL_ID]["name"] == "internal_search" + assert "config" in tools_dict[INTERNAL_TOOL_ID] + + def test_add_internal_search_tool_no_sources(self): + tools_dict = {} + retriever_config = {"source": {}} + + add_internal_search_tool(tools_dict, retriever_config) + + assert INTERNAL_TOOL_ID not in tools_dict + + def test_add_internal_search_tool_empty_config(self): + tools_dict = {} + add_internal_search_tool(tools_dict, {}) + assert INTERNAL_TOOL_ID not in tools_dict diff --git a/tests/agents/test_research_agent.py b/tests/agents/test_research_agent.py new file mode 100644 index 00000000..3cc1e6d3 --- /dev/null +++ b/tests/agents/test_research_agent.py @@ -0,0 +1,400 @@ +"""Tests for ResearchAgent — multi-step research with budget controls.""" + +import json +import time +from unittest.mock import Mock, patch + +import pytest +from application.agents.research_agent import ( + CitationManager, + ResearchAgent, + DEFAULT_MAX_STEPS, + DEFAULT_TIMEOUT_SECONDS, + DEFAULT_TOKEN_BUDGET, +) + + +# --------------------------------------------------------------------------- +# CitationManager +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestCitationManager: + + def test_add_returns_citation_number(self): + cm = CitationManager() + num = cm.add({"source": "s1", "title": "T1"}) + assert num == 1 + + def test_add_deduplicates(self): + cm = CitationManager() + n1 = cm.add({"source": "s1", "title": "T1"}) + n2 = cm.add({"source": "s1", "title": "T1"}) + assert n1 == n2 + assert len(cm.citations) == 1 + + def test_add_different_sources(self): + cm = CitationManager() + n1 = cm.add({"source": "s1", "title": "T1"}) + n2 = cm.add({"source": "s2", "title": "T2"}) + assert n1 != n2 + assert len(cm.citations) == 2 + + def test_add_docs_returns_mapping(self): + cm = CitationManager() + docs = [ + {"source": "s1", "title": "Doc A"}, + {"source": "s2", "title": "Doc B"}, + ] + text = cm.add_docs(docs) + assert "[1] Doc A" in text + assert "[2] Doc B" in text + + def test_format_references(self): + cm = CitationManager() + cm.add({"source": "http://example.com", "title": "Example", "filename": "ex.md"}) + refs = cm.format_references() + assert "[1]" in refs + assert "ex.md" in refs + + def test_format_references_empty(self): + cm = CitationManager() + assert "No sources" in cm.format_references() + + def test_get_all_docs(self): + cm = CitationManager() + cm.add({"source": "s1", "title": "T1"}) + cm.add({"source": "s2", "title": "T2"}) + docs = cm.get_all_docs() + assert len(docs) == 2 + + +# --------------------------------------------------------------------------- +# ResearchAgent Init & Budget +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestResearchAgentInit: + + def test_initialization( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = ResearchAgent(**agent_base_params) + assert isinstance(agent, ResearchAgent) + assert agent.max_steps == DEFAULT_MAX_STEPS + assert agent.timeout_seconds == DEFAULT_TIMEOUT_SECONDS + assert agent.token_budget == DEFAULT_TOKEN_BUDGET + assert agent.retriever_config == {} + + def test_custom_budget( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = ResearchAgent( + max_steps=3, + timeout_seconds=60, + token_budget=50_000, + **agent_base_params, + ) + assert agent.max_steps == 3 + assert agent.timeout_seconds == 60 + assert agent.token_budget == 50_000 + + def test_with_retriever_config( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + rc = {"source": {"active_docs": ["abc"]}} + agent = ResearchAgent(retriever_config=rc, **agent_base_params) + assert agent.retriever_config == rc + + +@pytest.mark.unit +class TestResearchAgentBudget: + + def _make_agent(self, agent_base_params, mock_llm_creator, mock_llm_handler_creator, **kwargs): + return ResearchAgent(**kwargs, **agent_base_params) + + def test_timeout_detection( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent( + agent_base_params, mock_llm_creator, mock_llm_handler_creator, + timeout_seconds=0, + ) + agent._start_time = time.monotonic() - 1 + assert agent._is_timed_out() is True + + def test_not_timed_out( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent( + agent_base_params, mock_llm_creator, mock_llm_handler_creator, + timeout_seconds=300, + ) + agent._start_time = time.monotonic() + assert agent._is_timed_out() is False + + def test_token_budget_tracking( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent( + agent_base_params, mock_llm_creator, mock_llm_handler_creator, + token_budget=1000, + ) + agent._track_tokens(500) + assert agent._budget_remaining() == 500 + assert agent._is_over_budget() is False + + agent._track_tokens(500) + assert agent._budget_remaining() == 0 + assert agent._is_over_budget() is True + + def test_snapshot_llm_tokens_returns_delta( + self, agent_base_params, mock_llm, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent( + agent_base_params, mock_llm_creator, mock_llm_handler_creator, + ) + mock_llm.token_usage = {"prompt_tokens": 100, "generated_tokens": 50} + + delta1 = agent._snapshot_llm_tokens() + assert delta1 == 150 + + # Simulate more tokens used + mock_llm.token_usage = {"prompt_tokens": 200, "generated_tokens": 100} + delta2 = agent._snapshot_llm_tokens() + assert delta2 == 150 # 300 - 150 + + def test_elapsed( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent( + agent_base_params, mock_llm_creator, mock_llm_handler_creator, + ) + agent._start_time = time.monotonic() - 1.5 + elapsed = agent._elapsed() + assert elapsed >= 1.0 + + +# --------------------------------------------------------------------------- +# ResearchAgent Phases +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestResearchAgentClarification: + + def test_is_follow_up_no_history( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = ResearchAgent(**agent_base_params) + assert agent._is_follow_up() is False + + def test_is_follow_up_with_clarification_metadata( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent_base_params["chat_history"] = [ + {"prompt": "What?", "response": "Clarify", "metadata": {"is_clarification": True}}, + ] + agent = ResearchAgent(**agent_base_params) + assert agent._is_follow_up() is True + + def test_is_follow_up_without_metadata( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent_base_params["chat_history"] = [ + {"prompt": "What?", "response": "Normal answer"}, + ] + agent = ResearchAgent(**agent_base_params) + assert agent._is_follow_up() is False + + def test_clarification_returns_none_on_no_clarification_needed( + self, agent_base_params, mock_llm, mock_llm_creator, mock_llm_handler_creator + ): + response = Mock() + response.choices = [Mock()] + response.choices[0].message = Mock() + response.choices[0].message.content = json.dumps( + {"needs_clarification": False, "reason": "Clear enough"} + ) + mock_llm.gen = Mock(return_value=response) + mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5} + + agent = ResearchAgent(**agent_base_params) + result = agent._clarification_phase("What is Python?") + assert result is None + + def test_clarification_returns_questions( + self, agent_base_params, mock_llm, mock_llm_creator, mock_llm_handler_creator + ): + clarification_json = json.dumps({ + "needs_clarification": True, + "questions": ["Which version?", "What context?"], + }) + # Return a plain string so _extract_text handles it directly + mock_llm.gen = Mock(return_value=clarification_json) + mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5} + + agent = ResearchAgent(**agent_base_params) + result = agent._clarification_phase("Tell me about it") + assert result is not None + assert "Which version?" in result + assert "What context?" in result + + +@pytest.mark.unit +class TestResearchAgentPlanning: + + def test_planning_returns_steps_and_complexity( + self, agent_base_params, mock_llm, mock_llm_creator, mock_llm_handler_creator + ): + plan_json = json.dumps({ + "complexity": "moderate", + "steps": [ + {"query": "sub-question 1", "rationale": "reason 1"}, + {"query": "sub-question 2", "rationale": "reason 2"}, + ], + }) + # Return plain string so _extract_text handles it directly + mock_llm.gen = Mock(return_value=plan_json) + mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5} + + agent = ResearchAgent(**agent_base_params) + steps, complexity = agent._planning_phase("Compare A and B") + + assert complexity == "moderate" + assert len(steps) == 2 + assert steps[0]["query"] == "sub-question 1" + + def test_planning_caps_steps_by_complexity( + self, agent_base_params, mock_llm, mock_llm_creator, mock_llm_handler_creator + ): + plan_json = json.dumps({ + "complexity": "simple", + "steps": [ + {"query": f"q{i}", "rationale": f"r{i}"} for i in range(10) + ], + }) + response = Mock() + response.choices = [Mock()] + response.choices[0].message = Mock() + response.choices[0].message.content = plan_json + mock_llm.gen = Mock(return_value=response) + mock_llm.token_usage = {"prompt_tokens": 10, "generated_tokens": 5} + + agent = ResearchAgent(**agent_base_params) + steps, complexity = agent._planning_phase("Simple question") + + assert complexity == "simple" + assert len(steps) <= 2 # COMPLEXITY_CAPS["simple"] == 2 + + def test_planning_fallback_on_error( + self, agent_base_params, mock_llm, mock_llm_creator, mock_llm_handler_creator + ): + mock_llm.gen = Mock(side_effect=Exception("LLM down")) + mock_llm.token_usage = {"prompt_tokens": 0, "generated_tokens": 0} + + agent = ResearchAgent(**agent_base_params) + steps, complexity = agent._planning_phase("Anything") + + assert complexity == "simple" + assert len(steps) == 1 + assert steps[0]["query"] == "Anything" + + +@pytest.mark.unit +class TestResearchAgentExtractText: + + def _make_agent(self, agent_base_params, mock_llm_creator, mock_llm_handler_creator): + return ResearchAgent(**agent_base_params) + + def test_extract_from_string( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + assert agent._extract_text("hello") == "hello" + + def test_extract_from_openai_response( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + response = Mock() + response.choices = [Mock()] + response.choices[0].message = Mock() + response.choices[0].message.content = "OpenAI content" + response.message = None + response.content = None + assert agent._extract_text(response) == "OpenAI content" + + def test_extract_from_anthropic_response( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + text_block = Mock() + text_block.text = "Anthropic content" + response = Mock() + response.content = [text_block] + response.message = None + response.choices = None + assert agent._extract_text(response) == "Anthropic content" + + def test_extract_from_none( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + assert agent._extract_text(None) == "" + + +@pytest.mark.unit +class TestResearchAgentParseJson: + + def _make_agent(self, agent_base_params, mock_llm_creator, mock_llm_handler_creator): + return ResearchAgent(**agent_base_params) + + def test_parse_plan_direct_json( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + text = '{"steps": [{"query": "q1"}], "complexity": "simple"}' + result = agent._parse_plan_json(text) + assert isinstance(result, dict) + assert len(result["steps"]) == 1 + + def test_parse_plan_from_code_fence( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + text = 'Here is the plan:\n```json\n{"steps": [{"query": "q1"}]}\n```' + result = agent._parse_plan_json(text) + assert isinstance(result, dict) + + def test_parse_plan_invalid_returns_empty( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + result = agent._parse_plan_json("not json at all") + assert result == [] + + def test_parse_clarification_json( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + text = '{"needs_clarification": false, "reason": "clear"}' + result = agent._parse_clarification_json(text) + assert result["needs_clarification"] is False + + def test_parse_clarification_json_from_code_fence( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + text = '```json\n{"needs_clarification": true, "questions": ["q1"]}\n```' + result = agent._parse_clarification_json(text) + assert result["needs_clarification"] is True + + def test_parse_clarification_json_invalid( + self, agent_base_params, mock_llm_creator, mock_llm_handler_creator + ): + agent = self._make_agent(agent_base_params, mock_llm_creator, mock_llm_handler_creator) + result = agent._parse_clarification_json("not json") + assert result is None diff --git a/tests/agents/test_think_tool.py b/tests/agents/test_think_tool.py new file mode 100644 index 00000000..45f9e727 --- /dev/null +++ b/tests/agents/test_think_tool.py @@ -0,0 +1,52 @@ +"""Tests for ThinkTool — the chain-of-thought pseudo-tool.""" + +import pytest +from application.agents.tools.think import ( + THINK_TOOL_ENTRY, + THINK_TOOL_ID, + ThinkTool, +) + + +@pytest.mark.unit +class TestThinkTool: + + def test_id_constant(self): + assert THINK_TOOL_ID == "think" + + def test_entry_has_reason_action(self): + actions = THINK_TOOL_ENTRY["actions"] + assert len(actions) == 1 + assert actions[0]["name"] == "reason" + assert actions[0]["active"] is True + + def test_execute_reason_returns_continue(self): + tool = ThinkTool() + result = tool.execute_action("reason", reasoning="step by step thinking") + assert result == "Continue." + + def test_execute_unknown_action_returns_continue(self): + tool = ThinkTool() + result = tool.execute_action("unknown_action") + assert result == "Continue." + + def test_get_actions_metadata(self): + tool = ThinkTool() + meta = tool.get_actions_metadata() + assert len(meta) == 1 + assert meta[0]["name"] == "reason" + props = meta[0]["parameters"]["properties"] + assert "reasoning" in props + assert props["reasoning"]["filled_by_llm"] is True + + def test_get_config_requirements_empty(self): + tool = ThinkTool() + assert tool.get_config_requirements() == {} + + def test_init_accepts_no_config(self): + tool = ThinkTool() + assert tool is not None + + def test_init_accepts_config(self): + tool = ThinkTool(config={"key": "value"}) + assert tool is not None diff --git a/tests/agents/test_tool_executor.py b/tests/agents/test_tool_executor.py new file mode 100644 index 00000000..d995c22b --- /dev/null +++ b/tests/agents/test_tool_executor.py @@ -0,0 +1,279 @@ +"""Tests for ToolExecutor — tool discovery, preparation, and execution.""" + +from unittest.mock import Mock + +import pytest +from application.agents.tool_executor import ToolExecutor + + +@pytest.mark.unit +class TestToolExecutorInit: + + def test_default_state(self): + executor = ToolExecutor() + assert executor.user_api_key is None + assert executor.user is None + assert executor.tool_calls == [] + assert executor._loaded_tools == {} + assert executor.conversation_id is None + + def test_init_with_params(self): + executor = ToolExecutor( + user_api_key="key", user="alice", decoded_token={"sub": "alice"} + ) + assert executor.user_api_key == "key" + assert executor.user == "alice" + + +@pytest.mark.unit +class TestToolExecutorGetTools: + + def test_get_tools_uses_api_key_when_present(self, mock_mongo_db): + executor = ToolExecutor(user_api_key="test_key", user="alice") + tools = executor.get_tools() + assert isinstance(tools, dict) + + def test_get_tools_uses_user_when_no_api_key(self, mock_mongo_db): + executor = ToolExecutor(user="alice") + tools = executor.get_tools() + assert isinstance(tools, dict) + + def test_get_tools_defaults_to_local(self, mock_mongo_db): + executor = ToolExecutor() + tools = executor.get_tools() + assert isinstance(tools, dict) + + +@pytest.mark.unit +class TestToolExecutorPrepare: + + def test_prepare_tools_for_llm_empty(self): + executor = ToolExecutor() + result = executor.prepare_tools_for_llm({}) + assert result == [] + + def test_prepare_tools_for_llm_non_api_tool(self): + executor = ToolExecutor() + tools_dict = { + "t1": { + "name": "test_tool", + "actions": [ + { + "name": "do_thing", + "description": "Does a thing", + "active": True, + "parameters": { + "properties": { + "query": { + "type": "string", + "description": "The query", + "filled_by_llm": True, + "required": True, + } + } + }, + } + ], + } + } + + result = executor.prepare_tools_for_llm(tools_dict) + assert len(result) == 1 + assert result[0]["type"] == "function" + assert result[0]["function"]["name"] == "do_thing_t1" + assert "query" in result[0]["function"]["parameters"]["properties"] + + def test_prepare_tools_skips_inactive_actions(self): + executor = ToolExecutor() + tools_dict = { + "t1": { + "name": "test_tool", + "actions": [ + {"name": "active_one", "description": "D", "active": True, "parameters": {"properties": {}}}, + {"name": "inactive_one", "description": "D", "active": False, "parameters": {"properties": {}}}, + ], + } + } + + result = executor.prepare_tools_for_llm(tools_dict) + assert len(result) == 1 + assert result[0]["function"]["name"] == "active_one_t1" + + def test_build_tool_parameters_filters_non_llm_fields(self): + executor = ToolExecutor() + action = { + "parameters": { + "properties": { + "query": { + "type": "string", + "description": "Search query", + "filled_by_llm": True, + "value": "default_val", + "required": True, + }, + "hidden": { + "type": "string", + "filled_by_llm": False, + }, + } + } + } + + result = executor._build_tool_parameters(action) + assert "query" in result["properties"] + assert "hidden" not in result["properties"] + assert "query" in result["required"] + # filled_by_llm, value, required stripped from schema + assert "filled_by_llm" not in result["properties"]["query"] + assert "value" not in result["properties"]["query"] + + +@pytest.mark.unit +class TestToolExecutorExecute: + + def _make_call(self, name="action_toolid", call_id="c1", arguments="{}"): + call = Mock() + call.name = name + call.id = call_id + call.arguments = arguments + return call + + def test_execute_parse_failure(self, monkeypatch): + executor = ToolExecutor() + + monkeypatch.setattr( + "application.agents.tool_executor.ToolActionParser", + lambda _cls: Mock(parse_args=Mock(return_value=(None, None, {}))), + ) + + call = self._make_call(name="bad") + gen = executor.execute({}, call, "MockLLM") + + events = [] + result = None + while True: + try: + events.append(next(gen)) + except StopIteration as e: + result = e.value + break + + assert result[0] == "Failed to parse tool call." + assert len(executor.tool_calls) == 1 + assert events[0]["data"]["status"] == "error" + + def test_execute_tool_not_found(self, monkeypatch): + executor = ToolExecutor() + + monkeypatch.setattr( + "application.agents.tool_executor.ToolActionParser", + lambda _cls: Mock(parse_args=Mock(return_value=("missing_id", "action", {}))), + ) + + call = self._make_call() + gen = executor.execute({}, call, "MockLLM") + + events = [] + result = None + while True: + try: + events.append(next(gen)) + except StopIteration as e: + result = e.value + break + + assert "not found" in result[0] + assert events[0]["data"]["status"] == "error" + + def test_execute_success(self, mock_tool_manager, monkeypatch): + executor = ToolExecutor(user="test_user") + + monkeypatch.setattr( + "application.agents.tool_executor.ToolActionParser", + lambda _cls: Mock(parse_args=Mock(return_value=("t1", "test_action", {"param1": "val"}))), + ) + + tools_dict = { + "t1": { + "name": "test_tool", + "config": {"key": "val"}, + "actions": [ + {"name": "test_action", "description": "Test", "parameters": {"properties": {}}}, + ], + } + } + + call = self._make_call(name="test_action_t1", call_id="c1") + gen = executor.execute(tools_dict, call, "MockLLM") + + events = [] + result = None + while True: + try: + events.append(next(gen)) + except StopIteration as e: + result = e.value + break + + assert result[0] == "Tool result" + assert result[1] == "c1" + + statuses = [e["data"]["status"] for e in events] + assert "pending" in statuses + assert "completed" in statuses + + def test_get_truncated_tool_calls(self): + executor = ToolExecutor() + executor.tool_calls = [ + { + "tool_name": "test", + "call_id": "1", + "action_name": "act", + "arguments": {}, + "result": "A" * 100, + } + ] + + truncated = executor.get_truncated_tool_calls() + assert len(truncated) == 1 + assert len(truncated[0]["result"]) <= 53 + assert truncated[0]["status"] == "completed" + + def test_tool_caching(self, mock_tool_manager, monkeypatch): + executor = ToolExecutor(user="test_user") + + monkeypatch.setattr( + "application.agents.tool_executor.ToolActionParser", + lambda _cls: Mock(parse_args=Mock(return_value=("t1", "test_action", {}))), + ) + + tools_dict = { + "t1": { + "name": "test_tool", + "config": {"key": "val"}, + "actions": [ + {"name": "test_action", "description": "Test", "parameters": {"properties": {}}}, + ], + } + } + + call = self._make_call(name="test_action_t1") + + # First execution — loads tool + gen = executor.execute(tools_dict, call, "MockLLM") + while True: + try: + next(gen) + except StopIteration: + break + + # Second execution — should use cache + gen = executor.execute(tools_dict, call, "MockLLM") + while True: + try: + next(gen) + except StopIteration: + break + + # load_tool called only once due to cache + assert mock_tool_manager.load_tool.call_count == 1 diff --git a/tests/agents/test_workflow_agent_types.py b/tests/agents/test_workflow_agent_types.py new file mode 100644 index 00000000..4b751d17 --- /dev/null +++ b/tests/agents/test_workflow_agent_types.py @@ -0,0 +1,475 @@ +"""Tests for new agent types (agentic, research) in the workflow builder.""" + +from types import SimpleNamespace +from typing import Any, Dict + +import pytest + +from application.agents.agentic_agent import AgenticAgent +from application.agents.classic_agent import ClassicAgent +from application.agents.research_agent import ResearchAgent +from application.agents.workflows.node_agent import ( + WorkflowNodeAgenticAgent, + WorkflowNodeAgentFactory, + WorkflowNodeClassicAgent, + WorkflowNodeResearchAgent, +) +from application.agents.workflows.schemas import ( + AgentNodeConfig, + AgentType, + NodeType, + Workflow, + WorkflowGraph, + WorkflowNode, +) +from application.agents.workflows.workflow_engine import WorkflowEngine + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class StubNodeAgent: + """Minimal agent stub that yields pre-defined events.""" + + def __init__(self, events): + self.events = events + + def gen(self, _prompt): + yield from self.events + + +def create_engine() -> WorkflowEngine: + graph = WorkflowGraph(workflow=Workflow(name="Test"), nodes=[], edges=[]) + agent = SimpleNamespace( + endpoint="stream", + llm_name="openai", + model_id="gpt-4o-mini", + api_key="test-key", + chat_history=[], + decoded_token={"sub": "user-1"}, + ) + return WorkflowEngine(graph, agent) + + +def create_agent_node( + node_id: str, + agent_type: str = "classic", + sources: list = None, + chunks: str = "2", + retriever: str = "", + output_variable: str = "", +) -> WorkflowNode: + config: Dict[str, Any] = { + "agent_type": agent_type, + "system_prompt": "You are a helpful assistant.", + "prompt_template": "", + "stream_to_user": False, + "tools": [], + "sources": sources or [], + "chunks": chunks, + "retriever": retriever, + } + if output_variable: + config["output_variable"] = output_variable + return WorkflowNode( + id=node_id, + workflow_id="workflow-1", + type=NodeType.AGENT, + title="Agent", + position={"x": 0, "y": 0}, + config=config, + ) + + +# --------------------------------------------------------------------------- +# AgentType enum +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestAgentTypeEnum: + + def test_agentic_value_exists(self): + assert AgentType.AGENTIC == "agentic" + + def test_research_value_exists(self): + assert AgentType.RESEARCH == "research" + + def test_classic_still_exists(self): + assert AgentType.CLASSIC == "classic" + + def test_react_still_exists(self): + assert AgentType.REACT == "react" + + +# --------------------------------------------------------------------------- +# AgentNodeConfig schema validation +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestAgentNodeConfigValidation: + + def test_accepts_agentic_agent_type(self): + config = AgentNodeConfig(agent_type="agentic") + assert config.agent_type == AgentType.AGENTIC + + def test_accepts_research_agent_type(self): + config = AgentNodeConfig(agent_type="research") + assert config.agent_type == AgentType.RESEARCH + + def test_rejects_unknown_agent_type(self): + with pytest.raises(Exception): + AgentNodeConfig(agent_type="nonexistent") + + def test_default_agent_type_is_classic(self): + config = AgentNodeConfig() + assert config.agent_type == AgentType.CLASSIC + + +# --------------------------------------------------------------------------- +# WorkflowNodeAgentFactory registry +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestWorkflowNodeAgentFactoryRegistry: + + def test_factory_has_agentic(self): + assert AgentType.AGENTIC in WorkflowNodeAgentFactory._agents + assert WorkflowNodeAgentFactory._agents[AgentType.AGENTIC] is WorkflowNodeAgenticAgent + + def test_factory_has_research(self): + assert AgentType.RESEARCH in WorkflowNodeAgentFactory._agents + assert WorkflowNodeAgentFactory._agents[AgentType.RESEARCH] is WorkflowNodeResearchAgent + + def test_factory_raises_for_unknown_type(self): + with pytest.raises(ValueError, match="Unsupported agent type"): + WorkflowNodeAgentFactory.create( + agent_type="nonexistent", + endpoint="stream", + llm_name="openai", + model_id="gpt-4o-mini", + api_key="key", + ) + + +# --------------------------------------------------------------------------- +# WorkflowNode agent classes (inheritance) +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestWorkflowNodeAgentClasses: + + def test_agentic_agent_inherits_correctly(self): + assert issubclass(WorkflowNodeAgenticAgent, AgenticAgent) + + def test_research_agent_inherits_correctly(self): + assert issubclass(WorkflowNodeResearchAgent, ResearchAgent) + + def test_classic_agent_inherits_correctly(self): + assert issubclass(WorkflowNodeClassicAgent, ClassicAgent) + + +# --------------------------------------------------------------------------- +# Workflow engine: agentic agent node execution +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestWorkflowEngineAgenticNode: + + def test_agentic_node_executes_and_saves_output(self, monkeypatch): + engine = create_engine() + node = create_agent_node( + node_id="agent_agentic", + agent_type="agentic", + output_variable="result", + ) + node_events = [{"answer": "agentic answer"}] + + captured: Dict[str, Any] = {} + + def capture_create(**kwargs): + captured.update(kwargs) + return StubNodeAgent(node_events) + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(capture_create), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + assert engine.state["node_agent_agentic_output"] == "agentic answer" + assert engine.state["result"] == "agentic answer" + + def test_agentic_node_passes_retriever_config(self, monkeypatch): + engine = create_engine() + node = create_agent_node( + node_id="agent_rc", + agent_type="agentic", + sources=["source-abc"], + chunks="4", + retriever="semantic", + ) + node_events = [{"answer": "ok"}] + + captured: Dict[str, Any] = {} + + def capture_create(**kwargs): + captured.update(kwargs) + return StubNodeAgent(node_events) + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(capture_create), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + rc = captured.get("retriever_config") + assert rc is not None + assert rc["source"] == {"active_docs": ["source-abc"]} + assert rc["retriever_name"] == "semantic" + assert rc["chunks"] == 4 + + def test_agentic_node_empty_sources_gives_empty_source_dict(self, monkeypatch): + engine = create_engine() + node = create_agent_node( + node_id="agent_nosrc", + agent_type="agentic", + sources=[], + ) + node_events = [{"answer": "ok"}] + + captured: Dict[str, Any] = {} + + def capture_create(**kwargs): + captured.update(kwargs) + return StubNodeAgent(node_events) + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(capture_create), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + rc = captured["retriever_config"] + assert rc["source"] == {} + + +# --------------------------------------------------------------------------- +# Workflow engine: research agent node execution +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestWorkflowEngineResearchNode: + + def test_research_node_executes_and_saves_output(self, monkeypatch): + engine = create_engine() + node = create_agent_node( + node_id="agent_research", + agent_type="research", + output_variable="report", + ) + node_events = [{"answer": "research report"}] + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(lambda **kwargs: StubNodeAgent(node_events)), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + assert engine.state["node_agent_research_output"] == "research report" + assert engine.state["report"] == "research report" + + def test_research_node_passes_retriever_config(self, monkeypatch): + engine = create_engine() + node = create_agent_node( + node_id="agent_rr", + agent_type="research", + sources=["doc-1", "doc-2"], + chunks="6", + ) + node_events = [{"answer": "ok"}] + + captured: Dict[str, Any] = {} + + def capture_create(**kwargs): + captured.update(kwargs) + return StubNodeAgent(node_events) + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(capture_create), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + rc = captured["retriever_config"] + assert rc["source"] == {"active_docs": ["doc-1", "doc-2"]} + assert rc["chunks"] == 6 + assert rc["decoded_token"] == {"sub": "user-1"} + + def test_research_node_handles_structured_output(self, monkeypatch): + engine = create_engine() + node = create_agent_node( + node_id="agent_rs", + agent_type="research", + output_variable="data", + ) + # Simulate structured JSON output from research agent + node_events = [ + {"answer": '{"findings": "important"}', "structured": True}, + ] + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(lambda **kwargs: StubNodeAgent(node_events)), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + # structured=True causes the engine to parse JSON + assert engine.state["data"] == {"findings": "important"} + + +# --------------------------------------------------------------------------- +# Workflow engine: classic node does NOT get retriever_config +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestWorkflowEngineClassicNodeNoRetrieverConfig: + + def test_classic_node_does_not_pass_retriever_config(self, monkeypatch): + engine = create_engine() + node = create_agent_node( + node_id="agent_classic", + agent_type="classic", + sources=["some-source"], + ) + node_events = [{"answer": "classic answer"}] + + captured: Dict[str, Any] = {} + + def capture_create(**kwargs): + captured.update(kwargs) + return StubNodeAgent(node_events) + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(capture_create), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + list(engine._execute_agent_node(node)) + + assert "retriever_config" not in captured + + +# --------------------------------------------------------------------------- +# Workflow engine: streaming events from new agent types +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestWorkflowEngineStreamingEvents: + + def test_agentic_node_streams_answer_events(self, monkeypatch): + engine = create_engine() + node = create_agent_node(node_id="agent_s1", agent_type="agentic") + # Modify config to enable streaming + node.config["stream_to_user"] = True + + node_events = [ + {"answer": "chunk 1"}, + {"answer": "chunk 2"}, + ] + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(lambda **kwargs: StubNodeAgent(node_events)), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + results = list(engine._execute_agent_node(node)) + answer_events = [r for r in results if "answer" in r] + assert len(answer_events) == 2 + + def test_research_node_passes_through_non_answer_events(self, monkeypatch): + """Research agents yield research_plan/research_progress events. + The workflow engine only forwards 'answer' events to the user.""" + engine = create_engine() + node = create_agent_node(node_id="agent_s2", agent_type="research") + node.config["stream_to_user"] = True + + node_events = [ + {"type": "research_plan", "data": {"steps": [], "complexity": "simple"}}, + {"type": "research_progress", "data": {"status": "planning"}}, + {"answer": "final report"}, + ] + + monkeypatch.setattr( + WorkflowNodeAgentFactory, + "create", + staticmethod(lambda **kwargs: StubNodeAgent(node_events)), + ) + monkeypatch.setattr( + "application.core.model_utils.get_api_key_for_provider", + lambda _provider: None, + ) + + results = list(engine._execute_agent_node(node)) + # Only answer events are streamed to user + answer_events = [r for r in results if "answer" in r] + assert len(answer_events) == 1 + assert answer_events[0]["answer"] == "final report" + + # State still captures the full text + assert engine.state["node_agent_s2_output"] == "final report" diff --git a/tests/integration/test_chat.py b/tests/integration/test_chat.py index 350faaba..2651d169 100644 --- a/tests/integration/test_chat.py +++ b/tests/integration/test_chat.py @@ -492,6 +492,486 @@ This is test documentation for integration tests. self.record_result(test_name, False, str(e)) return False + # ------------------------------------------------------------------------- + # Agentic / Research Agent Tests + # ------------------------------------------------------------------------- + + def _create_agent_with_type(self, agent_type: str) -> Optional[tuple]: + """Create a test agent with the given agent_type. Returns (agent_id, api_key) or None.""" + if not self.is_authenticated: + return None + + payload = { + "name": f"Chat Test {agent_type.title()} Agent {int(time.time())}", + "description": f"Integration test {agent_type} agent", + "prompt_id": "default", + "chunks": 2, + "retriever": "classic", + "agent_type": agent_type, + "status": "draft", + } + + try: + response = self.post("/api/create_agent", json=payload, timeout=10) + if response.status_code in [200, 201]: + result = response.json() + agent_id = result.get("id") + api_key = result.get("key") + if agent_id: + return (agent_id, api_key) + except Exception: + pass + + return None + + def test_stream_agentic_agent(self) -> bool: + """Test /stream endpoint with an agentic agent.""" + test_name = "Stream endpoint (agentic agent)" + + agent_result = self._create_agent_with_type("agentic") + if not agent_result: + if not self.require_auth(test_name): + return True + self.record_result(test_name, True, "Skipped (no agent)") + return True + + agent_id, _ = agent_result + self.print_header(f"Testing {test_name}") + + payload = { + "question": "What is DocsGPT?", + "history": "[]", + "agent_id": agent_id, + } + + try: + self.print_info(f"POST /stream with agentic agent_id={agent_id[:8]}...") + + response = requests.post( + f"{self.base_url}/stream", + json=payload, + headers=self.headers, + stream=True, + timeout=60, + ) + + self.print_info(f"Status Code: {response.status_code}") + + if response.status_code != 200: + self.print_error(f"Expected 200, got {response.status_code}") + self.record_result(test_name, False, f"Status {response.status_code}") + return False + + events = [] + full_response = "" + for line in response.iter_lines(): + if line: + line_str = line.decode("utf-8") + if line_str.startswith("data: "): + try: + data = json_module.loads(line_str[6:]) + events.append(data) + if data.get("type") in ["stream", "answer"]: + full_response += data.get("message", "") or data.get("answer", "") + elif data.get("type") == "end": + break + except json_module.JSONDecodeError: + pass + + self.print_success(f"Received {len(events)} events") + if full_response: + self.print_success(f"Answer preview: {full_response[:100]}...") + self.record_result(test_name, True, "Success") + return True + + except Exception as e: + self.print_error(f"Error: {str(e)}") + self.record_result(test_name, False, str(e)) + return False + + def test_answer_agentic_agent(self) -> bool: + """Test /api/answer endpoint with an agentic agent.""" + test_name = "Answer endpoint (agentic agent)" + + agent_result = self._create_agent_with_type("agentic") + if not agent_result: + if not self.require_auth(test_name): + return True + self.record_result(test_name, True, "Skipped (no agent)") + return True + + agent_id, _ = agent_result + self.print_header(f"Testing {test_name}") + + payload = { + "question": "What is DocsGPT?", + "history": "[]", + "agent_id": agent_id, + } + + try: + self.print_info(f"POST /api/answer with agentic agent_id={agent_id[:8]}...") + + response = self.post("/api/answer", json=payload, timeout=60) + + self.print_info(f"Status Code: {response.status_code}") + + if response.status_code != 200: + self.print_error(f"Expected 200, got {response.status_code}") + self.record_result(test_name, False, f"Status {response.status_code}") + return False + + result = response.json() + answer = result.get("answer", "") + self.print_success(f"Answer received: {answer[:100]}...") + self.record_result(test_name, True, "Success") + return True + + except Exception as e: + self.print_error(f"Error: {str(e)}") + self.record_result(test_name, False, str(e)) + return False + + def test_stream_research_agent(self) -> bool: + """Test /stream endpoint with a research agent. + + Research agents emit additional SSE event types: + - research_plan: the decomposed research steps + - research_progress: per-step status updates + """ + test_name = "Stream endpoint (research agent)" + + agent_result = self._create_agent_with_type("research") + if not agent_result: + if not self.require_auth(test_name): + return True + self.record_result(test_name, True, "Skipped (no agent)") + return True + + agent_id, _ = agent_result + self.print_header(f"Testing {test_name}") + + payload = { + "question": "What is DocsGPT?", + "history": "[]", + "agent_id": agent_id, + } + + try: + self.print_info(f"POST /stream with research agent_id={agent_id[:8]}...") + + response = requests.post( + f"{self.base_url}/stream", + json=payload, + headers=self.headers, + stream=True, + timeout=120, + ) + + self.print_info(f"Status Code: {response.status_code}") + + if response.status_code != 200: + self.print_error(f"Expected 200, got {response.status_code}") + self.record_result(test_name, False, f"Status {response.status_code}") + return False + + events = [] + full_response = "" + saw_plan = False + saw_progress = False + for line in response.iter_lines(): + if line: + line_str = line.decode("utf-8") + if line_str.startswith("data: "): + try: + data = json_module.loads(line_str[6:]) + events.append(data) + event_type = data.get("type", "") + if event_type in ["stream", "answer"]: + full_response += data.get("message", "") or data.get("answer", "") + elif event_type == "research_plan": + saw_plan = True + steps = data.get("data", {}).get("steps", []) + self.print_info(f"Research plan: {len(steps)} steps, complexity={data.get('data', {}).get('complexity')}") + elif event_type == "research_progress": + saw_progress = True + elif event_type == "end": + break + except json_module.JSONDecodeError: + pass + + self.print_success(f"Received {len(events)} events") + if saw_plan: + self.print_success("Received research_plan event") + if saw_progress: + self.print_success("Received research_progress events") + if full_response: + self.print_success(f"Report preview: {full_response[:100]}...") + + self.record_result(test_name, True, "Success") + return True + + except Exception as e: + self.print_error(f"Error: {str(e)}") + self.record_result(test_name, False, str(e)) + return False + + def test_answer_research_agent(self) -> bool: + """Test /api/answer endpoint with a research agent.""" + test_name = "Answer endpoint (research agent)" + + agent_result = self._create_agent_with_type("research") + if not agent_result: + if not self.require_auth(test_name): + return True + self.record_result(test_name, True, "Skipped (no agent)") + return True + + agent_id, _ = agent_result + self.print_header(f"Testing {test_name}") + + payload = { + "question": "What is DocsGPT?", + "history": "[]", + "agent_id": agent_id, + } + + try: + self.print_info(f"POST /api/answer with research agent_id={agent_id[:8]}...") + + response = self.post("/api/answer", json=payload, timeout=120) + + self.print_info(f"Status Code: {response.status_code}") + + if response.status_code != 200: + self.print_error(f"Expected 200, got {response.status_code}") + self.record_result(test_name, False, f"Status {response.status_code}") + return False + + result = response.json() + answer = result.get("answer", "") + self.print_success(f"Answer received: {answer[:100]}...") + self.record_result(test_name, True, "Success") + return True + + except Exception as e: + self.print_error(f"Error: {str(e)}") + self.record_result(test_name, False, str(e)) + return False + + # ------------------------------------------------------------------------- + # Workflow with Agentic Node Tests + # ------------------------------------------------------------------------- + + def _create_workflow_with_agentic_node(self) -> Optional[str]: + """Create a workflow with a single agentic agent node. + + Returns the workflow_id or None on failure. + """ + nodes = [ + { + "id": "start_1", + "type": "start", + "title": "Start", + "data": {}, + }, + { + "id": "agent_1", + "type": "agent", + "title": "Agentic Node", + "data": { + "agent_type": "agentic", + "system_prompt": "You are a helpful assistant.", + "prompt_template": "", + "stream_to_user": True, + "tools": [], + "sources": [], + "chunks": "2", + }, + }, + { + "id": "end_1", + "type": "end", + "title": "End", + "data": {}, + }, + ] + edges = [ + {"id": "edge_1", "source": "start_1", "target": "agent_1"}, + {"id": "edge_2", "source": "agent_1", "target": "end_1"}, + ] + + payload = { + "name": f"Agentic Workflow Test {int(time.time())}", + "nodes": nodes, + "edges": edges, + } + + try: + response = self.post("/api/workflows", json=payload, timeout=15) + if response.status_code in [200, 201]: + result = response.json() + return result.get("id") + except Exception: + pass + + return None + + def _create_workflow_agent(self, workflow_id: str) -> Optional[tuple]: + """Create an agent of type 'workflow' referencing a workflow. + + Returns (agent_id, api_key) or None. + """ + payload = { + "name": f"Workflow Agent Test {int(time.time())}", + "description": "Integration test workflow agent with agentic node", + "prompt_id": "default", + "chunks": 2, + "retriever": "classic", + "agent_type": "workflow", + "workflow": workflow_id, + "status": "draft", + } + + try: + response = self.post("/api/create_agent", json=payload, timeout=10) + if response.status_code in [200, 201]: + result = response.json() + agent_id = result.get("id") + api_key = result.get("key") + if agent_id: + return (agent_id, api_key) + except Exception: + pass + + return None + + def test_stream_workflow_with_agentic_node(self) -> bool: + """Test /stream with a workflow agent that contains an agentic node.""" + test_name = "Stream endpoint (workflow with agentic node)" + + if not self.require_auth(test_name): + return True + + workflow_id = self._create_workflow_with_agentic_node() + if not workflow_id: + self.print_warning("Could not create workflow") + self.record_result(test_name, True, "Skipped (workflow creation failed)") + return True + + agent_result = self._create_workflow_agent(workflow_id) + if not agent_result: + self.print_warning("Could not create workflow agent") + self.record_result(test_name, True, "Skipped (agent creation failed)") + return True + + agent_id, _ = agent_result + self.print_header(f"Testing {test_name}") + + payload = { + "question": "What is DocsGPT?", + "history": "[]", + "agent_id": agent_id, + } + + try: + self.print_info(f"POST /stream with workflow agent_id={agent_id[:8]}...") + self.print_info(f"Workflow ID: {workflow_id}") + + response = requests.post( + f"{self.base_url}/stream", + json=payload, + headers=self.headers, + stream=True, + timeout=60, + ) + + self.print_info(f"Status Code: {response.status_code}") + + if response.status_code != 200: + self.print_error(f"Expected 200, got {response.status_code}") + self.record_result(test_name, False, f"Status {response.status_code}") + return False + + events = [] + full_response = "" + for line in response.iter_lines(): + if line: + line_str = line.decode("utf-8") + if line_str.startswith("data: "): + try: + data = json_module.loads(line_str[6:]) + events.append(data) + if data.get("type") in ["stream", "answer"]: + full_response += data.get("message", "") or data.get("answer", "") + elif data.get("type") == "end": + break + except json_module.JSONDecodeError: + pass + + self.print_success(f"Received {len(events)} events") + if full_response: + self.print_success(f"Answer preview: {full_response[:100]}...") + + self.record_result(test_name, True, "Success") + return True + + except Exception as e: + self.print_error(f"Error: {str(e)}") + self.record_result(test_name, False, str(e)) + return False + + def test_answer_workflow_with_agentic_node(self) -> bool: + """Test /api/answer with a workflow agent that contains an agentic node.""" + test_name = "Answer endpoint (workflow with agentic node)" + + if not self.require_auth(test_name): + return True + + workflow_id = self._create_workflow_with_agentic_node() + if not workflow_id: + self.print_warning("Could not create workflow") + self.record_result(test_name, True, "Skipped (workflow creation failed)") + return True + + agent_result = self._create_workflow_agent(workflow_id) + if not agent_result: + self.print_warning("Could not create workflow agent") + self.record_result(test_name, True, "Skipped (agent creation failed)") + return True + + agent_id, _ = agent_result + self.print_header(f"Testing {test_name}") + + payload = { + "question": "What is DocsGPT?", + "history": "[]", + "agent_id": agent_id, + } + + try: + self.print_info(f"POST /api/answer with workflow agent_id={agent_id[:8]}...") + + response = self.post("/api/answer", json=payload, timeout=60) + + self.print_info(f"Status Code: {response.status_code}") + + if response.status_code != 200: + self.print_error(f"Expected 200, got {response.status_code}") + self.record_result(test_name, False, f"Status {response.status_code}") + return False + + result = response.json() + answer = result.get("answer", "") + self.print_success(f"Answer received: {answer[:100]}...") + self.record_result(test_name, True, "Success") + return True + + except Exception as e: + self.print_error(f"Error: {str(e)}") + self.record_result(test_name, False, str(e)) + return False + # ------------------------------------------------------------------------- # Validation Tests # ------------------------------------------------------------------------- @@ -916,6 +1396,27 @@ This is test documentation for integration tests. self.test_answer_endpoint_with_agent() time.sleep(1) + # Agentic agent tests + self.test_stream_agentic_agent() + time.sleep(1) + + self.test_answer_agentic_agent() + time.sleep(1) + + # Research agent tests + self.test_stream_research_agent() + time.sleep(2) + + self.test_answer_research_agent() + time.sleep(2) + + # Workflow with agentic node tests + self.test_stream_workflow_with_agentic_node() + time.sleep(2) + + self.test_answer_workflow_with_agentic_node() + time.sleep(2) + # API key tests self.test_stream_endpoint_with_api_key() time.sleep(1) From 820ee3a843a94966acaec8c7efa9c5231ec32283 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 25 Mar 2026 22:53:09 +0000 Subject: [PATCH 7/7] fix test --- tests/agents/test_research_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/agents/test_research_agent.py b/tests/agents/test_research_agent.py index 3cc1e6d3..30b29ebc 100644 --- a/tests/agents/test_research_agent.py +++ b/tests/agents/test_research_agent.py @@ -2,7 +2,7 @@ import json import time -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest from application.agents.research_agent import (