tiennm99/DocsGPT
1
0 Fork 0
mirror of https://github.com/tiennm99/DocsGPT.git synced 2026-10-12 22:14:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2327 from arc53/research-agent

Browse Source 
Research agent
This commit is contained in:
Alex authored and GitHub committed 2026-03-26 11:00:23 +00:00
2 parents ed34c2b929 820ee3a843
commit 175ed58d2e
50 files changed
+4467 -1270
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.gitignore
+1
View File
@@ -2,6 +2,7 @@
__pycache__/
*.py[cod]
*$py.class
results.txt
experiments/
experiments
application/agents/agent_creator.py
+5 -2
View File
@@ -1,7 +1,8 @@
import logging
from application.agents.agentic_agent import AgenticAgent
from application.agents.classic_agent import ClassicAgent
from application.agents.react_agent import ReActAgent
from application.agents.research_agent import ResearchAgent
from application.agents.workflow_agent import WorkflowAgent
logger = logging.getLogger(__name__)
@@ -10,7 +11,9 @@ logger = logging.getLogger(__name__)
class AgentCreator:
agents = {
"classic": ClassicAgent,
"react": ReActAgent,
"react": ClassicAgent, # backwards compat: react falls back to classic
"agentic": AgenticAgent,
"research": ResearchAgent,
"workflow": WorkflowAgent,
}
application/agents/agentic_agent.py
+63
View File
@@ -0,0 +1,63 @@
import logging
from typing import Dict, Generator, Optional
from application.agents.base import BaseAgent
from application.agents.tools.internal_search import (
INTERNAL_TOOL_ID,
add_internal_search_tool,
)
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]:
tools_dict = self.tool_executor.get_tools()
add_internal_search_tool(tools_dict, self.retriever_config)
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
application/agents/base.py
+61 -334
View File
@@ -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 (
application/agents/classic_agent.py
+1 -5
View File
@@ -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)
application/agents/react_agent.py
-238
View File
@@ -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)
application/agents/research_agent.py
+692
View File
@@ -0,0 +1,692 @@
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_ID,
add_internal_search_tool,
)
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()
CLARIFICATION_PROMPT = _load_prompt("clarification.txt")
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
self._last_token_snapshot: 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)
delta = current - self._last_token_snapshot
self._last_token_snapshot = current
return delta
# ------------------------------------------------------------------
# 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 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)
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()
add_internal_search_tool(tools_dict, self.retriever_config)
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 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)
# ------------------------------------------------------------------
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,
response_format={"type": "json_object"},
)
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:
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
# 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 ""
application/agents/tool_executor.py
+307
View File
@@ -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
]
application/agents/tools/internal_search.py
+436
View File
@@ -0,0 +1,436 @@
import json
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.
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=[],
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 _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."
try:
retriever = self._get_retriever()
docs = retriever.search(query)
except Exception as e:
logger.error(f"Internal search failed: {e}", exc_info=True)
return "Search failed: an internal error occurred."
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:
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 _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):
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."
),
"parameters": {
"properties": {
"query": {
"type": "string",
"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 {}
# Constants for building synthetic tools_dict entries
INTERNAL_TOOL_ID = "internal"
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": (
"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": 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 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",
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,
has_directory_structure: bool = False,
) -> 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,
"has_directory_structure": has_directory_structure,
}
application/agents/tools/think.py
+68
View File
@@ -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 {}
application/agents/workflows/node_agent.py
+16 -21
View File
@@ -2,9 +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.react_agent import ReActAgent
from application.agents.research_agent import ResearchAgent
from application.agents.workflows.schemas import AgentType
@@ -36,7 +37,8 @@ class ToolFilterMixin:
return filtered_tools
class WorkflowNodeClassicAgent(ToolFilterMixin, ClassicAgent):
class _WorkflowNodeMixin:
"""Common __init__ for all workflow node agents."""
def __init__(
self,
@@ -57,32 +59,25 @@ class WorkflowNodeClassicAgent(ToolFilterMixin, ClassicAgent):
self._allowed_tool_ids = tool_ids or []
class WorkflowNodeReActAgent(ToolFilterMixin, ReActAgent):
class WorkflowNodeClassicAgent(ToolFilterMixin, _WorkflowNodeMixin, ClassicAgent):
pass
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 WorkflowNodeAgenticAgent(ToolFilterMixin, _WorkflowNodeMixin, AgenticAgent):
pass
class WorkflowNodeResearchAgent(ToolFilterMixin, _WorkflowNodeMixin, ResearchAgent):
pass
class WorkflowNodeAgentFactory:
_agents: Dict[AgentType, Type[BaseAgent]] = {
AgentType.CLASSIC: WorkflowNodeClassicAgent,
AgentType.REACT: WorkflowNodeReActAgent,
AgentType.REACT: WorkflowNodeClassicAgent, # backwards compat
AgentType.AGENTIC: WorkflowNodeAgenticAgent,
AgentType.RESEARCH: WorkflowNodeResearchAgent,
}
@classmethod
application/agents/workflows/schemas.py
+2
View File
@@ -18,6 +18,8 @@ class NodeType(str, Enum):
class AgentType(str, Enum):
CLASSIC = "classic"
REACT = "react"
AGENTIC = "agentic"
RESEARCH = "research"
class ExecutionStatus(str, Enum):
application/agents/workflows/workflow_engine.py
+27 -12
View File
@@ -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] = []
application/api/answer/routes/answer.py
+1 -12
View File
@@ -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
application/api/answer/routes/base.py
+6 -1
View File
@@ -206,9 +206,12 @@ class BaseAnswerResource:
is_structured = False
schema_info = None
structured_chunks = []
query_metadata = {}
for line in agent.gen(query=question):
if "answer" in line:
if "metadata" in line:
query_metadata.update(line["metadata"])
elif "answer" in line:
response_full += str(line["answer"])
if line.get("structured"):
is_structured = True
@@ -295,6 +298,7 @@ class BaseAnswerResource:
is_shared_usage=is_shared_usage,
shared_token=shared_token,
attachment_ids=attachment_ids,
metadata=query_metadata if query_metadata else None,
)
# Persist compression metadata/summary if it exists and wasn't saved mid-execution
compression_meta = getattr(agent, "compression_metadata", None)
@@ -384,6 +388,7 @@ class BaseAnswerResource:
is_shared_usage=is_shared_usage,
shared_token=shared_token,
attachment_ids=attachment_ids,
metadata=query_metadata if query_metadata else None,
)
compression_meta = getattr(agent, "compression_metadata", None)
compression_saved = getattr(agent, "compression_saved", False)
application/api/answer/routes/stream.py
+1 -8
View File
@@ -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(
application/api/answer/services/conversation_service.py
+21 -12
View File
@@ -60,6 +60,7 @@ class ConversationService:
is_shared_usage: bool = False,
shared_token: Optional[str] = None,
attachment_ids: Optional[List[str]] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> str:
"""Save or update a conversation in the database"""
if decoded_token is None:
@@ -93,6 +94,11 @@ class ConversationService:
f"queries.{index}.timestamp": current_time,
f"queries.{index}.attachments": attachment_ids,
f"queries.{index}.model_id": model_id,
**(
{f"queries.{index}.metadata": metadata}
if metadata
else {}
),
}
},
)
@@ -124,6 +130,7 @@ class ConversationService:
"timestamp": current_time,
"attachments": attachment_ids,
"model_id": model_id,
**({"metadata": metadata} if metadata else {}),
}
}
},
@@ -156,22 +163,24 @@ class ConversationService:
if not completion or not completion.strip():
completion = question[:50] if question else "New Conversation"
query_doc = {
"prompt": question,
"response": response,
"thought": thought,
"sources": sources,
"tool_calls": tool_calls,
"timestamp": current_time,
"attachments": attachment_ids,
"model_id": model_id,
}
if metadata:
query_doc["metadata"] = metadata
conversation_data = {
"user": user_id,
"date": current_time,
"name": completion,
"queries": [
{
"prompt": question,
"response": response,
"thought": thought,
"sources": sources,
"tool_calls": tool_calls,
"timestamp": current_time,
"attachments": attachment_ids,
"model_id": model_id,
}
],
"queries": [query_doc],
}
if api_key:
application/api/answer/services/stream_processor.py
+150 -74
View File
@@ -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:
@@ -91,6 +101,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 +122,29 @@ 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()
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(
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:
@@ -124,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:
@@ -135,14 +177,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,
@@ -153,12 +189,15 @@ 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
# 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,17 +208,27 @@ class StreamProcessor:
f"+ {len(result.recent_queries)} recent messages"
)
# Build history from recent queries
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(
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"]}
{
"prompt": query["prompt"],
"response": query["response"],
**({"metadata": query["metadata"]} if "metadata" in query else {}),
}
for query in conversation.get("queries", [])
]
@@ -191,9 +240,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 +248,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 +277,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 +329,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 +354,6 @@ class StreamProcessor:
else:
data["sources"] = []
# Preserve model configuration from agent
data["default_model_id"] = data.get("default_model_id", "")
return data
@@ -376,7 +418,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 +429,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 +475,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 +558,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"
@@ -755,11 +777,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(
@@ -779,7 +811,30 @@ 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
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 = {
@@ -796,10 +851,31 @@ 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
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
application/api/user/agents/routes.py
+2
View File
@@ -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"]
application/api/user/base.py
+8 -2
View File
@@ -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"]),
application/app.py
+5
View File
@@ -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 = {
application/prompts/agentic/creative.txt
+16
View File
@@ -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.
application/prompts/agentic/default.txt
+15
View File
@@ -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.
application/prompts/agentic/strict.txt
+16
View File
@@ -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!
application/prompts/react_final_prompt.txt
-3
View File
@@ -1,3 +0,0 @@
Query: {query}
Observations: {observations}
Now, using the insights from the observations, formulate a well-structured and precise final answer.
application/prompts/react_planning_prompt.txt
-13
View File
@@ -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}
application/prompts/research/clarification.txt
+23
View File
@@ -0,0 +1,23 @@
You are a research assistant evaluating whether a user's question is clear enough to begin in-depth research.
Decide whether the question can be researched as-is, or whether you need to ask clarifying questions first.
Proceed WITHOUT clarification if:
- The question has a clear topic and intent
- There is enough context to form a research plan
- The question is broad but can be broken into sub-topics
- Documents/sources are available to search
Ask for clarification ONLY if:
- The question is critically ambiguous (e.g., "review the thing" with no context about what "thing" means)
- Key information is missing that would make research impossible (not just imperfect)
- The user seems to reference something specific but hasn't provided it
Err on the side of proceeding. Most questions are clear enough to start researching.
You MUST respond with ONLY a valid JSON object (no markdown, no code fences):
{
"needs_clarification": true or false,
"questions": ["question 1", "question 2"] (only if needs_clarification is true, max 3 questions),
"reason": "brief explanation of why clarification is or isn't needed"
}
application/prompts/research/planning.txt
+23
View File
@@ -0,0 +1,23 @@
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")
- "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"}
]
}
application/prompts/research/step.txt
+13
View File
@@ -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.
application/prompts/research/synthesis.txt
+22
View File
@@ -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.
application/seed/config/premade_agents.yaml
+1 -1
View File
@@ -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: |
frontend/src/agents/NewAgent.tsx
+2 -1
View File
@@ -138,7 +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: t('agents.form.agentTypes.react'), value: 'react' },
{ label: 'Agentic', value: 'agentic' },
{ label: 'Research', value: 'research' },
];
const isPublishable = () => {
frontend/src/agents/WorkflowBuilder.tsx
+7 -2
View File
@@ -55,7 +55,7 @@ import WorkflowPreview from './workflow/WorkflowPreview';
import type { Model } from '../models/types';
interface AgentNodeConfig {
agent_type: 'classic' | 'react';
agent_type: 'classic' | 'agentic' | 'research';
llm_name?: string;
model_id?: string;
system_prompt: string;
@@ -884,7 +884,12 @@ function WorkflowBuilderInner() {
<SelectItem value="classic">
Classic
</SelectItem>
<SelectItem value="react">ReAct</SelectItem>
<SelectItem value="agentic">
Agentic
</SelectItem>
<SelectItem value="research">
Research
</SelectItem>
</SelectContent>
</Select>
</div>
frontend/src/agents/workflow/WorkflowBuilder.tsx
+6 -3
View File
@@ -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' | 'agentic' | 'research';
llm_name?: string;
model_id?: string;
system_prompt: string;
@@ -1748,8 +1748,11 @@ function WorkflowBuilderInner() {
<SelectItem value="classic">
Classic
</SelectItem>
<SelectItem value="react">
ReAct
<SelectItem value="agentic">
Agentic
</SelectItem>
<SelectItem value="research">
Research
</SelectItem>
</SelectContent>
</Select>
frontend/src/assets/research.svg
+18
View File
@@ -0,0 +1,18 @@
<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
<path d="M11 2C15.968 2 20 6.032 20 11C20 13.1245 19.2637 15.0768 18.0319 16.6175L21.7071 20.2929C22.0976 20.6834 22.0976 21.3166 21.7071 21.7071C21.3166 22.0976 20.6834 22.0976 20.2929 21.7071L16.6175 18.0319C15.0768 19.2637 13.1245 20 11 20C6.032 20 2 15.968 2 11C2 6.032 6.032 2 11 2ZM11 4C7.13401 4 4 7.13401 4 11C4 14.866 7.13401 18 11 18C14.866 18 18 14.866 18 11C18 7.13401 14.866 4 11 4Z" fill="url(#paint0_research)"/>
<path d="M11 7C11.5523 7 12 7.44772 12 8V10H14C14.5523 10 15 10.4477 15 11C15 11.5523 14.5523 12 14 12H12V14C12 14.5523 11.5523 15 11 15C10.4477 15 10 14.5523 10 14V12H8C7.44772 12 7 11.5523 7 11C7 10.4477 7.44772 10 8 10H10V8C10 7.44772 10.4477 7 11 7Z" fill="url(#paint1_research)"/>
<defs>
<linearGradient id="paint0_research" x1="2" y1="12" x2="22" y2="12" gradientUnits="userSpaceOnUse">
<stop stop-color="#70FDF7"/>
<stop offset="0.325" stop-color="#747696"/>
<stop offset="0.68" stop-color="#BD5372"/>
<stop offset="1" stop-color="#F5A06C"/>
</linearGradient>
<linearGradient id="paint1_research" x1="7" y1="11" x2="15" y2="11" gradientUnits="userSpaceOnUse">
<stop stop-color="#70FDF7"/>
<stop offset="0.325" stop-color="#747696"/>
<stop offset="0.68" stop-color="#BD5372"/>
<stop offset="1" stop-color="#F5A06C"/>
</linearGradient>
</defs>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 1.3 KiB

frontend/src/conversation/ConversationBubble.tsx
+63 -3
View File
@@ -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 <a> tags we can style.
// Avoid matching inside code blocks or existing links.
processedContent = processedContent.replace(
/(?<!\[)\[(\d+)\](?!\()/g,
(_, num) => `[${num}](#cite-${num})`,
);
const contentSegments: Array<{
type: 'text' | 'mermaid';
@@ -290,6 +301,7 @@ const ConversationBubble = forwardRef<
return contentSegments;
};
bubble = (
<div
ref={ref}
@@ -320,7 +332,7 @@ const ConversationBubble = forwardRef<
<div className="fade-in mr-5 ml-3 max-w-[90vw] md:max-w-[70vw] lg:max-w-[50vw]">
<div className="grid grid-cols-2 gap-2 lg:grid-cols-4">
{sources?.slice(0, 3)?.map((source, index) => (
<div key={index} className="relative">
<div key={index} id={`source-${index}`} className="relative transition-all duration-300">
<div
className="bg-gray-1000 dark:bg-gun-metal h-28 cursor-pointer rounded-4xl p-4 hover:bg-[#F1F1F1] dark:hover:bg-[#2C2E3C]"
onMouseOver={() => setActiveTooltip(index)}
@@ -393,6 +405,7 @@ const ConversationBubble = forwardRef<
</div>
</div>
)}
{research && <ResearchProgress research={research} />}
{toolCalls && toolCalls.length > 0 && (
<ToolCalls toolCalls={toolCalls} />
)}
@@ -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 (
<button
type="button"
onClick={() => {
const el = document.getElementById(`source-${sourceIdx}`);
if (el) {
el.scrollIntoView({ behavior: 'smooth', block: 'center' });
el.classList.add('ring-2', 'ring-purple-500');
setTimeout(() => el.classList.remove('ring-2', 'ring-purple-500'), 2000);
}
}}
className="mx-0.5 inline-flex h-5 min-w-5 items-center justify-center rounded-full bg-purple-100 px-1.5 text-xs font-semibold text-purple-700 transition-colors hover:bg-purple-200 dark:bg-purple-900/40 dark:text-purple-300 dark:hover:bg-purple-900/60"
title={`Jump to source ${num}`}
>
{num}
</button>
);
}
return <a href={href} target="_blank" rel="noopener noreferrer">{children}</a>;
},
code(props) {
const {
children,
@@ -649,6 +686,29 @@ const ConversationBubble = forwardRef<
<div className="relative mr-2 block items-center justify-center">
<CopyButton textToCopy={message} />
</div>
{research && message && (
<div className="relative mr-2 block items-center justify-center">
<button
type="button"
onClick={() => {
const blob = new Blob([message], { type: 'text/markdown' });
const url = URL.createObjectURL(blob);
const link = document.createElement('a');
link.href = url;
link.download = `research-report.md`;
link.click();
URL.revokeObjectURL(url);
}}
className="bg-white-3000 dark:hover:bg-purple-taupe flex cursor-pointer items-center justify-center rounded-full p-2 hover:bg-[#EEEEEE] dark:bg-transparent"
aria-label="Export as Markdown"
title="Export as Markdown"
>
<svg className="h-5 w-5 stroke-gray-4000" fill="none" viewBox="0 0 24 24" strokeWidth={1.5}>
<path strokeLinecap="round" strokeLinejoin="round" d="M3 16.5v2.25A2.25 2.25 0 005.25 21h13.5A2.25 2.25 0 0021 18.75V16.5M16.5 12L12 16.5m0 0L7.5 12m4.5 4.5V3" />
</svg>
</button>
</div>
)}
<div className="relative mr-2 block items-center justify-center">
<SpeakButton text={message} />
</div>
frontend/src/conversation/ConversationMessages.tsx
+2 -1
View File
@@ -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}
frontend/src/conversation/ResearchProgress.tsx
+147
View File
@@ -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 = () => (
<svg className="h-3 w-3 text-green-500" fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={3}>
<path strokeLinecap="round" strokeLinejoin="round" d="M5 13l4 4L19 7" />
</svg>
);
const SmallSpinner = () => (
<svg className="h-3 w-3 animate-spin text-purple-500" fill="none" viewBox="0 0 24 24">
<circle className="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" strokeWidth="4" />
<path className="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4z" />
</svg>
);
const SmallPending = () => (
<div className="h-2.5 w-2.5 rounded-full border-[1.5px] border-gray-400 dark:border-gray-500" />
);
function StatusText({ status, elapsed }: { status: string; elapsed?: number }) {
const labels: Record<string, string> = {
planning: 'Planning research...',
researching: 'Researching...',
synthesizing: 'Writing report...',
complete: 'Complete',
};
const elapsed_str = elapsed ? ` \u00B7 ${Math.round(elapsed)}s` : '';
return (
<span className="text-xs text-gray-500 dark:text-gray-400">
{status === 'complete' ? (
<>
<span className="text-green-600 dark:text-green-400">{labels.complete}</span>
{elapsed_str}
</>
) : (
<>
{labels[status] || status}
{elapsed_str}
</>
)}
</span>
);
}
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 (
<div className="mb-4 flex w-full flex-col flex-wrap items-start self-start lg:flex-nowrap">
{/* Header row — matches Reasoning / Sources pattern */}
<div className="my-2 flex flex-row items-center gap-3">
<Avatar
className="h-[26px] w-[30px] text-xl"
avatar={
<img
src={ResearchIcon}
alt="Research"
className="h-full w-full object-fill"
/>
}
/>
<button
className="flex flex-row items-center gap-2"
onClick={() => setIsExpanded(!isExpanded)}
>
<p className="text-sm font-semibold">
{isExpanded ? 'Research' : summaryText}
</p>
<svg
className={`h-4 w-4 text-gray-500 transition-transform duration-200 dark:text-gray-400 ${isExpanded ? 'rotate-180' : ''}`}
fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={2}
>
<path strokeLinecap="round" strokeLinejoin="round" d="M19 9l-7 7-7-7" />
</svg>
</button>
{status && <StatusText status={status} elapsed={elapsed_seconds} />}
</div>
{/* Expanded: vertical timeline of steps */}
{isExpanded && plan && plan.length > 0 && (
<div className="fade-in mr-5 ml-[42px] max-w-[90vw] md:max-w-[70vw] lg:max-w-[50vw]">
<div className="space-y-0">
{plan.map((step, i) => {
const isLast = i === plan.length - 1;
return (
<div key={i} className="flex items-stretch gap-3">
{/* Timeline: dot + vertical line */}
<div className="flex flex-col items-center pt-1">
<div className="flex h-4 w-4 flex-shrink-0 items-center justify-center">
{step.status === 'complete' ? (
<SmallCheck />
) : step.status === 'researching' ? (
<SmallSpinner />
) : (
<SmallPending />
)}
</div>
{!isLast && (
<div className="mt-1 w-px flex-1 bg-gray-300 dark:bg-gray-600" />
)}
</div>
{/* Step content */}
<div className={`pb-3 ${isLast ? '' : ''}`}>
<p className={`text-sm ${
step.status === 'complete'
? 'text-gray-700 dark:text-gray-300'
: step.status === 'researching'
? 'font-medium text-purple-700 dark:text-purple-300'
: 'text-gray-500 dark:text-gray-500'
}`}>
{step.query}
</p>
</div>
</div>
);
})}
</div>
</div>
)}
</div>
);
}
frontend/src/conversation/conversationModels.ts
+15
View File
@@ -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 {
frontend/src/conversation/conversationSlice.ts
+106 -1
View File
@@ -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,16 @@ export const fetchAnswer = createAsyncThunk<
if (currentConversationId === state.conversation.conversationId) {
if (data.type === 'end') {
dispatch(conversationSlice.actions.setStatus('idle'));
// 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));
@@ -122,6 +138,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 +287,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 +369,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<Query> }>,
@@ -405,6 +508,8 @@ export const {
updateThought,
updateStreamingSource,
updateToolCall,
updateResearchPlan,
updateResearchProgress,
setConversation,
setStatus,
raiseError,
tests/agents/test_agent_creator.py
+2 -6
View File
@@ -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"
tests/agents/test_agentic_agent.py
+139
View File
@@ -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 == []
tests/agents/test_internal_search_tool.py
+250
View File
@@ -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
tests/agents/test_react_agent.py
-514
View File
@@ -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)
tests/agents/test_research_agent.py
+400
View File
@@ -0,0 +1,400 @@
"""Tests for ResearchAgent — multi-step research with budget controls."""
import json
import time
from unittest.mock import Mock
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
tests/agents/test_think_tool.py
+52
View File
@@ -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
tests/agents/test_tool_executor.py
+279
View File
@@ -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
tests/agents/test_workflow_agent_types.py
+475
View File
@@ -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"
tests/conftest.py
+1 -1
View File
@@ -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
tests/integration/test_chat.py
+501
View File
@@ -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)