Coverage for intelligence_toolkit/tests/unit/helpers/test_classes.py: 100%

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

2# Licensed under the MIT license. See LICENSE file in the project. 

3# 

4from unittest.mock import patch 

5 

6import pytest 

7 

8from intelligence_toolkit.AI.openai_configuration import OpenAIConfiguration 

9from intelligence_toolkit.helpers.classes import IntelligenceWorkflow 

10 

11 

12@pytest.fixture 

13def openai_config(): 

14 return OpenAIConfiguration({ 

15 "api_key": "test_key", 

16 "model": "gpt-4", 

17 "api_type": "OpenAI", 

18 }) 

19 

20 

21def test_intelligence_workflow_initialization_without_config(): 

22 workflow = IntelligenceWorkflow() 

23 assert workflow.ai_configuration is None 

24 assert workflow.embedder is None 

25 assert workflow.cache_embeddings is True 

26 

27 

28def test_intelligence_workflow_initialization_with_config(openai_config): 

29 workflow = IntelligenceWorkflow(openai_config) 

30 assert workflow.ai_configuration == openai_config 

31 assert workflow.embedder is None 

32 assert workflow.cache_embeddings is True 

33 

34 

35def test_set_ai_configuration(openai_config): 

36 workflow = IntelligenceWorkflow() 

37 

38 with patch.object(workflow, 'set_embedder') as mock_set_embedder: 

39 workflow.set_ai_configuration(openai_config) 

40 

41 assert workflow.ai_configuration == openai_config 

42 mock_set_embedder.assert_called_once() 

43 

44 

45def test_set_ai_configuration_with_existing_embedder(openai_config): 

46 workflow = IntelligenceWorkflow() 

47 workflow.embedder = "mock_embedder" 

48 

49 with patch.object(workflow, 'set_embedder') as mock_set_embedder: 

50 workflow.set_ai_configuration(openai_config) 

51 

52 assert workflow.ai_configuration == openai_config 

53 # Should not call set_embedder if embedder already exists 

54 mock_set_embedder.assert_not_called() 

55 

56 

57def test_set_embedder_local(openai_config): 

58 with patch("intelligence_toolkit.helpers.classes.LocalEmbedder") as mock_local: 

59 workflow = IntelligenceWorkflow(openai_config) 

60 workflow.set_embedder(local_embedding=True, cache_embeddings=False) 

61 

62 mock_local.assert_called_once() 

63 assert workflow.embedder is not None 

64 assert workflow.cache_embeddings is False 

65 

66 

67def test_set_embedder_openai(openai_config): 

68 with patch("intelligence_toolkit.helpers.classes.OpenAIEmbedder") as mock_openai, \ 

69 patch("intelligence_toolkit.AI.client.OpenAI"), \ 

70 patch("intelligence_toolkit.AI.client.AsyncOpenAI"): 

71 

72 workflow = IntelligenceWorkflow(openai_config) 

73 workflow.set_embedder(local_embedding=False, cache_embeddings=True) 

74 

75 mock_openai.assert_called_once_with(openai_config) 

76 assert workflow.embedder is not None 

77 assert workflow.cache_embeddings is True 

78 

79 

80def test_set_embedder_default_parameters(openai_config): 

81 with patch("intelligence_toolkit.helpers.classes.OpenAIEmbedder") as mock_openai, \ 

82 patch("intelligence_toolkit.AI.client.OpenAI"), \ 

83 patch("intelligence_toolkit.AI.client.AsyncOpenAI"): 

84 

85 workflow = IntelligenceWorkflow(openai_config) 

86 workflow.set_embedder() 

87 

88 # Default is OpenAI embedder with cache_embeddings=True 

89 mock_openai.assert_called_once() 

90 assert workflow.cache_embeddings is True