|
| 1 | +"""Tests for guardrail utility functions.""" |
| 2 | + |
| 3 | +import json |
| 4 | + |
| 5 | +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage |
| 6 | + |
| 7 | +from uipath_langchain.agent.guardrails.utils import ( |
| 8 | + _extract_tool_args_from_message, |
| 9 | + _extract_tool_output_data, |
| 10 | + _extract_tools_args_from_message, |
| 11 | + get_message_content, |
| 12 | +) |
| 13 | +from uipath_langchain.agent.react.types import AgentGuardrailsGraphState |
| 14 | + |
| 15 | + |
| 16 | +class TestExtractToolArgsFromMessage: |
| 17 | + """Tests for _extract_tool_args_from_message function.""" |
| 18 | + |
| 19 | + def test_extracts_args_from_matching_tool(self): |
| 20 | + """Should extract args from matching tool call.""" |
| 21 | + message = AIMessage( |
| 22 | + content="", |
| 23 | + tool_calls=[ |
| 24 | + { |
| 25 | + "name": "test_tool", |
| 26 | + "args": {"param1": "value1", "param2": 123}, |
| 27 | + "id": "call_1", |
| 28 | + } |
| 29 | + ], |
| 30 | + ) |
| 31 | + result = _extract_tool_args_from_message(message, "test_tool") |
| 32 | + assert result == {"param1": "value1", "param2": 123} |
| 33 | + |
| 34 | + def test_returns_empty_dict_for_non_matching_tool(self): |
| 35 | + """Should return empty dict when tool name doesn't match.""" |
| 36 | + message = AIMessage( |
| 37 | + content="", |
| 38 | + tool_calls=[ |
| 39 | + {"name": "other_tool", "args": {"data": "value"}, "id": "call_1"} |
| 40 | + ], |
| 41 | + ) |
| 42 | + result = _extract_tool_args_from_message(message, "test_tool") |
| 43 | + assert result == {} |
| 44 | + |
| 45 | + def test_returns_empty_dict_for_non_ai_message(self): |
| 46 | + """Should return empty dict when message is not AIMessage.""" |
| 47 | + message = HumanMessage(content="Test message") |
| 48 | + result = _extract_tool_args_from_message(message, "test_tool") |
| 49 | + assert result == {} |
| 50 | + |
| 51 | + def test_returns_first_matching_tool_when_multiple(self): |
| 52 | + """Should return args from first matching tool call.""" |
| 53 | + message = AIMessage( |
| 54 | + content="", |
| 55 | + tool_calls=[ |
| 56 | + {"name": "test_tool", "args": {"first": "call"}, "id": "call_1"}, |
| 57 | + {"name": "test_tool", "args": {"second": "call"}, "id": "call_2"}, |
| 58 | + ], |
| 59 | + ) |
| 60 | + result = _extract_tool_args_from_message(message, "test_tool") |
| 61 | + assert result == {"first": "call"} |
| 62 | + |
| 63 | + |
| 64 | +class TestExtractToolsArgsFromMessage: |
| 65 | + """Tests for _extract_tools_args_from_message function.""" |
| 66 | + |
| 67 | + def test_extracts_args_from_all_tool_calls(self): |
| 68 | + """Should extract args from all tool calls.""" |
| 69 | + message = AIMessage( |
| 70 | + content="", |
| 71 | + tool_calls=[ |
| 72 | + {"name": "tool1", "args": {"arg1": "val1"}, "id": "call_1"}, |
| 73 | + {"name": "tool2", "args": {"arg2": "val2"}, "id": "call_2"}, |
| 74 | + {"name": "tool3", "args": {"arg3": "val3"}, "id": "call_3"}, |
| 75 | + ], |
| 76 | + ) |
| 77 | + result = _extract_tools_args_from_message(message) |
| 78 | + assert result == [{"arg1": "val1"}, {"arg2": "val2"}, {"arg3": "val3"}] |
| 79 | + |
| 80 | + def test_returns_empty_list_for_non_ai_message(self): |
| 81 | + """Should return empty list when message is not AIMessage.""" |
| 82 | + message = HumanMessage(content="Test message") |
| 83 | + result = _extract_tools_args_from_message(message) |
| 84 | + assert result == [] |
| 85 | + |
| 86 | + def test_returns_empty_list_when_no_tool_calls(self): |
| 87 | + """Should return empty list when AIMessage has no tool calls.""" |
| 88 | + message = AIMessage(content="Test response") |
| 89 | + result = _extract_tools_args_from_message(message) |
| 90 | + assert result == [] |
| 91 | + |
| 92 | + |
| 93 | +class TestExtractToolOutputData: |
| 94 | + """Tests for _extract_tool_output_data function.""" |
| 95 | + |
| 96 | + def test_extracts_json_dict_content(self): |
| 97 | + """Should parse and return dict when content is JSON string.""" |
| 98 | + json_content = json.dumps({"result": "success", "data": {"value": 42}}) |
| 99 | + state = AgentGuardrailsGraphState( |
| 100 | + messages=[ToolMessage(content=json_content, tool_call_id="call_1")] |
| 101 | + ) |
| 102 | + result = _extract_tool_output_data(state) |
| 103 | + assert result == {"result": "success", "data": {"value": 42}} |
| 104 | + |
| 105 | + def test_wraps_non_json_string_in_output_field(self): |
| 106 | + """Should wrap non-JSON string content in 'output' field.""" |
| 107 | + state = AgentGuardrailsGraphState( |
| 108 | + messages=[ToolMessage(content="Plain text result", tool_call_id="call_1")] |
| 109 | + ) |
| 110 | + result = _extract_tool_output_data(state) |
| 111 | + assert result == {"output": "Plain text result"} |
| 112 | + |
| 113 | + def test_returns_empty_dict_for_empty_messages(self): |
| 114 | + """Should return empty dict when state has no messages.""" |
| 115 | + state = AgentGuardrailsGraphState(messages=[]) |
| 116 | + result = _extract_tool_output_data(state) |
| 117 | + assert result == {} |
| 118 | + |
| 119 | + def test_returns_empty_dict_for_non_tool_message(self): |
| 120 | + """Should return empty dict when last message is not ToolMessage.""" |
| 121 | + state = AgentGuardrailsGraphState( |
| 122 | + messages=[AIMessage(content="Not a tool message")] |
| 123 | + ) |
| 124 | + result = _extract_tool_output_data(state) |
| 125 | + assert result == {} |
| 126 | + |
| 127 | + |
| 128 | +class TestGetMessageContent: |
| 129 | + """Tests for get_message_content function.""" |
| 130 | + |
| 131 | + def test_extracts_string_content_from_human_message(self): |
| 132 | + """Should extract string content from HumanMessage.""" |
| 133 | + message = HumanMessage(content="Hello from human") |
| 134 | + result = get_message_content(message) |
| 135 | + assert result == "Hello from human" |
| 136 | + |
| 137 | + def test_extracts_content_from_ai_message(self): |
| 138 | + """Should extract content from AIMessage.""" |
| 139 | + message = AIMessage(content="AI response") |
| 140 | + result = get_message_content(message) |
| 141 | + assert result == "AI response" |
| 142 | + |
| 143 | + def test_extracts_content_from_tool_message(self): |
| 144 | + """Should extract content from ToolMessage.""" |
| 145 | + message = ToolMessage(content="Tool result", tool_call_id="call_1") |
| 146 | + result = get_message_content(message) |
| 147 | + assert result == "Tool result" |
| 148 | + |
| 149 | + def test_handles_empty_content(self): |
| 150 | + """Should handle empty content string.""" |
| 151 | + message = AIMessage(content="") |
| 152 | + result = get_message_content(message) |
| 153 | + assert result == "" |
0 commit comments