thanos-ctx 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.
- thanos_ctx-0.1.0/PKG-INFO +98 -0
- thanos_ctx-0.1.0/README.md +83 -0
- thanos_ctx-0.1.0/pyproject.toml +35 -0
- thanos_ctx-0.1.0/setup.cfg +4 -0
- thanos_ctx-0.1.0/tests/test_cli.py +71 -0
- thanos_ctx-0.1.0/tests/test_compressor.py +56 -0
- thanos_ctx-0.1.0/tests/test_guard.py +93 -0
- thanos_ctx-0.1.0/tests/test_pricing.py +41 -0
- thanos_ctx-0.1.0/thanos/__init__.py +20 -0
- thanos_ctx-0.1.0/thanos/budget.py +34 -0
- thanos_ctx-0.1.0/thanos/cli.py +167 -0
- thanos_ctx-0.1.0/thanos/compressor.py +50 -0
- thanos_ctx-0.1.0/thanos/data/models.json +6282 -0
- thanos_ctx-0.1.0/thanos/guard.py +106 -0
- thanos_ctx-0.1.0/thanos/pricing.py +107 -0
- thanos_ctx-0.1.0/thanos/tokenizer.py +43 -0
- thanos_ctx-0.1.0/thanos_ctx.egg-info/PKG-INFO +98 -0
- thanos_ctx-0.1.0/thanos_ctx.egg-info/SOURCES.txt +20 -0
- thanos_ctx-0.1.0/thanos_ctx.egg-info/dependency_links.txt +1 -0
- thanos_ctx-0.1.0/thanos_ctx.egg-info/entry_points.txt +2 -0
- thanos_ctx-0.1.0/thanos_ctx.egg-info/requires.txt +4 -0
- thanos_ctx-0.1.0/thanos_ctx.egg-info/top_level.txt +1 -0
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: thanos-ctx
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Lightweight context-window limits, pricing-tier awareness, and compression for LLM conversations
|
|
5
|
+
Author: Karthikeya Surabhi
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/freakomonk/thanos
|
|
8
|
+
Keywords: llm,context,tokens,pricing,prompt-caching,anthropic,openai
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Requires-Python: >=3.9
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
Provides-Extra: dev
|
|
13
|
+
Requires-Dist: pytest; extra == "dev"
|
|
14
|
+
Requires-Dist: ruff; extra == "dev"
|
|
15
|
+
|
|
16
|
+
# thanos (Python)
|
|
17
|
+
|
|
18
|
+
Keep an LLM conversation inside its context window and pricing tier, track
|
|
19
|
+
spend, and compress history automatically when a limit is hit.
|
|
20
|
+
|
|
21
|
+
```bash
|
|
22
|
+
pip install thanos-ctx
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
(Published as `thanos-ctx` — the bare `thanos` name on PyPI belongs to an
|
|
26
|
+
unrelated project. The import stays `from thanos import ...`.)
|
|
27
|
+
|
|
28
|
+
```python
|
|
29
|
+
from thanos import ContextGuard
|
|
30
|
+
|
|
31
|
+
guard = ContextGuard(model="claude-sonnet-4-5")
|
|
32
|
+
|
|
33
|
+
messages = [
|
|
34
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
35
|
+
{"role": "user", "content": "..."},
|
|
36
|
+
# ... growing conversation history
|
|
37
|
+
]
|
|
38
|
+
|
|
39
|
+
result = guard.prepare(messages)
|
|
40
|
+
|
|
41
|
+
# Send result.messages to the API instead of the original list.
|
|
42
|
+
response = client.messages.create(model="claude-sonnet-4-5", messages=result.messages, ...)
|
|
43
|
+
|
|
44
|
+
print(result.tokens, result.estimated_cost_usd, result.compressed)
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
## Budgets
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
from thanos import Budget, ContextGuard
|
|
51
|
+
|
|
52
|
+
budget = Budget(max_cost_usd=5.00)
|
|
53
|
+
guard = ContextGuard(model="claude-sonnet-4-5", budget=budget)
|
|
54
|
+
|
|
55
|
+
for turn in conversation:
|
|
56
|
+
result = guard.prepare(turn.messages) # auto-compresses once max_cost_usd is projected to be exceeded
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
## Custom compression
|
|
60
|
+
|
|
61
|
+
```python
|
|
62
|
+
from thanos import ContextGuard
|
|
63
|
+
|
|
64
|
+
def summarize_and_truncate(messages, target_tokens, pinned_prefix_len, tokenizer):
|
|
65
|
+
# Your own strategy — e.g. call an LLM to summarize the dropped turns
|
|
66
|
+
# into one system note, then fall back to truncation.
|
|
67
|
+
...
|
|
68
|
+
|
|
69
|
+
guard = ContextGuard(model="claude-sonnet-4-5", compressor=summarize_and_truncate)
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
## Reacting to limits
|
|
73
|
+
|
|
74
|
+
```python
|
|
75
|
+
def alert(info: dict) -> None:
|
|
76
|
+
print(f"compressed={info['compressed']} tokens={info['tokens']} cost=${info['estimated_cost_usd']:.4f}")
|
|
77
|
+
|
|
78
|
+
guard = ContextGuard(model="claude-sonnet-4-5", on_limit=alert)
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
Pass `raise_on_limit=True` to raise `thanos.LimitExceeded` instead of just
|
|
82
|
+
notifying `on_limit`.
|
|
83
|
+
|
|
84
|
+
## Adding or overriding a model
|
|
85
|
+
|
|
86
|
+
```python
|
|
87
|
+
guard = ContextGuard(
|
|
88
|
+
model="my-self-hosted-model",
|
|
89
|
+
models={
|
|
90
|
+
"my-self-hosted-model": {
|
|
91
|
+
"contextWindow": 128_000,
|
|
92
|
+
"tiers": [{"maxTokens": 128_000, "inputPerMTok": 0.0, "outputPerMTok": 0.0}],
|
|
93
|
+
}
|
|
94
|
+
},
|
|
95
|
+
)
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
See the [top-level README](../README.md) for the full design rationale.
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
# thanos (Python)
|
|
2
|
+
|
|
3
|
+
Keep an LLM conversation inside its context window and pricing tier, track
|
|
4
|
+
spend, and compress history automatically when a limit is hit.
|
|
5
|
+
|
|
6
|
+
```bash
|
|
7
|
+
pip install thanos-ctx
|
|
8
|
+
```
|
|
9
|
+
|
|
10
|
+
(Published as `thanos-ctx` — the bare `thanos` name on PyPI belongs to an
|
|
11
|
+
unrelated project. The import stays `from thanos import ...`.)
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
from thanos import ContextGuard
|
|
15
|
+
|
|
16
|
+
guard = ContextGuard(model="claude-sonnet-4-5")
|
|
17
|
+
|
|
18
|
+
messages = [
|
|
19
|
+
{"role": "system", "content": "You are a helpful assistant."},
|
|
20
|
+
{"role": "user", "content": "..."},
|
|
21
|
+
# ... growing conversation history
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
result = guard.prepare(messages)
|
|
25
|
+
|
|
26
|
+
# Send result.messages to the API instead of the original list.
|
|
27
|
+
response = client.messages.create(model="claude-sonnet-4-5", messages=result.messages, ...)
|
|
28
|
+
|
|
29
|
+
print(result.tokens, result.estimated_cost_usd, result.compressed)
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
## Budgets
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
from thanos import Budget, ContextGuard
|
|
36
|
+
|
|
37
|
+
budget = Budget(max_cost_usd=5.00)
|
|
38
|
+
guard = ContextGuard(model="claude-sonnet-4-5", budget=budget)
|
|
39
|
+
|
|
40
|
+
for turn in conversation:
|
|
41
|
+
result = guard.prepare(turn.messages) # auto-compresses once max_cost_usd is projected to be exceeded
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
## Custom compression
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
from thanos import ContextGuard
|
|
48
|
+
|
|
49
|
+
def summarize_and_truncate(messages, target_tokens, pinned_prefix_len, tokenizer):
|
|
50
|
+
# Your own strategy — e.g. call an LLM to summarize the dropped turns
|
|
51
|
+
# into one system note, then fall back to truncation.
|
|
52
|
+
...
|
|
53
|
+
|
|
54
|
+
guard = ContextGuard(model="claude-sonnet-4-5", compressor=summarize_and_truncate)
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
## Reacting to limits
|
|
58
|
+
|
|
59
|
+
```python
|
|
60
|
+
def alert(info: dict) -> None:
|
|
61
|
+
print(f"compressed={info['compressed']} tokens={info['tokens']} cost=${info['estimated_cost_usd']:.4f}")
|
|
62
|
+
|
|
63
|
+
guard = ContextGuard(model="claude-sonnet-4-5", on_limit=alert)
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
Pass `raise_on_limit=True` to raise `thanos.LimitExceeded` instead of just
|
|
67
|
+
notifying `on_limit`.
|
|
68
|
+
|
|
69
|
+
## Adding or overriding a model
|
|
70
|
+
|
|
71
|
+
```python
|
|
72
|
+
guard = ContextGuard(
|
|
73
|
+
model="my-self-hosted-model",
|
|
74
|
+
models={
|
|
75
|
+
"my-self-hosted-model": {
|
|
76
|
+
"contextWindow": 128_000,
|
|
77
|
+
"tiers": [{"maxTokens": 128_000, "inputPerMTok": 0.0, "outputPerMTok": 0.0}],
|
|
78
|
+
}
|
|
79
|
+
},
|
|
80
|
+
)
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
See the [top-level README](../README.md) for the full design rationale.
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "thanos-ctx"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Lightweight context-window limits, pricing-tier awareness, and compression for LLM conversations"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
requires-python = ">=3.9"
|
|
12
|
+
authors = [{ name = "Karthikeya Surabhi" }]
|
|
13
|
+
keywords = ["llm", "context", "tokens", "pricing", "prompt-caching", "anthropic", "openai"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
]
|
|
17
|
+
dependencies = []
|
|
18
|
+
|
|
19
|
+
[project.optional-dependencies]
|
|
20
|
+
dev = ["pytest", "ruff"]
|
|
21
|
+
|
|
22
|
+
[project.urls]
|
|
23
|
+
Homepage = "https://github.com/freakomonk/thanos"
|
|
24
|
+
|
|
25
|
+
[project.scripts]
|
|
26
|
+
thanos = "thanos.cli:main"
|
|
27
|
+
|
|
28
|
+
[tool.setuptools.packages.find]
|
|
29
|
+
include = ["thanos*"]
|
|
30
|
+
|
|
31
|
+
[tool.setuptools.package-data]
|
|
32
|
+
thanos = ["data/models.json"]
|
|
33
|
+
|
|
34
|
+
[tool.ruff]
|
|
35
|
+
target-version = "py39"
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
import json
|
|
2
|
+
|
|
3
|
+
from thanos.cli import main
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def _write_messages(tmp_path, messages):
|
|
7
|
+
path = tmp_path / "conversation.json"
|
|
8
|
+
path.write_text(json.dumps(messages))
|
|
9
|
+
return str(path)
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_check_prints_report_and_returns_zero_when_under_limit(tmp_path, capsys):
|
|
13
|
+
path = _write_messages(tmp_path, [{"role": "user", "content": "hello"}])
|
|
14
|
+
|
|
15
|
+
code = main(["check", path, "--model", "claude-haiku-4-5"])
|
|
16
|
+
|
|
17
|
+
out = capsys.readouterr().out
|
|
18
|
+
assert code == 0
|
|
19
|
+
assert "model:" in out
|
|
20
|
+
assert "claude-haiku-4-5" in out
|
|
21
|
+
assert "limit exceeded: False" in out
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def test_check_json_output_is_valid_json(tmp_path, capsys):
|
|
25
|
+
path = _write_messages(tmp_path, [{"role": "user", "content": "hello"}])
|
|
26
|
+
|
|
27
|
+
code = main(["check", path, "--model", "claude-haiku-4-5", "--json"])
|
|
28
|
+
|
|
29
|
+
out = capsys.readouterr().out
|
|
30
|
+
report = json.loads(out)
|
|
31
|
+
assert code == 0
|
|
32
|
+
assert report["model"] == "claude-haiku-4-5"
|
|
33
|
+
assert report["limit_exceeded"] is False
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def test_check_returns_nonzero_when_budget_exceeded(tmp_path, capsys):
|
|
37
|
+
path = _write_messages(tmp_path, [{"role": "user", "content": "x" * 4000}])
|
|
38
|
+
|
|
39
|
+
code = main(
|
|
40
|
+
["check", path, "--model", "claude-haiku-4-5", "--max-tokens", "10", "--json"]
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
report = json.loads(capsys.readouterr().out)
|
|
44
|
+
assert code == 1
|
|
45
|
+
assert report["limit_exceeded"] is True
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_check_accepts_object_with_messages_key(tmp_path, capsys):
|
|
49
|
+
path = tmp_path / "conversation.json"
|
|
50
|
+
path.write_text(json.dumps({"messages": [{"role": "user", "content": "hi"}]}))
|
|
51
|
+
|
|
52
|
+
code = main(["check", str(path), "--model", "claude-haiku-4-5"])
|
|
53
|
+
|
|
54
|
+
assert code == 0
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def test_models_lists_bundled_registry(capsys):
|
|
58
|
+
code = main(["models"])
|
|
59
|
+
|
|
60
|
+
out = capsys.readouterr().out
|
|
61
|
+
assert code == 0
|
|
62
|
+
assert "claude-sonnet-4-5" in out
|
|
63
|
+
assert "gpt-5" in out
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def test_models_json_output(capsys):
|
|
67
|
+
code = main(["models", "--json"])
|
|
68
|
+
|
|
69
|
+
registry = json.loads(capsys.readouterr().out)
|
|
70
|
+
assert code == 0
|
|
71
|
+
assert registry["claude-sonnet-4-5"]["provider"] == "anthropic"
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
from thanos import TruncateCompressor, estimate_tokens
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def _msg(role: str, text: str) -> dict:
|
|
5
|
+
return {"role": role, "content": text}
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_truncate_keeps_pinned_prefix_untouched():
|
|
9
|
+
messages = [
|
|
10
|
+
_msg("system", "You are a helpful assistant. " * 5),
|
|
11
|
+
_msg("user", "turn 1 " * 50),
|
|
12
|
+
_msg("assistant", "turn 1 reply " * 50),
|
|
13
|
+
_msg("user", "turn 2 " * 50),
|
|
14
|
+
_msg("assistant", "turn 2 reply " * 50),
|
|
15
|
+
_msg("user", "most recent turn " * 5),
|
|
16
|
+
]
|
|
17
|
+
compressor = TruncateCompressor()
|
|
18
|
+
target = estimate_tokens([messages[0], messages[-1]]) + 5
|
|
19
|
+
|
|
20
|
+
result = compressor(
|
|
21
|
+
messages, target, pinned_prefix_len=1, tokenizer=estimate_tokens
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
assert result[0] == messages[0]
|
|
25
|
+
assert estimate_tokens(result) <= target or len(result) == 1
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def test_truncate_drops_oldest_turns_first():
|
|
29
|
+
messages = [
|
|
30
|
+
_msg("system", "sys"),
|
|
31
|
+
_msg("user", "oldest"),
|
|
32
|
+
_msg("assistant", "oldest reply"),
|
|
33
|
+
_msg("user", "newest"),
|
|
34
|
+
]
|
|
35
|
+
compressor = TruncateCompressor()
|
|
36
|
+
# Target fits system + only the very last message.
|
|
37
|
+
target = estimate_tokens([messages[0], messages[-1]])
|
|
38
|
+
|
|
39
|
+
result = compressor(
|
|
40
|
+
messages, target, pinned_prefix_len=1, tokenizer=estimate_tokens
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
assert result[0]["role"] == "system"
|
|
44
|
+
assert result[-1] == messages[-1]
|
|
45
|
+
assert messages[1] not in result
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_truncate_never_touches_pinned_prefix_even_if_oversized():
|
|
49
|
+
messages = [_msg("system", "x" * 10_000), _msg("user", "hi")]
|
|
50
|
+
compressor = TruncateCompressor()
|
|
51
|
+
|
|
52
|
+
result = compressor(
|
|
53
|
+
messages, target_tokens=1, pinned_prefix_len=1, tokenizer=estimate_tokens
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
assert result == [messages[0]]
|
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
from thanos import Budget, ContextGuard
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def _msg(role: str, text: str) -> dict:
|
|
5
|
+
return {"role": role, "content": text}
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_prepare_returns_stats_without_compression_when_under_limit():
|
|
9
|
+
guard = ContextGuard(model="claude-haiku-4-5")
|
|
10
|
+
messages = [_msg("user", "hello there")]
|
|
11
|
+
|
|
12
|
+
result = guard.prepare(messages)
|
|
13
|
+
|
|
14
|
+
assert result.compressed is False
|
|
15
|
+
assert result.limit_exceeded is False
|
|
16
|
+
assert result.tokens > 0
|
|
17
|
+
assert result.messages == messages
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def test_prepare_compresses_when_over_tier_ceiling():
|
|
21
|
+
guard = ContextGuard(model="claude-haiku-4-5", stay_in_tier=0)
|
|
22
|
+
big_history = [_msg("system", "sys")] + [
|
|
23
|
+
_msg("user", "x" * 4000) for _ in range(300)
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
result = guard.prepare(big_history)
|
|
27
|
+
|
|
28
|
+
assert result.compressed is True
|
|
29
|
+
assert len(result.messages) < len(big_history)
|
|
30
|
+
assert result.messages[0]["role"] == "system"
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def test_on_limit_callback_fires_and_does_not_raise_by_default():
|
|
34
|
+
calls = []
|
|
35
|
+
guard = ContextGuard(
|
|
36
|
+
model="claude-haiku-4-5",
|
|
37
|
+
stay_in_tier=0,
|
|
38
|
+
on_limit=lambda info: calls.append(info),
|
|
39
|
+
)
|
|
40
|
+
big_history = [_msg("system", "sys")] + [
|
|
41
|
+
_msg("user", "x" * 4000) for _ in range(300)
|
|
42
|
+
]
|
|
43
|
+
|
|
44
|
+
guard.prepare(big_history)
|
|
45
|
+
|
|
46
|
+
assert len(calls) == 1
|
|
47
|
+
assert calls[0]["compressed"] is True
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def test_budget_ceiling_triggers_limit_and_raise_on_limit_raises():
|
|
51
|
+
from thanos import LimitExceeded
|
|
52
|
+
|
|
53
|
+
budget = Budget(max_tokens=50)
|
|
54
|
+
guard = ContextGuard(model="claude-haiku-4-5", budget=budget, raise_on_limit=True)
|
|
55
|
+
messages = [_msg("user", "x" * 4000)] # far more than 50 tokens
|
|
56
|
+
|
|
57
|
+
try:
|
|
58
|
+
guard.prepare(messages)
|
|
59
|
+
assert False, "expected LimitExceeded"
|
|
60
|
+
except LimitExceeded:
|
|
61
|
+
pass
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def test_custom_tokenizer_and_compressor_are_used():
|
|
65
|
+
from thanos.pricing import ModelProfile, PriceTier
|
|
66
|
+
|
|
67
|
+
calls = {"tokenizer": 0, "compressor": 0}
|
|
68
|
+
|
|
69
|
+
def fixed_tokenizer(messages):
|
|
70
|
+
calls["tokenizer"] += 1
|
|
71
|
+
return 10
|
|
72
|
+
|
|
73
|
+
def noop_compressor(messages, target_tokens, pinned_prefix_len, tokenizer):
|
|
74
|
+
calls["compressor"] += 1
|
|
75
|
+
return messages
|
|
76
|
+
|
|
77
|
+
tiny_model = ModelProfile(
|
|
78
|
+
model_id="tiny",
|
|
79
|
+
context_window=5,
|
|
80
|
+
tiers=[PriceTier(max_tokens=5, input_per_m_tok=1.0, output_per_m_tok=1.0)],
|
|
81
|
+
)
|
|
82
|
+
guard = ContextGuard(
|
|
83
|
+
model=tiny_model,
|
|
84
|
+
tokenizer=fixed_tokenizer,
|
|
85
|
+
compressor=noop_compressor,
|
|
86
|
+
stay_in_tier=None,
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
result = guard.prepare([_msg("user", "hi")])
|
|
90
|
+
|
|
91
|
+
assert calls["tokenizer"] >= 1
|
|
92
|
+
assert calls["compressor"] == 1
|
|
93
|
+
assert result.tokens == 10
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
from thanos import load_models
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def test_loads_bundled_models():
|
|
5
|
+
models = load_models()
|
|
6
|
+
assert "claude-sonnet-4-5" in models
|
|
7
|
+
sonnet = models["claude-sonnet-4-5"]
|
|
8
|
+
assert sonnet.context_window == 1_000_000
|
|
9
|
+
assert len(sonnet.tiers) == 2
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_tier_for_selects_cheaper_tier_below_threshold():
|
|
13
|
+
models = load_models()
|
|
14
|
+
sonnet = models["claude-sonnet-4-5"]
|
|
15
|
+
cheap_tier = sonnet.tier_for(100_000)
|
|
16
|
+
expensive_tier = sonnet.tier_for(500_000)
|
|
17
|
+
assert cheap_tier.input_per_m_tok < expensive_tier.input_per_m_tok
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def test_estimate_cost_usd_bills_cache_tokens_at_cache_rate():
|
|
21
|
+
models = load_models()
|
|
22
|
+
haiku = models["claude-haiku-4-5"]
|
|
23
|
+
# 1,000 total tokens, all served from cache read.
|
|
24
|
+
cost = haiku.estimate_cost_usd(total_input_tokens=1000, cache_read_tokens=1000)
|
|
25
|
+
expected = (1000 / 1_000_000) * haiku.cache.read_per_m_tok
|
|
26
|
+
assert cost == expected
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def test_extra_models_override_and_extend_registry():
|
|
30
|
+
models = load_models(
|
|
31
|
+
extra={
|
|
32
|
+
"my-custom-model": {
|
|
33
|
+
"contextWindow": 32000,
|
|
34
|
+
"tiers": [
|
|
35
|
+
{"maxTokens": 32000, "inputPerMTok": 0.5, "outputPerMTok": 1.5}
|
|
36
|
+
],
|
|
37
|
+
}
|
|
38
|
+
}
|
|
39
|
+
)
|
|
40
|
+
assert "my-custom-model" in models
|
|
41
|
+
assert models["my-custom-model"].context_window == 32000
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
from .budget import Budget
|
|
2
|
+
from .compressor import Compressor, TruncateCompressor
|
|
3
|
+
from .guard import ContextGuard, GuardResult, LimitExceeded
|
|
4
|
+
from .pricing import ModelProfile, load_models
|
|
5
|
+
from .tokenizer import Tokenizer, estimate_tokens
|
|
6
|
+
|
|
7
|
+
__all__ = [
|
|
8
|
+
"Budget",
|
|
9
|
+
"Compressor",
|
|
10
|
+
"ContextGuard",
|
|
11
|
+
"GuardResult",
|
|
12
|
+
"LimitExceeded",
|
|
13
|
+
"ModelProfile",
|
|
14
|
+
"Tokenizer",
|
|
15
|
+
"TruncateCompressor",
|
|
16
|
+
"estimate_tokens",
|
|
17
|
+
"load_models",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
__version__ = "0.1.0"
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@dataclass
|
|
7
|
+
class Budget:
|
|
8
|
+
"""Cumulative ceiling for a conversation/session.
|
|
9
|
+
|
|
10
|
+
Either limit is optional; set only the one(s) you care about.
|
|
11
|
+
`used_tokens`/`used_cost_usd` are running totals that `ContextGuard`
|
|
12
|
+
updates after each `prepare()` call when you pass `record=True`.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
max_tokens: int | None = None
|
|
16
|
+
max_cost_usd: float | None = None
|
|
17
|
+
used_tokens: int = 0
|
|
18
|
+
used_cost_usd: float = 0.0
|
|
19
|
+
|
|
20
|
+
def record(self, tokens: int, cost_usd: float) -> None:
|
|
21
|
+
self.used_tokens += tokens
|
|
22
|
+
self.used_cost_usd += cost_usd
|
|
23
|
+
|
|
24
|
+
def exceeded(self) -> bool:
|
|
25
|
+
if self.max_tokens is not None and self.used_tokens > self.max_tokens:
|
|
26
|
+
return True
|
|
27
|
+
return bool(
|
|
28
|
+
self.max_cost_usd is not None and self.used_cost_usd > self.max_cost_usd
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
def remaining_tokens(self, fallback: int) -> int:
|
|
32
|
+
if self.max_tokens is None:
|
|
33
|
+
return fallback
|
|
34
|
+
return max(0, self.max_tokens - self.used_tokens)
|