mirror of
https://github.com/arc53/DocsGPT.git
synced 2025-11-29 08:33:20 +00:00
feat: template-based prompt rendering with dynamic namespace injection (#2091)
* feat: template-based prompt rendering with dynamic namespace injection * refactor: improve template engine initialization with clearer formatting * refactor: streamline ReActAgent methods and improve content extraction logic feat: enhance error handling in NamespaceManager and TemplateEngine fix: update NewAgent component to ensure consistent form data submission test: modify tests for ReActAgent and prompt renderer to reflect method changes and improve coverage * feat: tools namespace + three-tier token budget * refactor: remove unused variable assignment in message building tests * Enhance prompt customization and tool pre-fetching functionality * ruff lint fix * refactor: cleaner error handling and reduce code clutter --------- Co-authored-by: Alex <a@tushynski.me>
This commit is contained in:
Siddhant Rai
@@ -250,3 +250,330 @@ class TestStreamProcessorAttachments:
|
||||
"attachments" not in processor.data
|
||||
or processor.data.get("attachments") is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestToolPreFetch:
|
||||
"""Tests for tool pre-fetching with saved parameter values from MongoDB"""
|
||||
|
||||
def test_cryptoprice_prefetch_with_saved_parameters(self, mock_mongo_db):
|
||||
"""Test that cryptoprice tool is pre-fetched with saved parameter values from MongoDB structure"""
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Setup MongoDB with cryptoprice tool configuration
|
||||
# NOTE: The collection is called "user_tools" not "tools"
|
||||
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
|
||||
tool_id = ObjectId()
|
||||
|
||||
tools_collection.insert_one(
|
||||
{
|
||||
"_id": tool_id,
|
||||
"name": "cryptoprice",
|
||||
"user": "user_123",
|
||||
"status": True, # Must be True for tool to be included
|
||||
"actions": [
|
||||
{
|
||||
"name": "cryptoprice_get",
|
||||
"description": "Get cryptocurrency price",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"symbol": {
|
||||
"type": "string",
|
||||
"description": "Crypto symbol",
|
||||
"value": "BTC" # Saved value in MongoDB
|
||||
},
|
||||
"currency": {
|
||||
"type": "string",
|
||||
"description": "Currency for price",
|
||||
"value": "USD" # Saved value in MongoDB
|
||||
}
|
||||
},
|
||||
"required": ["symbol", "currency"]
|
||||
}
|
||||
}
|
||||
],
|
||||
"config": {
|
||||
"token": ""
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"question": "What is the price of Bitcoin?",
|
||||
"tools": [str(tool_id)]
|
||||
}
|
||||
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
processor._required_tool_actions = {"cryptoprice": {"cryptoprice_get"}}
|
||||
|
||||
# Mock the ToolManager and tool instance
|
||||
with patch("application.agents.tools.tool_manager.ToolManager") as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance returned by load_tool
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
# Mock get_actions_metadata on the tool instance
|
||||
mock_tool.get_actions_metadata.return_value = [
|
||||
{
|
||||
"name": "cryptoprice_get",
|
||||
"description": "Get cryptocurrency price",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"symbol": {"type": "string", "description": "Crypto symbol"},
|
||||
"currency": {"type": "string", "description": "Currency for price"}
|
||||
},
|
||||
"required": ["symbol", "currency"]
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
# Mock execute_action on the tool instance to return price data
|
||||
mock_tool.execute_action.return_value = {
|
||||
"status_code": 200,
|
||||
"price": 45000.50,
|
||||
"message": "Price of BTC in USD retrieved successfully."
|
||||
}
|
||||
|
||||
# Execute pre-fetch
|
||||
tools_data = processor.pre_fetch_tools()
|
||||
|
||||
# Verify the tool was called
|
||||
assert mock_tool.execute_action.called
|
||||
|
||||
# Verify it was called with the saved parameters from MongoDB
|
||||
call_args = mock_tool.execute_action.call_args
|
||||
assert call_args is not None
|
||||
|
||||
# Check action name uses the full metadata name for execution
|
||||
assert call_args[0][0] == "cryptoprice_get"
|
||||
|
||||
# Check kwargs contain saved values
|
||||
kwargs = call_args[1]
|
||||
assert kwargs.get("symbol") == "BTC"
|
||||
assert kwargs.get("currency") == "USD"
|
||||
|
||||
# Verify tools_data structure
|
||||
assert "cryptoprice" in tools_data
|
||||
# Results are exposed under the full action name
|
||||
assert "cryptoprice_get" in tools_data["cryptoprice"]
|
||||
assert tools_data["cryptoprice"]["cryptoprice_get"]["price"] == 45000.50
|
||||
|
||||
def test_prefetch_with_missing_saved_values_uses_defaults(self, mock_mongo_db):
|
||||
"""Test that pre-fetch falls back to defaults when saved values are missing"""
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
|
||||
tool_id = ObjectId()
|
||||
|
||||
# Tool configuration without saved values
|
||||
tools_collection.insert_one(
|
||||
{
|
||||
"_id": tool_id,
|
||||
"name": "cryptoprice",
|
||||
"user": "user_123",
|
||||
"status": True,
|
||||
"actions": [
|
||||
{
|
||||
"name": "cryptoprice_get",
|
||||
"description": "Get cryptocurrency price",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"symbol": {
|
||||
"type": "string",
|
||||
"description": "Crypto symbol",
|
||||
"default": "ETH" # Only default, no saved value
|
||||
},
|
||||
"currency": {
|
||||
"type": "string",
|
||||
"description": "Currency",
|
||||
"default": "EUR"
|
||||
}
|
||||
},
|
||||
"required": ["symbol", "currency"]
|
||||
}
|
||||
}
|
||||
],
|
||||
"config": {}
|
||||
}
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"question": "Crypto price?",
|
||||
"tools": [str(tool_id)]
|
||||
}
|
||||
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
processor._required_tool_actions = {"cryptoprice": {"cryptoprice_get"}}
|
||||
|
||||
with patch("application.agents.tools.tool_manager.ToolManager") as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
mock_tool.get_actions_metadata.return_value = [
|
||||
{
|
||||
"name": "cryptoprice_get",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"symbol": {"type": "string", "default": "ETH"},
|
||||
"currency": {"type": "string", "default": "EUR"}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
mock_tool.execute_action.return_value = {
|
||||
"status_code": 200,
|
||||
"price": 2500.00
|
||||
}
|
||||
|
||||
processor.pre_fetch_tools()
|
||||
|
||||
# Should use default values when saved values are missing
|
||||
call_args = mock_tool.execute_action.call_args
|
||||
if call_args:
|
||||
kwargs = call_args[1]
|
||||
# Either uses defaults or skips if no values available
|
||||
assert kwargs.get("symbol") in ["ETH", None]
|
||||
assert kwargs.get("currency") in ["EUR", None]
|
||||
|
||||
def test_prefetch_with_tool_id_reference(self, mock_mongo_db):
|
||||
"""Test that tools can be referenced by MongoDB ObjectId in templates"""
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
|
||||
tool_id = ObjectId()
|
||||
|
||||
# Create a tool in the database
|
||||
tools_collection.insert_one(
|
||||
{
|
||||
"_id": tool_id,
|
||||
"name": "memory",
|
||||
"user": "user_123",
|
||||
"status": True,
|
||||
"actions": [
|
||||
{
|
||||
"name": "memory_ls",
|
||||
"description": "List files",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
}
|
||||
],
|
||||
"config": {},
|
||||
}
|
||||
)
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
|
||||
# Mock the filtering to require this specific tool by ID
|
||||
processor._required_tool_actions = {
|
||||
str(tool_id): {"memory_ls"} # Reference by ObjectId string
|
||||
}
|
||||
|
||||
with patch(
|
||||
"application.agents.tools.tool_manager.ToolManager"
|
||||
) as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
mock_tool.get_actions_metadata.return_value = [
|
||||
{"name": "memory_ls", "description": "List files", "parameters": {"properties": {}}}
|
||||
]
|
||||
mock_tool.execute_action.return_value = "Directory: /\n- file.txt"
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
# Tool data should be available under both name and ID
|
||||
assert result is not None
|
||||
assert "memory" in result
|
||||
assert str(tool_id) in result
|
||||
# Both should point to the same data
|
||||
assert result["memory"] == result[str(tool_id)]
|
||||
assert "memory_ls" in result[str(tool_id)]
|
||||
|
||||
def test_prefetch_with_multiple_same_name_tools(self, mock_mongo_db):
|
||||
"""Test that multiple tools with the same name can be distinguished by ID"""
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
|
||||
|
||||
# Create two memory tools with different IDs
|
||||
tool_id_1 = ObjectId()
|
||||
tool_id_2 = ObjectId()
|
||||
|
||||
tools_collection.insert_many([
|
||||
{
|
||||
"_id": tool_id_1,
|
||||
"name": "memory",
|
||||
"user": "user_123",
|
||||
"status": True,
|
||||
"actions": [{"name": "memory_ls", "parameters": {"properties": {}}}],
|
||||
"config": {"path": "/home"},
|
||||
},
|
||||
{
|
||||
"_id": tool_id_2,
|
||||
"name": "memory",
|
||||
"user": "user_123",
|
||||
"status": True,
|
||||
"actions": [{"name": "memory_ls", "parameters": {"properties": {}}}],
|
||||
"config": {"path": "/work"},
|
||||
}
|
||||
])
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
|
||||
# Mock the filtering to require only the second tool by ID
|
||||
processor._required_tool_actions = {
|
||||
str(tool_id_2): {"memory_ls"} # Only reference the second one
|
||||
}
|
||||
|
||||
with patch(
|
||||
"application.agents.tools.tool_manager.ToolManager"
|
||||
) as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
mock_tool.get_actions_metadata.return_value = [
|
||||
{"name": "memory_ls", "parameters": {"properties": {}}}
|
||||
]
|
||||
mock_tool.execute_action.return_value = "Work directory"
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
# Only the second tool should be fetched (referenced by ID)
|
||||
assert result is not None
|
||||
assert str(tool_id_2) in result
|
||||
# Since filtering is enabled and only tool_id_2 is referenced,
|
||||
# only tool_id_2 should be pre-fetched
|
||||
# The "memory" key will still exist because we store under both name and ID
|
||||
assert "memory" in result
|
||||
|
||||
Reference in New Issue
Block a user