Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/hyde/generators/multiple.py: 96%

1"""Multiple HyDE generator implementation.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any 

6 

7from lexigram.ai.rag.hyde.base import AbstractHyDEGenerator 

8from lexigram.ai.rag.hyde.protocols import EmbeddingClientProtocol 

9from lexigram.ai.rag.hyde.types import HyDEResult, HyDEStrategy, HypotheticalDocument 

10from lexigram.contracts import ( 

11 LLMClientProtocol, 

12) 

13 

14 

15class MultipleHyDEGenerator(AbstractHyDEGenerator): 

16 """Generator for multiple hypothetical documents.""" 

17 

18 def __init__( 

19 self, 

20 llm_client: LLMClientProtocol, 

21 embedding_client: EmbeddingClientProtocol | None = None, 

22 temperature: float = 0.9, 

23 max_tokens: int = 150, 

24 default_num_documents: int = 3, 

25 ): 

26 """Initialize multiple HyDE generator. 

27 

28 Args: 

29 llm_client: Client for generating hypothetical documents 

30 embedding_client: Optional client for generating embeddings 

31 temperature: Higher temperature for diversity 

32 max_tokens: Maximum tokens per document 

33 default_num_documents: Default number of documents 

34 """ 

35 super().__init__(llm_client, embedding_client) 

36 self.temperature = temperature 

37 self.max_tokens = max_tokens 

38 self.default_num_documents = default_num_documents 

39 

40 async def generate( 

41 self, 

42 query: str, 

43 num_documents: int | None = None, 

44 **kwargs: Any, 

45 ) -> HyDEResult: 

46 """Generate multiple hypothetical documents. 

47 

48 Args: 

49 query: User query 

50 num_documents: Number of documents (default: default_num_documents) 

51 **kwargs: Additional parameters (context, domain) 

52 

53 Returns: 

54 HyDE result with multiple hypothetical documents 

55 """ 

56 if num_documents is None: 

57 num_documents = self.default_num_documents 

58 

59 # Generate multiple documents 

60 documents = [] 

61 for i in range(num_documents): 

62 content = await self._generate_single_document( 

63 query, 

64 temperature=self.temperature, 

65 max_tokens=self.max_tokens, 

66 **kwargs, 

67 ) 

68 

69 doc = HypotheticalDocument( 

70 content=content, 

71 query=query, 

72 confidence=1.0 / (i + 1), # Decrease confidence for later docs 

73 metadata={ 

74 "index": i, 

75 "temperature": self.temperature, 

76 "max_tokens": self.max_tokens, 

77 }, 

78 ) 

79 documents.append(doc) 

80 

81 # Generate embeddings if client available 

82 aggregated_embedding = None 

83 if self.embedding_client: 

84 embeddings = await self._embed_documents(documents) 

85 # Average embeddings 

86 aggregated_embedding = self._aggregate_embeddings(embeddings) 

87 

88 return HyDEResult( 

89 query=query, 

90 hypothetical_docs=documents, 

91 strategy=HyDEStrategy.MULTIPLE, 

92 aggregated_embedding=aggregated_embedding, 

93 metadata={ 

94 "temperature": self.temperature, 

95 "max_tokens": self.max_tokens, 

96 "num_documents": num_documents, 

97 }, 

98 )