Coverage for src/lilbee/retrieval/query/history_window.py: 100%

1"""Token-budget history windowing for chat conversations.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from typing import TYPE_CHECKING 

7 

8from lilbee.data.extract.chunk import CHARS_PER_TOKEN 

9 

10if TYPE_CHECKING: 

11 from lilbee.retrieval.query.searcher import ChatMessage 

12 

13 

14def estimate_text_tokens(text: str) -> int: 

15 """Cheap char/4 token estimate for a string.""" 

16 return max(1, len(text) // CHARS_PER_TOKEN) 

17 

18 

19def chars_for_tokens(tokens: int) -> int: 

20 """Rough char budget for a token budget: the inverse of estimate_text_tokens. 

21 

22 The chars-per-token ratio itself is owned by :mod:`lilbee.data.extract.chunk`, the 

23 one place it is defined; this is the token->char direction of it. 

24 """ 

25 return tokens * CHARS_PER_TOKEN 

26 

27 

28def estimate_tokens(message: ChatMessage) -> int: 

29 """Cheap char/4 token estimate for one message.""" 

30 return estimate_text_tokens(message["content"]) 

31 

32 

33def windowed_history( 

34 messages: list[ChatMessage], 

35 *, 

36 max_tokens: int, 

37 estimator: Callable[[ChatMessage], int] = estimate_tokens, 

38) -> list[ChatMessage]: 

39 """Return the suffix of *messages* whose token cost fits in *max_tokens*. 

40 

41 Drops messages from the front in pairs so the window starts at a user 

42 message; never strands an orphan assistant reply with no preceding user 

43 turn for the model to anchor to. The newest pair is always kept even 

44 if it exceeds the budget on its own (caller decides what to do then). 

45 

46 A non-positive *max_tokens* disables windowing and returns everything, 

47 rather than windowing hardest. No production caller can reach it today 

48 (the context target has a floor), but a caller deriving a budget that 

49 goes non-positive gets the full history, not an empty one. 

50 """ 

51 if max_tokens <= 0 or not messages: 

52 return list(messages) 

53 sizes = [estimator(m) for m in messages] 

54 total = sum(sizes) 

55 if total <= max_tokens: 

56 return list(messages) 

57 start = 0 

58 # ``len(messages) - 2`` keeps the newest user/assistant pair even when it 

59 # exceeds the budget on its own. The caller decides what to do if the 

60 # final pair is over-sized (typically: send it anyway and let the chat 

61 # server error if it must, rather than send nothing at all). 

62 while start < len(messages) - 2 and total > max_tokens: 

63 # Drop the front pair (user + assistant). If the front isn't a user 

64 # message (malformed input), drop one to realign. 

65 drop = 2 if messages[start]["role"] == "user" else 1 

66 for i in range(start, min(start + drop, len(messages))): 

67 total -= sizes[i] 

68 start += drop 

69 return list(messages[start:])