Coverage for intelligence_toolkit/tests/unit/AI/test_utils.py: 100%

1# Copyright (c) 2024 Microsoft Corporation. All rights reserved. 

2import pytest 

3 

4from intelligence_toolkit.AI.utils import ( 

5 get_token_count, 

6 hash_text, 

7 prepare_messages, 

8 prepare_validation, 

9 try_parse_json_object, 

10) 

11 

12 

13def test_get_token_count(): 

14 text = "example text" 

15 result = get_token_count(text) 

16 expected = 4 

17 

18 assert result == expected 

19 

20 

21def test_get_token_count_with_model(): 

22 text = "example text" 

23 result = get_token_count(text, None, "gpt-4.1-mini") 

24 expected = 4 

25 assert result == expected 

26 

27 

28def test_get_token_count_with_encoding(): 

29 text = "example text" 

30 result = get_token_count(text, "o200k_base") 

31 expected = 4 

32 assert result == expected 

33 

34 

35def test_hash_text(): 

36 text = "example test \n\n example test two" 

37 hash_returned = hash_text(text) 

38 assert ( 

39 hash_returned 

40 == "5aaa5f424725140ddfeb0f5747e7543a52dc18ebde854b83e554f155ba2abfab" 

41 ) 

42 

43 

44def test_try_parse_json_object_ok(): 

45 obj_test = '{"key": "value"}' 

46 result = try_parse_json_object(obj_test) 

47 assert isinstance(result, dict) 

48 

49 

50def test_try_parse_json_object_exception(): 

51 obj_test = {"key": "value"} 

52 with pytest.raises(Exception): 

53 try_parse_json_object(obj_test) 

54 

55 

56def test_prepare_validation(): 

57 messages = [{"role": "user", "content": "Write me a poem about a cat"}] 

58 report = "The cat poem: The cat is a furry animal that likes to play with yarn. It is a pet that is loved by many people. The cat is a furry animal that likes to play with yarn. It is a pet that is loved by many people." 

59 message = prepare_validation(messages, report) 

60 expected = "Write me a poem about a cat" 

61 assert expected in message[0]["content"] 

62 

63 

64def test_prepare_messages(): 

65 variables = {"animal": "cat", "place": "tree"} 

66 system_message = "I can write you a poem about a {animal} in a {place}" 

67 message = prepare_messages(system_message, variables) 

68 expected = "I can write you a poem about a cat in a tree" 

69 assert message[0]["content"] == expected 

70 

71 

72def test_prepare_messages_user(): 

73 variables = {"animal": "cat", "place": "tree", "rescue": "fireman"} 

74 system_message = "I can write you a poem about a {animal} in a {place}" 

75 user_message = "Make the {rescue} save the cat" 

76 message = prepare_messages(system_message, variables, user_message) 

77 assert message[0]["content"] == "I can write you a poem about a cat in a tree" 

78 assert message[1]["content"] == "Make the fireman save the cat" 

79 

80 

81def test_generate_messages(): 

82 from intelligence_toolkit.AI.utils import generate_messages 

83 

84 user_prompt = "Write a story about a {animal}" 

85 system_prompt = "You are a helpful assistant." 

86 variables = {"animal": "dog"} 

87 safety_prompt = "Keep it family friendly." 

88 

89 messages = generate_messages(user_prompt, system_prompt, variables, safety_prompt) 

90 

91 assert len(messages) == 1 

92 assert messages[0]["role"] == "system" 

93 assert "dog" in messages[0]["content"] 

94 assert "helpful assistant" in messages[0]["content"] 

95 assert "family friendly" in messages[0]["content"] 

96 

97 

98def test_try_parse_json_object_invalid_json(): 

99 obj_test = '{"key": invalid}' 

100 with pytest.raises(Exception): 

101 try_parse_json_object(obj_test) 

102 

103 

104def test_try_parse_json_object_not_dict(): 

105 obj_test = '["item1", "item2"]' 

106 with pytest.raises(TypeError): 

107 try_parse_json_object(obj_test)