tokensor 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.
- tokensor-0.1.0/.env.example +2 -0
- tokensor-0.1.0/.gitignore +12 -0
- tokensor-0.1.0/PKG-INFO +100 -0
- tokensor-0.1.0/README.md +66 -0
- tokensor-0.1.0/pyproject.toml +53 -0
- tokensor-0.1.0/src/tokensor/__init__.py +72 -0
- tokensor-0.1.0/src/tokensor/_adapters/__init__.py +0 -0
- tokensor-0.1.0/src/tokensor/_adapters/_anthropic.py +120 -0
- tokensor-0.1.0/src/tokensor/_adapters/_openai.py +124 -0
- tokensor-0.1.0/src/tokensor/_client.py +131 -0
- tokensor-0.1.0/src/tokensor/_config.py +30 -0
- tokensor-0.1.0/src/tokensor/_hash.py +32 -0
- tokensor-0.1.0/src/tokensor/_models.py +29 -0
- tokensor-0.1.0/tests/__init__.py +0 -0
- tokensor-0.1.0/tests/test_anthropic.py +183 -0
- tokensor-0.1.0/tests/test_client.py +138 -0
- tokensor-0.1.0/tests/test_config.py +47 -0
- tokensor-0.1.0/tests/test_hash.py +34 -0
- tokensor-0.1.0/tests/test_models.py +45 -0
- tokensor-0.1.0/tests/test_openai.py +143 -0
- tokensor-0.1.0/tests/test_wrap.py +89 -0
tokensor-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tokensor
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: LLM token observability SDK — wraps OpenAI and Anthropic clients to capture metadata
|
|
5
|
+
Project-URL: Homepage, https://tokensor.com
|
|
6
|
+
Project-URL: Repository, https://github.com/tokensor/tokensor
|
|
7
|
+
Project-URL: Bug Tracker, https://github.com/tokensor/tokensor/issues
|
|
8
|
+
Author-email: Tokensor <hello@tokensor.com>
|
|
9
|
+
License: MIT
|
|
10
|
+
Keywords: anthropic,cost,llm,observability,openai,tokens
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
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 :: Software Development :: Libraries :: Python Modules
|
|
20
|
+
Requires-Python: >=3.9
|
|
21
|
+
Requires-Dist: httpx>=0.28
|
|
22
|
+
Provides-Extra: anthropic
|
|
23
|
+
Requires-Dist: anthropic>=0.30; extra == 'anthropic'
|
|
24
|
+
Provides-Extra: dev
|
|
25
|
+
Requires-Dist: anthropic>=0.30; extra == 'dev'
|
|
26
|
+
Requires-Dist: openai>=1.0; extra == 'dev'
|
|
27
|
+
Requires-Dist: pytest-asyncio>=0.24; extra == 'dev'
|
|
28
|
+
Requires-Dist: pytest-cov>=6.0; extra == 'dev'
|
|
29
|
+
Requires-Dist: pytest>=8.3; extra == 'dev'
|
|
30
|
+
Requires-Dist: respx>=0.21; extra == 'dev'
|
|
31
|
+
Provides-Extra: openai
|
|
32
|
+
Requires-Dist: openai>=1.0; extra == 'openai'
|
|
33
|
+
Description-Content-Type: text/markdown
|
|
34
|
+
|
|
35
|
+
# tokensor
|
|
36
|
+
|
|
37
|
+
LLM token observability for Python. Wraps your OpenAI or Anthropic client to capture call metadata (model, tokens, cost) and forward it to your Tokensor dashboard — one line of code, zero latency overhead.
|
|
38
|
+
|
|
39
|
+
## Install
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
pip install tokensor
|
|
43
|
+
# With OpenAI support:
|
|
44
|
+
pip install "tokensor[openai]"
|
|
45
|
+
# With Anthropic support:
|
|
46
|
+
pip install "tokensor[anthropic]"
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
## Quickstart
|
|
50
|
+
|
|
51
|
+
Get your API key from [app.tokensor.com](https://app.tokensor.com) → Settings → API Keys.
|
|
52
|
+
|
|
53
|
+
**OpenAI:**
|
|
54
|
+
|
|
55
|
+
```python
|
|
56
|
+
import os
|
|
57
|
+
import openai
|
|
58
|
+
import tokensor
|
|
59
|
+
|
|
60
|
+
client = tokensor.wrap(
|
|
61
|
+
openai.OpenAI(),
|
|
62
|
+
api_key=os.environ["TOKENSOR_API_KEY"],
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
# Use exactly like openai.OpenAI()
|
|
66
|
+
response = client.chat.completions.create(
|
|
67
|
+
model="gpt-4o-mini",
|
|
68
|
+
messages=[{"role": "user", "content": "Hello"}],
|
|
69
|
+
)
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
**Anthropic:**
|
|
73
|
+
|
|
74
|
+
```python
|
|
75
|
+
import os
|
|
76
|
+
import anthropic
|
|
77
|
+
import tokensor
|
|
78
|
+
|
|
79
|
+
client = tokensor.wrap(
|
|
80
|
+
anthropic.Anthropic(),
|
|
81
|
+
api_key=os.environ["TOKENSOR_API_KEY"],
|
|
82
|
+
)
|
|
83
|
+
|
|
84
|
+
message = client.messages.create(
|
|
85
|
+
model="claude-3-5-haiku-latest",
|
|
86
|
+
max_tokens=1024,
|
|
87
|
+
messages=[{"role": "user", "content": "Hello"}],
|
|
88
|
+
)
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
## Environment variables
|
|
92
|
+
|
|
93
|
+
| Variable | Required | Description |
|
|
94
|
+
|----------|----------|-------------|
|
|
95
|
+
| `TOKENSOR_API_KEY` | yes | Your Tokensor API key |
|
|
96
|
+
| `TOKENSOR_HOST` | no | Override API host (default: `https://api.tokensor.com`) |
|
|
97
|
+
|
|
98
|
+
## License
|
|
99
|
+
|
|
100
|
+
MIT
|
tokensor-0.1.0/README.md
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
# tokensor
|
|
2
|
+
|
|
3
|
+
LLM token observability for Python. Wraps your OpenAI or Anthropic client to capture call metadata (model, tokens, cost) and forward it to your Tokensor dashboard — one line of code, zero latency overhead.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install tokensor
|
|
9
|
+
# With OpenAI support:
|
|
10
|
+
pip install "tokensor[openai]"
|
|
11
|
+
# With Anthropic support:
|
|
12
|
+
pip install "tokensor[anthropic]"
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
## Quickstart
|
|
16
|
+
|
|
17
|
+
Get your API key from [app.tokensor.com](https://app.tokensor.com) → Settings → API Keys.
|
|
18
|
+
|
|
19
|
+
**OpenAI:**
|
|
20
|
+
|
|
21
|
+
```python
|
|
22
|
+
import os
|
|
23
|
+
import openai
|
|
24
|
+
import tokensor
|
|
25
|
+
|
|
26
|
+
client = tokensor.wrap(
|
|
27
|
+
openai.OpenAI(),
|
|
28
|
+
api_key=os.environ["TOKENSOR_API_KEY"],
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
# Use exactly like openai.OpenAI()
|
|
32
|
+
response = client.chat.completions.create(
|
|
33
|
+
model="gpt-4o-mini",
|
|
34
|
+
messages=[{"role": "user", "content": "Hello"}],
|
|
35
|
+
)
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
**Anthropic:**
|
|
39
|
+
|
|
40
|
+
```python
|
|
41
|
+
import os
|
|
42
|
+
import anthropic
|
|
43
|
+
import tokensor
|
|
44
|
+
|
|
45
|
+
client = tokensor.wrap(
|
|
46
|
+
anthropic.Anthropic(),
|
|
47
|
+
api_key=os.environ["TOKENSOR_API_KEY"],
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
message = client.messages.create(
|
|
51
|
+
model="claude-3-5-haiku-latest",
|
|
52
|
+
max_tokens=1024,
|
|
53
|
+
messages=[{"role": "user", "content": "Hello"}],
|
|
54
|
+
)
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
## Environment variables
|
|
58
|
+
|
|
59
|
+
| Variable | Required | Description |
|
|
60
|
+
|----------|----------|-------------|
|
|
61
|
+
| `TOKENSOR_API_KEY` | yes | Your Tokensor API key |
|
|
62
|
+
| `TOKENSOR_HOST` | no | Override API host (default: `https://api.tokensor.com`) |
|
|
63
|
+
|
|
64
|
+
## License
|
|
65
|
+
|
|
66
|
+
MIT
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "tokensor"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "LLM token observability SDK — wraps OpenAI and Anthropic clients to capture metadata"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = { text = "MIT" }
|
|
7
|
+
requires-python = ">=3.9"
|
|
8
|
+
authors = [
|
|
9
|
+
{ name = "Tokensor", email = "hello@tokensor.com" },
|
|
10
|
+
]
|
|
11
|
+
keywords = ["llm", "observability", "openai", "anthropic", "tokens", "cost"]
|
|
12
|
+
classifiers = [
|
|
13
|
+
"Development Status :: 3 - Alpha",
|
|
14
|
+
"Intended Audience :: Developers",
|
|
15
|
+
"License :: OSI Approved :: MIT License",
|
|
16
|
+
"Programming Language :: Python :: 3",
|
|
17
|
+
"Programming Language :: Python :: 3.9",
|
|
18
|
+
"Programming Language :: Python :: 3.10",
|
|
19
|
+
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Programming Language :: Python :: 3.12",
|
|
21
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
22
|
+
]
|
|
23
|
+
dependencies = [
|
|
24
|
+
"httpx>=0.28",
|
|
25
|
+
]
|
|
26
|
+
|
|
27
|
+
[project.urls]
|
|
28
|
+
Homepage = "https://tokensor.com"
|
|
29
|
+
Repository = "https://github.com/tokensor/tokensor"
|
|
30
|
+
"Bug Tracker" = "https://github.com/tokensor/tokensor/issues"
|
|
31
|
+
|
|
32
|
+
[project.optional-dependencies]
|
|
33
|
+
openai = ["openai>=1.0"]
|
|
34
|
+
anthropic = ["anthropic>=0.30"]
|
|
35
|
+
dev = [
|
|
36
|
+
"pytest>=8.3",
|
|
37
|
+
"pytest-asyncio>=0.24",
|
|
38
|
+
"pytest-cov>=6.0",
|
|
39
|
+
"openai>=1.0",
|
|
40
|
+
"anthropic>=0.30",
|
|
41
|
+
"respx>=0.21",
|
|
42
|
+
]
|
|
43
|
+
|
|
44
|
+
[tool.pytest.ini_options]
|
|
45
|
+
asyncio_mode = "auto"
|
|
46
|
+
testpaths = ["tests"]
|
|
47
|
+
|
|
48
|
+
[build-system]
|
|
49
|
+
requires = ["hatchling"]
|
|
50
|
+
build-backend = "hatchling.build"
|
|
51
|
+
|
|
52
|
+
[tool.hatch.build.targets.wheel]
|
|
53
|
+
packages = ["src/tokensor"]
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
from typing import Any
|
|
3
|
+
|
|
4
|
+
from tokensor._config import load_config
|
|
5
|
+
from tokensor._client import TokensorClient
|
|
6
|
+
from tokensor._adapters._openai import WrappedOpenAI, WrappedAsyncOpenAI
|
|
7
|
+
from tokensor._adapters._anthropic import WrappedAnthropic, WrappedAsyncAnthropic
|
|
8
|
+
|
|
9
|
+
_OPENAI_NAMES = {"OpenAI", "AzureOpenAI"}
|
|
10
|
+
_ASYNC_OPENAI_NAMES = {"AsyncOpenAI", "AsyncAzureOpenAI"}
|
|
11
|
+
_ANTHROPIC_NAMES = {"Anthropic"}
|
|
12
|
+
_ASYNC_ANTHROPIC_NAMES = {"AsyncAnthropic"}
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def wrap(
|
|
16
|
+
client: Any,
|
|
17
|
+
*,
|
|
18
|
+
api_key: str | None = None,
|
|
19
|
+
host: str | None = None,
|
|
20
|
+
provider: str | None = None,
|
|
21
|
+
_async: bool = False,
|
|
22
|
+
) -> Any:
|
|
23
|
+
"""Wrap an OpenAI or Anthropic client to capture LLM call metadata.
|
|
24
|
+
|
|
25
|
+
Args:
|
|
26
|
+
client: An openai.OpenAI, openai.AsyncOpenAI, anthropic.Anthropic,
|
|
27
|
+
or anthropic.AsyncAnthropic instance.
|
|
28
|
+
api_key: Tokensor API key. Falls back to TOKENSOR_API_KEY env var.
|
|
29
|
+
host: Tokensor API host. Falls back to TOKENSOR_HOST env var or default.
|
|
30
|
+
provider: "openai" or "anthropic". Auto-detected from class name if omitted.
|
|
31
|
+
_async: Force async wrapper. Usually auto-detected from class name.
|
|
32
|
+
|
|
33
|
+
Returns:
|
|
34
|
+
A wrapped client with the same interface as the original.
|
|
35
|
+
"""
|
|
36
|
+
config = load_config(api_key=api_key, host=host)
|
|
37
|
+
tokensor_client = TokensorClient(config)
|
|
38
|
+
|
|
39
|
+
class_name = type(client).__name__
|
|
40
|
+
|
|
41
|
+
resolved_provider = provider
|
|
42
|
+
resolved_async = _async
|
|
43
|
+
|
|
44
|
+
if resolved_provider is None:
|
|
45
|
+
if class_name in _OPENAI_NAMES:
|
|
46
|
+
resolved_provider = "openai"
|
|
47
|
+
elif class_name in _ASYNC_OPENAI_NAMES:
|
|
48
|
+
resolved_provider = "openai"
|
|
49
|
+
resolved_async = True
|
|
50
|
+
elif class_name in _ANTHROPIC_NAMES:
|
|
51
|
+
resolved_provider = "anthropic"
|
|
52
|
+
elif class_name in _ASYNC_ANTHROPIC_NAMES:
|
|
53
|
+
resolved_provider = "anthropic"
|
|
54
|
+
resolved_async = True
|
|
55
|
+
else:
|
|
56
|
+
tokensor_client.shutdown()
|
|
57
|
+
raise ValueError(
|
|
58
|
+
f"Cannot detect provider from class '{class_name}'. "
|
|
59
|
+
"Pass provider='openai' or provider='anthropic' explicitly."
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
if resolved_provider == "openai":
|
|
63
|
+
if resolved_async:
|
|
64
|
+
return WrappedAsyncOpenAI(client, tokensor_client)
|
|
65
|
+
return WrappedOpenAI(client, tokensor_client)
|
|
66
|
+
elif resolved_provider == "anthropic":
|
|
67
|
+
if resolved_async:
|
|
68
|
+
return WrappedAsyncAnthropic(client, tokensor_client)
|
|
69
|
+
return WrappedAnthropic(client, tokensor_client)
|
|
70
|
+
else:
|
|
71
|
+
tokensor_client.shutdown()
|
|
72
|
+
raise ValueError(f"Unknown provider '{resolved_provider}'. Use 'openai' or 'anthropic'.")
|
|
File without changes
|
|
@@ -0,0 +1,120 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
import time
|
|
3
|
+
from datetime import datetime, timezone
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
from tokensor._client import TokensorClient
|
|
7
|
+
from tokensor._hash import endpoint_hash
|
|
8
|
+
from tokensor._models import EventPayload
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def _record_message(
|
|
12
|
+
tokensor_client: TokensorClient,
|
|
13
|
+
h: str,
|
|
14
|
+
model: str,
|
|
15
|
+
messages: list | None,
|
|
16
|
+
max_tokens: int | None,
|
|
17
|
+
called_at: datetime,
|
|
18
|
+
start: float,
|
|
19
|
+
response: Any | None,
|
|
20
|
+
exc: Exception | None,
|
|
21
|
+
) -> None:
|
|
22
|
+
if response is not None:
|
|
23
|
+
status_code: int | None = 200
|
|
24
|
+
input_tokens = getattr(response.usage, "input_tokens", 0) or 0
|
|
25
|
+
output_tokens = getattr(response.usage, "output_tokens", 0) or 0
|
|
26
|
+
cache_creation_tokens = getattr(response.usage, "cache_creation_input_tokens", 0) or 0
|
|
27
|
+
cache_read_tokens = getattr(response.usage, "cache_read_input_tokens", 0) or 0
|
|
28
|
+
finish_reason = getattr(response, "stop_reason", None)
|
|
29
|
+
else:
|
|
30
|
+
status_code = getattr(exc, "status_code", None)
|
|
31
|
+
input_tokens = 0
|
|
32
|
+
output_tokens = 0
|
|
33
|
+
cache_creation_tokens = 0
|
|
34
|
+
cache_read_tokens = 0
|
|
35
|
+
finish_reason = None
|
|
36
|
+
|
|
37
|
+
latency_ms = int((time.monotonic() - start) * 1000)
|
|
38
|
+
event = EventPayload(
|
|
39
|
+
called_at=called_at,
|
|
40
|
+
model=model,
|
|
41
|
+
endpoint_hash=h,
|
|
42
|
+
input_tokens=input_tokens,
|
|
43
|
+
output_tokens=output_tokens,
|
|
44
|
+
cache_creation_tokens=cache_creation_tokens,
|
|
45
|
+
cache_read_tokens=cache_read_tokens,
|
|
46
|
+
message_count=len(messages) if messages is not None else None,
|
|
47
|
+
max_tokens=max_tokens,
|
|
48
|
+
finish_reason=finish_reason,
|
|
49
|
+
latency_ms=latency_ms,
|
|
50
|
+
status_code=status_code,
|
|
51
|
+
)
|
|
52
|
+
tokensor_client.enqueue(event)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class _Messages:
|
|
56
|
+
def __init__(self, original_messages, tokensor_client: TokensorClient) -> None:
|
|
57
|
+
self._messages = original_messages
|
|
58
|
+
self._tokensor = tokensor_client
|
|
59
|
+
|
|
60
|
+
def create(self, model: str, messages: list | None = None, max_tokens: int | None = None, **kwargs) -> Any:
|
|
61
|
+
h = endpoint_hash(depth=2)
|
|
62
|
+
start = time.monotonic()
|
|
63
|
+
called_at = datetime.now(timezone.utc)
|
|
64
|
+
response, exc_to_raise = None, None
|
|
65
|
+
try:
|
|
66
|
+
response = self._messages.create(
|
|
67
|
+
model=model, messages=messages, max_tokens=max_tokens, **kwargs
|
|
68
|
+
)
|
|
69
|
+
except Exception as e:
|
|
70
|
+
exc_to_raise = e
|
|
71
|
+
_record_message(self._tokensor, h, model, messages, max_tokens, called_at, start, response, exc_to_raise)
|
|
72
|
+
if exc_to_raise is not None:
|
|
73
|
+
raise exc_to_raise
|
|
74
|
+
return response
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
class _AsyncMessages:
|
|
78
|
+
def __init__(self, original_messages, tokensor_client: TokensorClient) -> None:
|
|
79
|
+
self._messages = original_messages
|
|
80
|
+
self._tokensor = tokensor_client
|
|
81
|
+
|
|
82
|
+
async def create(self, model: str, messages: list | None = None, max_tokens: int | None = None, **kwargs) -> Any:
|
|
83
|
+
h = endpoint_hash(depth=2)
|
|
84
|
+
start = time.monotonic()
|
|
85
|
+
called_at = datetime.now(timezone.utc)
|
|
86
|
+
response, exc_to_raise = None, None
|
|
87
|
+
try:
|
|
88
|
+
response = await self._messages.create(
|
|
89
|
+
model=model, messages=messages, max_tokens=max_tokens, **kwargs
|
|
90
|
+
)
|
|
91
|
+
except Exception as e:
|
|
92
|
+
exc_to_raise = e
|
|
93
|
+
_record_message(self._tokensor, h, model, messages, max_tokens, called_at, start, response, exc_to_raise)
|
|
94
|
+
if exc_to_raise is not None:
|
|
95
|
+
raise exc_to_raise
|
|
96
|
+
return response
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
class WrappedAnthropic:
|
|
100
|
+
"""Wraps anthropic.Anthropic to capture call metadata."""
|
|
101
|
+
|
|
102
|
+
def __init__(self, client: Any, tokensor_client: TokensorClient) -> None:
|
|
103
|
+
self._client = client
|
|
104
|
+
self._tokensor = tokensor_client
|
|
105
|
+
self.messages = _Messages(client.messages, tokensor_client)
|
|
106
|
+
|
|
107
|
+
def __getattr__(self, name: str) -> Any:
|
|
108
|
+
return getattr(self._client, name)
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
class WrappedAsyncAnthropic:
|
|
112
|
+
"""Wraps anthropic.AsyncAnthropic to capture call metadata."""
|
|
113
|
+
|
|
114
|
+
def __init__(self, client: Any, tokensor_client: TokensorClient) -> None:
|
|
115
|
+
self._client = client
|
|
116
|
+
self._tokensor = tokensor_client
|
|
117
|
+
self.messages = _AsyncMessages(client.messages, tokensor_client)
|
|
118
|
+
|
|
119
|
+
def __getattr__(self, name: str) -> Any:
|
|
120
|
+
return getattr(self._client, name)
|
|
@@ -0,0 +1,124 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
import time
|
|
3
|
+
from datetime import datetime, timezone
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
from tokensor._client import TokensorClient
|
|
7
|
+
from tokensor._hash import endpoint_hash
|
|
8
|
+
from tokensor._models import EventPayload
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def _record_completion(
|
|
12
|
+
tokensor_client: TokensorClient,
|
|
13
|
+
h: str,
|
|
14
|
+
model: str,
|
|
15
|
+
messages: list | None,
|
|
16
|
+
max_tokens: int | None,
|
|
17
|
+
called_at: datetime,
|
|
18
|
+
start: float,
|
|
19
|
+
response: Any | None,
|
|
20
|
+
exc: Exception | None,
|
|
21
|
+
) -> None:
|
|
22
|
+
if response is not None:
|
|
23
|
+
status_code: int | None = 200
|
|
24
|
+
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
|
25
|
+
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
|
26
|
+
finish_reason = (
|
|
27
|
+
getattr(response.choices[0], "finish_reason", None) if response.choices else None
|
|
28
|
+
)
|
|
29
|
+
else:
|
|
30
|
+
status_code = getattr(exc, "status_code", None)
|
|
31
|
+
input_tokens = 0
|
|
32
|
+
output_tokens = 0
|
|
33
|
+
finish_reason = None
|
|
34
|
+
|
|
35
|
+
latency_ms = int((time.monotonic() - start) * 1000)
|
|
36
|
+
event = EventPayload(
|
|
37
|
+
called_at=called_at,
|
|
38
|
+
model=model,
|
|
39
|
+
endpoint_hash=h,
|
|
40
|
+
input_tokens=input_tokens,
|
|
41
|
+
output_tokens=output_tokens,
|
|
42
|
+
message_count=len(messages) if messages is not None else None,
|
|
43
|
+
max_tokens=max_tokens,
|
|
44
|
+
finish_reason=finish_reason,
|
|
45
|
+
latency_ms=latency_ms,
|
|
46
|
+
status_code=status_code,
|
|
47
|
+
)
|
|
48
|
+
tokensor_client.enqueue(event)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
class _ChatCompletions:
|
|
52
|
+
def __init__(self, original_chat, tokensor_client: TokensorClient) -> None:
|
|
53
|
+
self._chat = original_chat
|
|
54
|
+
self._tokensor = tokensor_client
|
|
55
|
+
|
|
56
|
+
@property
|
|
57
|
+
def completions(self):
|
|
58
|
+
return self
|
|
59
|
+
|
|
60
|
+
def create(self, model: str, messages: list | None = None, max_tokens: int | None = None, **kwargs) -> Any:
|
|
61
|
+
h = endpoint_hash(depth=2)
|
|
62
|
+
start = time.monotonic()
|
|
63
|
+
called_at = datetime.now(timezone.utc)
|
|
64
|
+
response, exc_to_raise = None, None
|
|
65
|
+
try:
|
|
66
|
+
response = self._chat.completions.create(
|
|
67
|
+
model=model, messages=messages, max_tokens=max_tokens, **kwargs
|
|
68
|
+
)
|
|
69
|
+
except Exception as e:
|
|
70
|
+
exc_to_raise = e
|
|
71
|
+
_record_completion(self._tokensor, h, model, messages, max_tokens, called_at, start, response, exc_to_raise)
|
|
72
|
+
if exc_to_raise is not None:
|
|
73
|
+
raise exc_to_raise
|
|
74
|
+
return response
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
class _AsyncChatCompletions:
|
|
78
|
+
def __init__(self, original_chat, tokensor_client: TokensorClient) -> None:
|
|
79
|
+
self._chat = original_chat
|
|
80
|
+
self._tokensor = tokensor_client
|
|
81
|
+
|
|
82
|
+
@property
|
|
83
|
+
def completions(self):
|
|
84
|
+
return self
|
|
85
|
+
|
|
86
|
+
async def create(self, model: str, messages: list | None = None, max_tokens: int | None = None, **kwargs) -> Any:
|
|
87
|
+
h = endpoint_hash(depth=2)
|
|
88
|
+
start = time.monotonic()
|
|
89
|
+
called_at = datetime.now(timezone.utc)
|
|
90
|
+
response, exc_to_raise = None, None
|
|
91
|
+
try:
|
|
92
|
+
response = await self._chat.completions.create(
|
|
93
|
+
model=model, messages=messages, max_tokens=max_tokens, **kwargs
|
|
94
|
+
)
|
|
95
|
+
except Exception as e:
|
|
96
|
+
exc_to_raise = e
|
|
97
|
+
_record_completion(self._tokensor, h, model, messages, max_tokens, called_at, start, response, exc_to_raise)
|
|
98
|
+
if exc_to_raise is not None:
|
|
99
|
+
raise exc_to_raise
|
|
100
|
+
return response
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
class WrappedOpenAI:
|
|
104
|
+
"""Wraps openai.OpenAI to capture call metadata."""
|
|
105
|
+
|
|
106
|
+
def __init__(self, client: Any, tokensor_client: TokensorClient) -> None:
|
|
107
|
+
self._client = client
|
|
108
|
+
self._tokensor = tokensor_client
|
|
109
|
+
self.chat = _ChatCompletions(client.chat, tokensor_client)
|
|
110
|
+
|
|
111
|
+
def __getattr__(self, name: str) -> Any:
|
|
112
|
+
return getattr(self._client, name)
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
class WrappedAsyncOpenAI:
|
|
116
|
+
"""Wraps openai.AsyncOpenAI to capture call metadata."""
|
|
117
|
+
|
|
118
|
+
def __init__(self, client: Any, tokensor_client: TokensorClient) -> None:
|
|
119
|
+
self._client = client
|
|
120
|
+
self._tokensor = tokensor_client
|
|
121
|
+
self.chat = _AsyncChatCompletions(client.chat, tokensor_client)
|
|
122
|
+
|
|
123
|
+
def __getattr__(self, name: str) -> Any:
|
|
124
|
+
return getattr(self._client, name)
|
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
import functools
|
|
3
|
+
import logging
|
|
4
|
+
import queue
|
|
5
|
+
import threading
|
|
6
|
+
import time
|
|
7
|
+
from typing import Callable, TypeVar
|
|
8
|
+
|
|
9
|
+
import httpx
|
|
10
|
+
|
|
11
|
+
from tokensor._config import TokensorConfig
|
|
12
|
+
from tokensor._models import EventPayload, events_to_batch_payload
|
|
13
|
+
|
|
14
|
+
logger = logging.getLogger("tokensor")
|
|
15
|
+
|
|
16
|
+
_SENTINEL = object()
|
|
17
|
+
_MAX_RETRIES = 3
|
|
18
|
+
_RETRY_DELAYS = (0.5, 1.0, 2.0)
|
|
19
|
+
|
|
20
|
+
F = TypeVar("F", bound=Callable)
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def _retryable(max_retries: int = _MAX_RETRIES, delays: tuple = _RETRY_DELAYS) -> Callable[[F], F]:
|
|
24
|
+
"""Decorator: retry on httpx errors or 5xx responses; discard permanently on 4xx."""
|
|
25
|
+
def decorator(fn: F) -> F:
|
|
26
|
+
@functools.wraps(fn)
|
|
27
|
+
def wrapper(self, events: list[EventPayload]) -> None:
|
|
28
|
+
for attempt in range(max_retries):
|
|
29
|
+
try:
|
|
30
|
+
result = fn(self, events)
|
|
31
|
+
if result is not None and result < 500:
|
|
32
|
+
return # success or permanent 4xx (already logged inside fn)
|
|
33
|
+
if result is not None:
|
|
34
|
+
# 5xx — fall through to retry
|
|
35
|
+
logger.warning(
|
|
36
|
+
"Tokensor: server error %d for %d events (attempt %d/%d)",
|
|
37
|
+
result, len(events), attempt + 1, max_retries,
|
|
38
|
+
)
|
|
39
|
+
except Exception as exc:
|
|
40
|
+
logger.warning(
|
|
41
|
+
"Tokensor: network error sending %d events (attempt %d/%d): %s",
|
|
42
|
+
len(events), attempt + 1, max_retries, exc,
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
if attempt < max_retries - 1:
|
|
46
|
+
time.sleep(delays[attempt])
|
|
47
|
+
|
|
48
|
+
logger.error(
|
|
49
|
+
"Tokensor: permanently dropping %d events after %d retries",
|
|
50
|
+
len(events), max_retries,
|
|
51
|
+
)
|
|
52
|
+
return wrapper # type: ignore[return-value]
|
|
53
|
+
return decorator
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class TokensorClient:
|
|
57
|
+
"""HTTP sender with a background batch queue thread."""
|
|
58
|
+
|
|
59
|
+
def __init__(self, config: TokensorConfig) -> None:
|
|
60
|
+
self._config = config
|
|
61
|
+
self._queue: queue.Queue = queue.Queue()
|
|
62
|
+
self._stop_event = threading.Event()
|
|
63
|
+
self._thread = threading.Thread(target=self._run, daemon=True)
|
|
64
|
+
self._thread.start()
|
|
65
|
+
|
|
66
|
+
def __enter__(self) -> "TokensorClient":
|
|
67
|
+
return self
|
|
68
|
+
|
|
69
|
+
def __exit__(self, *_) -> None:
|
|
70
|
+
self.shutdown()
|
|
71
|
+
|
|
72
|
+
def enqueue(self, event: EventPayload) -> None:
|
|
73
|
+
"""Add event to the batch queue. Returns immediately."""
|
|
74
|
+
self._queue.put_nowait(event)
|
|
75
|
+
|
|
76
|
+
def shutdown(self) -> None:
|
|
77
|
+
"""Signal the background thread to stop, flush remaining events, and wait."""
|
|
78
|
+
self._stop_event.set()
|
|
79
|
+
self._queue.put_nowait(_SENTINEL) # Wake the thread if it's blocked on get()
|
|
80
|
+
self._thread.join(timeout=10)
|
|
81
|
+
|
|
82
|
+
def _run(self) -> None:
|
|
83
|
+
batch: list[EventPayload] = []
|
|
84
|
+
deadline = time.monotonic() + self._config.flush_interval
|
|
85
|
+
|
|
86
|
+
while True:
|
|
87
|
+
timeout = max(0.001, deadline - time.monotonic())
|
|
88
|
+
try:
|
|
89
|
+
item = self._queue.get(timeout=timeout)
|
|
90
|
+
except queue.Empty:
|
|
91
|
+
# flush_interval elapsed — send whatever we have
|
|
92
|
+
if batch:
|
|
93
|
+
self._send(batch)
|
|
94
|
+
batch = []
|
|
95
|
+
deadline = time.monotonic() + self._config.flush_interval
|
|
96
|
+
continue
|
|
97
|
+
|
|
98
|
+
if item is _SENTINEL:
|
|
99
|
+
# Drain any remaining items enqueued before shutdown
|
|
100
|
+
try:
|
|
101
|
+
while True:
|
|
102
|
+
item2 = self._queue.get_nowait()
|
|
103
|
+
if item2 is not _SENTINEL:
|
|
104
|
+
batch.append(item2)
|
|
105
|
+
except queue.Empty:
|
|
106
|
+
pass
|
|
107
|
+
if batch:
|
|
108
|
+
self._send(batch)
|
|
109
|
+
break
|
|
110
|
+
|
|
111
|
+
batch.append(item)
|
|
112
|
+
if len(batch) >= self._config.batch_size:
|
|
113
|
+
self._send(batch)
|
|
114
|
+
batch = []
|
|
115
|
+
deadline = time.monotonic() + self._config.flush_interval
|
|
116
|
+
|
|
117
|
+
@_retryable()
|
|
118
|
+
def _send(self, events: list[EventPayload]) -> int | None:
|
|
119
|
+
"""POST a batch to the API. Returns status_code; decorator handles retries."""
|
|
120
|
+
resp = httpx.post(
|
|
121
|
+
f"{self._config.host}/v1/events/batch",
|
|
122
|
+
json=events_to_batch_payload(events),
|
|
123
|
+
headers={"X-Tokensor-Key": self._config.api_key},
|
|
124
|
+
timeout=5.0,
|
|
125
|
+
)
|
|
126
|
+
if 400 <= resp.status_code < 500:
|
|
127
|
+
logger.error(
|
|
128
|
+
"Tokensor: dropping %d events — server rejected with %d",
|
|
129
|
+
len(events), resp.status_code,
|
|
130
|
+
)
|
|
131
|
+
return resp.status_code
|