tokensave 0.1.1__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.
- tokensave-0.1.1/PKG-INFO +29 -0
- tokensave-0.1.1/README.md +19 -0
- tokensave-0.1.1/pyproject.toml +24 -0
- tokensave-0.1.1/setup.cfg +4 -0
- tokensave-0.1.1/tokensave/__init__.py +85 -0
- tokensave-0.1.1/tokensave/cli.py +77 -0
- tokensave-0.1.1/tokensave/compressors/__init__.py +50 -0
- tokensave-0.1.1/tokensave/detectors/__init__.py +143 -0
- tokensave-0.1.1/tokensave/proxy.py +318 -0
- tokensave-0.1.1/tokensave.egg-info/PKG-INFO +29 -0
- tokensave-0.1.1/tokensave.egg-info/SOURCES.txt +13 -0
- tokensave-0.1.1/tokensave.egg-info/dependency_links.txt +1 -0
- tokensave-0.1.1/tokensave.egg-info/entry_points.txt +2 -0
- tokensave-0.1.1/tokensave.egg-info/requires.txt +1 -0
- tokensave-0.1.1/tokensave.egg-info/top_level.txt +1 -0
tokensave-0.1.1/PKG-INFO
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tokensave
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: Context optimization for LLM API calls. One import, fewer tokens.
|
|
5
|
+
Author: raydatalab
|
|
6
|
+
License: Apache 2.0
|
|
7
|
+
Requires-Python: >=3.10
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Requires-Dist: headroom-ai>=0.30.0
|
|
10
|
+
|
|
11
|
+
# TokenSave ⚡
|
|
12
|
+
|
|
13
|
+
Context optimization for LLM API calls. One import, fewer tokens.
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
pip install tokensave
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
from tokensave import OpenAI
|
|
21
|
+
|
|
22
|
+
client = OpenAI(api_key="...")
|
|
23
|
+
# Open box. Save tokens. That's it.
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
- No daemon, no config, no proxy.
|
|
27
|
+
- `tokensave proxy` for non-Python tools (power-user).
|
|
28
|
+
|
|
29
|
+
[github.com/raydatalab/tokensave](https://github.com/raydatalab/tokensave) · Apache 2.0
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
# TokenSave ⚡
|
|
2
|
+
|
|
3
|
+
Context optimization for LLM API calls. One import, fewer tokens.
|
|
4
|
+
|
|
5
|
+
```bash
|
|
6
|
+
pip install tokensave
|
|
7
|
+
```
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
from tokensave import OpenAI
|
|
11
|
+
|
|
12
|
+
client = OpenAI(api_key="...")
|
|
13
|
+
# Open box. Save tokens. That's it.
|
|
14
|
+
```
|
|
15
|
+
|
|
16
|
+
- No daemon, no config, no proxy.
|
|
17
|
+
- `tokensave proxy` for non-Python tools (power-user).
|
|
18
|
+
|
|
19
|
+
[github.com/raydatalab/tokensave](https://github.com/raydatalab/tokensave) · Apache 2.0
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=64", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "tokensave"
|
|
7
|
+
version = "0.1.1"
|
|
8
|
+
description = "Context optimization for LLM API calls. One import, fewer tokens."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = {text = "Apache 2.0"}
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
authors = [
|
|
13
|
+
{name = "raydatalab"}
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
dependencies = [
|
|
17
|
+
"headroom-ai>=0.30.0",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
[project.scripts]
|
|
21
|
+
tokensave = "tokensave.cli:main"
|
|
22
|
+
|
|
23
|
+
[tool.setuptools.packages.find]
|
|
24
|
+
include = ["tokensave*"]
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
"""
|
|
2
|
+
TokenSave — Context optimization for LLM API calls.
|
|
3
|
+
|
|
4
|
+
Safe defaults: pip install does nothing but make the package available.
|
|
5
|
+
Use `from tokensave import OpenAI` as a drop-in for `from openai import OpenAI`.
|
|
6
|
+
If headroom is installed, messages are compressed before sending. If not,
|
|
7
|
+
requests pass through unchanged. Never breaks your existing workflow.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
__version__ = "0.1.1"
|
|
11
|
+
|
|
12
|
+
import logging
|
|
13
|
+
|
|
14
|
+
logger = logging.getLogger("tokensave")
|
|
15
|
+
|
|
16
|
+
# Check headroom availability once at import time
|
|
17
|
+
_HAS_HEADROOM = False
|
|
18
|
+
try:
|
|
19
|
+
import headroom # noqa: F401
|
|
20
|
+
_HAS_HEADROOM = True
|
|
21
|
+
except ImportError:
|
|
22
|
+
pass
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class OpenAI:
|
|
26
|
+
"""
|
|
27
|
+
Drop-in replacement for `openai.OpenAI` with optional compression.
|
|
28
|
+
|
|
29
|
+
Usage:
|
|
30
|
+
# Instead of:
|
|
31
|
+
# from openai import OpenAI
|
|
32
|
+
# client = OpenAI(api_key="...")
|
|
33
|
+
|
|
34
|
+
from tokensave import OpenAI
|
|
35
|
+
client = OpenAI(api_key="...")
|
|
36
|
+
|
|
37
|
+
When headroom is installed ('pip install tokensave[compress]'), messages
|
|
38
|
+
are compressed before being sent to the API. When headroom is not available,
|
|
39
|
+
requests pass through unchanged — no errors, no config changes, no breakage.
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
def __new__(cls, *args, **kwargs):
|
|
43
|
+
try:
|
|
44
|
+
from openai import OpenAI as _RealOpenAI
|
|
45
|
+
client = _RealOpenAI(*args, **kwargs)
|
|
46
|
+
except ImportError:
|
|
47
|
+
raise ImportError(
|
|
48
|
+
"tokensave requires 'openai' package. Install with: pip install openai"
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
if _HAS_HEADROOM:
|
|
52
|
+
return _CompressingOpenAI(client)
|
|
53
|
+
return client
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class _CompressingOpenAI:
|
|
57
|
+
"""Wraps an openai.OpenAI client with Headroom compression."""
|
|
58
|
+
|
|
59
|
+
def __init__(self, client):
|
|
60
|
+
self._client = client
|
|
61
|
+
self._compressor = None
|
|
62
|
+
self._init_compressor()
|
|
63
|
+
|
|
64
|
+
def _init_compressor(self):
|
|
65
|
+
try:
|
|
66
|
+
from headroom import compress as _hr_compress
|
|
67
|
+
self._hr_compress = _hr_compress
|
|
68
|
+
logger.info("tokensave: Headroom compression active")
|
|
69
|
+
except Exception as e:
|
|
70
|
+
logger.warning(f"tokensave: Headroom init failed ({e}), running passthrough")
|
|
71
|
+
self._hr_compress = None
|
|
72
|
+
|
|
73
|
+
def _compress_messages(self, messages):
|
|
74
|
+
if not self._hr_compress:
|
|
75
|
+
return messages
|
|
76
|
+
try:
|
|
77
|
+
result = self._hr_compress(messages)
|
|
78
|
+
return result.messages
|
|
79
|
+
except Exception as e:
|
|
80
|
+
logger.warning(f"tokensave: compression failed ({e}), sending uncompressed")
|
|
81
|
+
return messages
|
|
82
|
+
|
|
83
|
+
def __getattr__(self, name):
|
|
84
|
+
"""Pass through any attribute access to the underlying client."""
|
|
85
|
+
return getattr(self._client, name)
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
"""
|
|
3
|
+
TokenSave — Context optimization for LLM API calls.
|
|
4
|
+
|
|
5
|
+
Core usage:
|
|
6
|
+
from tokensave import OpenAI # drop-in for openai.OpenAI
|
|
7
|
+
|
|
8
|
+
CLI:
|
|
9
|
+
tokensave setup Detect environment, check dependencies
|
|
10
|
+
tokensave proxy Start/stop compression proxy (power-user)
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
import argparse
|
|
14
|
+
import logging
|
|
15
|
+
import sys
|
|
16
|
+
|
|
17
|
+
from . import __version__
|
|
18
|
+
from .detectors import detect_all, summary
|
|
19
|
+
from .compressors import check_headroom, list_compressors, get_headroom_version
|
|
20
|
+
|
|
21
|
+
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
22
|
+
logger = logging.getLogger("tokensave")
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def cmd_setup(args):
|
|
26
|
+
"""Detect environment and check dependencies. Read-only."""
|
|
27
|
+
print("⚡ TokenSave\n")
|
|
28
|
+
env = detect_all()
|
|
29
|
+
print(summary(env))
|
|
30
|
+
|
|
31
|
+
if check_headroom():
|
|
32
|
+
print(f"\n ✓ headroom {get_headroom_version()} — compression ready")
|
|
33
|
+
else:
|
|
34
|
+
print("\n ○ headroom not installed")
|
|
35
|
+
print(" pip install 'tokensave[compress]'")
|
|
36
|
+
|
|
37
|
+
print("\n → from tokensave import OpenAI")
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def cmd_proxy(args):
|
|
41
|
+
"""Start/stop compression proxy."""
|
|
42
|
+
from .proxy import TokenSaveProxy
|
|
43
|
+
proxy = TokenSaveProxy()
|
|
44
|
+
|
|
45
|
+
if args.action == "start":
|
|
46
|
+
if proxy.start():
|
|
47
|
+
print("✅ Proxy on 127.0.0.1:18787")
|
|
48
|
+
else:
|
|
49
|
+
print("⚠ Port 18787 in use")
|
|
50
|
+
elif args.action == "stop":
|
|
51
|
+
proxy.stop()
|
|
52
|
+
print("🛑 Proxy stopped")
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def main():
|
|
56
|
+
parser = argparse.ArgumentParser(
|
|
57
|
+
description="TokenSave — Context optimization for LLM API calls"
|
|
58
|
+
)
|
|
59
|
+
parser.add_argument("--version", action="version", version=f"TokenSave {__version__}")
|
|
60
|
+
|
|
61
|
+
sub = parser.add_subparsers(dest="command")
|
|
62
|
+
sub.add_parser("setup", help="Detect environment, check dependencies")
|
|
63
|
+
|
|
64
|
+
p = sub.add_parser("proxy", help="Start/stop compression proxy (power-user)")
|
|
65
|
+
p.add_argument("action", choices=["start", "stop"])
|
|
66
|
+
|
|
67
|
+
args = parser.parse_args()
|
|
68
|
+
if args.command == "setup":
|
|
69
|
+
cmd_setup(args)
|
|
70
|
+
elif args.command == "proxy":
|
|
71
|
+
cmd_proxy(args)
|
|
72
|
+
else:
|
|
73
|
+
parser.print_help()
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
if __name__ == "__main__":
|
|
77
|
+
main()
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
"""Compressor backends. Primary: headroom. Future: llmlingua, selective-context."""
|
|
2
|
+
|
|
3
|
+
import importlib
|
|
4
|
+
import logging
|
|
5
|
+
|
|
6
|
+
logger = logging.getLogger("tokensave")
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def check_headroom() -> bool:
|
|
10
|
+
"""Check if headroom is installed and importable."""
|
|
11
|
+
try:
|
|
12
|
+
import headroom # noqa: F401
|
|
13
|
+
return True
|
|
14
|
+
except ImportError:
|
|
15
|
+
return False
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def get_headroom_version() -> str:
|
|
19
|
+
"""Get headroom version string."""
|
|
20
|
+
try:
|
|
21
|
+
import headroom
|
|
22
|
+
return getattr(headroom, "__version__", "unknown")
|
|
23
|
+
except ImportError:
|
|
24
|
+
return "not installed"
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def wrap_command(tool: str, compress_refs: bool = True, keep_active: bool = True) -> list[str]:
|
|
28
|
+
"""Return the headroom wrap command args for a given tool."""
|
|
29
|
+
cmd = ["headroom", "wrap", tool]
|
|
30
|
+
if compress_refs:
|
|
31
|
+
cmd.append("--compress-refs")
|
|
32
|
+
if keep_active:
|
|
33
|
+
cmd.append("--keep-active")
|
|
34
|
+
return cmd
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
AVAILABLE_COMPRESSORS = {"headroom": check_headroom()}
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def list_compressors() -> list[dict]:
|
|
41
|
+
"""List available compressors and their status."""
|
|
42
|
+
return [
|
|
43
|
+
{
|
|
44
|
+
"name": "headroom",
|
|
45
|
+
"available": check_headroom(),
|
|
46
|
+
"version": get_headroom_version(),
|
|
47
|
+
"description": "Content-aware context compression. 60-95% token reduction.",
|
|
48
|
+
"default": True,
|
|
49
|
+
},
|
|
50
|
+
]
|
|
@@ -0,0 +1,143 @@
|
|
|
1
|
+
"""Detect what LLM tools the user has installed and which API provider they use."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import shutil
|
|
5
|
+
import subprocess
|
|
6
|
+
import json
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
|
|
9
|
+
HOME = Path.home()
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def detect_claude_code() -> dict:
|
|
13
|
+
"""Detect Claude Code installation and configuration."""
|
|
14
|
+
result = {"installed": False, "mode": None, "config_path": None}
|
|
15
|
+
|
|
16
|
+
# Check if claude command exists
|
|
17
|
+
claude_path = shutil.which("claude")
|
|
18
|
+
if claude_path:
|
|
19
|
+
result["installed"] = True
|
|
20
|
+
result["binary"] = str(claude_path)
|
|
21
|
+
|
|
22
|
+
# Check API key / auth mode
|
|
23
|
+
if os.environ.get("ANTHROPIC_API_KEY"):
|
|
24
|
+
result["mode"] = "api"
|
|
25
|
+
result["api_key_prefix"] = os.environ["ANTHROPIC_API_KEY"][:8] + "..."
|
|
26
|
+
elif os.environ.get("CLAUDE_API_KEY"):
|
|
27
|
+
result["mode"] = "api"
|
|
28
|
+
else:
|
|
29
|
+
result["mode"] = "subscription"
|
|
30
|
+
|
|
31
|
+
# Check config
|
|
32
|
+
config_paths = [
|
|
33
|
+
HOME / ".claude" / "settings.json",
|
|
34
|
+
HOME / ".claude" / "claude.json",
|
|
35
|
+
HOME / ".config" / "claude" / "settings.json",
|
|
36
|
+
]
|
|
37
|
+
for p in config_paths:
|
|
38
|
+
if p.exists():
|
|
39
|
+
result["config_path"] = str(p)
|
|
40
|
+
try:
|
|
41
|
+
data = json.loads(p.read_text())
|
|
42
|
+
result["config"] = data
|
|
43
|
+
except (json.JSONDecodeError, OSError):
|
|
44
|
+
pass
|
|
45
|
+
break
|
|
46
|
+
|
|
47
|
+
return result
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def detect_sillytavern() -> dict:
|
|
51
|
+
"""Detect SillyTavern installation."""
|
|
52
|
+
result = {"installed": False, "path": None}
|
|
53
|
+
|
|
54
|
+
# Common install locations
|
|
55
|
+
search_paths = [
|
|
56
|
+
HOME / "SillyTavern",
|
|
57
|
+
HOME / "sillytavern",
|
|
58
|
+
HOME / "ST",
|
|
59
|
+
HOME / "SillyTavern-Launcher",
|
|
60
|
+
]
|
|
61
|
+
for p in search_paths:
|
|
62
|
+
if p.exists() and (p / "config.yaml").exists():
|
|
63
|
+
result["installed"] = True
|
|
64
|
+
result["path"] = str(p)
|
|
65
|
+
break
|
|
66
|
+
|
|
67
|
+
return result
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def detect_openai_key() -> bool:
|
|
71
|
+
"""Check if user has OpenAI API key configured."""
|
|
72
|
+
return bool(os.environ.get("OPENAI_API_KEY"))
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def detect_anthropic_key() -> bool:
|
|
76
|
+
"""Check if user has Anthropic API key configured."""
|
|
77
|
+
return bool(os.environ.get("ANTHROPIC_API_KEY"))
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def get_api_base_urls() -> dict:
|
|
81
|
+
"""Get current API base URLs from environment (may have been set by other tools)."""
|
|
82
|
+
return {
|
|
83
|
+
"openai": os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1"),
|
|
84
|
+
"anthropic": os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com"),
|
|
85
|
+
}
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def detect_all() -> dict:
|
|
89
|
+
"""Run all detectors and return a complete profile."""
|
|
90
|
+
env = {
|
|
91
|
+
"claude_code": detect_claude_code(),
|
|
92
|
+
"sillytavern": detect_sillytavern(),
|
|
93
|
+
"api_keys": {
|
|
94
|
+
"openai": detect_openai_key(),
|
|
95
|
+
"anthropic": detect_anthropic_key(),
|
|
96
|
+
},
|
|
97
|
+
"base_urls": get_api_base_urls(),
|
|
98
|
+
"shell": os.environ.get("SHELL", "unknown"),
|
|
99
|
+
"platform": os.uname().sysname if hasattr(os, "uname") else "unknown",
|
|
100
|
+
}
|
|
101
|
+
|
|
102
|
+
# Determine primary use case
|
|
103
|
+
if env["claude_code"]["installed"]:
|
|
104
|
+
env["primary_tool"] = "claude_code"
|
|
105
|
+
elif env["sillytavern"]["installed"]:
|
|
106
|
+
env["primary_tool"] = "sillytavern"
|
|
107
|
+
elif env["api_keys"]["openai"] or env["api_keys"]["anthropic"]:
|
|
108
|
+
env["primary_tool"] = "api_user"
|
|
109
|
+
else:
|
|
110
|
+
env["primary_tool"] = "unknown"
|
|
111
|
+
|
|
112
|
+
return env
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def summary(env: dict) -> str:
|
|
116
|
+
"""Return a human-readable summary of the detected environment."""
|
|
117
|
+
lines = ["🔍 Environment Detection:\n"]
|
|
118
|
+
|
|
119
|
+
cc = env["claude_code"]
|
|
120
|
+
if cc["installed"]:
|
|
121
|
+
mode_str = "API key" if cc["mode"] == "api" else "subscription"
|
|
122
|
+
lines.append(f" ✓ Claude Code detected ({mode_str} mode)")
|
|
123
|
+
else:
|
|
124
|
+
lines.append(" ✗ Claude Code not detected")
|
|
125
|
+
|
|
126
|
+
st = env["sillytavern"]
|
|
127
|
+
if st["installed"]:
|
|
128
|
+
lines.append(f" ✓ SillyTavern detected at {st['path']}")
|
|
129
|
+
else:
|
|
130
|
+
lines.append(" ✗ SillyTavern not detected")
|
|
131
|
+
|
|
132
|
+
keys = env["api_keys"]
|
|
133
|
+
if keys["openai"]:
|
|
134
|
+
lines.append(" ✓ OpenAI API key found")
|
|
135
|
+
if keys["anthropic"]:
|
|
136
|
+
lines.append(" ✓ Anthropic API key found")
|
|
137
|
+
|
|
138
|
+
if not any(keys.values()):
|
|
139
|
+
if not cc["installed"]:
|
|
140
|
+
lines.append(" ⚠ No API keys detected. TokenSave works best with API-based access.")
|
|
141
|
+
|
|
142
|
+
lines.append(f"\n Primary tool: {env['primary_tool']}")
|
|
143
|
+
return "\n".join(lines)
|
|
@@ -0,0 +1,318 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
"""
|
|
3
|
+
Transparent local proxy that compresses LLM API calls on the fly.
|
|
4
|
+
Speaks both OpenAI-compatible (/v1/chat/completions) and Anthropic-compatible (/v1/messages) formats.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import json
|
|
8
|
+
import logging
|
|
9
|
+
import os
|
|
10
|
+
import threading
|
|
11
|
+
import time
|
|
12
|
+
import urllib.request
|
|
13
|
+
import urllib.error
|
|
14
|
+
from http.server import HTTPServer, BaseHTTPRequestHandler
|
|
15
|
+
from dataclasses import dataclass, field
|
|
16
|
+
from pathlib import Path
|
|
17
|
+
|
|
18
|
+
logger = logging.getLogger("tokensave.proxy")
|
|
19
|
+
|
|
20
|
+
STATS_DIR = Path.home() / ".tokensave"
|
|
21
|
+
STATS_FILE = STATS_DIR / "stats.json"
|
|
22
|
+
PROXY_PORT = 18787
|
|
23
|
+
HEADROOM_PORT = 8787 # Default headroom proxy port, if running
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
@dataclass
|
|
27
|
+
class ProxyStats:
|
|
28
|
+
"""Tracks token savings across proxy sessions."""
|
|
29
|
+
total_requests: int = 0
|
|
30
|
+
total_input_tokens_before: int = 0
|
|
31
|
+
total_input_tokens_after: int = 0
|
|
32
|
+
total_output_tokens: int = 0
|
|
33
|
+
estimated_cost_saved: float = 0.0
|
|
34
|
+
start_time: float = field(default_factory=time.time)
|
|
35
|
+
|
|
36
|
+
def save(self):
|
|
37
|
+
STATS_DIR.mkdir(parents=True, exist_ok=True)
|
|
38
|
+
data = {
|
|
39
|
+
"total_requests": self.total_requests,
|
|
40
|
+
"total_input_tokens_before": self.total_input_tokens_before,
|
|
41
|
+
"total_input_tokens_after": self.total_input_tokens_after,
|
|
42
|
+
"total_output_tokens": self.total_output_tokens,
|
|
43
|
+
"estimated_cost_saved": round(self.estimated_cost_saved, 4),
|
|
44
|
+
"start_time": self.start_time,
|
|
45
|
+
}
|
|
46
|
+
STATS_FILE.write_text(json.dumps(data, indent=2))
|
|
47
|
+
|
|
48
|
+
@classmethod
|
|
49
|
+
def load(cls):
|
|
50
|
+
if STATS_FILE.exists():
|
|
51
|
+
try:
|
|
52
|
+
data = json.loads(STATS_FILE.read_text())
|
|
53
|
+
return cls(**data)
|
|
54
|
+
except (json.JSONDecodeError, TypeError):
|
|
55
|
+
pass
|
|
56
|
+
return cls()
|
|
57
|
+
|
|
58
|
+
@property
|
|
59
|
+
def compression_ratio(self) -> float:
|
|
60
|
+
if self.total_input_tokens_before > 0:
|
|
61
|
+
return 1 - (self.total_input_tokens_after / self.total_input_tokens_before)
|
|
62
|
+
return 0.0
|
|
63
|
+
|
|
64
|
+
@property
|
|
65
|
+
def runtime_hours(self) -> float:
|
|
66
|
+
return (time.time() - self.start_time) / 3600
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
stats = ProxyStats.load()
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class ProxyHandler(BaseHTTPRequestHandler):
|
|
73
|
+
"""HTTP handler that proxies LLM API requests through compression."""
|
|
74
|
+
|
|
75
|
+
def do_OPTIONS(self):
|
|
76
|
+
self._cors_headers()
|
|
77
|
+
self.send_response(200)
|
|
78
|
+
self.end_headers()
|
|
79
|
+
|
|
80
|
+
def do_POST(self):
|
|
81
|
+
content_len = int(self.headers.get("Content-Length", 0))
|
|
82
|
+
body = self.rfile.read(content_len)
|
|
83
|
+
|
|
84
|
+
# Determine target from path
|
|
85
|
+
path = self.path
|
|
86
|
+
|
|
87
|
+
if "/v1/chat/completions" in path:
|
|
88
|
+
self._handle_openai_chat(body)
|
|
89
|
+
elif "/v1/messages" in path:
|
|
90
|
+
self._handle_anthropic_messages(body)
|
|
91
|
+
else:
|
|
92
|
+
self._forward_raw(body)
|
|
93
|
+
|
|
94
|
+
def _cors_headers(self):
|
|
95
|
+
self.send_header("Access-Control-Allow-Origin", "*")
|
|
96
|
+
self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
|
|
97
|
+
self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization, x-api-key")
|
|
98
|
+
|
|
99
|
+
def _compress_messages(self, messages: list) -> list:
|
|
100
|
+
"""Compress messages using headroom proxy if available, otherwise return as-is."""
|
|
101
|
+
try:
|
|
102
|
+
import headroom
|
|
103
|
+
from headroom import Headroom
|
|
104
|
+
|
|
105
|
+
# Use headroom's compressor directly
|
|
106
|
+
hr = Headroom()
|
|
107
|
+
compressed = []
|
|
108
|
+
for msg in messages:
|
|
109
|
+
if isinstance(msg.get("content"), str) and len(msg["content"]) > 200:
|
|
110
|
+
compressed_content = hr.compress(msg["content"])
|
|
111
|
+
compressed.append({**msg, "content": compressed_content})
|
|
112
|
+
else:
|
|
113
|
+
compressed.append(msg)
|
|
114
|
+
return compressed
|
|
115
|
+
except ImportError:
|
|
116
|
+
# Headroom not directly importable for compression; skip
|
|
117
|
+
return messages
|
|
118
|
+
except Exception as e:
|
|
119
|
+
logger.warning(f"Compression failed for message batch: {e}")
|
|
120
|
+
return messages
|
|
121
|
+
|
|
122
|
+
def _estimate_tokens(self, text: str) -> int:
|
|
123
|
+
"""Rough token estimate (4 chars ≈ 1 token)."""
|
|
124
|
+
return len(text) // 4
|
|
125
|
+
|
|
126
|
+
def _handle_openai_chat(self, body: bytes):
|
|
127
|
+
try:
|
|
128
|
+
req = json.loads(body)
|
|
129
|
+
messages = req.get("messages", [])
|
|
130
|
+
|
|
131
|
+
# Count tokens before
|
|
132
|
+
text_before = json.dumps(messages)
|
|
133
|
+
tokens_before = self._estimate_tokens(text_before)
|
|
134
|
+
|
|
135
|
+
# Compress
|
|
136
|
+
compressed_messages = self._compress_messages(messages)
|
|
137
|
+
req["messages"] = compressed_messages
|
|
138
|
+
|
|
139
|
+
# Count after
|
|
140
|
+
text_after = json.dumps(compressed_messages)
|
|
141
|
+
tokens_after = self._estimate_tokens(text_after)
|
|
142
|
+
|
|
143
|
+
# Forward to real API
|
|
144
|
+
real_url = os.environ.get(
|
|
145
|
+
"TOKENSAVE_OPENAI_BASE_URL",
|
|
146
|
+
os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1"),
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
api_key = os.environ.get("OPENAI_API_KEY", "")
|
|
150
|
+
target_url = real_url.rstrip("/") + "/chat/completions"
|
|
151
|
+
|
|
152
|
+
response_data = self._forward_request(
|
|
153
|
+
target_url, json.dumps(req).encode(),
|
|
154
|
+
{"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
|
|
155
|
+
)
|
|
156
|
+
|
|
157
|
+
# Update stats
|
|
158
|
+
stats.total_requests += 1
|
|
159
|
+
stats.total_input_tokens_before += tokens_before
|
|
160
|
+
stats.total_input_tokens_after += tokens_after
|
|
161
|
+
|
|
162
|
+
if response_data:
|
|
163
|
+
resp_json = json.loads(response_data)
|
|
164
|
+
if "usage" in resp_json:
|
|
165
|
+
stats.total_output_tokens += resp_json["usage"].get("completion_tokens", 0)
|
|
166
|
+
|
|
167
|
+
# Rough cost estimate (GPT-4o pricing)
|
|
168
|
+
savings = (tokens_before - tokens_after) / 1_000_000 * 2.50
|
|
169
|
+
stats.estimated_cost_saved += savings
|
|
170
|
+
stats.save()
|
|
171
|
+
|
|
172
|
+
self._cors_headers()
|
|
173
|
+
self.send_response(200)
|
|
174
|
+
self.send_header("Content-Type", "application/json")
|
|
175
|
+
self.end_headers()
|
|
176
|
+
self.wfile.write(response_data)
|
|
177
|
+
else:
|
|
178
|
+
self._send_error(502, "Upstream API failed")
|
|
179
|
+
|
|
180
|
+
except Exception as e:
|
|
181
|
+
logger.error(f"OpenAI proxy error: {e}")
|
|
182
|
+
self._send_error(500, str(e))
|
|
183
|
+
|
|
184
|
+
def _handle_anthropic_messages(self, body: bytes):
|
|
185
|
+
try:
|
|
186
|
+
req = json.loads(body)
|
|
187
|
+
messages = req.get("messages", [])
|
|
188
|
+
|
|
189
|
+
text_before = json.dumps(messages)
|
|
190
|
+
tokens_before = self._estimate_tokens(text_before)
|
|
191
|
+
|
|
192
|
+
compressed_messages = self._compress_messages(messages)
|
|
193
|
+
req["messages"] = compressed_messages
|
|
194
|
+
|
|
195
|
+
text_after = json.dumps(compressed_messages)
|
|
196
|
+
tokens_after = self._estimate_tokens(text_after)
|
|
197
|
+
|
|
198
|
+
real_url = os.environ.get(
|
|
199
|
+
"TOKENSAVE_ANTHROPIC_BASE_URL",
|
|
200
|
+
os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com"),
|
|
201
|
+
)
|
|
202
|
+
|
|
203
|
+
api_key = os.environ.get("ANTHROPIC_API_KEY", "")
|
|
204
|
+
target_url = real_url.rstrip("/") + "/messages"
|
|
205
|
+
|
|
206
|
+
response_data = self._forward_request(
|
|
207
|
+
target_url, json.dumps(req).encode(),
|
|
208
|
+
{
|
|
209
|
+
"x-api-key": api_key,
|
|
210
|
+
"anthropic-version": self.headers.get("anthropic-version", "2023-06-01"),
|
|
211
|
+
"Content-Type": "application/json",
|
|
212
|
+
}
|
|
213
|
+
)
|
|
214
|
+
|
|
215
|
+
stats.total_requests += 1
|
|
216
|
+
stats.total_input_tokens_before += tokens_before
|
|
217
|
+
stats.total_input_tokens_after += tokens_after
|
|
218
|
+
|
|
219
|
+
if response_data:
|
|
220
|
+
resp_json = json.loads(response_data)
|
|
221
|
+
if "usage" in resp_json:
|
|
222
|
+
stats.total_output_tokens += resp_json["usage"].get("output_tokens", 0)
|
|
223
|
+
input_tokens = resp_json["usage"].get("input_tokens", 0)
|
|
224
|
+
# Use actual reported input tokens for savings calc
|
|
225
|
+
savings = (tokens_before - tokens_after) / 1_000_000 * 3.0
|
|
226
|
+
stats.estimated_cost_saved += savings
|
|
227
|
+
stats.save()
|
|
228
|
+
|
|
229
|
+
self._cors_headers()
|
|
230
|
+
self.send_response(200)
|
|
231
|
+
self.send_header("Content-Type", "application/json")
|
|
232
|
+
self.end_headers()
|
|
233
|
+
self.wfile.write(response_data)
|
|
234
|
+
else:
|
|
235
|
+
self._send_error(502, "Upstream API failed")
|
|
236
|
+
|
|
237
|
+
except Exception as e:
|
|
238
|
+
logger.error(f"Anthropic proxy error: {e}")
|
|
239
|
+
self._send_error(500, str(e))
|
|
240
|
+
|
|
241
|
+
def _forward_request(self, url: str, data: bytes, headers: dict) -> bytes | None:
|
|
242
|
+
try:
|
|
243
|
+
req = urllib.request.Request(url, data=data, headers=headers, method="POST")
|
|
244
|
+
with urllib.request.urlopen(req, timeout=60) as resp:
|
|
245
|
+
return resp.read()
|
|
246
|
+
except urllib.error.HTTPError as e:
|
|
247
|
+
logger.error(f"Upstream HTTP {e.code}: {e.read().decode()[:200]}")
|
|
248
|
+
return None
|
|
249
|
+
except Exception as e:
|
|
250
|
+
logger.error(f"Forward error: {e}")
|
|
251
|
+
return None
|
|
252
|
+
|
|
253
|
+
def _forward_raw(self, body: bytes):
|
|
254
|
+
"""Pass through for non-chat endpoints."""
|
|
255
|
+
self._cors_headers()
|
|
256
|
+
self.send_response(200)
|
|
257
|
+
self.send_header("Content-Type", "application/json")
|
|
258
|
+
self.end_headers()
|
|
259
|
+
self.wfile.write(body)
|
|
260
|
+
|
|
261
|
+
def _send_error(self, code: int, message: str):
|
|
262
|
+
self._cors_headers()
|
|
263
|
+
self.send_response(code)
|
|
264
|
+
self.send_header("Content-Type", "application/json")
|
|
265
|
+
self.end_headers()
|
|
266
|
+
self.wfile.write(json.dumps({"error": {"message": message}}).encode())
|
|
267
|
+
|
|
268
|
+
def log_message(self, format, *args):
|
|
269
|
+
logger.debug(f"{self.address_string()} - {format % args}")
|
|
270
|
+
|
|
271
|
+
|
|
272
|
+
class TokenSaveProxy:
|
|
273
|
+
"""Manages the transparent proxy lifecycle."""
|
|
274
|
+
|
|
275
|
+
def __init__(self, port: int = PROXY_PORT):
|
|
276
|
+
self.port = port
|
|
277
|
+
self.server = None
|
|
278
|
+
self.thread = None
|
|
279
|
+
self.running = False
|
|
280
|
+
|
|
281
|
+
def start(self, daemon: bool = True) -> bool:
|
|
282
|
+
if self.running:
|
|
283
|
+
logger.info("Proxy already running")
|
|
284
|
+
return True
|
|
285
|
+
|
|
286
|
+
try:
|
|
287
|
+
self.server = HTTPServer(("127.0.0.1", self.port), ProxyHandler)
|
|
288
|
+
self.thread = threading.Thread(target=self.server.serve_forever, daemon=daemon)
|
|
289
|
+
self.thread.start()
|
|
290
|
+
self.running = True
|
|
291
|
+
logger.info(f"TokenSave proxy started on 127.0.0.1:{self.port}")
|
|
292
|
+
return True
|
|
293
|
+
except OSError as e:
|
|
294
|
+
logger.error(f"Failed to start proxy on port {self.port}: {e}")
|
|
295
|
+
return False
|
|
296
|
+
|
|
297
|
+
def stop(self):
|
|
298
|
+
if self.server:
|
|
299
|
+
self.server.shutdown()
|
|
300
|
+
self.server.server_close()
|
|
301
|
+
self.running = False
|
|
302
|
+
logger.info("TokenSave proxy stopped")
|
|
303
|
+
|
|
304
|
+
def is_running(self) -> bool:
|
|
305
|
+
return self.running
|
|
306
|
+
|
|
307
|
+
@property
|
|
308
|
+
def status_line(self) -> str:
|
|
309
|
+
if self.running:
|
|
310
|
+
r = stats.total_requests
|
|
311
|
+
ratio = stats.compression_ratio * 100
|
|
312
|
+
saved = stats.estimated_cost_saved
|
|
313
|
+
return (
|
|
314
|
+
f" ✓ Proxy running on 127.0.0.1:{self.port}\n"
|
|
315
|
+
f" 📊 {r} requests processed | {ratio:.1f}% compression"
|
|
316
|
+
f" | ~${saved:.2f} saved"
|
|
317
|
+
)
|
|
318
|
+
return " ✗ Proxy not running"
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tokensave
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: Context optimization for LLM API calls. One import, fewer tokens.
|
|
5
|
+
Author: raydatalab
|
|
6
|
+
License: Apache 2.0
|
|
7
|
+
Requires-Python: >=3.10
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Requires-Dist: headroom-ai>=0.30.0
|
|
10
|
+
|
|
11
|
+
# TokenSave ⚡
|
|
12
|
+
|
|
13
|
+
Context optimization for LLM API calls. One import, fewer tokens.
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
pip install tokensave
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
from tokensave import OpenAI
|
|
21
|
+
|
|
22
|
+
client = OpenAI(api_key="...")
|
|
23
|
+
# Open box. Save tokens. That's it.
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
- No daemon, no config, no proxy.
|
|
27
|
+
- `tokensave proxy` for non-Python tools (power-user).
|
|
28
|
+
|
|
29
|
+
[github.com/raydatalab/tokensave](https://github.com/raydatalab/tokensave) · Apache 2.0
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
tokensave/__init__.py
|
|
4
|
+
tokensave/cli.py
|
|
5
|
+
tokensave/proxy.py
|
|
6
|
+
tokensave.egg-info/PKG-INFO
|
|
7
|
+
tokensave.egg-info/SOURCES.txt
|
|
8
|
+
tokensave.egg-info/dependency_links.txt
|
|
9
|
+
tokensave.egg-info/entry_points.txt
|
|
10
|
+
tokensave.egg-info/requires.txt
|
|
11
|
+
tokensave.egg-info/top_level.txt
|
|
12
|
+
tokensave/compressors/__init__.py
|
|
13
|
+
tokensave/detectors/__init__.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
headroom-ai>=0.30.0
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
tokensave
|