Coverage for intelligence_toolkit/AI/openai_embedder.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.base_embedder import BaseEmbedder 

6from intelligence_toolkit.AI.client import OpenAIClient 

7from intelligence_toolkit.AI.defaults import DEFAULT_CONCURRENT_COROUTINES 

8from intelligence_toolkit.AI.openai_configuration import OpenAIConfiguration 

9from intelligence_toolkit.helpers.constants import CACHE_PATH 

10 

11 

12class OpenAIEmbedder(BaseEmbedder): 

13 def __init__( 

14 self, 

15 configuration: OpenAIConfiguration, 

16 db_name: str = "embeddings", 

17 db_path=CACHE_PATH, 

18 concurrent_coroutines: int | None = DEFAULT_CONCURRENT_COROUTINES, 

19 ): 

20 super().__init__( 

21 db_name, db_path, configuration.max_tokens, concurrent_coroutines 

22 ) 

23 self.configuration = configuration 

24 self.openai_client = OpenAIClient(configuration) 

25 

26 def _generate_embedding(self, text: str) -> list[float]: 

27 return self.openai_client.generate_embedding( 

28 text, model=self.configuration.embedding_model 

29 ) 

30 

31 async def _generate_embedding_async(self, text: str) -> list[float]: 

32 return await self.openai_client.generate_embedding_async( 

33 text, model=self.configuration.embedding_model 

34 )