diff --git a/test/backend/agents/test_create_agent_info.py b/test/backend/agents/test_create_agent_info.py index 44ce32832..1a133700b 100644 --- a/test/backend/agents/test_create_agent_info.py +++ b/test/backend/agents/test_create_agent_info.py @@ -182,7 +182,6 @@ def _create_stub_module(name: str, **attrs): prepare_prompt_templates, _get_skills_for_template, _get_skill_script_tools, - _print_prompt_with_token_count, ) # Import constants for testing @@ -433,88 +432,6 @@ def test_get_skill_script_tools_tool_descriptions(self): assert "skill" in desc.lower() -class TestPrintPromptWithTokenCount: - """Tests for the _print_prompt_with_token_count function""" - - def test_print_prompt_with_token_count_success(self): - """Test successful token counting with tiktoken available""" - import tiktoken - - with patch('backend.agents.create_agent_info.logger') as mock_logger: - mock_encoding = MagicMock() - mock_encoding.encode.return_value = ["token1", "token2", "token3"] - with patch.object(tiktoken, 'get_encoding', return_value=mock_encoding): - _print_prompt_with_token_count("test prompt content", agent_id=123, stage="TEST") - - mock_encoding.encode.assert_called_once_with("test prompt content") - mock_logger.info.assert_called() - - # Check that log messages contain expected content - log_calls = mock_logger.info.call_args_list - log_text = " ".join([str(call) for call in log_calls]) - assert "TEST" in log_text - assert "123" in log_text - assert "3" in log_text # Token count - - def test_print_prompt_with_token_count_tiktoken_failure(self): - """Test graceful handling when tiktoken fails""" - import tiktoken - - with patch('backend.agents.create_agent_info.logger') as mock_logger: - with patch.object(tiktoken, 'get_encoding', side_effect=Exception("tiktoken not available")): - _print_prompt_with_token_count("test prompt", agent_id=456, stage="FALLBACK") - - # Should log a warning and then log the prompt - mock_logger.warning.assert_called_once() - assert "Failed to count tokens: tiktoken not available" in mock_logger.warning.call_args[0][0] - - # Should still log the prompt - mock_logger.info.assert_called() - - def test_print_prompt_with_token_count_default_stage(self): - """Test with default stage parameter""" - import tiktoken - - with patch('backend.agents.create_agent_info.logger') as mock_logger: - mock_encoding = MagicMock() - mock_encoding.encode.return_value = ["a", "b"] - with patch.object(tiktoken, 'get_encoding', return_value=mock_encoding): - _print_prompt_with_token_count("short prompt") - - log_calls = mock_logger.info.call_args_list - log_text = " ".join([str(call) for call in log_calls]) - assert "PROMPT" in log_text # Default stage - - def test_print_prompt_with_token_count_empty_prompt(self): - """Test with empty prompt""" - import tiktoken - - with patch('backend.agents.create_agent_info.logger') as mock_logger: - mock_encoding = MagicMock() - mock_encoding.encode.return_value = [] - with patch.object(tiktoken, 'get_encoding', return_value=mock_encoding): - _print_prompt_with_token_count("", agent_id=1, stage="EMPTY") - - mock_encoding.encode.assert_called_once_with("") - # Should log token count of 0 - log_calls = mock_logger.info.call_args_list - log_text = " ".join([str(call) for call in log_calls]) - assert "0" in log_text - - def test_print_prompt_with_token_count_none_agent_id(self): - """Test with None agent_id""" - import tiktoken - - with patch('backend.agents.create_agent_info.logger') as mock_logger: - mock_encoding = MagicMock() - mock_encoding.encode.return_value = ["token"] - with patch.object(tiktoken, 'get_encoding', return_value=mock_encoding): - _print_prompt_with_token_count("prompt", agent_id=None, stage="NO_ID") - - # Should not raise an error - mock_encoding.encode.assert_called_once_with("prompt") - - class TestDiscoverLangchainTools: """Tests for the discover_langchain_tools function""" diff --git a/test/sdk/core/utils/test_prompt_template_utils.py b/test/sdk/core/utils/test_prompt_template_utils.py index 8c1788c39..c0a3ad634 100644 --- a/test/sdk/core/utils/test_prompt_template_utils.py +++ b/test/sdk/core/utils/test_prompt_template_utils.py @@ -86,22 +86,6 @@ def test_get_prompt_template_unsupported_type(self): assert "Unsupported template type" in str(excinfo.value) assert "unsupported_type" in str(excinfo.value) - @patch('builtins.open', new_callable=mock_open, read_data='system_prompt: "Test prompt"') - @patch('yaml.safe_load') - def test_get_prompt_template_with_kwargs(self, mock_yaml_load, mock_file): - """Test get_prompt_template with additional kwargs (should be logged but not used)""" - mock_yaml_load.return_value = {"system_prompt": "Test prompt"} - - with patch('sdk.nexent.core.utils.prompt_template_utils.logger') as mock_logger: - result = get_prompt_template(template_type='analyze_image', language='en', extra_param='value') - - # Verify kwargs were logged - log_calls = [str(call) for call in mock_logger.info.call_args_list] - assert any("extra_param" in str(call) or "kwargs" in str(call) for call in log_calls) - - # Verify function still works - assert result == {"system_prompt": "Test prompt"} - @patch('builtins.open', side_effect=FileNotFoundError("File not found")) def test_get_prompt_template_file_not_found(self, mock_file): """Test get_prompt_template when template file is not found""" @@ -119,21 +103,6 @@ def test_get_prompt_template_yaml_error(self, mock_yaml_load, mock_file): assert "YAML parse error" in str(excinfo.value) - @patch('builtins.open', new_callable=mock_open, read_data='system_prompt: "Test prompt"') - @patch('yaml.safe_load') - @patch('sdk.nexent.core.utils.prompt_template_utils.logger') - def test_get_prompt_template_logging(self, mock_logger, mock_yaml_load, mock_file): - """Test that get_prompt_template logs correctly""" - mock_yaml_load.return_value = {"system_prompt": "Test prompt"} - - get_prompt_template(template_type='analyze_image', language='en') - - # Verify logger was called - mock_logger.info.assert_called_once() - log_call = str(mock_logger.info.call_args) - assert "analyze_image" in log_call - assert "en" in log_call - @patch('builtins.open', new_callable=mock_open, read_data='system_prompt: "Test prompt"') @patch('yaml.safe_load') def test_get_prompt_template_path_construction(self, mock_yaml_load, mock_file):