Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-session/src/lexigram/ai/session/context/pruner.py: 21%

1"""Relevance-based conversation history pruner.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING 

6 

7if TYPE_CHECKING: 

8 from lexigram.contracts.ai.llm import ChatMessageProtocol 

9 

10 

11def _tokenize(text: str) -> list[str]: 

12 """Split text into lowercase word tokens. 

13 

14 Args: 

15 text: Input string. 

16 

17 Returns: 

18 List of lowercase, whitespace-delimited tokens. 

19 """ 

20 return [w.lower() for w in text.split() if w] 

21 

22 

23class RelevanceContextPruner: 

24 """Keyword-overlap relevance pruner for conversation history. 

25 

26 Implements ``ContextPrunerProtocol``. When the history exceeds 

27 ``max_turns``, system messages are always preserved and non-system 

28 turns are scored by keyword overlap with the current query plus a 

29 recency bonus, then the top-scoring turns are retained in their 

30 original chronological order. 

31 """ 

32 

33 async def prune( 

34 self, 

35 history: list[ChatMessageProtocol], 

36 current_query: str, 

37 max_turns: int, 

38 ) -> list[ChatMessageProtocol]: 

39 """Prune history to at most ``max_turns``, keeping relevant turns. 

40 

41 Scoring formula per non-system turn:: 

42 

43 score = overlap_count + recency_ratio * 0.3 

44 

45 where ``overlap_count`` is the number of tokens shared between the 

46 turn content and ``current_query``, and ``recency_ratio`` is the 

47 turn's normalised position in the non-system history (0 = oldest, 

48 1 = newest). 

49 

50 Args: 

51 history: Full conversation history. 

52 current_query: The current user query for relevance scoring. 

53 max_turns: Maximum number of turns to retain. 

54 

55 Returns: 

56 Pruned history in chronological order. 

57 """ 

58 if len(history) <= max_turns: 

59 return list(history) 

60 

61 system_msgs = [m for m in history if str(m.role).lower() == "system"] 

62 non_system = [m for m in history if str(m.role).lower() != "system"] 

63 

64 system_count = len(system_msgs) 

65 available_slots = max_turns - system_count 

66 

67 if available_slots <= 0: 

68 return system_msgs[:max_turns] 

69 

70 query_tokens = set(_tokenize(current_query)) 

71 total = len(non_system) 

72 normalization = max(total - 1, 1) 

73 

74 scored: list[tuple[float, int, ChatMessageProtocol]] = [] 

75 for i, msg in enumerate(non_system): 

76 content = str(msg.content) if msg.content is not None else "" 

77 msg_tokens = set(_tokenize(content)) 

78 overlap = len(query_tokens & msg_tokens) 

79 recency = i / normalization 

80 score = float(overlap) + recency * 0.3 

81 scored.append((score, i, msg)) 

82 

83 scored.sort(key=lambda x: x[0], reverse=True) 

84 selected_ids = {id(item[2]) for item in scored[:available_slots]} 

85 

86 return [ 

87 msg 

88 for msg in history 

89 if str(msg.role).lower() == "system" or id(msg) in selected_ids 

90 ] 

91 

92 

93__all__ = ["RelevanceContextPruner"]