tokengrip 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.
- tokengrip-0.1.0/.gitignore +11 -0
- tokengrip-0.1.0/PKG-INFO +151 -0
- tokengrip-0.1.0/README.md +129 -0
- tokengrip-0.1.0/pyproject.toml +30 -0
- tokengrip-0.1.0/src/tokengrip/__init__.py +26 -0
- tokengrip-0.1.0/src/tokengrip/client.py +204 -0
- tokengrip-0.1.0/src/tokengrip/types.py +39 -0
- tokengrip-0.1.0/src/tokengrip/wrappers/__init__.py +1 -0
- tokengrip-0.1.0/src/tokengrip/wrappers/anthropic.py +423 -0
- tokengrip-0.1.0/src/tokengrip/wrappers/openai.py +462 -0
- tokengrip-0.1.0/tests/test_anthropic.py +329 -0
- tokengrip-0.1.0/tests/test_client.py +237 -0
- tokengrip-0.1.0/tests/test_openai.py +304 -0
tokengrip-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tokengrip
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Track and optimize your AI spending with Tokengrip
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Keywords: ai,anthropic,cost,llm,openai,tokengrip,tracking
|
|
7
|
+
Classifier: Development Status :: 3 - Alpha
|
|
8
|
+
Classifier: Intended Audience :: Developers
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
17
|
+
Requires-Python: >=3.9
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
|
|
20
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# tokengrip
|
|
24
|
+
|
|
25
|
+
Track and optimize your AI spending with [Tokengrip](https://tokengrip.com).
|
|
26
|
+
|
|
27
|
+
## Installation
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
pip install tokengrip
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Quick Start
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
from tokengrip import Tokengrip, wrap_openai
|
|
37
|
+
import openai
|
|
38
|
+
|
|
39
|
+
tg = Tokengrip(api_key="tq_live_...")
|
|
40
|
+
client = wrap_openai(openai.OpenAI(), tg)
|
|
41
|
+
|
|
42
|
+
# Usage is automatically tracked
|
|
43
|
+
response = client.chat.completions.create(
|
|
44
|
+
model="gpt-4.1",
|
|
45
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
# Don't forget to flush on exit
|
|
49
|
+
tg.shutdown()
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
## Wrappers
|
|
53
|
+
|
|
54
|
+
### OpenAI
|
|
55
|
+
|
|
56
|
+
```python
|
|
57
|
+
from tokengrip import Tokengrip, wrap_openai
|
|
58
|
+
import openai
|
|
59
|
+
|
|
60
|
+
tg = Tokengrip(api_key="tq_live_...", project_id="my-project")
|
|
61
|
+
client = wrap_openai(openai.OpenAI(), tg)
|
|
62
|
+
|
|
63
|
+
# Chat Completions API
|
|
64
|
+
response = client.chat.completions.create(
|
|
65
|
+
model="gpt-4.1",
|
|
66
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
67
|
+
)
|
|
68
|
+
|
|
69
|
+
# Responses API
|
|
70
|
+
response = client.responses.create(
|
|
71
|
+
model="gpt-4.1",
|
|
72
|
+
input="Hello!",
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
# Streaming works too
|
|
76
|
+
stream = client.chat.completions.create(
|
|
77
|
+
model="gpt-4.1",
|
|
78
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
79
|
+
stream=True,
|
|
80
|
+
)
|
|
81
|
+
for chunk in stream:
|
|
82
|
+
print(chunk.choices[0].delta.content or "", end="")
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
### Anthropic
|
|
86
|
+
|
|
87
|
+
```python
|
|
88
|
+
from tokengrip import Tokengrip, wrap_anthropic
|
|
89
|
+
import anthropic
|
|
90
|
+
|
|
91
|
+
tg = Tokengrip(api_key="tq_live_...", project_id="my-project")
|
|
92
|
+
client = wrap_anthropic(anthropic.Anthropic(), tg)
|
|
93
|
+
|
|
94
|
+
response = client.messages.create(
|
|
95
|
+
model="claude-sonnet-4-5-20250929",
|
|
96
|
+
max_tokens=1024,
|
|
97
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
98
|
+
)
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
### Async Clients
|
|
102
|
+
|
|
103
|
+
Both sync and async clients are supported transparently:
|
|
104
|
+
|
|
105
|
+
```python
|
|
106
|
+
import openai
|
|
107
|
+
from tokengrip import Tokengrip, wrap_openai
|
|
108
|
+
|
|
109
|
+
tg = Tokengrip(api_key="tq_live_...")
|
|
110
|
+
client = wrap_openai(openai.AsyncOpenAI(), tg)
|
|
111
|
+
|
|
112
|
+
response = await client.chat.completions.create(
|
|
113
|
+
model="gpt-4.1",
|
|
114
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
115
|
+
)
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
## Options
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
tg = Tokengrip(
|
|
122
|
+
api_key="tq_live_...", # Required
|
|
123
|
+
base_url="https://...", # Default: https://tokengrip.com
|
|
124
|
+
project_id="my-project", # Applied to all records
|
|
125
|
+
agent_name="my-agent", # Applied to all records
|
|
126
|
+
debug=True, # Log tracking activity
|
|
127
|
+
batch_size=10, # Records per batch
|
|
128
|
+
flush_interval_ms=5000, # Auto-flush interval
|
|
129
|
+
)
|
|
130
|
+
|
|
131
|
+
# Wrapper-specific options
|
|
132
|
+
client = wrap_openai(openai.OpenAI(), tg, {
|
|
133
|
+
"project_id": "override", # Override per-wrapper
|
|
134
|
+
"agent_name": "override",
|
|
135
|
+
"task_type": "summarization",
|
|
136
|
+
"capture_content": True, # Capture truncated prompt/response
|
|
137
|
+
})
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
## Manual Tracking
|
|
141
|
+
|
|
142
|
+
```python
|
|
143
|
+
tg.track(
|
|
144
|
+
provider="openai",
|
|
145
|
+
model="gpt-4.1",
|
|
146
|
+
inputTokens=500,
|
|
147
|
+
outputTokens=150,
|
|
148
|
+
latencyMs=1200,
|
|
149
|
+
projectId="my-project",
|
|
150
|
+
)
|
|
151
|
+
```
|
|
@@ -0,0 +1,129 @@
|
|
|
1
|
+
# tokengrip
|
|
2
|
+
|
|
3
|
+
Track and optimize your AI spending with [Tokengrip](https://tokengrip.com).
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install tokengrip
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Quick Start
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
from tokengrip import Tokengrip, wrap_openai
|
|
15
|
+
import openai
|
|
16
|
+
|
|
17
|
+
tg = Tokengrip(api_key="tq_live_...")
|
|
18
|
+
client = wrap_openai(openai.OpenAI(), tg)
|
|
19
|
+
|
|
20
|
+
# Usage is automatically tracked
|
|
21
|
+
response = client.chat.completions.create(
|
|
22
|
+
model="gpt-4.1",
|
|
23
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
24
|
+
)
|
|
25
|
+
|
|
26
|
+
# Don't forget to flush on exit
|
|
27
|
+
tg.shutdown()
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
## Wrappers
|
|
31
|
+
|
|
32
|
+
### OpenAI
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
from tokengrip import Tokengrip, wrap_openai
|
|
36
|
+
import openai
|
|
37
|
+
|
|
38
|
+
tg = Tokengrip(api_key="tq_live_...", project_id="my-project")
|
|
39
|
+
client = wrap_openai(openai.OpenAI(), tg)
|
|
40
|
+
|
|
41
|
+
# Chat Completions API
|
|
42
|
+
response = client.chat.completions.create(
|
|
43
|
+
model="gpt-4.1",
|
|
44
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
# Responses API
|
|
48
|
+
response = client.responses.create(
|
|
49
|
+
model="gpt-4.1",
|
|
50
|
+
input="Hello!",
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
# Streaming works too
|
|
54
|
+
stream = client.chat.completions.create(
|
|
55
|
+
model="gpt-4.1",
|
|
56
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
57
|
+
stream=True,
|
|
58
|
+
)
|
|
59
|
+
for chunk in stream:
|
|
60
|
+
print(chunk.choices[0].delta.content or "", end="")
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
### Anthropic
|
|
64
|
+
|
|
65
|
+
```python
|
|
66
|
+
from tokengrip import Tokengrip, wrap_anthropic
|
|
67
|
+
import anthropic
|
|
68
|
+
|
|
69
|
+
tg = Tokengrip(api_key="tq_live_...", project_id="my-project")
|
|
70
|
+
client = wrap_anthropic(anthropic.Anthropic(), tg)
|
|
71
|
+
|
|
72
|
+
response = client.messages.create(
|
|
73
|
+
model="claude-sonnet-4-5-20250929",
|
|
74
|
+
max_tokens=1024,
|
|
75
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
76
|
+
)
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
### Async Clients
|
|
80
|
+
|
|
81
|
+
Both sync and async clients are supported transparently:
|
|
82
|
+
|
|
83
|
+
```python
|
|
84
|
+
import openai
|
|
85
|
+
from tokengrip import Tokengrip, wrap_openai
|
|
86
|
+
|
|
87
|
+
tg = Tokengrip(api_key="tq_live_...")
|
|
88
|
+
client = wrap_openai(openai.AsyncOpenAI(), tg)
|
|
89
|
+
|
|
90
|
+
response = await client.chat.completions.create(
|
|
91
|
+
model="gpt-4.1",
|
|
92
|
+
messages=[{"role": "user", "content": "Hello!"}],
|
|
93
|
+
)
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
## Options
|
|
97
|
+
|
|
98
|
+
```python
|
|
99
|
+
tg = Tokengrip(
|
|
100
|
+
api_key="tq_live_...", # Required
|
|
101
|
+
base_url="https://...", # Default: https://tokengrip.com
|
|
102
|
+
project_id="my-project", # Applied to all records
|
|
103
|
+
agent_name="my-agent", # Applied to all records
|
|
104
|
+
debug=True, # Log tracking activity
|
|
105
|
+
batch_size=10, # Records per batch
|
|
106
|
+
flush_interval_ms=5000, # Auto-flush interval
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
# Wrapper-specific options
|
|
110
|
+
client = wrap_openai(openai.OpenAI(), tg, {
|
|
111
|
+
"project_id": "override", # Override per-wrapper
|
|
112
|
+
"agent_name": "override",
|
|
113
|
+
"task_type": "summarization",
|
|
114
|
+
"capture_content": True, # Capture truncated prompt/response
|
|
115
|
+
})
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
## Manual Tracking
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
tg.track(
|
|
122
|
+
provider="openai",
|
|
123
|
+
model="gpt-4.1",
|
|
124
|
+
inputTokens=500,
|
|
125
|
+
outputTokens=150,
|
|
126
|
+
latencyMs=1200,
|
|
127
|
+
projectId="my-project",
|
|
128
|
+
)
|
|
129
|
+
```
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "tokengrip"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Track and optimize your AI spending with Tokengrip"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
requires-python = ">=3.9"
|
|
12
|
+
keywords = ["ai", "llm", "cost", "tracking", "openai", "anthropic", "tokengrip"]
|
|
13
|
+
classifiers = [
|
|
14
|
+
"Development Status :: 3 - Alpha",
|
|
15
|
+
"Intended Audience :: Developers",
|
|
16
|
+
"License :: OSI Approved :: MIT License",
|
|
17
|
+
"Programming Language :: Python :: 3",
|
|
18
|
+
"Programming Language :: Python :: 3.9",
|
|
19
|
+
"Programming Language :: Python :: 3.10",
|
|
20
|
+
"Programming Language :: Python :: 3.11",
|
|
21
|
+
"Programming Language :: Python :: 3.12",
|
|
22
|
+
"Programming Language :: Python :: 3.13",
|
|
23
|
+
"Topic :: Software Development :: Libraries",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
[project.optional-dependencies]
|
|
27
|
+
dev = ["pytest>=7.0", "pytest-asyncio>=0.21"]
|
|
28
|
+
|
|
29
|
+
[tool.hatch.build.targets.wheel]
|
|
30
|
+
packages = ["src/tokengrip"]
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
"""tokengrip - Track and optimize your AI spending."""
|
|
2
|
+
|
|
3
|
+
from tokengrip.client import Tokengrip
|
|
4
|
+
from tokengrip.types import TrackUsagePayload, WrapOptions
|
|
5
|
+
|
|
6
|
+
__all__ = [
|
|
7
|
+
"Tokengrip",
|
|
8
|
+
"wrap_openai",
|
|
9
|
+
"wrap_anthropic",
|
|
10
|
+
"TrackUsagePayload",
|
|
11
|
+
"WrapOptions",
|
|
12
|
+
]
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def wrap_openai(client, tg, opts=None): # type: ignore[no-untyped-def]
|
|
16
|
+
"""Wrap an OpenAI client to automatically track token usage."""
|
|
17
|
+
from tokengrip.wrappers.openai import wrap_openai as _wrap
|
|
18
|
+
|
|
19
|
+
return _wrap(client, tg, opts)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def wrap_anthropic(client, tg, opts=None): # type: ignore[no-untyped-def]
|
|
23
|
+
"""Wrap an Anthropic client to automatically track token usage."""
|
|
24
|
+
from tokengrip.wrappers.anthropic import wrap_anthropic as _wrap
|
|
25
|
+
|
|
26
|
+
return _wrap(client, tg, opts)
|
|
@@ -0,0 +1,204 @@
|
|
|
1
|
+
"""Core Tokengrip client with batching and auto-flush."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import atexit
|
|
6
|
+
import json
|
|
7
|
+
import logging
|
|
8
|
+
import threading
|
|
9
|
+
import urllib.error
|
|
10
|
+
import urllib.request
|
|
11
|
+
from typing import Any, Dict, List, Optional
|
|
12
|
+
|
|
13
|
+
from tokengrip.types import TrackUsagePayload
|
|
14
|
+
|
|
15
|
+
logger = logging.getLogger("tokengrip")
|
|
16
|
+
|
|
17
|
+
_SDK_VERSION = "0.1.0"
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class Tokengrip:
|
|
21
|
+
"""Buffer usage records and send them in batches to the Tokengrip API.
|
|
22
|
+
|
|
23
|
+
Example::
|
|
24
|
+
|
|
25
|
+
tg = Tokengrip(api_key="tq_live_abc123", project_id="my-project")
|
|
26
|
+
tg.track(provider="openai", model="gpt-4.1",
|
|
27
|
+
inputTokens=500, outputTokens=150)
|
|
28
|
+
tg.shutdown()
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
def __init__(
|
|
32
|
+
self,
|
|
33
|
+
api_key: str,
|
|
34
|
+
base_url: str = "https://tokengrip.com",
|
|
35
|
+
project_id: str = "",
|
|
36
|
+
agent_name: str = "",
|
|
37
|
+
debug: bool = False,
|
|
38
|
+
batch_size: int = 10,
|
|
39
|
+
flush_interval_ms: int = 5000,
|
|
40
|
+
) -> None:
|
|
41
|
+
if not api_key:
|
|
42
|
+
raise ValueError(
|
|
43
|
+
"Tokengrip: api_key is required. "
|
|
44
|
+
"Get yours at https://tokengrip.com/settings/api-keys"
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
self._api_key = api_key
|
|
48
|
+
self._base_url = base_url.rstrip("/")
|
|
49
|
+
self._project_id = project_id
|
|
50
|
+
self._agent_name = agent_name
|
|
51
|
+
self._debug = debug
|
|
52
|
+
self._batch_size = batch_size
|
|
53
|
+
self._flush_interval_s = flush_interval_ms / 1000.0
|
|
54
|
+
|
|
55
|
+
self._buffer: List[Dict[str, Any]] = []
|
|
56
|
+
self._lock = threading.Lock()
|
|
57
|
+
self._is_flushing = False
|
|
58
|
+
self._is_shutdown = False
|
|
59
|
+
|
|
60
|
+
self._timer: Optional[threading.Timer] = None
|
|
61
|
+
if self._flush_interval_s > 0:
|
|
62
|
+
self._schedule_flush()
|
|
63
|
+
|
|
64
|
+
atexit.register(self.shutdown)
|
|
65
|
+
|
|
66
|
+
def _schedule_flush(self) -> None:
|
|
67
|
+
if self._is_shutdown:
|
|
68
|
+
return
|
|
69
|
+
self._timer = threading.Timer(self._flush_interval_s, self._periodic_flush)
|
|
70
|
+
self._timer.daemon = True
|
|
71
|
+
self._timer.start()
|
|
72
|
+
|
|
73
|
+
def _periodic_flush(self) -> None:
|
|
74
|
+
self.flush()
|
|
75
|
+
self._schedule_flush()
|
|
76
|
+
|
|
77
|
+
def track(
|
|
78
|
+
self,
|
|
79
|
+
payload: Optional[TrackUsagePayload] = None,
|
|
80
|
+
**kwargs: Any,
|
|
81
|
+
) -> None:
|
|
82
|
+
"""Buffer a single usage record.
|
|
83
|
+
|
|
84
|
+
Accepts a ``TrackUsagePayload`` dict, keyword arguments, or both.
|
|
85
|
+
Never raises — errors are logged at debug level when ``debug=True``.
|
|
86
|
+
"""
|
|
87
|
+
if self._is_shutdown:
|
|
88
|
+
self._log("warn", "track() called after shutdown — record dropped")
|
|
89
|
+
return
|
|
90
|
+
|
|
91
|
+
try:
|
|
92
|
+
record: Dict[str, Any] = dict(payload) if payload else {}
|
|
93
|
+
record.update(kwargs)
|
|
94
|
+
|
|
95
|
+
# Compute totalTokens if absent
|
|
96
|
+
input_t = record.get("inputTokens", 0)
|
|
97
|
+
output_t = record.get("outputTokens", 0)
|
|
98
|
+
if "totalTokens" not in record:
|
|
99
|
+
record["totalTokens"] = input_t + output_t
|
|
100
|
+
|
|
101
|
+
# Apply defaults from config
|
|
102
|
+
if not record.get("projectId") and self._project_id:
|
|
103
|
+
record["projectId"] = self._project_id
|
|
104
|
+
if not record.get("agentName") and self._agent_name:
|
|
105
|
+
record["agentName"] = self._agent_name
|
|
106
|
+
|
|
107
|
+
# Strip falsy optional fields for a clean payload
|
|
108
|
+
for key in ("projectId", "agentName", "taskType", "metadata"):
|
|
109
|
+
if key in record and not record[key]:
|
|
110
|
+
del record[key]
|
|
111
|
+
if "latencyMs" in record and record["latencyMs"] is None:
|
|
112
|
+
del record["latencyMs"]
|
|
113
|
+
|
|
114
|
+
with self._lock:
|
|
115
|
+
self._buffer.append(record)
|
|
116
|
+
buffer_len = len(self._buffer)
|
|
117
|
+
|
|
118
|
+
self._log(
|
|
119
|
+
"debug",
|
|
120
|
+
f"Buffered: {record.get('provider')}/{record.get('model')} "
|
|
121
|
+
f"({input_t}+{output_t} tokens)",
|
|
122
|
+
)
|
|
123
|
+
|
|
124
|
+
if buffer_len >= self._batch_size:
|
|
125
|
+
self.flush()
|
|
126
|
+
|
|
127
|
+
except Exception as e:
|
|
128
|
+
self._log("warn", f"Failed to buffer record: {e}")
|
|
129
|
+
|
|
130
|
+
def flush(self) -> None:
|
|
131
|
+
"""Force-flush all buffered records to the API. Thread-safe."""
|
|
132
|
+
with self._lock:
|
|
133
|
+
if not self._buffer or self._is_flushing:
|
|
134
|
+
return
|
|
135
|
+
self._is_flushing = True
|
|
136
|
+
records = self._buffer[:]
|
|
137
|
+
self._buffer.clear()
|
|
138
|
+
|
|
139
|
+
try:
|
|
140
|
+
self._send_batch(records)
|
|
141
|
+
except Exception as e:
|
|
142
|
+
self._log("warn", f"Flush failed: {e}")
|
|
143
|
+
finally:
|
|
144
|
+
with self._lock:
|
|
145
|
+
self._is_flushing = False
|
|
146
|
+
|
|
147
|
+
def shutdown(self) -> None:
|
|
148
|
+
"""Flush remaining records and stop the auto-flush timer."""
|
|
149
|
+
if self._is_shutdown:
|
|
150
|
+
return
|
|
151
|
+
self._is_shutdown = True
|
|
152
|
+
|
|
153
|
+
if self._timer is not None:
|
|
154
|
+
self._timer.cancel()
|
|
155
|
+
self._timer = None
|
|
156
|
+
|
|
157
|
+
with self._lock:
|
|
158
|
+
self._is_flushing = False
|
|
159
|
+
|
|
160
|
+
self.flush()
|
|
161
|
+
self._log("debug", "Tokengrip client shut down")
|
|
162
|
+
|
|
163
|
+
def _send_batch(self, records: List[Dict[str, Any]]) -> None:
|
|
164
|
+
if not records:
|
|
165
|
+
return
|
|
166
|
+
|
|
167
|
+
url = f"{self._base_url}/api/v1/track"
|
|
168
|
+
self._log("debug", f"Sending {len(records)} record(s) to {url}")
|
|
169
|
+
|
|
170
|
+
data = json.dumps(records).encode("utf-8")
|
|
171
|
+
req = urllib.request.Request(
|
|
172
|
+
url,
|
|
173
|
+
data=data,
|
|
174
|
+
method="POST",
|
|
175
|
+
headers={
|
|
176
|
+
"Content-Type": "application/json",
|
|
177
|
+
"Authorization": f"Bearer {self._api_key}",
|
|
178
|
+
"User-Agent": f"tokengrip-python/{_SDK_VERSION}",
|
|
179
|
+
},
|
|
180
|
+
)
|
|
181
|
+
|
|
182
|
+
try:
|
|
183
|
+
with urllib.request.urlopen(req, timeout=10) as resp:
|
|
184
|
+
self._log(
|
|
185
|
+
"debug", f"Sent {len(records)} record(s), status {resp.status}"
|
|
186
|
+
)
|
|
187
|
+
except urllib.error.HTTPError as e:
|
|
188
|
+
body = ""
|
|
189
|
+
try:
|
|
190
|
+
body = e.read().decode("utf-8", errors="replace")
|
|
191
|
+
except Exception:
|
|
192
|
+
pass
|
|
193
|
+
self._log("warn", f"API returned {e.code}: {body}")
|
|
194
|
+
except Exception as e:
|
|
195
|
+
self._log("warn", f"Failed to send batch: {e}")
|
|
196
|
+
|
|
197
|
+
def _log(self, level: str, message: str) -> None:
|
|
198
|
+
if not self._debug:
|
|
199
|
+
return
|
|
200
|
+
prefix = "[tokengrip]"
|
|
201
|
+
if level == "warn":
|
|
202
|
+
logger.warning("%s %s", prefix, message)
|
|
203
|
+
else:
|
|
204
|
+
logger.debug("%s %s", prefix, message)
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""Type definitions for the Tokengrip SDK."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any, Dict, Optional
|
|
6
|
+
|
|
7
|
+
try:
|
|
8
|
+
from typing import TypedDict
|
|
9
|
+
except ImportError:
|
|
10
|
+
from typing_extensions import TypedDict # type: ignore[assignment]
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class TrackUsagePayload(TypedDict, total=False):
|
|
14
|
+
"""Payload sent to the Tokengrip tracking API.
|
|
15
|
+
|
|
16
|
+
Field names use camelCase to match the server's Zod schema.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
provider: str
|
|
20
|
+
model: str
|
|
21
|
+
inputTokens: int
|
|
22
|
+
outputTokens: int
|
|
23
|
+
totalTokens: int
|
|
24
|
+
latencyMs: int
|
|
25
|
+
projectId: str
|
|
26
|
+
agentName: str
|
|
27
|
+
taskType: str
|
|
28
|
+
metadata: Dict[str, Any]
|
|
29
|
+
promptContent: str
|
|
30
|
+
responseContent: str
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class WrapOptions(TypedDict, total=False):
|
|
34
|
+
"""Options for AI client wrappers."""
|
|
35
|
+
|
|
36
|
+
project_id: str
|
|
37
|
+
agent_name: str
|
|
38
|
+
task_type: str
|
|
39
|
+
capture_content: bool
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""AI client wrappers for automatic usage tracking."""
|