Coverage for intelligence_toolkit/helpers/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# 

4 

5from intelligence_toolkit.AI.local_embedder import LocalEmbedder 

6from intelligence_toolkit.AI.openai_configuration import OpenAIConfiguration 

7from intelligence_toolkit.AI.openai_embedder import OpenAIEmbedder 

8 

9 

10class IntelligenceWorkflow: 

11 embedder = None 

12 cache_embeddings = True 

13 

14 # Base class for all AI workflows 

15 def __init__(self, ai_configuration: OpenAIConfiguration | None = None) -> None: 

16 self.ai_configuration = ai_configuration 

17 

18 def set_ai_configuration(self, ai_configuration: OpenAIConfiguration) -> None: 

19 self.ai_configuration = ai_configuration 

20 if not self.embedder: 

21 self.set_embedder() 

22 

23 def set_embedder( 

24 self, local_embedding: bool = False, cache_embeddings: bool = True 

25 ): 

26 if local_embedding: 

27 self.embedder = LocalEmbedder() 

28 else: 

29 self.embedder = OpenAIEmbedder(self.ai_configuration) 

30 self.cache_embeddings = cache_embeddings