ctx-compact 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ctx_compact-0.1.0/.gitignore +3 -0
- ctx_compact-0.1.0/PKG-INFO +141 -0
- ctx_compact-0.1.0/README.md +124 -0
- ctx_compact-0.1.0/ctx_compact.py +380 -0
- ctx_compact-0.1.0/py.typed +0 -0
- ctx_compact-0.1.0/pyproject.toml +27 -0
- ctx_compact-0.1.0/test_ctx_compact.py +315 -0
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: ctx-compact
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Trim a plain OpenAI-shaped message array down to a token budget without ever orphaning a tool result.
|
|
5
|
+
Project-URL: Homepage, https://github.com/pjdurden/ctx-compact
|
|
6
|
+
Project-URL: Source, https://github.com/pjdurden/ctx-compact
|
|
7
|
+
Author: Prajjwal Chittori
|
|
8
|
+
License: MIT
|
|
9
|
+
Keywords: compaction,context-window,conversation,llm,openai,token-budget,tool-calls
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Typing :: Typed
|
|
15
|
+
Requires-Python: >=3.9
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# ctx-compact
|
|
19
|
+
|
|
20
|
+
Trim a plain OpenAI-shaped message array down to a token budget without ever orphaning a tool result. This is the Python port of the [ctx-compact](https://www.npmjs.com/package/ctx-compact) npm package.
|
|
21
|
+
|
|
22
|
+
## The problem
|
|
23
|
+
|
|
24
|
+
Every agent framework ships its own conversation compaction (LangGraph, Inspect, MS Agent Framework, the Claude SDK all have one), and each is welded to that framework's message type. If you are working with a plain list of `{role, content, ...}` messages, people tend to hand-roll a "drop the oldest N messages" loop. That works until an assistant message with `tool_calls` gets dropped but its matching `tool` result messages do not (or the reverse). Most providers reject that shape outright, so the trim silently turns into an API error on the next call. `ctx-compact` is a small, framework-neutral compactor that keeps tool-call and tool-result messages paired and dropped or kept as a unit.
|
|
25
|
+
|
|
26
|
+
## Install
|
|
27
|
+
|
|
28
|
+
```
|
|
29
|
+
pip install ctx-compact
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
## Usage
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
from ctx_compact import compact, compact_with_summary
|
|
36
|
+
|
|
37
|
+
messages = [
|
|
38
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
39
|
+
{"role": "user", "content": "What is the weather in Denver?"},
|
|
40
|
+
{
|
|
41
|
+
"role": "assistant",
|
|
42
|
+
"content": None,
|
|
43
|
+
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city":"Denver"}'}}],
|
|
44
|
+
},
|
|
45
|
+
{"role": "tool", "tool_call_id": "call_1", "content": '{"tempF":72}'},
|
|
46
|
+
{"role": "assistant", "content": "It is 72F in Denver."},
|
|
47
|
+
{"role": "user", "content": "What about Austin?"},
|
|
48
|
+
{
|
|
49
|
+
"role": "assistant",
|
|
50
|
+
"content": None,
|
|
51
|
+
"tool_calls": [{"id": "call_2", "type": "function", "function": {"name": "get_weather", "arguments": '{"city":"Austin"}'}}],
|
|
52
|
+
},
|
|
53
|
+
{"role": "tool", "tool_call_id": "call_2", "content": '{"tempF":88}'},
|
|
54
|
+
{"role": "assistant", "content": "It is 88F in Austin."},
|
|
55
|
+
{"role": "user", "content": "And tomorrow in Denver?"},
|
|
56
|
+
]
|
|
57
|
+
|
|
58
|
+
result = compact(messages, max_tokens=150, keep_head=1, keep_tail=2)
|
|
59
|
+
# result.tokens_before -> 195
|
|
60
|
+
# result.tokens_after -> 124
|
|
61
|
+
# result.fits -> True (124 <= 150)
|
|
62
|
+
# result.dropped -> the 3 oldest droppable messages: the first "What is
|
|
63
|
+
# the weather in Denver?" turn and its whole
|
|
64
|
+
# assistant/tool_calls + tool group
|
|
65
|
+
|
|
66
|
+
# Synchronous variant: summarize whatever got dropped and splice a note back in.
|
|
67
|
+
with_summary = compact_with_summary(
|
|
68
|
+
messages,
|
|
69
|
+
max_tokens=150,
|
|
70
|
+
keep_head=1,
|
|
71
|
+
keep_tail=2,
|
|
72
|
+
summarize=lambda dropped: f"Earlier in this conversation: {len(dropped)} messages were removed.",
|
|
73
|
+
)
|
|
74
|
+
# with_summary.summary -> 'Earlier in this conversation: 3 messages were removed.'
|
|
75
|
+
# with_summary.tokens_after -> 145 (the 124 kept after dropping, plus the inserted summary message)
|
|
76
|
+
# with_summary.fits -> True (145 <= 150)
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
The example above is exact output from running this code against this package.
|
|
80
|
+
|
|
81
|
+
## API
|
|
82
|
+
|
|
83
|
+
### `compact(messages, *, max_tokens, count_tokens=None, keep_head=1, keep_tail=4) -> CompactResult`
|
|
84
|
+
|
|
85
|
+
Drops messages from the middle of `messages` until the estimated token count fits the budget.
|
|
86
|
+
|
|
87
|
+
- `messages`: list of dicts shaped like `{role, content, tool_calls?, tool_call_id?, name?}`. `role` is one of `'system' | 'user' | 'assistant' | 'tool'`.
|
|
88
|
+
- `max_tokens` (required, keyword-only, `float`) - the budget. Raises `TypeError` if missing or not a positive number.
|
|
89
|
+
- `count_tokens` - `(message) -> int`. Default: `math.ceil(len(json.dumps(message, separators=(",", ":"), ensure_ascii=False)) / 4)`, with the length measured in UTF-16 code units (matching JavaScript's `String.length`), not Python codepoints.
|
|
90
|
+
- `keep_head` - number of leading messages always kept. Default `1`.
|
|
91
|
+
- `keep_tail` - number of trailing messages always kept. Default `4`.
|
|
92
|
+
|
|
93
|
+
Returns a frozen dataclass:
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
@dataclass(frozen=True)
|
|
97
|
+
class CompactResult:
|
|
98
|
+
messages: list # the compacted list
|
|
99
|
+
dropped: list # what was removed, in original order
|
|
100
|
+
tokens_before: int # summed estimated tokens of the input list
|
|
101
|
+
tokens_after: int # summed estimated tokens of the output list
|
|
102
|
+
fits: bool # tokens_after <= max_tokens
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
If the input already fits, it is returned unchanged with `dropped: []`.
|
|
106
|
+
|
|
107
|
+
### `compact_with_summary(messages, *, max_tokens, count_tokens=None, keep_head=1, keep_tail=4, summarize=None, summary_role="user") -> CompactResultWithSummary`
|
|
108
|
+
|
|
109
|
+
Same as `compact`, then, if anything was dropped and `summarize` is provided, calls `summarize(dropped)` and inserts the returned string as a message `{"role": summary_role, "content": <summary>}` immediately after the head-kept messages.
|
|
110
|
+
|
|
111
|
+
- `summarize` - `(dropped: list) -> str`. If omitted, behaves exactly like `compact` and returns `summary: None`.
|
|
112
|
+
- `summary_role` - default `'user'`. Some providers reject a second `system` message, which is why the default is `'user'` rather than `'system'`.
|
|
113
|
+
|
|
114
|
+
The inserted summary message counts toward the budget: after insertion the result is re-checked, and if it no longer fits, additional whole groups are dropped (oldest first) to make room. `fits` is reported `False` if it is still over budget after that.
|
|
115
|
+
|
|
116
|
+
Returns a frozen dataclass with the same fields as `CompactResult` plus `summary: Optional[str]`.
|
|
117
|
+
|
|
118
|
+
**Difference from the JavaScript version:** the JS `compactWithSummary` is `async` and does `await options.summarize(dropped)`, since JS summarizers are typically an async LLM call. This Python port is **synchronous**: `summarize` is a plain callable, `dropped -> str`, called directly with no `await`. If your summarizer needs to be async in Python, run it yourself (e.g. via `asyncio.run` or your event loop) before calling `compact_with_summary`, and pass a synchronous wrapper.
|
|
119
|
+
|
|
120
|
+
### `estimate_tokens(message, count_tokens=None) -> int`
|
|
121
|
+
|
|
122
|
+
Runs `count_tokens` (or the default heuristic) against a single message. Exported so callers can reuse the same estimator `compact`/`compact_with_summary` use, e.g. to pre-check a message before appending it.
|
|
123
|
+
|
|
124
|
+
## How it works
|
|
125
|
+
|
|
126
|
+
1. **Group first.** Before anything is dropped, the whole list is split into groups: an assistant message carrying `tool_calls` plus every immediately-following `tool` message whose `tool_call_id` matches one of that assistant's `tool_calls[].id` forms one group. Every other message (including a `tool` message with no matching assistant) is its own group. Groups are always dropped or kept whole, so a tool result is never left without its assistant call, or vice versa.
|
|
127
|
+
2. **Snap keep_head/keep_tail to group boundaries.** `keep_head` and `keep_tail` are counted in messages, but if the boundary would land inside a group, it expands outward to keep that whole group.
|
|
128
|
+
3. **Drop oldest-first.** Whatever is left in the middle is droppable. Groups are dropped oldest first until the running token total fits `max_tokens` or nothing droppable is left.
|
|
129
|
+
4. Token counts are an estimate (`~length / 4` by default, over the compact JSON serialization of the message, or your own `count_tokens`), not a real tokenizer. There is no LLM call, no tokenizer library, and no streaming. If your `count_tokens` is inaccurate, `fits` will be inaccurate too. Pass a `count_tokens` backed by your provider's real tokenizer if you need exact numbers.
|
|
130
|
+
5. `compact_with_summary` never re-summarizes after dropping additional groups to make room for the summary itself; it just drops more of the already-dropped-eligible messages. If you need every dropped message reflected in the summary text, make sure `max_tokens` leaves enough headroom for the summary you expect `summarize` to produce.
|
|
131
|
+
6. The default estimator uses `json.dumps(message, separators=(",", ":"), ensure_ascii=False)`, matching the length JavaScript's `JSON.stringify` plus `.length` produces. Two of Python's `json.dumps` defaults disagree with `JSON.stringify` and both are overridden here:
|
|
132
|
+
- `separators` defaults to `", "` and `": "` (with spaces) in Python; `JSON.stringify` never inserts spaces. Passing plain `json.dumps(message)` would inflate every count.
|
|
133
|
+
- `ensure_ascii` defaults to `True` in Python, which `\uXXXX`-escapes every non-ASCII codepoint (accents, CJK, emoji); `JSON.stringify` never escapes non-ASCII text. Leaving `ensure_ascii` at its default would silently inflate the token count, and therefore over-compact, for any non-English or emoji-bearing content.
|
|
134
|
+
|
|
135
|
+
A third, more subtle difference is corrected for internally rather than in the `json.dumps` call: JavaScript's `String.length` counts UTF-16 code units, so a character outside the Basic Multilingual Plane (most emoji, e.g. U+1F389) counts as a 2-unit surrogate pair there, while Python's `len()` counts it as a single codepoint. The estimator measures the serialized JSON's length in UTF-16 code units (not `len()` codepoints) so astral-plane characters do not silently disagree with the JS package's token numbers for the same message.
|
|
136
|
+
|
|
137
|
+
See the JavaScript version at the repo root (`../index.js`, `../README.md`) for the original implementation this port matches.
|
|
138
|
+
|
|
139
|
+
## License
|
|
140
|
+
|
|
141
|
+
MIT
|
|
@@ -0,0 +1,124 @@
|
|
|
1
|
+
# ctx-compact
|
|
2
|
+
|
|
3
|
+
Trim a plain OpenAI-shaped message array down to a token budget without ever orphaning a tool result. This is the Python port of the [ctx-compact](https://www.npmjs.com/package/ctx-compact) npm package.
|
|
4
|
+
|
|
5
|
+
## The problem
|
|
6
|
+
|
|
7
|
+
Every agent framework ships its own conversation compaction (LangGraph, Inspect, MS Agent Framework, the Claude SDK all have one), and each is welded to that framework's message type. If you are working with a plain list of `{role, content, ...}` messages, people tend to hand-roll a "drop the oldest N messages" loop. That works until an assistant message with `tool_calls` gets dropped but its matching `tool` result messages do not (or the reverse). Most providers reject that shape outright, so the trim silently turns into an API error on the next call. `ctx-compact` is a small, framework-neutral compactor that keeps tool-call and tool-result messages paired and dropped or kept as a unit.
|
|
8
|
+
|
|
9
|
+
## Install
|
|
10
|
+
|
|
11
|
+
```
|
|
12
|
+
pip install ctx-compact
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
## Usage
|
|
16
|
+
|
|
17
|
+
```python
|
|
18
|
+
from ctx_compact import compact, compact_with_summary
|
|
19
|
+
|
|
20
|
+
messages = [
|
|
21
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
22
|
+
{"role": "user", "content": "What is the weather in Denver?"},
|
|
23
|
+
{
|
|
24
|
+
"role": "assistant",
|
|
25
|
+
"content": None,
|
|
26
|
+
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city":"Denver"}'}}],
|
|
27
|
+
},
|
|
28
|
+
{"role": "tool", "tool_call_id": "call_1", "content": '{"tempF":72}'},
|
|
29
|
+
{"role": "assistant", "content": "It is 72F in Denver."},
|
|
30
|
+
{"role": "user", "content": "What about Austin?"},
|
|
31
|
+
{
|
|
32
|
+
"role": "assistant",
|
|
33
|
+
"content": None,
|
|
34
|
+
"tool_calls": [{"id": "call_2", "type": "function", "function": {"name": "get_weather", "arguments": '{"city":"Austin"}'}}],
|
|
35
|
+
},
|
|
36
|
+
{"role": "tool", "tool_call_id": "call_2", "content": '{"tempF":88}'},
|
|
37
|
+
{"role": "assistant", "content": "It is 88F in Austin."},
|
|
38
|
+
{"role": "user", "content": "And tomorrow in Denver?"},
|
|
39
|
+
]
|
|
40
|
+
|
|
41
|
+
result = compact(messages, max_tokens=150, keep_head=1, keep_tail=2)
|
|
42
|
+
# result.tokens_before -> 195
|
|
43
|
+
# result.tokens_after -> 124
|
|
44
|
+
# result.fits -> True (124 <= 150)
|
|
45
|
+
# result.dropped -> the 3 oldest droppable messages: the first "What is
|
|
46
|
+
# the weather in Denver?" turn and its whole
|
|
47
|
+
# assistant/tool_calls + tool group
|
|
48
|
+
|
|
49
|
+
# Synchronous variant: summarize whatever got dropped and splice a note back in.
|
|
50
|
+
with_summary = compact_with_summary(
|
|
51
|
+
messages,
|
|
52
|
+
max_tokens=150,
|
|
53
|
+
keep_head=1,
|
|
54
|
+
keep_tail=2,
|
|
55
|
+
summarize=lambda dropped: f"Earlier in this conversation: {len(dropped)} messages were removed.",
|
|
56
|
+
)
|
|
57
|
+
# with_summary.summary -> 'Earlier in this conversation: 3 messages were removed.'
|
|
58
|
+
# with_summary.tokens_after -> 145 (the 124 kept after dropping, plus the inserted summary message)
|
|
59
|
+
# with_summary.fits -> True (145 <= 150)
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
The example above is exact output from running this code against this package.
|
|
63
|
+
|
|
64
|
+
## API
|
|
65
|
+
|
|
66
|
+
### `compact(messages, *, max_tokens, count_tokens=None, keep_head=1, keep_tail=4) -> CompactResult`
|
|
67
|
+
|
|
68
|
+
Drops messages from the middle of `messages` until the estimated token count fits the budget.
|
|
69
|
+
|
|
70
|
+
- `messages`: list of dicts shaped like `{role, content, tool_calls?, tool_call_id?, name?}`. `role` is one of `'system' | 'user' | 'assistant' | 'tool'`.
|
|
71
|
+
- `max_tokens` (required, keyword-only, `float`) - the budget. Raises `TypeError` if missing or not a positive number.
|
|
72
|
+
- `count_tokens` - `(message) -> int`. Default: `math.ceil(len(json.dumps(message, separators=(",", ":"), ensure_ascii=False)) / 4)`, with the length measured in UTF-16 code units (matching JavaScript's `String.length`), not Python codepoints.
|
|
73
|
+
- `keep_head` - number of leading messages always kept. Default `1`.
|
|
74
|
+
- `keep_tail` - number of trailing messages always kept. Default `4`.
|
|
75
|
+
|
|
76
|
+
Returns a frozen dataclass:
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
@dataclass(frozen=True)
|
|
80
|
+
class CompactResult:
|
|
81
|
+
messages: list # the compacted list
|
|
82
|
+
dropped: list # what was removed, in original order
|
|
83
|
+
tokens_before: int # summed estimated tokens of the input list
|
|
84
|
+
tokens_after: int # summed estimated tokens of the output list
|
|
85
|
+
fits: bool # tokens_after <= max_tokens
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
If the input already fits, it is returned unchanged with `dropped: []`.
|
|
89
|
+
|
|
90
|
+
### `compact_with_summary(messages, *, max_tokens, count_tokens=None, keep_head=1, keep_tail=4, summarize=None, summary_role="user") -> CompactResultWithSummary`
|
|
91
|
+
|
|
92
|
+
Same as `compact`, then, if anything was dropped and `summarize` is provided, calls `summarize(dropped)` and inserts the returned string as a message `{"role": summary_role, "content": <summary>}` immediately after the head-kept messages.
|
|
93
|
+
|
|
94
|
+
- `summarize` - `(dropped: list) -> str`. If omitted, behaves exactly like `compact` and returns `summary: None`.
|
|
95
|
+
- `summary_role` - default `'user'`. Some providers reject a second `system` message, which is why the default is `'user'` rather than `'system'`.
|
|
96
|
+
|
|
97
|
+
The inserted summary message counts toward the budget: after insertion the result is re-checked, and if it no longer fits, additional whole groups are dropped (oldest first) to make room. `fits` is reported `False` if it is still over budget after that.
|
|
98
|
+
|
|
99
|
+
Returns a frozen dataclass with the same fields as `CompactResult` plus `summary: Optional[str]`.
|
|
100
|
+
|
|
101
|
+
**Difference from the JavaScript version:** the JS `compactWithSummary` is `async` and does `await options.summarize(dropped)`, since JS summarizers are typically an async LLM call. This Python port is **synchronous**: `summarize` is a plain callable, `dropped -> str`, called directly with no `await`. If your summarizer needs to be async in Python, run it yourself (e.g. via `asyncio.run` or your event loop) before calling `compact_with_summary`, and pass a synchronous wrapper.
|
|
102
|
+
|
|
103
|
+
### `estimate_tokens(message, count_tokens=None) -> int`
|
|
104
|
+
|
|
105
|
+
Runs `count_tokens` (or the default heuristic) against a single message. Exported so callers can reuse the same estimator `compact`/`compact_with_summary` use, e.g. to pre-check a message before appending it.
|
|
106
|
+
|
|
107
|
+
## How it works
|
|
108
|
+
|
|
109
|
+
1. **Group first.** Before anything is dropped, the whole list is split into groups: an assistant message carrying `tool_calls` plus every immediately-following `tool` message whose `tool_call_id` matches one of that assistant's `tool_calls[].id` forms one group. Every other message (including a `tool` message with no matching assistant) is its own group. Groups are always dropped or kept whole, so a tool result is never left without its assistant call, or vice versa.
|
|
110
|
+
2. **Snap keep_head/keep_tail to group boundaries.** `keep_head` and `keep_tail` are counted in messages, but if the boundary would land inside a group, it expands outward to keep that whole group.
|
|
111
|
+
3. **Drop oldest-first.** Whatever is left in the middle is droppable. Groups are dropped oldest first until the running token total fits `max_tokens` or nothing droppable is left.
|
|
112
|
+
4. Token counts are an estimate (`~length / 4` by default, over the compact JSON serialization of the message, or your own `count_tokens`), not a real tokenizer. There is no LLM call, no tokenizer library, and no streaming. If your `count_tokens` is inaccurate, `fits` will be inaccurate too. Pass a `count_tokens` backed by your provider's real tokenizer if you need exact numbers.
|
|
113
|
+
5. `compact_with_summary` never re-summarizes after dropping additional groups to make room for the summary itself; it just drops more of the already-dropped-eligible messages. If you need every dropped message reflected in the summary text, make sure `max_tokens` leaves enough headroom for the summary you expect `summarize` to produce.
|
|
114
|
+
6. The default estimator uses `json.dumps(message, separators=(",", ":"), ensure_ascii=False)`, matching the length JavaScript's `JSON.stringify` plus `.length` produces. Two of Python's `json.dumps` defaults disagree with `JSON.stringify` and both are overridden here:
|
|
115
|
+
- `separators` defaults to `", "` and `": "` (with spaces) in Python; `JSON.stringify` never inserts spaces. Passing plain `json.dumps(message)` would inflate every count.
|
|
116
|
+
- `ensure_ascii` defaults to `True` in Python, which `\uXXXX`-escapes every non-ASCII codepoint (accents, CJK, emoji); `JSON.stringify` never escapes non-ASCII text. Leaving `ensure_ascii` at its default would silently inflate the token count, and therefore over-compact, for any non-English or emoji-bearing content.
|
|
117
|
+
|
|
118
|
+
A third, more subtle difference is corrected for internally rather than in the `json.dumps` call: JavaScript's `String.length` counts UTF-16 code units, so a character outside the Basic Multilingual Plane (most emoji, e.g. U+1F389) counts as a 2-unit surrogate pair there, while Python's `len()` counts it as a single codepoint. The estimator measures the serialized JSON's length in UTF-16 code units (not `len()` codepoints) so astral-plane characters do not silently disagree with the JS package's token numbers for the same message.
|
|
119
|
+
|
|
120
|
+
See the JavaScript version at the repo root (`../index.js`, `../README.md`) for the original implementation this port matches.
|
|
121
|
+
|
|
122
|
+
## License
|
|
123
|
+
|
|
124
|
+
MIT
|
|
@@ -0,0 +1,380 @@
|
|
|
1
|
+
"""
|
|
2
|
+
ctx-compact - framework-neutral conversation compaction for plain
|
|
3
|
+
OpenAI-shaped message arrays. See README.md for the full picture.
|
|
4
|
+
|
|
5
|
+
This is a Python port of the ctx-compact npm package. The JavaScript
|
|
6
|
+
source (../index.js in the same repo) is the specification; this module
|
|
7
|
+
matches its behavior exactly, with one intentional API difference:
|
|
8
|
+
`compact_with_summary` is synchronous here, while the JS `compactWithSummary`
|
|
9
|
+
is async. See README.md for details.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
import math
|
|
14
|
+
from dataclasses import dataclass
|
|
15
|
+
from typing import Any, Callable, Dict, List, Optional
|
|
16
|
+
|
|
17
|
+
Message = Dict[str, Any]
|
|
18
|
+
CountTokensFn = Callable[[Message], int]
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _utf16_length(s: str) -> int:
|
|
22
|
+
"""Number of UTF-16 code units in s, matching JavaScript's
|
|
23
|
+
String.prototype.length. JS strings are UTF-16 internally, so any
|
|
24
|
+
codepoint outside the Basic Multilingual Plane (astral-plane
|
|
25
|
+
characters, which includes most emoji, e.g. U+1F389) counts as a
|
|
26
|
+
2-unit surrogate pair there, even though it is a single Python
|
|
27
|
+
string character (Python's len() counts Unicode codepoints, not
|
|
28
|
+
UTF-16 units). Encoding to UTF-16LE and counting 2-byte units
|
|
29
|
+
reproduces JS's notion of length exactly, including that surrogate
|
|
30
|
+
pair doubling.
|
|
31
|
+
"""
|
|
32
|
+
return len(s.encode("utf-16-le")) // 2
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def _default_count_tokens(message: Message) -> int:
|
|
36
|
+
"""Default token estimator: ~4 chars per token on the JSON-serialized
|
|
37
|
+
message. Good enough for budgeting, not a real tokenizer.
|
|
38
|
+
|
|
39
|
+
Three ways Python's json.dumps/str differ from JavaScript's
|
|
40
|
+
JSON.stringify/String.length must be overridden or corrected for so
|
|
41
|
+
the lengths (and therefore token counts) match:
|
|
42
|
+
|
|
43
|
+
- separators: json.dumps defaults to ", " and ": " (with spaces);
|
|
44
|
+
JSON.stringify never inserts spaces. We pass separators=(",", ":").
|
|
45
|
+
- ensure_ascii: json.dumps defaults to True, which \\uXXXX-escapes
|
|
46
|
+
every non-ASCII codepoint. JSON.stringify never escapes non-ASCII
|
|
47
|
+
text, it emits it as literal characters. We pass
|
|
48
|
+
ensure_ascii=False so accented text, CJK, and emoji are not
|
|
49
|
+
inflated by escaping.
|
|
50
|
+
- string length semantics: Python's len() on a str counts Unicode
|
|
51
|
+
codepoints; JavaScript's .length (which is what JSON.stringify's
|
|
52
|
+
result length uses) counts UTF-16 code units, so an astral-plane
|
|
53
|
+
character such as an emoji counts as 2 in JS but 1 in plain Python
|
|
54
|
+
len(). We use _utf16_length() instead of len() so that difference
|
|
55
|
+
does not creep back in for emoji-bearing content.
|
|
56
|
+
"""
|
|
57
|
+
serialized = json.dumps(message, separators=(",", ":"), ensure_ascii=False)
|
|
58
|
+
return math.ceil(_utf16_length(serialized) / 4)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def estimate_tokens(message: Message, count_tokens: Optional[CountTokensFn] = None) -> int:
|
|
62
|
+
"""Estimate the token cost of a single message.
|
|
63
|
+
|
|
64
|
+
Args:
|
|
65
|
+
message: A message object.
|
|
66
|
+
count_tokens: Custom estimator. Defaults to a ~4-chars-per-token
|
|
67
|
+
heuristic over the JSON-serialized message.
|
|
68
|
+
|
|
69
|
+
Returns:
|
|
70
|
+
Estimated token count.
|
|
71
|
+
"""
|
|
72
|
+
fn = count_tokens if count_tokens is not None else _default_count_tokens
|
|
73
|
+
return fn(message)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _validate_max_tokens(max_tokens: Any) -> float:
|
|
77
|
+
if (
|
|
78
|
+
isinstance(max_tokens, bool)
|
|
79
|
+
or not isinstance(max_tokens, (int, float))
|
|
80
|
+
or not math.isfinite(max_tokens)
|
|
81
|
+
or max_tokens <= 0
|
|
82
|
+
):
|
|
83
|
+
raise TypeError("max_tokens must be a positive number")
|
|
84
|
+
return max_tokens
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def _build_groups(messages: List[Message]) -> List[List[Message]]:
|
|
88
|
+
"""Split messages into groups. An assistant message carrying
|
|
89
|
+
`tool_calls` forms one group with every immediately-following `tool`
|
|
90
|
+
message whose `tool_call_id` matches one of that assistant's
|
|
91
|
+
`tool_calls[].id`. Every other message (including an orphan `tool`
|
|
92
|
+
message) is its own group. Groups are dropped or kept whole.
|
|
93
|
+
"""
|
|
94
|
+
groups: List[List[Message]] = []
|
|
95
|
+
i = 0
|
|
96
|
+
n = len(messages)
|
|
97
|
+
while i < n:
|
|
98
|
+
msg = messages[i]
|
|
99
|
+
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
|
|
100
|
+
if msg.get("role") == "assistant" and isinstance(tool_calls, list) and len(tool_calls) > 0:
|
|
101
|
+
ids = {tc.get("id") for tc in tool_calls}
|
|
102
|
+
group_messages = [msg]
|
|
103
|
+
j = i + 1
|
|
104
|
+
while (
|
|
105
|
+
j < n
|
|
106
|
+
and messages[j].get("role") == "tool"
|
|
107
|
+
and messages[j].get("tool_call_id") in ids
|
|
108
|
+
):
|
|
109
|
+
group_messages.append(messages[j])
|
|
110
|
+
j += 1
|
|
111
|
+
groups.append(group_messages)
|
|
112
|
+
i = j
|
|
113
|
+
else:
|
|
114
|
+
groups.append([msg])
|
|
115
|
+
i += 1
|
|
116
|
+
return groups
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def _count_groups_covering_front(groups: List[List[Message]], keep_count: int) -> int:
|
|
120
|
+
"""Walk groups from the front, accumulating message counts, until at
|
|
121
|
+
least `keep_count` messages are covered. Returns the number of whole
|
|
122
|
+
groups needed to do that (0 if keep_count <= 0).
|
|
123
|
+
"""
|
|
124
|
+
if keep_count <= 0:
|
|
125
|
+
return 0
|
|
126
|
+
covered = 0
|
|
127
|
+
groups_used = 0
|
|
128
|
+
for g in groups:
|
|
129
|
+
if covered >= keep_count:
|
|
130
|
+
break
|
|
131
|
+
covered += len(g)
|
|
132
|
+
groups_used += 1
|
|
133
|
+
return groups_used
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
def _count_groups_covering_back(groups: List[List[Message]], keep_count: int) -> int:
|
|
137
|
+
"""Same as _count_groups_covering_front but walking from the back."""
|
|
138
|
+
if keep_count <= 0:
|
|
139
|
+
return 0
|
|
140
|
+
covered = 0
|
|
141
|
+
groups_used = 0
|
|
142
|
+
for g in reversed(groups):
|
|
143
|
+
if covered >= keep_count:
|
|
144
|
+
break
|
|
145
|
+
covered += len(g)
|
|
146
|
+
groups_used += 1
|
|
147
|
+
return groups_used
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
@dataclass
|
|
151
|
+
class _CoreResult:
|
|
152
|
+
"""Internal result shared by compact() and compact_with_summary()."""
|
|
153
|
+
|
|
154
|
+
messages: List[Message]
|
|
155
|
+
dropped: List[Message]
|
|
156
|
+
tokens_before: int
|
|
157
|
+
tokens_after: int
|
|
158
|
+
fits: bool
|
|
159
|
+
head_kept_count: int
|
|
160
|
+
remaining_groups: List[List[Message]]
|
|
161
|
+
count_tokens_fn: CountTokensFn
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def _compact_core(
|
|
165
|
+
messages: List[Message],
|
|
166
|
+
max_tokens: Any,
|
|
167
|
+
count_tokens: Optional[CountTokensFn],
|
|
168
|
+
keep_head: int,
|
|
169
|
+
keep_tail: int,
|
|
170
|
+
) -> _CoreResult:
|
|
171
|
+
max_tokens = _validate_max_tokens(max_tokens)
|
|
172
|
+
count_tokens_fn = count_tokens if count_tokens is not None else _default_count_tokens
|
|
173
|
+
|
|
174
|
+
tokens_before = sum(estimate_tokens(m, count_tokens_fn) for m in messages)
|
|
175
|
+
|
|
176
|
+
if tokens_before <= max_tokens:
|
|
177
|
+
return _CoreResult(
|
|
178
|
+
messages=list(messages),
|
|
179
|
+
dropped=[],
|
|
180
|
+
tokens_before=tokens_before,
|
|
181
|
+
tokens_after=tokens_before,
|
|
182
|
+
fits=True,
|
|
183
|
+
head_kept_count=0,
|
|
184
|
+
remaining_groups=[],
|
|
185
|
+
count_tokens_fn=count_tokens_fn,
|
|
186
|
+
)
|
|
187
|
+
|
|
188
|
+
groups = _build_groups(messages)
|
|
189
|
+
head_group_count = _count_groups_covering_front(groups, keep_head)
|
|
190
|
+
tail_group_count = _count_groups_covering_back(groups, keep_tail)
|
|
191
|
+
|
|
192
|
+
middle_start = head_group_count
|
|
193
|
+
middle_end = len(groups) - tail_group_count
|
|
194
|
+
|
|
195
|
+
if middle_start >= middle_end:
|
|
196
|
+
# keep_head + keep_tail (snapped to group boundaries) cover the
|
|
197
|
+
# whole array. Nothing is droppable.
|
|
198
|
+
return _CoreResult(
|
|
199
|
+
messages=list(messages),
|
|
200
|
+
dropped=[],
|
|
201
|
+
tokens_before=tokens_before,
|
|
202
|
+
tokens_after=tokens_before,
|
|
203
|
+
fits=False,
|
|
204
|
+
head_kept_count=0,
|
|
205
|
+
remaining_groups=[],
|
|
206
|
+
count_tokens_fn=count_tokens_fn,
|
|
207
|
+
)
|
|
208
|
+
|
|
209
|
+
head_groups = groups[:head_group_count]
|
|
210
|
+
middle_groups = groups[middle_start:middle_end]
|
|
211
|
+
tail_groups = groups[middle_end:]
|
|
212
|
+
|
|
213
|
+
current_tokens = tokens_before
|
|
214
|
+
dropped_groups: List[List[Message]] = []
|
|
215
|
+
kept_middle_groups: List[List[Message]] = []
|
|
216
|
+
for g in middle_groups:
|
|
217
|
+
if current_tokens > max_tokens:
|
|
218
|
+
group_tokens = sum(estimate_tokens(m, count_tokens_fn) for m in g)
|
|
219
|
+
current_tokens -= group_tokens
|
|
220
|
+
dropped_groups.append(g)
|
|
221
|
+
else:
|
|
222
|
+
kept_middle_groups.append(g)
|
|
223
|
+
|
|
224
|
+
kept_messages = (
|
|
225
|
+
[m for g in head_groups for m in g]
|
|
226
|
+
+ [m for g in kept_middle_groups for m in g]
|
|
227
|
+
+ [m for g in tail_groups for m in g]
|
|
228
|
+
)
|
|
229
|
+
dropped = [m for g in dropped_groups for m in g]
|
|
230
|
+
head_kept_count = sum(len(g) for g in head_groups)
|
|
231
|
+
|
|
232
|
+
return _CoreResult(
|
|
233
|
+
messages=kept_messages,
|
|
234
|
+
dropped=dropped,
|
|
235
|
+
tokens_before=tokens_before,
|
|
236
|
+
tokens_after=current_tokens,
|
|
237
|
+
fits=current_tokens <= max_tokens,
|
|
238
|
+
head_kept_count=head_kept_count,
|
|
239
|
+
remaining_groups=kept_middle_groups,
|
|
240
|
+
count_tokens_fn=count_tokens_fn,
|
|
241
|
+
)
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
@dataclass(frozen=True)
|
|
245
|
+
class CompactResult:
|
|
246
|
+
"""Result of compact()."""
|
|
247
|
+
|
|
248
|
+
messages: List[Message]
|
|
249
|
+
dropped: List[Message]
|
|
250
|
+
tokens_before: int
|
|
251
|
+
tokens_after: int
|
|
252
|
+
fits: bool
|
|
253
|
+
|
|
254
|
+
|
|
255
|
+
@dataclass(frozen=True)
|
|
256
|
+
class CompactResultWithSummary:
|
|
257
|
+
"""Result of compact_with_summary()."""
|
|
258
|
+
|
|
259
|
+
messages: List[Message]
|
|
260
|
+
dropped: List[Message]
|
|
261
|
+
tokens_before: int
|
|
262
|
+
tokens_after: int
|
|
263
|
+
fits: bool
|
|
264
|
+
summary: Optional[str]
|
|
265
|
+
|
|
266
|
+
|
|
267
|
+
def compact(
|
|
268
|
+
messages: List[Message],
|
|
269
|
+
*,
|
|
270
|
+
max_tokens: Any,
|
|
271
|
+
count_tokens: Optional[CountTokensFn] = None,
|
|
272
|
+
keep_head: int = 1,
|
|
273
|
+
keep_tail: int = 4,
|
|
274
|
+
) -> CompactResult:
|
|
275
|
+
"""Drop messages from the middle of a conversation until the
|
|
276
|
+
estimated token count fits the budget, never splitting an
|
|
277
|
+
assistant/tool_calls group from its matching tool results.
|
|
278
|
+
|
|
279
|
+
Args:
|
|
280
|
+
messages: Plain OpenAI-shaped messages: dicts with
|
|
281
|
+
`role`, `content`, and optionally `tool_calls`, `tool_call_id`,
|
|
282
|
+
`name`.
|
|
283
|
+
max_tokens: Required token budget. Must be a positive number.
|
|
284
|
+
Raises TypeError if missing or not a positive number.
|
|
285
|
+
count_tokens: Custom token estimator. Defaults to
|
|
286
|
+
`math.ceil(utf16_length(json.dumps(message, separators=(",", ":"), ensure_ascii=False)) / 4)`,
|
|
287
|
+
where `utf16_length` counts UTF-16 code units (matching
|
|
288
|
+
JavaScript's `String.length`), not Python codepoints.
|
|
289
|
+
keep_head: Leading messages always kept (snapped outward to group
|
|
290
|
+
boundaries). Default 1.
|
|
291
|
+
keep_tail: Trailing messages always kept (snapped outward to group
|
|
292
|
+
boundaries). Default 4.
|
|
293
|
+
|
|
294
|
+
Returns:
|
|
295
|
+
CompactResult(messages, dropped, tokens_before, tokens_after, fits)
|
|
296
|
+
"""
|
|
297
|
+
result = _compact_core(messages, max_tokens, count_tokens, keep_head, keep_tail)
|
|
298
|
+
return CompactResult(
|
|
299
|
+
messages=result.messages,
|
|
300
|
+
dropped=result.dropped,
|
|
301
|
+
tokens_before=result.tokens_before,
|
|
302
|
+
tokens_after=result.tokens_after,
|
|
303
|
+
fits=result.fits,
|
|
304
|
+
)
|
|
305
|
+
|
|
306
|
+
|
|
307
|
+
def compact_with_summary(
|
|
308
|
+
messages: List[Message],
|
|
309
|
+
*,
|
|
310
|
+
max_tokens: Any,
|
|
311
|
+
count_tokens: Optional[CountTokensFn] = None,
|
|
312
|
+
keep_head: int = 1,
|
|
313
|
+
keep_tail: int = 4,
|
|
314
|
+
summarize: Optional[Callable[[List[Message]], str]] = None,
|
|
315
|
+
summary_role: str = "user",
|
|
316
|
+
) -> CompactResultWithSummary:
|
|
317
|
+
"""Same as compact(), then, if anything was dropped and `summarize` is
|
|
318
|
+
provided, calls `summarize(dropped)` and inserts the resulting string
|
|
319
|
+
as a message right after the head-kept messages. The inserted summary
|
|
320
|
+
counts toward the budget: if it pushes the result back over
|
|
321
|
+
max_tokens, additional groups are dropped (oldest first) to make room.
|
|
322
|
+
|
|
323
|
+
Note: unlike the JavaScript compactWithSummary (which is async and
|
|
324
|
+
awaits `summarize`), this Python port is synchronous. `summarize` is a
|
|
325
|
+
plain callable taking the dropped list and returning a string.
|
|
326
|
+
|
|
327
|
+
Args:
|
|
328
|
+
messages: Plain OpenAI-shaped messages.
|
|
329
|
+
max_tokens: Required token budget. Must be a positive number.
|
|
330
|
+
count_tokens: Custom token estimator.
|
|
331
|
+
keep_head: Leading messages always kept. Default 1.
|
|
332
|
+
keep_tail: Trailing messages always kept. Default 4.
|
|
333
|
+
summarize: Produces a summary of the dropped messages. If
|
|
334
|
+
omitted, behaves exactly like compact().
|
|
335
|
+
summary_role: Role for the inserted summary message. Defaults to
|
|
336
|
+
'user' since some providers reject a second 'system' message.
|
|
337
|
+
|
|
338
|
+
Returns:
|
|
339
|
+
CompactResultWithSummary(messages, dropped, tokens_before,
|
|
340
|
+
tokens_after, fits, summary)
|
|
341
|
+
"""
|
|
342
|
+
base = _compact_core(messages, max_tokens, count_tokens, keep_head, keep_tail)
|
|
343
|
+
|
|
344
|
+
if len(base.dropped) == 0 or not callable(summarize):
|
|
345
|
+
return CompactResultWithSummary(
|
|
346
|
+
messages=base.messages,
|
|
347
|
+
dropped=base.dropped,
|
|
348
|
+
tokens_before=base.tokens_before,
|
|
349
|
+
tokens_after=base.tokens_after,
|
|
350
|
+
fits=base.fits,
|
|
351
|
+
summary=None,
|
|
352
|
+
)
|
|
353
|
+
|
|
354
|
+
max_tokens = _validate_max_tokens(max_tokens)
|
|
355
|
+
count_tokens_fn = base.count_tokens_fn
|
|
356
|
+
summary_text = summarize(base.dropped)
|
|
357
|
+
summary_message: Message = {"role": summary_role, "content": summary_text}
|
|
358
|
+
|
|
359
|
+
kept_messages = list(base.messages)
|
|
360
|
+
kept_messages.insert(base.head_kept_count, summary_message)
|
|
361
|
+
current_tokens = base.tokens_after + estimate_tokens(summary_message, count_tokens_fn)
|
|
362
|
+
|
|
363
|
+
extra_dropped: List[Message] = []
|
|
364
|
+
remaining = list(base.remaining_groups)
|
|
365
|
+
while current_tokens > max_tokens and remaining:
|
|
366
|
+
g = remaining.pop(0)
|
|
367
|
+
group_tokens = sum(estimate_tokens(m, count_tokens_fn) for m in g)
|
|
368
|
+
current_tokens -= group_tokens
|
|
369
|
+
extra_dropped.extend(g)
|
|
370
|
+
to_remove_ids = {id(m) for m in g}
|
|
371
|
+
kept_messages = [m for m in kept_messages if id(m) not in to_remove_ids]
|
|
372
|
+
|
|
373
|
+
return CompactResultWithSummary(
|
|
374
|
+
messages=kept_messages,
|
|
375
|
+
dropped=list(base.dropped) + extra_dropped,
|
|
376
|
+
tokens_before=base.tokens_before,
|
|
377
|
+
tokens_after=current_tokens,
|
|
378
|
+
fits=current_tokens <= max_tokens,
|
|
379
|
+
summary=summary_text,
|
|
380
|
+
)
|
|
File without changes
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "ctx-compact"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Trim a plain OpenAI-shaped message array down to a token budget without ever orphaning a tool result."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = {text = "MIT"}
|
|
12
|
+
authors = [{name = "Prajjwal Chittori"}]
|
|
13
|
+
keywords = ["llm", "context-window", "token-budget", "conversation", "compaction", "openai", "tool-calls"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 4 - Beta",
|
|
16
|
+
"Intended Audience :: Developers",
|
|
17
|
+
"License :: OSI Approved :: MIT License",
|
|
18
|
+
"Programming Language :: Python :: 3",
|
|
19
|
+
"Typing :: Typed",
|
|
20
|
+
]
|
|
21
|
+
|
|
22
|
+
[project.urls]
|
|
23
|
+
Homepage = "https://github.com/pjdurden/ctx-compact"
|
|
24
|
+
Source = "https://github.com/pjdurden/ctx-compact"
|
|
25
|
+
|
|
26
|
+
[tool.hatch.build.targets.wheel]
|
|
27
|
+
include = ["ctx_compact.py", "py.typed"]
|
|
@@ -0,0 +1,315 @@
|
|
|
1
|
+
"""Port of ../test/ctx-compact.test.js, one test for one, plus a fuzz sweep."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import math
|
|
5
|
+
import unittest
|
|
6
|
+
|
|
7
|
+
from ctx_compact import (
|
|
8
|
+
CompactResult,
|
|
9
|
+
CompactResultWithSummary,
|
|
10
|
+
compact,
|
|
11
|
+
compact_with_summary,
|
|
12
|
+
estimate_tokens,
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def const_tokens(_message):
|
|
17
|
+
return 1
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class CompactTests(unittest.TestCase):
|
|
21
|
+
def test_input_already_under_budget_is_returned_unchanged(self):
|
|
22
|
+
messages = [
|
|
23
|
+
{"role": "system", "content": "sys"},
|
|
24
|
+
{"role": "user", "content": "hi"},
|
|
25
|
+
{"role": "assistant", "content": "hello"},
|
|
26
|
+
]
|
|
27
|
+
result = compact(messages, max_tokens=100_000)
|
|
28
|
+
self.assertEqual(result.messages, messages)
|
|
29
|
+
self.assertEqual(result.dropped, [])
|
|
30
|
+
self.assertTrue(result.fits)
|
|
31
|
+
self.assertEqual(result.tokens_before, result.tokens_after)
|
|
32
|
+
|
|
33
|
+
def test_throws_type_error_when_max_tokens_missing_or_not_a_positive_number(self):
|
|
34
|
+
messages = [{"role": "user", "content": "hi"}]
|
|
35
|
+
with self.assertRaises(TypeError):
|
|
36
|
+
compact(messages)
|
|
37
|
+
with self.assertRaises(TypeError):
|
|
38
|
+
compact(messages, max_tokens=0)
|
|
39
|
+
with self.assertRaises(TypeError):
|
|
40
|
+
compact(messages, max_tokens=-5)
|
|
41
|
+
with self.assertRaises(TypeError):
|
|
42
|
+
compact(messages, max_tokens="lots")
|
|
43
|
+
|
|
44
|
+
def test_drops_an_assistant_tool_result_group_as_a_whole_unit_oldest_first(self):
|
|
45
|
+
msg_assistant_tools = {
|
|
46
|
+
"role": "assistant",
|
|
47
|
+
"content": None,
|
|
48
|
+
"tool_calls": [{"id": "a"}, {"id": "b"}],
|
|
49
|
+
}
|
|
50
|
+
msg_tool_a = {"role": "tool", "tool_call_id": "a", "content": "ra"}
|
|
51
|
+
msg_tool_b = {"role": "tool", "tool_call_id": "b", "content": "rb"}
|
|
52
|
+
messages = [
|
|
53
|
+
{"role": "system", "content": "sys"}, # 0 - head
|
|
54
|
+
{"role": "user", "content": "u1"}, # 1 - droppable
|
|
55
|
+
msg_assistant_tools, # 2 - droppable group start
|
|
56
|
+
msg_tool_a, # 3
|
|
57
|
+
msg_tool_b, # 4
|
|
58
|
+
{"role": "user", "content": "u2"}, # 5 - kept (not needed to drop further)
|
|
59
|
+
{"role": "assistant", "content": "reply2"}, # 6 - kept
|
|
60
|
+
{"role": "user", "content": "u3"}, # 7 - tail
|
|
61
|
+
{"role": "assistant", "content": "final"}, # 8 - tail
|
|
62
|
+
]
|
|
63
|
+
|
|
64
|
+
result = compact(messages, max_tokens=6, count_tokens=const_tokens, keep_head=1, keep_tail=2)
|
|
65
|
+
|
|
66
|
+
self.assertEqual(result.tokens_before, 9)
|
|
67
|
+
self.assertEqual(result.tokens_after, 5)
|
|
68
|
+
self.assertTrue(result.fits)
|
|
69
|
+
|
|
70
|
+
# The tool group is dropped whole: no orphan tool message survives.
|
|
71
|
+
self.assertFalse(any(m["role"] == "tool" for m in result.messages))
|
|
72
|
+
self.assertEqual(
|
|
73
|
+
[m["role"] for m in result.messages],
|
|
74
|
+
["system", "user", "assistant", "user", "assistant"],
|
|
75
|
+
)
|
|
76
|
+
self.assertEqual(result.dropped, [messages[1], msg_assistant_tools, msg_tool_a, msg_tool_b])
|
|
77
|
+
|
|
78
|
+
def test_keep_tail_boundary_landing_mid_group_snaps_outward(self):
|
|
79
|
+
msg_assistant_tools = {
|
|
80
|
+
"role": "assistant",
|
|
81
|
+
"content": None,
|
|
82
|
+
"tool_calls": [{"id": "x"}, {"id": "y"}],
|
|
83
|
+
}
|
|
84
|
+
msg_tool_x = {"role": "tool", "tool_call_id": "x", "content": "rx"}
|
|
85
|
+
msg_tool_y = {"role": "tool", "tool_call_id": "y", "content": "ry"}
|
|
86
|
+
messages = [
|
|
87
|
+
{"role": "system", "content": "sys"}, # 0
|
|
88
|
+
{"role": "user", "content": "u1"}, # 1 - only droppable message
|
|
89
|
+
msg_assistant_tools, # 2
|
|
90
|
+
msg_tool_x, # 3
|
|
91
|
+
msg_tool_y, # 4
|
|
92
|
+
{"role": "user", "content": "u2"}, # 5
|
|
93
|
+
]
|
|
94
|
+
|
|
95
|
+
# keep_tail: 2 messages would naively land inside the tool group
|
|
96
|
+
# (msg_tool_y, u2), splitting it. It must snap outward to keep the
|
|
97
|
+
# whole group instead.
|
|
98
|
+
result = compact(messages, max_tokens=5, count_tokens=const_tokens, keep_head=1, keep_tail=2)
|
|
99
|
+
|
|
100
|
+
self.assertEqual(len(result.dropped), 1)
|
|
101
|
+
self.assertEqual(result.dropped, [messages[1]])
|
|
102
|
+
self.assertIn(msg_assistant_tools, result.messages)
|
|
103
|
+
self.assertIn(msg_tool_x, result.messages)
|
|
104
|
+
self.assertIn(msg_tool_y, result.messages)
|
|
105
|
+
self.assertTrue(result.fits)
|
|
106
|
+
self.assertEqual(result.tokens_after, 5)
|
|
107
|
+
|
|
108
|
+
def test_reports_fits_false_when_keep_head_and_keep_tail_protect_everything(self):
|
|
109
|
+
messages = [
|
|
110
|
+
{"role": "system", "content": "s"},
|
|
111
|
+
{"role": "user", "content": "u"},
|
|
112
|
+
{"role": "assistant", "content": "a"},
|
|
113
|
+
]
|
|
114
|
+
result = compact(messages, max_tokens=5, count_tokens=lambda m: 10, keep_head=2, keep_tail=2)
|
|
115
|
+
|
|
116
|
+
self.assertFalse(result.fits)
|
|
117
|
+
self.assertEqual(result.dropped, [])
|
|
118
|
+
self.assertEqual(result.messages, messages)
|
|
119
|
+
self.assertEqual(result.tokens_after, result.tokens_before)
|
|
120
|
+
|
|
121
|
+
def test_compact_with_summary_inserts_summary_and_drops_more_if_it_no_longer_fits(self):
|
|
122
|
+
messages = [
|
|
123
|
+
{"role": "system", "content": "sys"}, # 0 - head
|
|
124
|
+
{"role": "user", "content": "u1"}, # 1 - droppable
|
|
125
|
+
{"role": "user", "content": "u2"}, # 2 - droppable
|
|
126
|
+
{"role": "user", "content": "u3"}, # 3 - tail
|
|
127
|
+
{"role": "assistant", "content": "a"}, # 4 - tail
|
|
128
|
+
]
|
|
129
|
+
|
|
130
|
+
result = compact_with_summary(
|
|
131
|
+
messages,
|
|
132
|
+
max_tokens=4,
|
|
133
|
+
count_tokens=const_tokens,
|
|
134
|
+
keep_head=1,
|
|
135
|
+
keep_tail=2,
|
|
136
|
+
summarize=lambda dropped: f"dropped {len(dropped)} messages",
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
# Inserting the 1-token summary would have pushed tokens_after back
|
|
140
|
+
# to 5, over the budget of 4, so one more group had to go.
|
|
141
|
+
self.assertEqual(result.dropped, [messages[1], messages[2]])
|
|
142
|
+
self.assertEqual(result.summary, "dropped 1 messages")
|
|
143
|
+
self.assertEqual(result.messages[0], messages[0])
|
|
144
|
+
self.assertEqual(result.messages[1]["role"], "user")
|
|
145
|
+
self.assertEqual(result.messages[1]["content"], "dropped 1 messages")
|
|
146
|
+
self.assertEqual(result.messages[2:], [messages[3], messages[4]])
|
|
147
|
+
self.assertEqual(result.tokens_after, 4)
|
|
148
|
+
self.assertTrue(result.fits)
|
|
149
|
+
|
|
150
|
+
def test_compact_with_summary_returns_summary_none_when_summarize_not_provided(self):
|
|
151
|
+
messages = [
|
|
152
|
+
{"role": "system", "content": "sys"},
|
|
153
|
+
{"role": "user", "content": "u1"},
|
|
154
|
+
{"role": "user", "content": "u2"},
|
|
155
|
+
{"role": "user", "content": "u3"},
|
|
156
|
+
]
|
|
157
|
+
plain = compact(messages, max_tokens=2, count_tokens=const_tokens, keep_head=1, keep_tail=1)
|
|
158
|
+
with_summary = compact_with_summary(
|
|
159
|
+
messages, max_tokens=2, count_tokens=const_tokens, keep_head=1, keep_tail=1
|
|
160
|
+
)
|
|
161
|
+
|
|
162
|
+
self.assertIsNone(with_summary.summary)
|
|
163
|
+
self.assertEqual(with_summary.messages, plain.messages)
|
|
164
|
+
self.assertEqual(with_summary.dropped, plain.dropped)
|
|
165
|
+
|
|
166
|
+
def test_compact_with_summary_does_not_call_summarize_when_nothing_dropped(self):
|
|
167
|
+
called = {"value": False}
|
|
168
|
+
messages = [{"role": "user", "content": "hi"}]
|
|
169
|
+
|
|
170
|
+
def summarize(_dropped):
|
|
171
|
+
called["value"] = True
|
|
172
|
+
return "should not happen"
|
|
173
|
+
|
|
174
|
+
result = compact_with_summary(messages, max_tokens=1000, summarize=summarize)
|
|
175
|
+
|
|
176
|
+
self.assertFalse(called["value"])
|
|
177
|
+
self.assertIsNone(result.summary)
|
|
178
|
+
self.assertEqual(result.dropped, [])
|
|
179
|
+
|
|
180
|
+
def test_estimate_tokens_default_heuristic_and_custom_estimator(self):
|
|
181
|
+
message = {"role": "user", "content": "hi"}
|
|
182
|
+
default_count = math.ceil(
|
|
183
|
+
len(json.dumps(message, separators=(",", ":"), ensure_ascii=False)) / 4
|
|
184
|
+
)
|
|
185
|
+
self.assertEqual(estimate_tokens(message), default_count)
|
|
186
|
+
self.assertEqual(estimate_tokens(message, lambda m: 42), 42)
|
|
187
|
+
|
|
188
|
+
# --- Unicode parity with JSON.stringify --------------------------------
|
|
189
|
+
#
|
|
190
|
+
# Python's json.dumps escapes non-ASCII by default (ensure_ascii=True),
|
|
191
|
+
# which JavaScript's JSON.stringify never does. The following expected
|
|
192
|
+
# token counts were computed directly from Node with the JS package's
|
|
193
|
+
# own formula (Math.ceil(JSON.stringify(message).length / 4)), so these
|
|
194
|
+
# assert literal parity with the JS original, not just "some number".
|
|
195
|
+
|
|
196
|
+
def test_estimate_tokens_matches_js_for_accented_latin_text(self):
|
|
197
|
+
message = {"role": "user", "content": "héllo wörld café"}
|
|
198
|
+
# node -e "console.log(JSON.stringify({role:'user',content:'héllo wörld café'}).length)" -> 44
|
|
199
|
+
self.assertEqual(estimate_tokens(message), 11)
|
|
200
|
+
|
|
201
|
+
def test_estimate_tokens_matches_js_for_cjk_text(self):
|
|
202
|
+
message = {"role": "user", "content": "日本語のテスト"}
|
|
203
|
+
# node -e "..." -> JSON.stringify length 35
|
|
204
|
+
self.assertEqual(estimate_tokens(message), 9)
|
|
205
|
+
|
|
206
|
+
def test_estimate_tokens_matches_js_for_astral_plane_emoji(self):
|
|
207
|
+
# U+1F389 PARTY POPPER is outside the Basic Multilingual Plane, so
|
|
208
|
+
# JS counts it as a 2-unit UTF-16 surrogate pair; plain Python
|
|
209
|
+
# len() would count it as 1 codepoint. The estimator must reproduce
|
|
210
|
+
# the JS (UTF-16) length to keep token counts identical.
|
|
211
|
+
message = {"role": "user", "content": "party \U0001F389 time"}
|
|
212
|
+
# node -e "..." -> JSON.stringify length 41
|
|
213
|
+
self.assertEqual(estimate_tokens(message), 11)
|
|
214
|
+
|
|
215
|
+
def test_estimate_tokens_matches_js_for_mixed_unicode_content(self):
|
|
216
|
+
message = {
|
|
217
|
+
"role": "user",
|
|
218
|
+
"content": "héllo wörld — café 日本語 emoji \U0001F389 test",
|
|
219
|
+
}
|
|
220
|
+
# node -e "..." -> JSON.stringify length 64, matching the exact
|
|
221
|
+
# case reported by the differential review.
|
|
222
|
+
self.assertEqual(estimate_tokens(message), 16)
|
|
223
|
+
|
|
224
|
+
def test_estimate_tokens_ascii_only_message_unaffected_by_the_fix(self):
|
|
225
|
+
message = {"role": "user", "content": "plain ascii text"}
|
|
226
|
+
# node -e "..." -> JSON.stringify length 44; ensure_ascii=False and
|
|
227
|
+
# the UTF-16 length calculation must not change ASCII-only counts.
|
|
228
|
+
self.assertEqual(estimate_tokens(message), 11)
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
class FuzzSweepTests(unittest.TestCase):
|
|
232
|
+
"""Equivalent of the JS package's 350-configuration fuzz sweep: two
|
|
233
|
+
separate tool-call groups back to back, swept across max_tokens 1-14,
|
|
234
|
+
keep_head 0-4, keep_tail 0-4. Asserts no orphan tool message ever
|
|
235
|
+
survives and that messages + dropped always accounts for every input
|
|
236
|
+
message.
|
|
237
|
+
"""
|
|
238
|
+
|
|
239
|
+
@staticmethod
|
|
240
|
+
def _build_conversation():
|
|
241
|
+
return [
|
|
242
|
+
{"role": "system", "content": "sys"}, # 0
|
|
243
|
+
{"role": "user", "content": "u1"}, # 1
|
|
244
|
+
{ # 2 - group 1 start
|
|
245
|
+
"role": "assistant",
|
|
246
|
+
"content": None,
|
|
247
|
+
"tool_calls": [{"id": "a"}],
|
|
248
|
+
},
|
|
249
|
+
{"role": "tool", "tool_call_id": "a", "content": "ra"}, # 3 - group 1
|
|
250
|
+
{ # 4 - group 2 start, immediately after group 1
|
|
251
|
+
"role": "assistant",
|
|
252
|
+
"content": None,
|
|
253
|
+
"tool_calls": [{"id": "b"}],
|
|
254
|
+
},
|
|
255
|
+
{"role": "tool", "tool_call_id": "b", "content": "rb"}, # 5 - group 2
|
|
256
|
+
{"role": "user", "content": "u2"}, # 6
|
|
257
|
+
{"role": "assistant", "content": "final"}, # 7
|
|
258
|
+
]
|
|
259
|
+
|
|
260
|
+
@staticmethod
|
|
261
|
+
def _assistant_ids_by_tool_call_id(messages):
|
|
262
|
+
owner = {}
|
|
263
|
+
for m in messages:
|
|
264
|
+
if m.get("role") == "assistant" and isinstance(m.get("tool_calls"), list):
|
|
265
|
+
for tc in m["tool_calls"]:
|
|
266
|
+
owner[tc["id"]] = id(m)
|
|
267
|
+
return owner
|
|
268
|
+
|
|
269
|
+
def test_no_orphan_tool_messages_across_the_full_sweep(self):
|
|
270
|
+
messages = self._build_conversation()
|
|
271
|
+
owner_by_tool_call_id = self._assistant_ids_by_tool_call_id(messages)
|
|
272
|
+
|
|
273
|
+
configs_checked = 0
|
|
274
|
+
for max_tokens in range(1, 15): # 1..14
|
|
275
|
+
for keep_head in range(0, 5): # 0..4
|
|
276
|
+
for keep_tail in range(0, 5): # 0..4
|
|
277
|
+
configs_checked += 1
|
|
278
|
+
result = compact(
|
|
279
|
+
messages,
|
|
280
|
+
max_tokens=max_tokens,
|
|
281
|
+
count_tokens=const_tokens,
|
|
282
|
+
keep_head=keep_head,
|
|
283
|
+
keep_tail=keep_tail,
|
|
284
|
+
)
|
|
285
|
+
|
|
286
|
+
# Every message is accounted for exactly once.
|
|
287
|
+
self.assertEqual(
|
|
288
|
+
len(result.messages) + len(result.dropped),
|
|
289
|
+
len(messages),
|
|
290
|
+
msg=f"config max_tokens={max_tokens} keep_head={keep_head} keep_tail={keep_tail}",
|
|
291
|
+
)
|
|
292
|
+
|
|
293
|
+
# No orphan tool message: every surviving tool message's
|
|
294
|
+
# owning assistant message also survives.
|
|
295
|
+
result_ids = {id(m) for m in result.messages}
|
|
296
|
+
for m in result.messages:
|
|
297
|
+
if m.get("role") == "tool":
|
|
298
|
+
owner_id = owner_by_tool_call_id.get(m.get("tool_call_id"))
|
|
299
|
+
self.assertIsNotNone(
|
|
300
|
+
owner_id,
|
|
301
|
+
msg=f"tool message with no owning assistant survived, "
|
|
302
|
+
f"config max_tokens={max_tokens} keep_head={keep_head} keep_tail={keep_tail}",
|
|
303
|
+
)
|
|
304
|
+
self.assertIn(
|
|
305
|
+
owner_id,
|
|
306
|
+
result_ids,
|
|
307
|
+
msg=f"orphan tool message survived without its assistant, "
|
|
308
|
+
f"config max_tokens={max_tokens} keep_head={keep_head} keep_tail={keep_tail}",
|
|
309
|
+
)
|
|
310
|
+
|
|
311
|
+
self.assertEqual(configs_checked, 350)
|
|
312
|
+
|
|
313
|
+
|
|
314
|
+
if __name__ == "__main__":
|
|
315
|
+
unittest.main()
|