Coverage for src/qdrant_loader/core/embedding_service.py: 100%

1import logging 

2 

3 

4class EmbeddingService: 

5 def __init__(self, client, model_name): 

6 self.client = client 

7 self.model_name = model_name 

8 self.logger = logging.getLogger(__name__) 

9 

10 async def get_embedding(self, text: str) -> list[float]: 

11 """Get embedding for a single text. 

12 

13 Args: 

14 text: The text to get embedding for 

15 

16 Returns: 

17 List of floats representing the embedding vector 

18 """ 

19 self.logger.debug( 

20 "Getting embedding for text", 

21 extra={"text_length": len(text), "model": self.model_name}, 

22 ) 

23 

24 try: 

25 # Get embedding from OpenAI 

26 response = await self.client.embeddings.create( 

27 model=self.model_name, input=text 

28 ) 

29 

30 embedding = response.data[0].embedding 

31 self.logger.debug( 

32 "Successfully generated embedding", 

33 extra={"text_length": len(text), "embedding_size": len(embedding)}, 

34 ) 

35 return embedding 

36 except Exception as e: 

37 self.logger.error( 

38 f"Error generating embedding: {str(e)}", 

39 extra={ 

40 "text_length": len(text), 

41 "error": str(e), 

42 "error_type": type(e).__name__, 

43 }, 

44 ) 

45 raise 

46 

47 async def get_embeddings(self, texts: list[str]) -> list[list[float]]: 

48 """Get embeddings for multiple texts. 

49 

50 Args: 

51 texts: List of texts to get embeddings for 

52 

53 Returns: 

54 List of embedding vectors 

55 """ 

56 self.logger.debug( 

57 "Getting embeddings for texts", 

58 extra={"text_count": len(texts), "model": self.model_name}, 

59 ) 

60 

61 try: 

62 # Get embeddings from OpenAI 

63 response = await self.client.embeddings.create( 

64 model=self.model_name, input=texts 

65 ) 

66 

67 embeddings = [data.embedding for data in response.data] 

68 self.logger.debug( 

69 "Successfully generated embeddings", 

70 extra={ 

71 "text_count": len(texts), 

72 "embedding_count": len(embeddings), 

73 "embedding_size": len(embeddings[0]) if embeddings else 0, 

74 }, 

75 ) 

76 return embeddings 

77 except Exception as e: 

78 self.logger.error( 

79 f"Error generating embeddings: {str(e)}", 

80 extra={ 

81 "text_count": len(texts), 

82 "error": str(e), 

83 "error_type": type(e).__name__, 

84 }, 

85 ) 

86 raise