Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/context_compression/abstractive.py: 84%

1"""Abstractive compression strategies.""" 

2 

3from __future__ import annotations 

4 

5from datetime import UTC, datetime 

6 

7from lexigram.ai.rag.context_compression.base import AbstractCompressor 

8from lexigram.ai.rag.context_compression.types import ( 

9 CompressionResult, 

10 CompressionStrategy, 

11) 

12from lexigram.contracts import ( 

13 ChatMessage, 

14 LLMClientProtocol, 

15) 

16 

17 

18class AbstractiveCompressor(AbstractCompressor): 

19 """LLM-based abstractive compression. 

20 

21 Uses an LLM to generate a compressed summary that preserves 

22 the most important information relevant to the query. 

23 

24 Example: 

25 >>> compressor = AbstractiveCompressor( 

26 ... llm_client=llm, 

27 ... max_tokens=150, 

28 ... temperature=0.3 

29 ... ) 

30 >>> result = await compressor.compress( 

31 ... context=long_article, 

32 ... query="What are the main findings?" 

33 ... ) 

34 """ 

35 

36 def __init__( 

37 self, 

38 llm_client: LLMClientProtocol, 

39 max_tokens: int = 200, 

40 temperature: float = 0.3, 

41 ): 

42 """Initialize abstractive compressor. 

43 

44 Args: 

45 llm_client: LLM client for summarization. 

46 max_tokens: Maximum tokens for compressed output. 

47 temperature: Temperature for generation. 

48 """ 

49 self.llm_client = llm_client 

50 self.max_tokens = max_tokens 

51 self.temperature = temperature 

52 

53 async def compress( 

54 self, 

55 context: str | list[str], 

56 query: str | None = None, 

57 **kwargs, 

58 ) -> CompressionResult: 

59 """Compress using LLM-based summarization.""" 

60 original_text = self._normalize_context(context) 

61 original_tokens = self._estimate_tokens(original_text) 

62 

63 # Build summarization prompt 

64 prompt = self._build_prompt(original_text, query) 

65 

66 # Generate summary 

67 result = await self.llm_client.complete( 

68 messages=[ 

69 ChatMessage( 

70 role="system", 

71 content="You are a helpful assistant that creates concise, information-dense summaries.", 

72 ), 

73 ChatMessage( 

74 role="user", 

75 content=prompt, 

76 ), 

77 ], 

78 temperature=self.temperature, 

79 max_tokens=self.max_tokens, 

80 ) 

81 if result.is_err(): 

82 raise result.unwrap_err() 

83 response = result.unwrap() 

84 

85 # Extract compressed text 

86 if hasattr(response, "content"): 

87 compressed_text = response.content 

88 elif hasattr(response, "choices") and response.choices: 

89 compressed_text = response.choices[0].message.content 

90 elif isinstance(response, dict) and "content" in response: 

91 compressed_text = response["content"] 

92 else: 

93 compressed_text = str(response) 

94 

95 compressed_tokens = self._estimate_tokens(compressed_text) 

96 compression_ratio = ( 

97 compressed_tokens / original_tokens if original_tokens > 0 else 1.0 

98 ) 

99 

100 return CompressionResult( 

101 original_text=original_text, 

102 compressed_text=compressed_text, 

103 original_tokens=original_tokens, 

104 compressed_tokens=compressed_tokens, 

105 compression_ratio=compression_ratio, 

106 strategy=CompressionStrategy.ABSTRACTIVE, 

107 metadata={ 

108 "max_tokens": self.max_tokens, 

109 "temperature": self.temperature, 

110 "query_used": query is not None, 

111 "timestamp": datetime.now(UTC).isoformat(), 

112 }, 

113 ) 

114 

115 def _build_prompt(self, text: str, query: str | None) -> str: 

116 """Build summarization prompt.""" 

117 prompt = f"Text to compress:\n\n{text}\n\n" 

118 

119 if query: 

120 prompt += f"Focus on information relevant to: {query}\n\n" 

121 

122 prompt += ( 

123 f"Provide a concise summary in no more than {self.max_tokens // 4} words " 

124 f"that preserves the most important information" 

125 ) 

126 

127 if query: 

128 prompt += " relevant to the query" 

129 

130 prompt += ":" 

131 

132 return prompt