Coverage for src/lilbee/providers/fleet/windowing.py: 100%

1"""Fit a chat message list to a served context window by dropping oldest turns.""" 

2 

3from __future__ import annotations 

4 

5import json 

6import math 

7from dataclasses import dataclass 

8from typing import Any 

9 

10from lilbee.providers.base import CHAT_MESSAGE_OVERHEAD_TOKENS 

11 

12_SYSTEM_ROLE = "system" 

13_TOOL_ROLE = "tool" 

14# Conservative chars-per-token: below the ~4 English average so the estimate 

15# over-counts tokens and the window errs toward dropping more, never overflowing. 

16_CHARS_PER_TOKEN = 3 

17# Per-message token overhead for the role markers and chat-template wrappers the 

18# server adds around each message. Owned by base; retrieval budgets the same number. 

19 

20 

21def estimate_tokens(text: str) -> int: 

22 """Conservative token estimate for a text fragment.""" 

23 return math.ceil(len(text) / _CHARS_PER_TOKEN) 

24 

25 

26def _message_tokens(message: dict[str, Any]) -> int: 

27 """Estimated tokens a wire message contributes (content + tool-call JSON + overhead).""" 

28 total = CHAT_MESSAGE_OVERHEAD_TOKENS 

29 content = message.get("content") 

30 if isinstance(content, str): 

31 total += estimate_tokens(content) 

32 elif content: 

33 total += estimate_tokens(json.dumps(content)) 

34 tool_calls = message.get("tool_calls") 

35 if tool_calls: 

36 total += estimate_tokens(json.dumps(tool_calls)) 

37 return total 

38 

39 

40def _tools_tokens(tools: list[dict[str, Any]] | None) -> int: 

41 """Estimated tokens the tool schemas contribute to the rendered prompt.""" 

42 if not tools: 

43 return 0 

44 return estimate_tokens(json.dumps(tools)) 

45 

46 

47@dataclass(frozen=True) 

48class WindowResult: 

49 """Outcome of fitting messages to a budget.""" 

50 

51 messages: list[dict[str, Any]] # system + kept suffix (best-effort even on overflow) 

52 fits: bool 

53 prompt_tokens: int # estimated tokens of ``messages`` plus the tools passed in 

54 dropped: int # number of conversation messages dropped 

55 

56 

57def window_messages( 

58 messages: list[dict[str, Any]], 

59 tools: list[dict[str, Any]] | None, 

60 budget: int, 

61) -> WindowResult: 

62 """Drop oldest conversation turns until the estimated prompt fits ``budget``. 

63 

64 System messages and the most recent turn are always kept; tool-call/result 

65 pairs drop together (a kept suffix never starts with an orphan ``tool`` 

66 message whose originating call was dropped). ``fits`` is False when even the 

67 system messages, tools, and the final message exceed the budget; the caller 

68 turns that into a context-overflow error. 

69 """ 

70 system = [m for m in messages if m.get("role") == _SYSTEM_ROLE] 

71 convo = [m for m in messages if m.get("role") != _SYSTEM_ROLE] 

72 fixed = sum(_message_tokens(m) for m in system) + _tools_tokens(tools) 

73 

74 if not convo: 

75 return WindowResult(list(system), fixed <= budget, fixed, 0) 

76 

77 # Keep conversation messages newest-first while they fit; the most recent is 

78 # always kept (an empty guard) so the current turn survives. 

79 kept_rev: list[dict[str, Any]] = [] 

80 used = fixed 

81 for msg in reversed(convo): 

82 cost = _message_tokens(msg) 

83 if kept_rev and used + cost > budget: 

84 break 

85 kept_rev.append(msg) 

86 used += cost 

87 kept = list(reversed(kept_rev)) 

88 

89 # A kept suffix must not begin with an orphan tool result (its call dropped). 

90 while kept and kept[0].get("role") == _TOOL_ROLE: 

91 kept = kept[1:] 

92 

93 prompt_tokens = fixed + sum(_message_tokens(m) for m in kept) 

94 fits = prompt_tokens <= budget and bool(kept) 

95 return WindowResult(system + kept, fits, prompt_tokens, len(convo) - len(kept))