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 

10_SYSTEM_ROLE = "system" 

11_TOOL_ROLE = "tool" 

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

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

14_CHARS_PER_TOKEN = 3 

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

16# server adds around each message. 

17_PER_MESSAGE_OVERHEAD = 8 

18 

19 

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

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

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

23 

24 

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

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

27 total = _PER_MESSAGE_OVERHEAD 

28 content = message.get("content") 

29 if isinstance(content, str): 

30 total += estimate_tokens(content) 

31 elif content: 

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

33 tool_calls = message.get("tool_calls") 

34 if tool_calls: 

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

36 return total 

37 

38 

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

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

41 if not tools: 

42 return 0 

43 return estimate_tokens(json.dumps(tools)) 

44 

45 

46@dataclass(frozen=True) 

47class WindowResult: 

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

49 

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

51 fits: bool 

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

53 dropped: int # number of conversation messages dropped 

54 

55 

56def window_messages( 

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

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

59 budget: int, 

60) -> WindowResult: 

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

62 

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

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

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

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

67 turns that into a context-overflow error. 

68 """ 

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

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

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

72 

73 if not convo: 

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

75 

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

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

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

79 used = fixed 

80 for msg in reversed(convo): 

81 cost = _message_tokens(msg) 

82 if kept_rev and used + cost > budget: 

83 break 

84 kept_rev.append(msg) 

85 used += cost 

86 kept = list(reversed(kept_rev)) 

87 

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

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

90 kept = kept[1:] 

91 

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

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

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