test_user_interfaces.py
1 # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 2 # 3 # SPDX-License-Identifier: Apache-2.0 4 5 from unittest.mock import MagicMock, patch 6 7 import pytest 8 9 from haystack.human_in_the_loop import ConfirmationUIResult, RichConsoleUI, SimpleConsoleUI 10 from haystack.tools import create_tool_from_function 11 12 13 def multiply_tool(x: int) -> int: 14 return x * 2 15 16 17 @pytest.fixture 18 def tool(): 19 return create_tool_from_function( 20 function=multiply_tool, name="test_tool", description="A test tool that multiplies input by 2." 21 ) 22 23 24 class TestRichConsoleUI: 25 @pytest.mark.parametrize("choice", ["y"]) 26 def test_process_choice_confirm(self, tool, choice): 27 ui = RichConsoleUI(console=MagicMock()) 28 29 with patch("haystack.human_in_the_loop.user_interfaces.Prompt.ask", side_effect=[choice, "feedback"]): 30 result = ui.get_user_confirmation(tool.name, tool.description, {"x": 1}) 31 32 assert isinstance(result, ConfirmationUIResult) 33 assert result.action == "confirm" 34 assert result.new_tool_params is None 35 assert result.feedback is None 36 37 @pytest.mark.parametrize("choice", ["m"]) 38 def test_process_choice_modify(self, tool, choice): 39 ui = RichConsoleUI(console=MagicMock()) 40 41 with patch("haystack.human_in_the_loop.user_interfaces.Prompt.ask", side_effect=["m", "2"]): 42 result = ui.get_user_confirmation(tool.name, tool.description, {"x": 1}) 43 44 assert isinstance(result, ConfirmationUIResult) 45 assert result.action == "modify" 46 assert result.new_tool_params == {"x": 2} 47 48 def test_process_choice_modify_dict_param(self, tool): 49 ui = RichConsoleUI(console=MagicMock()) 50 51 with patch("haystack.human_in_the_loop.user_interfaces.Prompt.ask", side_effect=["m", '{"key": "value"}']): 52 result = ui.get_user_confirmation(tool.name, tool.description, {"param1": {"old_key": "old_value"}}) 53 54 assert isinstance(result, ConfirmationUIResult) 55 assert result.action == "modify" 56 assert result.new_tool_params == {"param1": {"key": "value"}} 57 58 def test_process_choice_modify_dict_param_invalid_json(self, tool): 59 ui = RichConsoleUI(console=MagicMock()) 60 61 with patch( 62 "haystack.human_in_the_loop.user_interfaces.Prompt.ask", 63 side_effect=["m", "invalid_json", '{"key": "value"}'], 64 ): 65 result = ui.get_user_confirmation(tool.name, tool.description, {"param1": {"old_key": "old_value"}}) 66 67 assert isinstance(result, ConfirmationUIResult) 68 assert result.action == "modify" 69 assert result.new_tool_params == {"param1": {"key": "value"}} 70 71 @pytest.mark.parametrize("choice", ["n"]) 72 def test_process_choice_reject(self, tool, choice): 73 ui = RichConsoleUI(console=MagicMock()) 74 75 with patch("haystack.human_in_the_loop.user_interfaces.Prompt.ask", side_effect=["n", "Changed my mind"]): 76 result = ui.get_user_confirmation(tool.name, tool.description, {"x": 1}) 77 78 assert isinstance(result, ConfirmationUIResult) 79 assert result.action == "reject" 80 assert result.feedback == "Changed my mind" 81 82 def test_to_dict(self): 83 ui = RichConsoleUI() 84 data = ui.to_dict() 85 assert data["type"] == ("haystack.human_in_the_loop.user_interfaces.RichConsoleUI") 86 assert data["init_parameters"]["console"] is None 87 88 def test_from_dict(self): 89 ui = RichConsoleUI() 90 data = ui.to_dict() 91 new_ui = RichConsoleUI.from_dict(data) 92 assert isinstance(new_ui, RichConsoleUI) 93 94 95 class TestSimpleConsoleUI: 96 @pytest.mark.parametrize("choice", ["y", "yes", "Y", "YES"]) 97 def test_process_choice_confirm(self, tool, choice): 98 ui = SimpleConsoleUI() 99 100 with patch("builtins.input", side_effect=[choice]): 101 result = ui.get_user_confirmation(tool.name, tool.description, {"y": "abc"}) 102 103 assert isinstance(result, ConfirmationUIResult) 104 assert result.action == "confirm" 105 106 @pytest.mark.parametrize("choice", ["m", "modify", "M", "MODIFY"]) 107 def test_process_choice_modify(self, tool, choice): 108 ui = SimpleConsoleUI() 109 110 with patch("builtins.input", side_effect=[choice, "new_value"]): 111 result = ui.get_user_confirmation(tool.name, tool.description, {"y": "abc"}) 112 113 assert isinstance(result, ConfirmationUIResult) 114 assert result.action == "modify" 115 assert result.new_tool_params == {"y": "new_value"} 116 117 def test_process_choice_modify_dict_param(self, tool): 118 ui = SimpleConsoleUI() 119 120 with patch("builtins.input", side_effect=["m", '{"key": "value"}']): 121 result = ui.get_user_confirmation(tool.name, tool.description, {"param1": {"old_key": "old_value"}}) 122 123 assert isinstance(result, ConfirmationUIResult) 124 assert result.action == "modify" 125 assert result.new_tool_params == {"param1": {"key": "value"}} 126 127 def test_process_choice_modify_dict_param_invalid_json(self, tool): 128 ui = SimpleConsoleUI() 129 130 with patch("builtins.input", side_effect=["m", "invalid_json", '{"key": "value"}']): 131 result = ui.get_user_confirmation(tool.name, tool.description, {"param1": {"old_key": "old_value"}}) 132 133 assert isinstance(result, ConfirmationUIResult) 134 assert result.action == "modify" 135 assert result.new_tool_params == {"param1": {"key": "value"}} 136 137 @pytest.mark.parametrize("choice", ["n", "no", "N", "NO"]) 138 def test_process_choice_reject(self, tool, choice): 139 ui = SimpleConsoleUI() 140 141 with patch("builtins.input", side_effect=[choice, "Changed my mind"]): 142 result = ui.get_user_confirmation(tool.name, tool.description, {"param1": "value1"}) 143 144 assert isinstance(result, ConfirmationUIResult) 145 assert result.action == "reject" 146 assert result.feedback == "Changed my mind" 147 148 def test_process_choice_no_tool_params_confirm(self, tool): 149 ui = SimpleConsoleUI() 150 151 with patch("builtins.input", side_effect=["y"]): 152 result = ui.get_user_confirmation(tool.name, tool.description, {}) 153 154 assert isinstance(result, ConfirmationUIResult) 155 assert result.action == "confirm" 156 assert result.new_tool_params is None 157 assert result.feedback is None 158 159 def test_process_choice_no_tool_params_modify(self, tool): 160 ui = SimpleConsoleUI() 161 162 with patch("builtins.input", side_effect=["m"]): 163 result = ui.get_user_confirmation(tool.name, tool.description, {}) 164 165 assert isinstance(result, ConfirmationUIResult) 166 assert result.action == "modify" 167 assert result.new_tool_params == {} 168 assert result.feedback is None 169 170 def test_process_choice_no_tool_params_reject(self, tool): 171 ui = SimpleConsoleUI() 172 173 with patch("builtins.input", side_effect=["n", "Changed my mind"]): 174 result = ui.get_user_confirmation(tool.name, tool.description, {}) 175 176 assert isinstance(result, ConfirmationUIResult) 177 assert result.action == "reject" 178 assert result.new_tool_params is None 179 assert result.feedback == "Changed my mind" 180 181 def test_to_dict(self): 182 ui = SimpleConsoleUI() 183 data = ui.to_dict() 184 assert data["type"] == ("haystack.human_in_the_loop.user_interfaces.SimpleConsoleUI") 185 assert data["init_parameters"] == {} 186 187 def test_from_dict(self): 188 ui = SimpleConsoleUI() 189 data = ui.to_dict() 190 new_ui = SimpleConsoleUI.from_dict(data) 191 assert isinstance(new_ui, SimpleConsoleUI) 192 193 def test_get_user_confirmation_invalid_input_then_valid(self, tool): 194 ui = SimpleConsoleUI() 195 196 with patch("builtins.input", side_effect=["invalid", "y"]): 197 result = ui.get_user_confirmation(tool.name, tool.description, {"x": 1}) 198 199 assert isinstance(result, ConfirmationUIResult) 200 assert result.action == "confirm" 201 assert result.new_tool_params is None 202 assert result.feedback is None