cheapskate 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.
- cheapskate-0.1.0/LICENSE +21 -0
- cheapskate-0.1.0/PKG-INFO +207 -0
- cheapskate-0.1.0/README.md +174 -0
- cheapskate-0.1.0/cheapskate/__init__.py +22 -0
- cheapskate-0.1.0/cheapskate/callback.py +115 -0
- cheapskate-0.1.0/cheapskate/compressor.py +229 -0
- cheapskate-0.1.0/cheapskate/pruner.py +318 -0
- cheapskate-0.1.0/cheapskate/router.py +429 -0
- cheapskate-0.1.0/cheapskate/token_counter.py +189 -0
- cheapskate-0.1.0/pyproject.toml +65 -0
cheapskate-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 CheapSkate Contributors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,207 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cheapskate
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Drop-in LangChain/LangGraph middleware that cuts frontier LLM API costs via tool pruning, prompt squeezing, and hybrid SLM routing.
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
License-File: LICENSE
|
|
7
|
+
Keywords: langchain,langgraph,llm,cost-optimization,routing,tool-pruning,prompt-compression
|
|
8
|
+
Author: CheapSkate Contributors
|
|
9
|
+
Author-email: maintainers@cheapskate.dev
|
|
10
|
+
Requires-Python: >=3.9,<4.0
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Operating System :: OS Independent
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Classifier: Typing :: Typed
|
|
21
|
+
Requires-Dist: langchain-core (>=0.3.0,<1.0.0)
|
|
22
|
+
Requires-Dist: langchain-groq (>=0.2.0,<1.0.0)
|
|
23
|
+
Requires-Dist: langchain-openai (>=0.2.0,<1.0.0)
|
|
24
|
+
Requires-Dist: loguru (>=0.7.0)
|
|
25
|
+
Requires-Dist: numpy (>=1.24.0)
|
|
26
|
+
Requires-Dist: tiktoken (>=0.7.0)
|
|
27
|
+
Project-URL: Bug Tracker, https://github.com/cheapskate-ai/cheapskate/issues
|
|
28
|
+
Project-URL: Documentation, https://github.com/cheapskate-ai/cheapskate#readme
|
|
29
|
+
Project-URL: Homepage, https://github.com/cheapskate-ai/cheapskate
|
|
30
|
+
Project-URL: Repository, https://github.com/cheapskate-ai/cheapskate
|
|
31
|
+
Description-Content-Type: text/markdown
|
|
32
|
+
|
|
33
|
+
# cheapskate
|
|
34
|
+
|
|
35
|
+
**Drop-in cost control for LangChain and LangGraph agents.**
|
|
36
|
+
|
|
37
|
+
`cheapskate` sits in front of expensive frontier models (GPT-4o, Claude-class APIs) and automatically applies three levers that cut token waste and route easy work to cheap SLMs — without rewriting your agent graph.
|
|
38
|
+
|
|
39
|
+
Built for teams that already ship tool-calling agents and want lower bills without giving up quality on hard tasks.
|
|
40
|
+
|
|
41
|
+
---
|
|
42
|
+
|
|
43
|
+
## Why it exists
|
|
44
|
+
|
|
45
|
+
Agent stacks get expensive for boring reasons:
|
|
46
|
+
|
|
47
|
+
- every turn re-sends bloated tool schemas
|
|
48
|
+
- prompts accumulate filler and repeated boilerplate
|
|
49
|
+
- simple classify / summarize / translate calls still hit GPT-4o prices
|
|
50
|
+
|
|
51
|
+
`cheapskate` intercepts those calls as a LangChain `BaseChatModel`, so it works as a middleware layer inside existing chains and stateful LangGraph workflows.
|
|
52
|
+
|
|
53
|
+
---
|
|
54
|
+
|
|
55
|
+
## Core features
|
|
56
|
+
|
|
57
|
+
### 1. Advanced tool pruning
|
|
58
|
+
Keeps only the tool schemas that matter for the current turn.
|
|
59
|
+
|
|
60
|
+
- Groups tools into **namespaces** (weather, finance, ops, …)
|
|
61
|
+
- Scores relevance with **context-aware vector similarity** over full message history, including prior tool results / scratchpad text
|
|
62
|
+
- Uses numpy cosine similarity — not naive keyword matching alone
|
|
63
|
+
- Optional **namespace classifier hook** (plug in an SLM or custom ranker)
|
|
64
|
+
- Hard guarantee: tools listed in `always_keep` are never pruned
|
|
65
|
+
|
|
66
|
+
**Effect:** fewer tools in the prompt → fewer input tokens on every frontier call.
|
|
67
|
+
|
|
68
|
+
### 2. Prompt squeezing
|
|
69
|
+
Compresses conversational fluff while protecting structured payloads.
|
|
70
|
+
|
|
71
|
+
- Operates on a **shallow copy** of messages so LangGraph checkpointers keep an uncorrupted history
|
|
72
|
+
- Scrubs high-frequency boilerplate and filler phrases
|
|
73
|
+
- Uses **token-density** heuristics (via `tiktoken`) to truncate low-signal older turns
|
|
74
|
+
- Leaves valid JSON blocks inside message text intact
|
|
75
|
+
|
|
76
|
+
**Effect:** smaller context windows without silently breaking tool arguments or stored state.
|
|
77
|
+
|
|
78
|
+
### 3. Hybrid SLM routing
|
|
79
|
+
Sends easy, deterministic work to a cheap secondary model (Groq / Together-style Llama-class SLMs).
|
|
80
|
+
|
|
81
|
+
- Lightweight heuristics decide when a task is “simple enough”
|
|
82
|
+
- Hard / ambiguous / tool-heavy work stays on the primary frontier model
|
|
83
|
+
- **Failsafe path:** if the SLM errors, times out, or returns invalid structured output, cheapskate logs the failure and immediately retries on the primary model with the **original unmodified payload**
|
|
84
|
+
|
|
85
|
+
**Effect:** most of the spend reduction on repetitive turns, without single-point-of-failure routing.
|
|
86
|
+
|
|
87
|
+
### 4. Observability built in
|
|
88
|
+
`CheapSkateCallbackHandler` + router metrics track:
|
|
89
|
+
|
|
90
|
+
- baseline vs optimized token counts
|
|
91
|
+
- prune / compression / routing decisions
|
|
92
|
+
- failover events
|
|
93
|
+
|
|
94
|
+
---
|
|
95
|
+
|
|
96
|
+
## Quick start
|
|
97
|
+
|
|
98
|
+
```bash
|
|
99
|
+
pip install cheapskate
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
```python
|
|
103
|
+
from langchain_openai import ChatOpenAI
|
|
104
|
+
from langchain_groq import ChatGroq
|
|
105
|
+
from cheapskate import CheapSkateRouter, ToolNamespace
|
|
106
|
+
|
|
107
|
+
primary = ChatOpenAI(model="gpt-4o")
|
|
108
|
+
secondary = ChatGroq(model="llama-3.2-3b-preview")
|
|
109
|
+
|
|
110
|
+
router = CheapSkateRouter(
|
|
111
|
+
primary_model=primary,
|
|
112
|
+
secondary_model=secondary,
|
|
113
|
+
always_keep_tools=["search_docs", "calculator"],
|
|
114
|
+
tool_namespaces=[
|
|
115
|
+
ToolNamespace(
|
|
116
|
+
name="weather",
|
|
117
|
+
tools=("get_weather", "get_forecast"),
|
|
118
|
+
description="weather forecast temperature",
|
|
119
|
+
),
|
|
120
|
+
ToolNamespace(
|
|
121
|
+
name="finance",
|
|
122
|
+
tools=("get_stock_price", "list_portfolios"),
|
|
123
|
+
description="stocks portfolios equity",
|
|
124
|
+
),
|
|
125
|
+
],
|
|
126
|
+
)
|
|
127
|
+
|
|
128
|
+
response = router.invoke(
|
|
129
|
+
[{"role": "user", "content": "Summarize this in one sentence: rain is likely Saturday."}]
|
|
130
|
+
)
|
|
131
|
+
print(response.content)
|
|
132
|
+
print(router.metrics.summary)
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
Drop `CheapSkateRouter` anywhere you currently pass a chat model in LangChain / LangGraph.
|
|
136
|
+
|
|
137
|
+
---
|
|
138
|
+
|
|
139
|
+
## Architecture
|
|
140
|
+
|
|
141
|
+
```text
|
|
142
|
+
Agent / LangGraph node
|
|
143
|
+
│
|
|
144
|
+
▼
|
|
145
|
+
┌───────────────────────────┐
|
|
146
|
+
│ CheapSkateRouter │
|
|
147
|
+
│ 1. prune tool schemas │
|
|
148
|
+
│ 2. squeeze prompt tokens │
|
|
149
|
+
│ 3. route simple → SLM │
|
|
150
|
+
│ 4. failover → primary │
|
|
151
|
+
└─────────────┬─────────────┘
|
|
152
|
+
│
|
|
153
|
+
┌────────┴────────┐
|
|
154
|
+
▼ ▼
|
|
155
|
+
cheap SLM frontier LLM
|
|
156
|
+
(Groq/Together) (GPT-4o / Claude)
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
---
|
|
160
|
+
|
|
161
|
+
## Benchmarking without API keys
|
|
162
|
+
|
|
163
|
+
You can measure pruning, compression, routing, and **synthetic cost impact** offline:
|
|
164
|
+
|
|
165
|
+
```bash
|
|
166
|
+
poetry install --with dev
|
|
167
|
+
poetry run python benchmarks/offline_bench.py
|
|
168
|
+
```
|
|
169
|
+
|
|
170
|
+
This uses mock chat models + `tiktoken` + published list-price assumptions. No OpenAI / Groq keys required. Results are written to `benchmarks/offline_results.json`.
|
|
171
|
+
|
|
172
|
+
Live quality-vs-cost A/B against production traffic still needs real keys — the offline bench validates the middleware mechanics and estimates spend deltas.
|
|
173
|
+
|
|
174
|
+
---
|
|
175
|
+
|
|
176
|
+
## What is production-ready today
|
|
177
|
+
|
|
178
|
+
| Capability | Status |
|
|
179
|
+
|---|---|
|
|
180
|
+
| tiktoken token accounting | Complete |
|
|
181
|
+
| Tool pruning + `always_keep` | Complete |
|
|
182
|
+
| Prompt compression on shallow copies | Complete |
|
|
183
|
+
| Hybrid routing heuristics | Complete |
|
|
184
|
+
| SLM → primary failsafe | Complete |
|
|
185
|
+
| Sync + async `_generate` / `_agenerate` | Complete |
|
|
186
|
+
| Metrics / callback handler | Complete |
|
|
187
|
+
| Default hosted SLM namespace classifier | Hook only (bring your own) |
|
|
188
|
+
| Neural embedding provider | Local hash embeddings by default |
|
|
189
|
+
|
|
190
|
+
The common “60–80% savings” range is a **target envelope** for tool-heavy agents with lots of prompt fluff and easy turns. Real savings depend on tool cardinality, history length, and how often work is SLM-eligible. Run the offline bench, then validate on your traffic.
|
|
191
|
+
|
|
192
|
+
---
|
|
193
|
+
|
|
194
|
+
## Development
|
|
195
|
+
|
|
196
|
+
```bash
|
|
197
|
+
poetry install --with dev
|
|
198
|
+
poetry run pytest -q
|
|
199
|
+
./publish_pipeline.sh # needs Artifactory/PyPI credentials to publish
|
|
200
|
+
```
|
|
201
|
+
|
|
202
|
+
---
|
|
203
|
+
|
|
204
|
+
## License
|
|
205
|
+
|
|
206
|
+
MIT
|
|
207
|
+
|
|
@@ -0,0 +1,174 @@
|
|
|
1
|
+
# cheapskate
|
|
2
|
+
|
|
3
|
+
**Drop-in cost control for LangChain and LangGraph agents.**
|
|
4
|
+
|
|
5
|
+
`cheapskate` sits in front of expensive frontier models (GPT-4o, Claude-class APIs) and automatically applies three levers that cut token waste and route easy work to cheap SLMs — without rewriting your agent graph.
|
|
6
|
+
|
|
7
|
+
Built for teams that already ship tool-calling agents and want lower bills without giving up quality on hard tasks.
|
|
8
|
+
|
|
9
|
+
---
|
|
10
|
+
|
|
11
|
+
## Why it exists
|
|
12
|
+
|
|
13
|
+
Agent stacks get expensive for boring reasons:
|
|
14
|
+
|
|
15
|
+
- every turn re-sends bloated tool schemas
|
|
16
|
+
- prompts accumulate filler and repeated boilerplate
|
|
17
|
+
- simple classify / summarize / translate calls still hit GPT-4o prices
|
|
18
|
+
|
|
19
|
+
`cheapskate` intercepts those calls as a LangChain `BaseChatModel`, so it works as a middleware layer inside existing chains and stateful LangGraph workflows.
|
|
20
|
+
|
|
21
|
+
---
|
|
22
|
+
|
|
23
|
+
## Core features
|
|
24
|
+
|
|
25
|
+
### 1. Advanced tool pruning
|
|
26
|
+
Keeps only the tool schemas that matter for the current turn.
|
|
27
|
+
|
|
28
|
+
- Groups tools into **namespaces** (weather, finance, ops, …)
|
|
29
|
+
- Scores relevance with **context-aware vector similarity** over full message history, including prior tool results / scratchpad text
|
|
30
|
+
- Uses numpy cosine similarity — not naive keyword matching alone
|
|
31
|
+
- Optional **namespace classifier hook** (plug in an SLM or custom ranker)
|
|
32
|
+
- Hard guarantee: tools listed in `always_keep` are never pruned
|
|
33
|
+
|
|
34
|
+
**Effect:** fewer tools in the prompt → fewer input tokens on every frontier call.
|
|
35
|
+
|
|
36
|
+
### 2. Prompt squeezing
|
|
37
|
+
Compresses conversational fluff while protecting structured payloads.
|
|
38
|
+
|
|
39
|
+
- Operates on a **shallow copy** of messages so LangGraph checkpointers keep an uncorrupted history
|
|
40
|
+
- Scrubs high-frequency boilerplate and filler phrases
|
|
41
|
+
- Uses **token-density** heuristics (via `tiktoken`) to truncate low-signal older turns
|
|
42
|
+
- Leaves valid JSON blocks inside message text intact
|
|
43
|
+
|
|
44
|
+
**Effect:** smaller context windows without silently breaking tool arguments or stored state.
|
|
45
|
+
|
|
46
|
+
### 3. Hybrid SLM routing
|
|
47
|
+
Sends easy, deterministic work to a cheap secondary model (Groq / Together-style Llama-class SLMs).
|
|
48
|
+
|
|
49
|
+
- Lightweight heuristics decide when a task is “simple enough”
|
|
50
|
+
- Hard / ambiguous / tool-heavy work stays on the primary frontier model
|
|
51
|
+
- **Failsafe path:** if the SLM errors, times out, or returns invalid structured output, cheapskate logs the failure and immediately retries on the primary model with the **original unmodified payload**
|
|
52
|
+
|
|
53
|
+
**Effect:** most of the spend reduction on repetitive turns, without single-point-of-failure routing.
|
|
54
|
+
|
|
55
|
+
### 4. Observability built in
|
|
56
|
+
`CheapSkateCallbackHandler` + router metrics track:
|
|
57
|
+
|
|
58
|
+
- baseline vs optimized token counts
|
|
59
|
+
- prune / compression / routing decisions
|
|
60
|
+
- failover events
|
|
61
|
+
|
|
62
|
+
---
|
|
63
|
+
|
|
64
|
+
## Quick start
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
pip install cheapskate
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
```python
|
|
71
|
+
from langchain_openai import ChatOpenAI
|
|
72
|
+
from langchain_groq import ChatGroq
|
|
73
|
+
from cheapskate import CheapSkateRouter, ToolNamespace
|
|
74
|
+
|
|
75
|
+
primary = ChatOpenAI(model="gpt-4o")
|
|
76
|
+
secondary = ChatGroq(model="llama-3.2-3b-preview")
|
|
77
|
+
|
|
78
|
+
router = CheapSkateRouter(
|
|
79
|
+
primary_model=primary,
|
|
80
|
+
secondary_model=secondary,
|
|
81
|
+
always_keep_tools=["search_docs", "calculator"],
|
|
82
|
+
tool_namespaces=[
|
|
83
|
+
ToolNamespace(
|
|
84
|
+
name="weather",
|
|
85
|
+
tools=("get_weather", "get_forecast"),
|
|
86
|
+
description="weather forecast temperature",
|
|
87
|
+
),
|
|
88
|
+
ToolNamespace(
|
|
89
|
+
name="finance",
|
|
90
|
+
tools=("get_stock_price", "list_portfolios"),
|
|
91
|
+
description="stocks portfolios equity",
|
|
92
|
+
),
|
|
93
|
+
],
|
|
94
|
+
)
|
|
95
|
+
|
|
96
|
+
response = router.invoke(
|
|
97
|
+
[{"role": "user", "content": "Summarize this in one sentence: rain is likely Saturday."}]
|
|
98
|
+
)
|
|
99
|
+
print(response.content)
|
|
100
|
+
print(router.metrics.summary)
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
Drop `CheapSkateRouter` anywhere you currently pass a chat model in LangChain / LangGraph.
|
|
104
|
+
|
|
105
|
+
---
|
|
106
|
+
|
|
107
|
+
## Architecture
|
|
108
|
+
|
|
109
|
+
```text
|
|
110
|
+
Agent / LangGraph node
|
|
111
|
+
│
|
|
112
|
+
▼
|
|
113
|
+
┌───────────────────────────┐
|
|
114
|
+
│ CheapSkateRouter │
|
|
115
|
+
│ 1. prune tool schemas │
|
|
116
|
+
│ 2. squeeze prompt tokens │
|
|
117
|
+
│ 3. route simple → SLM │
|
|
118
|
+
│ 4. failover → primary │
|
|
119
|
+
└─────────────┬─────────────┘
|
|
120
|
+
│
|
|
121
|
+
┌────────┴────────┐
|
|
122
|
+
▼ ▼
|
|
123
|
+
cheap SLM frontier LLM
|
|
124
|
+
(Groq/Together) (GPT-4o / Claude)
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
---
|
|
128
|
+
|
|
129
|
+
## Benchmarking without API keys
|
|
130
|
+
|
|
131
|
+
You can measure pruning, compression, routing, and **synthetic cost impact** offline:
|
|
132
|
+
|
|
133
|
+
```bash
|
|
134
|
+
poetry install --with dev
|
|
135
|
+
poetry run python benchmarks/offline_bench.py
|
|
136
|
+
```
|
|
137
|
+
|
|
138
|
+
This uses mock chat models + `tiktoken` + published list-price assumptions. No OpenAI / Groq keys required. Results are written to `benchmarks/offline_results.json`.
|
|
139
|
+
|
|
140
|
+
Live quality-vs-cost A/B against production traffic still needs real keys — the offline bench validates the middleware mechanics and estimates spend deltas.
|
|
141
|
+
|
|
142
|
+
---
|
|
143
|
+
|
|
144
|
+
## What is production-ready today
|
|
145
|
+
|
|
146
|
+
| Capability | Status |
|
|
147
|
+
|---|---|
|
|
148
|
+
| tiktoken token accounting | Complete |
|
|
149
|
+
| Tool pruning + `always_keep` | Complete |
|
|
150
|
+
| Prompt compression on shallow copies | Complete |
|
|
151
|
+
| Hybrid routing heuristics | Complete |
|
|
152
|
+
| SLM → primary failsafe | Complete |
|
|
153
|
+
| Sync + async `_generate` / `_agenerate` | Complete |
|
|
154
|
+
| Metrics / callback handler | Complete |
|
|
155
|
+
| Default hosted SLM namespace classifier | Hook only (bring your own) |
|
|
156
|
+
| Neural embedding provider | Local hash embeddings by default |
|
|
157
|
+
|
|
158
|
+
The common “60–80% savings” range is a **target envelope** for tool-heavy agents with lots of prompt fluff and easy turns. Real savings depend on tool cardinality, history length, and how often work is SLM-eligible. Run the offline bench, then validate on your traffic.
|
|
159
|
+
|
|
160
|
+
---
|
|
161
|
+
|
|
162
|
+
## Development
|
|
163
|
+
|
|
164
|
+
```bash
|
|
165
|
+
poetry install --with dev
|
|
166
|
+
poetry run pytest -q
|
|
167
|
+
./publish_pipeline.sh # needs Artifactory/PyPI credentials to publish
|
|
168
|
+
```
|
|
169
|
+
|
|
170
|
+
---
|
|
171
|
+
|
|
172
|
+
## License
|
|
173
|
+
|
|
174
|
+
MIT
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
from cheapskate.callback import CheapSkateCallbackHandler
|
|
2
|
+
from cheapskate.compressor import PromptCompressor
|
|
3
|
+
from cheapskate.pruner import ToolNamespace, ToolPruner
|
|
4
|
+
from cheapskate.router import CheapSkateRouter
|
|
5
|
+
from cheapskate.token_counter import TokenCounter
|
|
6
|
+
|
|
7
|
+
__all__ = [
|
|
8
|
+
"CheapSkateRouter",
|
|
9
|
+
"CheapSkateCallbackHandler",
|
|
10
|
+
"TokenCounter",
|
|
11
|
+
"ToolPruner",
|
|
12
|
+
"ToolNamespace",
|
|
13
|
+
"PromptCompressor",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
__version__ = "0.1.0"
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def package_info() -> str:
|
|
20
|
+
info = f"cheapskate {__version__}"
|
|
21
|
+
print(info)
|
|
22
|
+
return info
|
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Any, Dict, List, Optional, Union
|
|
4
|
+
from uuid import UUID
|
|
5
|
+
|
|
6
|
+
from langchain_core.callbacks import BaseCallbackHandler
|
|
7
|
+
from langchain_core.messages import BaseMessage
|
|
8
|
+
from langchain_core.outputs import LLMResult
|
|
9
|
+
from loguru import logger
|
|
10
|
+
|
|
11
|
+
from cheapskate.token_counter import TokenCounter
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class CheapSkateCallbackHandler(BaseCallbackHandler):
|
|
15
|
+
def __init__(self, token_counter: Optional[TokenCounter] = None) -> None:
|
|
16
|
+
super().__init__()
|
|
17
|
+
self.token_counter = token_counter or TokenCounter()
|
|
18
|
+
self.routing_events: List[Dict[str, Any]] = []
|
|
19
|
+
self.prune_events: List[Dict[str, Any]] = []
|
|
20
|
+
self.compression_events: List[Dict[str, Any]] = []
|
|
21
|
+
self.fallback_events: List[Dict[str, Any]] = []
|
|
22
|
+
self._pending_baseline: Dict[UUID, int] = {}
|
|
23
|
+
|
|
24
|
+
def on_chat_model_start(
|
|
25
|
+
self,
|
|
26
|
+
serialized: Dict[str, Any],
|
|
27
|
+
messages: List[List[BaseMessage]],
|
|
28
|
+
*,
|
|
29
|
+
run_id: UUID,
|
|
30
|
+
parent_run_id: Optional[UUID] = None,
|
|
31
|
+
tags: Optional[List[str]] = None,
|
|
32
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
33
|
+
**kwargs: Any,
|
|
34
|
+
) -> None:
|
|
35
|
+
try:
|
|
36
|
+
flat_messages = [message for batch in messages for message in batch]
|
|
37
|
+
baseline = self.token_counter.record_baseline(flat_messages)
|
|
38
|
+
self._pending_baseline[run_id] = baseline
|
|
39
|
+
logger.debug("CheapSkate tracked chat start run_id={} baseline_tokens={}", run_id, baseline)
|
|
40
|
+
except Exception as exc:
|
|
41
|
+
logger.warning("CheapSkateCallbackHandler failed during chat start: {}", exc)
|
|
42
|
+
|
|
43
|
+
def on_llm_end(
|
|
44
|
+
self,
|
|
45
|
+
response: LLMResult,
|
|
46
|
+
*,
|
|
47
|
+
run_id: UUID,
|
|
48
|
+
parent_run_id: Optional[UUID] = None,
|
|
49
|
+
**kwargs: Any,
|
|
50
|
+
) -> None:
|
|
51
|
+
try:
|
|
52
|
+
self._pending_baseline.pop(run_id, None)
|
|
53
|
+
except Exception as exc:
|
|
54
|
+
logger.warning("CheapSkateCallbackHandler failed during llm end: {}", exc)
|
|
55
|
+
|
|
56
|
+
def on_llm_error(
|
|
57
|
+
self,
|
|
58
|
+
error: BaseException,
|
|
59
|
+
*,
|
|
60
|
+
run_id: UUID,
|
|
61
|
+
parent_run_id: Optional[UUID] = None,
|
|
62
|
+
**kwargs: Any,
|
|
63
|
+
) -> None:
|
|
64
|
+
logger.error("CheapSkate observed LLM error for run_id={}: {}", run_id, error)
|
|
65
|
+
self._pending_baseline.pop(run_id, None)
|
|
66
|
+
|
|
67
|
+
def record_routing(self, destination: str, reason: str, metadata: Optional[Dict[str, Any]] = None) -> None:
|
|
68
|
+
event = {"destination": destination, "reason": reason, "metadata": metadata or {}}
|
|
69
|
+
self.routing_events.append(event)
|
|
70
|
+
logger.info("CheapSkate routed to {} because {}", destination, reason)
|
|
71
|
+
|
|
72
|
+
def record_prune(
|
|
73
|
+
self,
|
|
74
|
+
retained: List[str],
|
|
75
|
+
dropped: List[str],
|
|
76
|
+
namespaces: List[str],
|
|
77
|
+
) -> None:
|
|
78
|
+
event = {
|
|
79
|
+
"retained": list(retained),
|
|
80
|
+
"dropped": list(dropped),
|
|
81
|
+
"namespaces": list(namespaces),
|
|
82
|
+
}
|
|
83
|
+
self.prune_events.append(event)
|
|
84
|
+
|
|
85
|
+
def record_compression(self, before_tokens: int, after_tokens: int) -> None:
|
|
86
|
+
event = {
|
|
87
|
+
"before_tokens": before_tokens,
|
|
88
|
+
"after_tokens": after_tokens,
|
|
89
|
+
"tokens_saved": max(0, before_tokens - after_tokens),
|
|
90
|
+
}
|
|
91
|
+
self.compression_events.append(event)
|
|
92
|
+
|
|
93
|
+
def record_fallback(self, error: Union[str, BaseException], recovered: bool = True) -> None:
|
|
94
|
+
event = {"error": str(error), "recovered": recovered}
|
|
95
|
+
self.fallback_events.append(event)
|
|
96
|
+
logger.warning("CheapSkate fallback engaged: {}", error)
|
|
97
|
+
|
|
98
|
+
@property
|
|
99
|
+
def summary(self) -> Dict[str, Any]:
|
|
100
|
+
return {
|
|
101
|
+
"token_stats": self.token_counter.stats,
|
|
102
|
+
"routing_events": len(self.routing_events),
|
|
103
|
+
"prune_events": len(self.prune_events),
|
|
104
|
+
"compression_events": len(self.compression_events),
|
|
105
|
+
"fallback_events": len(self.fallback_events),
|
|
106
|
+
"last_route": self.routing_events[-1] if self.routing_events else None,
|
|
107
|
+
}
|
|
108
|
+
|
|
109
|
+
def reset(self) -> None:
|
|
110
|
+
self.routing_events.clear()
|
|
111
|
+
self.prune_events.clear()
|
|
112
|
+
self.compression_events.clear()
|
|
113
|
+
self.fallback_events.clear()
|
|
114
|
+
self._pending_baseline.clear()
|
|
115
|
+
self.token_counter.reset()
|