tokenator 0.1.16__py3-none-any.whl → 0.2.1__py3-none-any.whl
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.
- tokenator/__init__.py +8 -1
- tokenator/base_wrapper.py +4 -1
- tokenator/gemini/__init__.py +5 -0
- tokenator/gemini/client_gemini.py +230 -0
- tokenator/gemini/stream_interceptors.py +77 -0
- tokenator/usage.py +35 -24
- tokenator/utils.py +7 -4
- {tokenator-0.1.16.dist-info → tokenator-0.2.1.dist-info}/METADATA +56 -4
- {tokenator-0.1.16.dist-info → tokenator-0.2.1.dist-info}/RECORD +11 -8
- {tokenator-0.1.16.dist-info → tokenator-0.2.1.dist-info}/WHEEL +1 -1
- {tokenator-0.1.16.dist-info → tokenator-0.2.1.dist-info}/LICENSE +0 -0
tokenator/__init__.py
CHANGED
@@ -3,11 +3,18 @@
|
|
3
3
|
import logging
|
4
4
|
from .openai.client_openai import tokenator_openai
|
5
5
|
from .anthropic.client_anthropic import tokenator_anthropic
|
6
|
+
from .gemini.client_gemini import tokenator_gemini
|
6
7
|
from . import usage
|
7
8
|
from .utils import get_default_db_path
|
8
9
|
from .usage import TokenUsageService
|
9
10
|
|
10
11
|
usage = TokenUsageService() # noqa: F811
|
11
|
-
__all__ = [
|
12
|
+
__all__ = [
|
13
|
+
"tokenator_openai",
|
14
|
+
"tokenator_anthropic",
|
15
|
+
"tokenator_gemini",
|
16
|
+
"usage",
|
17
|
+
"get_default_db_path",
|
18
|
+
]
|
12
19
|
|
13
20
|
logger = logging.getLogger(__name__)
|
tokenator/base_wrapper.py
CHANGED
@@ -112,7 +112,10 @@ class BaseWrapper:
|
|
112
112
|
try:
|
113
113
|
self._log_usage_impl(token_usage_stats, session, execution_id)
|
114
114
|
session.commit()
|
115
|
-
logger.debug(
|
115
|
+
logger.debug(
|
116
|
+
"Successfully committed token usage for execution_id: %s",
|
117
|
+
execution_id,
|
118
|
+
)
|
116
119
|
except Exception as e:
|
117
120
|
logger.error("Failed to log token usage: %s", str(e))
|
118
121
|
session.rollback()
|
@@ -0,0 +1,230 @@
|
|
1
|
+
"""Gemini client wrapper with token usage tracking."""
|
2
|
+
|
3
|
+
from typing import Any, Optional, Iterator, AsyncIterator
|
4
|
+
import logging
|
5
|
+
|
6
|
+
from google import genai
|
7
|
+
from google.genai.types import GenerateContentResponse
|
8
|
+
|
9
|
+
from ..models import (
|
10
|
+
TokenMetrics,
|
11
|
+
TokenUsageStats,
|
12
|
+
)
|
13
|
+
from ..base_wrapper import BaseWrapper, ResponseType
|
14
|
+
from .stream_interceptors import (
|
15
|
+
GeminiAsyncStreamInterceptor,
|
16
|
+
GeminiSyncStreamInterceptor,
|
17
|
+
)
|
18
|
+
from ..state import is_tokenator_enabled
|
19
|
+
|
20
|
+
logger = logging.getLogger(__name__)
|
21
|
+
|
22
|
+
|
23
|
+
def _create_usage_callback(execution_id, log_usage_fn):
|
24
|
+
"""Creates a callback function for processing usage statistics from stream chunks."""
|
25
|
+
|
26
|
+
def usage_callback(chunks):
|
27
|
+
if not chunks:
|
28
|
+
return
|
29
|
+
|
30
|
+
# Skip if tokenator is disabled
|
31
|
+
if not is_tokenator_enabled:
|
32
|
+
logger.debug("Tokenator is disabled - skipping stream usage logging")
|
33
|
+
return
|
34
|
+
|
35
|
+
logger.debug("Processing stream usage for execution_id: %s", execution_id)
|
36
|
+
|
37
|
+
# Build usage_data from the first chunk's model
|
38
|
+
usage_data = TokenUsageStats(
|
39
|
+
model=chunks[0].model_version,
|
40
|
+
usage=TokenMetrics(),
|
41
|
+
)
|
42
|
+
|
43
|
+
# Only take usage from the last chunk as it contains complete usage info
|
44
|
+
last_chunk = chunks[-1]
|
45
|
+
if last_chunk.usage_metadata:
|
46
|
+
usage_data.usage.prompt_tokens = (
|
47
|
+
last_chunk.usage_metadata.prompt_token_count
|
48
|
+
)
|
49
|
+
usage_data.usage.completion_tokens = (
|
50
|
+
last_chunk.usage_metadata.candidates_token_count or 0
|
51
|
+
)
|
52
|
+
usage_data.usage.total_tokens = last_chunk.usage_metadata.total_token_count
|
53
|
+
log_usage_fn(usage_data, execution_id=execution_id)
|
54
|
+
|
55
|
+
return usage_callback
|
56
|
+
|
57
|
+
|
58
|
+
class BaseGeminiWrapper(BaseWrapper):
|
59
|
+
def __init__(self, client, db_path=None, provider: str = "gemini"):
|
60
|
+
super().__init__(client, db_path)
|
61
|
+
self.provider = provider
|
62
|
+
self._async_wrapper = None
|
63
|
+
|
64
|
+
def _process_response_usage(
|
65
|
+
self, response: ResponseType
|
66
|
+
) -> Optional[TokenUsageStats]:
|
67
|
+
"""Process and log usage statistics from a response."""
|
68
|
+
try:
|
69
|
+
if isinstance(response, GenerateContentResponse):
|
70
|
+
if response.usage_metadata is None:
|
71
|
+
return None
|
72
|
+
usage = TokenMetrics(
|
73
|
+
prompt_tokens=response.usage_metadata.prompt_token_count,
|
74
|
+
completion_tokens=response.usage_metadata.candidates_token_count,
|
75
|
+
total_tokens=response.usage_metadata.total_token_count,
|
76
|
+
)
|
77
|
+
return TokenUsageStats(model=response.model_version, usage=usage)
|
78
|
+
|
79
|
+
elif isinstance(response, dict):
|
80
|
+
usage_dict = response.get("usage_metadata")
|
81
|
+
if not usage_dict:
|
82
|
+
return None
|
83
|
+
usage = TokenMetrics(
|
84
|
+
prompt_tokens=usage_dict.get("prompt_token_count", 0),
|
85
|
+
completion_tokens=usage_dict.get("candidates_token_count", 0),
|
86
|
+
total_tokens=usage_dict.get("total_token_count", 0),
|
87
|
+
)
|
88
|
+
return TokenUsageStats(
|
89
|
+
model=response.get("model", "unknown"), usage=usage
|
90
|
+
)
|
91
|
+
except Exception as e:
|
92
|
+
logger.warning("Failed to process usage stats: %s", str(e))
|
93
|
+
return None
|
94
|
+
return None
|
95
|
+
|
96
|
+
@property
|
97
|
+
def chat(self):
|
98
|
+
return self
|
99
|
+
|
100
|
+
@property
|
101
|
+
def chats(self):
|
102
|
+
return self
|
103
|
+
|
104
|
+
@property
|
105
|
+
def models(self):
|
106
|
+
return self
|
107
|
+
|
108
|
+
@property
|
109
|
+
def aio(self):
|
110
|
+
if self._async_wrapper is None:
|
111
|
+
self._async_wrapper = AsyncGeminiWrapper(self)
|
112
|
+
return self._async_wrapper
|
113
|
+
|
114
|
+
def count_tokens(self, *args: Any, **kwargs: Any):
|
115
|
+
return self.client.models.count_tokens(*args, **kwargs)
|
116
|
+
|
117
|
+
|
118
|
+
class AsyncGeminiWrapper:
|
119
|
+
"""Async wrapper for Gemini client to match the official SDK structure."""
|
120
|
+
|
121
|
+
def __init__(self, wrapper: BaseGeminiWrapper):
|
122
|
+
self.wrapper = wrapper
|
123
|
+
self._models = None
|
124
|
+
|
125
|
+
@property
|
126
|
+
def models(self):
|
127
|
+
if self._models is None:
|
128
|
+
self._models = AsyncModelsWrapper(self.wrapper)
|
129
|
+
return self._models
|
130
|
+
|
131
|
+
|
132
|
+
class AsyncModelsWrapper:
|
133
|
+
"""Async wrapper for models to match the official SDK structure."""
|
134
|
+
|
135
|
+
def __init__(self, wrapper: BaseGeminiWrapper):
|
136
|
+
self.wrapper = wrapper
|
137
|
+
|
138
|
+
async def generate_content(
|
139
|
+
self, *args: Any, **kwargs: Any
|
140
|
+
) -> GenerateContentResponse:
|
141
|
+
"""Async method for generate_content."""
|
142
|
+
execution_id = kwargs.pop("execution_id", None)
|
143
|
+
return await self.wrapper.generate_content_async(
|
144
|
+
*args, execution_id=execution_id, **kwargs
|
145
|
+
)
|
146
|
+
|
147
|
+
async def generate_content_stream(
|
148
|
+
self, *args: Any, **kwargs: Any
|
149
|
+
) -> AsyncIterator[GenerateContentResponse]:
|
150
|
+
"""Async method for generate_content_stream."""
|
151
|
+
execution_id = kwargs.pop("execution_id", None)
|
152
|
+
return await self.wrapper.generate_content_stream_async(
|
153
|
+
*args, execution_id=execution_id, **kwargs
|
154
|
+
)
|
155
|
+
|
156
|
+
|
157
|
+
class GeminiWrapper(BaseGeminiWrapper):
|
158
|
+
def generate_content(
|
159
|
+
self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
|
160
|
+
) -> GenerateContentResponse:
|
161
|
+
"""Generate content and log token usage."""
|
162
|
+
logger.debug("Generating content with args: %s, kwargs: %s", args, kwargs)
|
163
|
+
|
164
|
+
response = self.client.models.generate_content(*args, **kwargs)
|
165
|
+
usage_data = self._process_response_usage(response)
|
166
|
+
if usage_data:
|
167
|
+
self._log_usage(usage_data, execution_id=execution_id)
|
168
|
+
|
169
|
+
return response
|
170
|
+
|
171
|
+
def generate_content_stream(
|
172
|
+
self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
|
173
|
+
) -> Iterator[GenerateContentResponse]:
|
174
|
+
"""Generate content with streaming and log token usage."""
|
175
|
+
logger.debug(
|
176
|
+
"Generating content stream with args: %s, kwargs: %s", args, kwargs
|
177
|
+
)
|
178
|
+
|
179
|
+
base_stream = self.client.models.generate_content_stream(*args, **kwargs)
|
180
|
+
return GeminiSyncStreamInterceptor(
|
181
|
+
base_stream=base_stream,
|
182
|
+
usage_callback=_create_usage_callback(execution_id, self._log_usage),
|
183
|
+
)
|
184
|
+
|
185
|
+
async def generate_content_async(
|
186
|
+
self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
|
187
|
+
) -> GenerateContentResponse:
|
188
|
+
"""Generate content asynchronously and log token usage."""
|
189
|
+
logger.debug("Generating content async with args: %s, kwargs: %s", args, kwargs)
|
190
|
+
|
191
|
+
response = await self.client.aio.models.generate_content(*args, **kwargs)
|
192
|
+
usage_data = self._process_response_usage(response)
|
193
|
+
if usage_data:
|
194
|
+
self._log_usage(usage_data, execution_id=execution_id)
|
195
|
+
|
196
|
+
return response
|
197
|
+
|
198
|
+
async def generate_content_stream_async(
|
199
|
+
self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
|
200
|
+
) -> AsyncIterator[GenerateContentResponse]:
|
201
|
+
"""Generate content with async streaming and log token usage."""
|
202
|
+
logger.debug(
|
203
|
+
"Generating content stream async with args: %s, kwargs: %s", args, kwargs
|
204
|
+
)
|
205
|
+
|
206
|
+
base_stream = await self.client.aio.models.generate_content_stream(
|
207
|
+
*args, **kwargs
|
208
|
+
)
|
209
|
+
return GeminiAsyncStreamInterceptor(
|
210
|
+
base_stream=base_stream,
|
211
|
+
usage_callback=_create_usage_callback(execution_id, self._log_usage),
|
212
|
+
)
|
213
|
+
|
214
|
+
|
215
|
+
def tokenator_gemini(
|
216
|
+
client: genai.Client,
|
217
|
+
db_path: Optional[str] = None,
|
218
|
+
provider: str = "gemini",
|
219
|
+
) -> GeminiWrapper:
|
220
|
+
"""Create a token-tracking wrapper for a Gemini client.
|
221
|
+
|
222
|
+
Args:
|
223
|
+
client: Gemini client instance
|
224
|
+
db_path: Optional path to SQLite database for token tracking
|
225
|
+
provider: Provider name, defaults to "gemini"
|
226
|
+
"""
|
227
|
+
if not isinstance(client, genai.Client):
|
228
|
+
raise ValueError("Client must be an instance of genai.Client")
|
229
|
+
|
230
|
+
return GeminiWrapper(client=client, db_path=db_path, provider=provider)
|
@@ -0,0 +1,77 @@
|
|
1
|
+
"""Stream interceptors for Gemini responses."""
|
2
|
+
|
3
|
+
import logging
|
4
|
+
from typing import AsyncIterator, Callable, List, Optional, TypeVar, Iterator
|
5
|
+
|
6
|
+
|
7
|
+
logger = logging.getLogger(__name__)
|
8
|
+
|
9
|
+
_T = TypeVar("_T") # GenerateContentResponse
|
10
|
+
|
11
|
+
|
12
|
+
class GeminiAsyncStreamInterceptor(AsyncIterator[_T]):
|
13
|
+
"""
|
14
|
+
A wrapper around Gemini async stream that intercepts each chunk to handle usage or
|
15
|
+
logging logic.
|
16
|
+
"""
|
17
|
+
|
18
|
+
def __init__(
|
19
|
+
self,
|
20
|
+
base_stream: AsyncIterator[_T],
|
21
|
+
usage_callback: Optional[Callable[[List[_T]], None]] = None,
|
22
|
+
):
|
23
|
+
self._base_stream = base_stream
|
24
|
+
self._usage_callback = usage_callback
|
25
|
+
self._chunks: List[_T] = []
|
26
|
+
|
27
|
+
def __aiter__(self) -> AsyncIterator[_T]:
|
28
|
+
"""Return self as async iterator."""
|
29
|
+
return self
|
30
|
+
|
31
|
+
async def __anext__(self) -> _T:
|
32
|
+
"""Get next chunk and track it."""
|
33
|
+
try:
|
34
|
+
chunk = await self._base_stream.__anext__()
|
35
|
+
except StopAsyncIteration:
|
36
|
+
# Once the base stream is fully consumed, we can do final usage/logging.
|
37
|
+
if self._usage_callback and self._chunks:
|
38
|
+
self._usage_callback(self._chunks)
|
39
|
+
raise
|
40
|
+
|
41
|
+
# Intercept each chunk
|
42
|
+
self._chunks.append(chunk)
|
43
|
+
return chunk
|
44
|
+
|
45
|
+
|
46
|
+
class GeminiSyncStreamInterceptor(Iterator[_T]):
|
47
|
+
"""
|
48
|
+
A wrapper around Gemini sync stream that intercepts each chunk to handle usage or
|
49
|
+
logging logic.
|
50
|
+
"""
|
51
|
+
|
52
|
+
def __init__(
|
53
|
+
self,
|
54
|
+
base_stream: Iterator[_T],
|
55
|
+
usage_callback: Optional[Callable[[List[_T]], None]] = None,
|
56
|
+
):
|
57
|
+
self._base_stream = base_stream
|
58
|
+
self._usage_callback = usage_callback
|
59
|
+
self._chunks: List[_T] = []
|
60
|
+
|
61
|
+
def __iter__(self) -> Iterator[_T]:
|
62
|
+
"""Return self as iterator."""
|
63
|
+
return self
|
64
|
+
|
65
|
+
def __next__(self) -> _T:
|
66
|
+
"""Get next chunk and track it."""
|
67
|
+
try:
|
68
|
+
chunk = next(self._base_stream)
|
69
|
+
except StopIteration:
|
70
|
+
# Once the base stream is fully consumed, we can do final usage/logging.
|
71
|
+
if self._usage_callback and self._chunks:
|
72
|
+
self._usage_callback(self._chunks)
|
73
|
+
raise
|
74
|
+
|
75
|
+
# Intercept each chunk
|
76
|
+
self._chunks.append(chunk)
|
77
|
+
return chunk
|
tokenator/usage.py
CHANGED
@@ -28,7 +28,7 @@ class TokenUsageService:
|
|
28
28
|
logger.info("Tokenator is disabled. Database access is unavailable.")
|
29
29
|
self.MODEL_COSTS = self._get_model_costs()
|
30
30
|
except Exception as e:
|
31
|
-
logger.
|
31
|
+
logger.warning(f"Error in __init__: {e}")
|
32
32
|
self.MODEL_COSTS = {}
|
33
33
|
|
34
34
|
def _get_model_costs(self) -> Dict[str, TokenRate]:
|
@@ -54,13 +54,14 @@ class TokenUsageService:
|
|
54
54
|
prompt_audio=info.get("input_cost_per_audio_token"),
|
55
55
|
completion_audio=info.get("output_cost_per_audio_token"),
|
56
56
|
prompt_cached_input=info.get("cache_read_input_token_cost") or 0,
|
57
|
-
prompt_cached_creation=info.get("
|
57
|
+
prompt_cached_creation=info.get("cache_creation_input_token_cost")
|
58
|
+
or 0,
|
58
59
|
)
|
59
60
|
model_costs[model] = rate
|
60
61
|
|
61
62
|
return model_costs
|
62
63
|
except Exception as e:
|
63
|
-
logger.
|
64
|
+
logger.warning(f"Error in _get_model_costs: {e}")
|
64
65
|
return {}
|
65
66
|
|
66
67
|
def _calculate_cost(
|
@@ -96,7 +97,9 @@ class TokenUsageService:
|
|
96
97
|
elif f"{usage.provider}/{usage.model}" in self.MODEL_COSTS:
|
97
98
|
model_key = f"{usage.provider}/{usage.model}"
|
98
99
|
else:
|
99
|
-
matched_keys = [
|
100
|
+
matched_keys = [
|
101
|
+
k for k in self.MODEL_COSTS.keys() if usage.model in k
|
102
|
+
]
|
100
103
|
if matched_keys:
|
101
104
|
model_key = matched_keys[0]
|
102
105
|
logger.warning(
|
@@ -164,7 +167,9 @@ class TokenUsageService:
|
|
164
167
|
)
|
165
168
|
|
166
169
|
prompt_cost = prompt_text_tokens * model_rates.prompt
|
167
|
-
completion_cost =
|
170
|
+
completion_cost = (
|
171
|
+
completion_text_tokens * model_rates.completion
|
172
|
+
)
|
168
173
|
model_cost += prompt_cost + completion_cost
|
169
174
|
|
170
175
|
# Audio token costs
|
@@ -355,7 +360,7 @@ class TokenUsageService:
|
|
355
360
|
},
|
356
361
|
)
|
357
362
|
except Exception as e:
|
358
|
-
logger.
|
363
|
+
logger.warning(f"Error in _calculate_cost: {e}")
|
359
364
|
return TokenUsageReport()
|
360
365
|
|
361
366
|
def _query_usage(
|
@@ -377,7 +382,7 @@ class TokenUsageService:
|
|
377
382
|
)
|
378
383
|
|
379
384
|
if provider:
|
380
|
-
query = query.filter(TokenUsage.provider
|
385
|
+
query = query.filter(TokenUsage.provider.ilike(provider))
|
381
386
|
if model:
|
382
387
|
query = query.filter(TokenUsage.model == model)
|
383
388
|
|
@@ -385,12 +390,12 @@ class TokenUsageService:
|
|
385
390
|
|
386
391
|
return self._calculate_cost(usages, provider or "all")
|
387
392
|
except Exception as e:
|
388
|
-
logger.
|
393
|
+
logger.warning(f"Error querying usage: {e}")
|
389
394
|
return TokenUsageReport()
|
390
395
|
finally:
|
391
396
|
session.close()
|
392
397
|
except Exception as e:
|
393
|
-
logger.
|
398
|
+
logger.warning(f"Unexpected error in _query_usage: {e}")
|
394
399
|
return TokenUsageReport()
|
395
400
|
|
396
401
|
def last_hour(
|
@@ -406,7 +411,7 @@ class TokenUsageService:
|
|
406
411
|
start = end - timedelta(hours=1)
|
407
412
|
return self._query_usage(start, end, provider, model)
|
408
413
|
except Exception as e:
|
409
|
-
logger.
|
414
|
+
logger.warning(f"Error in last_hour: {e}")
|
410
415
|
return TokenUsageReport()
|
411
416
|
|
412
417
|
def last_day(
|
@@ -422,7 +427,7 @@ class TokenUsageService:
|
|
422
427
|
start = end - timedelta(days=1)
|
423
428
|
return self._query_usage(start, end, provider, model)
|
424
429
|
except Exception as e:
|
425
|
-
logger.
|
430
|
+
logger.warning(f"Error in last_day: {e}")
|
426
431
|
return TokenUsageReport()
|
427
432
|
|
428
433
|
def last_week(
|
@@ -438,7 +443,7 @@ class TokenUsageService:
|
|
438
443
|
start = end - timedelta(weeks=1)
|
439
444
|
return self._query_usage(start, end, provider, model)
|
440
445
|
except Exception as e:
|
441
|
-
logger.
|
446
|
+
logger.warning(f"Error in last_week: {e}")
|
442
447
|
return TokenUsageReport()
|
443
448
|
|
444
449
|
def last_month(
|
@@ -454,7 +459,7 @@ class TokenUsageService:
|
|
454
459
|
start = end - timedelta(days=30)
|
455
460
|
return self._query_usage(start, end, provider, model)
|
456
461
|
except Exception as e:
|
457
|
-
logger.
|
462
|
+
logger.warning(f"Error in last_month: {e}")
|
458
463
|
return TokenUsageReport()
|
459
464
|
|
460
465
|
def between(
|
@@ -499,7 +504,7 @@ class TokenUsageService:
|
|
499
504
|
|
500
505
|
return self._query_usage(start, end, provider, model)
|
501
506
|
except Exception as e:
|
502
|
-
logger.
|
507
|
+
logger.warning(f"Error in between: {e}")
|
503
508
|
return TokenUsageReport()
|
504
509
|
|
505
510
|
def for_execution(self, execution_id: str) -> TokenUsageReport:
|
@@ -514,12 +519,12 @@ class TokenUsageService:
|
|
514
519
|
)
|
515
520
|
return self._calculate_cost(query.all())
|
516
521
|
except Exception as e:
|
517
|
-
logger.
|
522
|
+
logger.warning(f"Error querying for_execution: {e}")
|
518
523
|
return TokenUsageReport()
|
519
524
|
finally:
|
520
525
|
session.close()
|
521
526
|
except Exception as e:
|
522
|
-
logger.
|
527
|
+
logger.warning(f"Unexpected error in for_execution: {e}")
|
523
528
|
return TokenUsageReport()
|
524
529
|
|
525
530
|
def last_execution(self) -> TokenUsageReport:
|
@@ -530,18 +535,20 @@ class TokenUsageService:
|
|
530
535
|
session = get_session()()
|
531
536
|
try:
|
532
537
|
query = (
|
533
|
-
session.query(TokenUsage)
|
538
|
+
session.query(TokenUsage)
|
539
|
+
.order_by(TokenUsage.created_at.desc())
|
540
|
+
.first()
|
534
541
|
)
|
535
542
|
if query:
|
536
543
|
return self.for_execution(query.execution_id)
|
537
544
|
return TokenUsageReport()
|
538
545
|
except Exception as e:
|
539
|
-
logger.
|
546
|
+
logger.warning(f"Error querying last_execution: {e}")
|
540
547
|
return TokenUsageReport()
|
541
548
|
finally:
|
542
549
|
session.close()
|
543
550
|
except Exception as e:
|
544
|
-
logger.
|
551
|
+
logger.warning(f"Unexpected error in last_execution: {e}")
|
545
552
|
return TokenUsageReport()
|
546
553
|
|
547
554
|
def all_time(self) -> TokenUsageReport:
|
@@ -549,22 +556,26 @@ class TokenUsageService:
|
|
549
556
|
if not state.is_tokenator_enabled:
|
550
557
|
return TokenUsageReport()
|
551
558
|
|
552
|
-
logger.warning(
|
559
|
+
logger.warning(
|
560
|
+
"Getting cost analysis for all time. This may take a while..."
|
561
|
+
)
|
553
562
|
session = get_session()()
|
554
563
|
try:
|
555
564
|
query = session.query(TokenUsage)
|
556
565
|
return self._calculate_cost(query.all())
|
557
566
|
except Exception as e:
|
558
|
-
logger.
|
567
|
+
logger.warning(f"Error querying all_time usage: {e}")
|
559
568
|
return TokenUsageReport()
|
560
569
|
finally:
|
561
570
|
session.close()
|
562
571
|
except Exception as e:
|
563
|
-
logger.
|
572
|
+
logger.warning(f"Unexpected error in all_time: {e}")
|
564
573
|
return TokenUsageReport()
|
565
574
|
|
566
575
|
def wipe(self):
|
567
|
-
logger.warning(
|
576
|
+
logger.warning(
|
577
|
+
"All your usage data is about to be wiped, are you sure you want to do this? You have 5 seconds to cancel this operation."
|
578
|
+
)
|
568
579
|
for i in range(5, 0, -1):
|
569
580
|
logger.warning(str(i))
|
570
581
|
time.sleep(1)
|
@@ -574,6 +585,6 @@ class TokenUsageService:
|
|
574
585
|
session.commit()
|
575
586
|
logger.warning("All usage data has been deleted.")
|
576
587
|
except Exception as e:
|
577
|
-
logger.
|
588
|
+
logger.warning(f"Error wiping data: {e}")
|
578
589
|
finally:
|
579
590
|
session.close()
|
tokenator/utils.py
CHANGED
@@ -8,19 +8,22 @@ from pathlib import Path
|
|
8
8
|
|
9
9
|
logger = logging.getLogger(__name__)
|
10
10
|
|
11
|
+
|
11
12
|
def is_notebook() -> bool:
|
12
13
|
try:
|
13
|
-
from IPython import get_ipython
|
14
|
+
from IPython import get_ipython # type: ignore
|
15
|
+
|
14
16
|
shell = get_ipython().__class__.__name__
|
15
|
-
if shell ==
|
16
|
-
return True
|
17
|
-
elif shell ==
|
17
|
+
if shell == "ZMQInteractiveShell":
|
18
|
+
return True # Jupyter notebook or qtconsole
|
19
|
+
elif shell == "TerminalInteractiveShell":
|
18
20
|
return False # Terminal running IPython
|
19
21
|
else:
|
20
22
|
return False # Other type (?)
|
21
23
|
except NameError:
|
22
24
|
return False
|
23
25
|
|
26
|
+
|
24
27
|
def is_colab() -> bool:
|
25
28
|
"""Check if running in Google Colab."""
|
26
29
|
try:
|
@@ -1,10 +1,10 @@
|
|
1
1
|
Metadata-Version: 2.3
|
2
2
|
Name: tokenator
|
3
|
-
Version: 0.1
|
4
|
-
Summary: Token usage tracking
|
3
|
+
Version: 0.2.1
|
4
|
+
Summary: Token usage and cost tracking for LLMs
|
5
5
|
License: MIT
|
6
6
|
Author: Ujjwal Maheshwari
|
7
|
-
Author-email:
|
7
|
+
Author-email: ujjwalm29@gmail.com
|
8
8
|
Requires-Python: >=3.9,<4.0
|
9
9
|
Classifier: License :: OSI Approved :: MIT License
|
10
10
|
Classifier: Programming Language :: Python :: 3
|
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
|
|
15
15
|
Classifier: Programming Language :: Python :: 3.13
|
16
16
|
Requires-Dist: alembic (>=1.13.0,<2.0.0)
|
17
17
|
Requires-Dist: anthropic (>=0.43.0,<0.44.0)
|
18
|
+
Requires-Dist: google-genai (>=1.3.0,<2.0.0)
|
18
19
|
Requires-Dist: ipython
|
19
20
|
Requires-Dist: openai (>=1.59.0,<2.0.0)
|
20
21
|
Requires-Dist: requests (>=2.32.3,<3.0.0)
|
@@ -33,6 +34,9 @@ Afraid not, tokenator is here! With tokenator's easy to use functions, you can s
|
|
33
34
|
|
34
35
|
Get started with just 3 lines of code!
|
35
36
|
|
37
|
+
Tokenator supports the official SDKs from openai, anthropic and google-genai(the new one).
|
38
|
+
LLM providers which use the openai SDK like perplexity, deepseek and xAI are also supported.
|
39
|
+
|
36
40
|
## Installation
|
37
41
|
|
38
42
|
```bash
|
@@ -178,6 +182,54 @@ print(usage.last_execution().model_dump_json(indent=4))
|
|
178
182
|
"""
|
179
183
|
```
|
180
184
|
|
185
|
+
### Google (Gemini - through AI studio)
|
186
|
+
|
187
|
+
```python
|
188
|
+
from google import genai
|
189
|
+
from tokenator import tokenator_gemini
|
190
|
+
|
191
|
+
gemini_client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
|
192
|
+
|
193
|
+
# Wrap it with Tokenator
|
194
|
+
client = tokenator_gemini(gemini_client)
|
195
|
+
|
196
|
+
# Use it exactly like the google-genai client
|
197
|
+
response = models.generate_content(
|
198
|
+
model="gemini-2.0-flash",
|
199
|
+
contents="hello how are you",
|
200
|
+
)
|
201
|
+
|
202
|
+
print(response)
|
203
|
+
|
204
|
+
print(usage.last_execution().model_dump_json(indent=4))
|
205
|
+
"""
|
206
|
+
{
|
207
|
+
"total_cost": 0.0001,
|
208
|
+
"total_tokens": 23,
|
209
|
+
"prompt_tokens": 10,
|
210
|
+
"completion_tokens": 13,
|
211
|
+
"providers": [
|
212
|
+
{
|
213
|
+
"total_cost": 0.0001,
|
214
|
+
"total_tokens": 23,
|
215
|
+
"prompt_tokens": 10,
|
216
|
+
"completion_tokens": 13,
|
217
|
+
"provider": "gemini",
|
218
|
+
"models": [
|
219
|
+
{
|
220
|
+
"total_cost": 0.0004,
|
221
|
+
"total_tokens": 79,
|
222
|
+
"prompt_tokens": 52,
|
223
|
+
"completion_tokens": 27,
|
224
|
+
"model": "gemini-2.0-flash"
|
225
|
+
}
|
226
|
+
]
|
227
|
+
}
|
228
|
+
]
|
229
|
+
}
|
230
|
+
"""
|
231
|
+
```
|
232
|
+
|
181
233
|
### xAI
|
182
234
|
|
183
235
|
You can use xAI models through the `openai` SDK and track usage using `provider` parameter in `tokenator`.
|
@@ -226,7 +278,7 @@ client = tokenator_openai(perplexity_client, db_path=temp_db, provider="perplexi
|
|
226
278
|
|
227
279
|
# Use it exactly like the OpenAI client but with perplexity models
|
228
280
|
response = client.chat.completions.create(
|
229
|
-
model="
|
281
|
+
model="sonar",
|
230
282
|
messages=[{"role": "user", "content": "Hello!"}]
|
231
283
|
)
|
232
284
|
|
@@ -1,8 +1,11 @@
|
|
1
|
-
tokenator/__init__.py,sha256=
|
1
|
+
tokenator/__init__.py,sha256=NB2UOm5oDxj4KLabed4PTSGGzkXvEYUSolOo44ei7XQ,559
|
2
2
|
tokenator/anthropic/client_anthropic.py,sha256=2oxTLb5-sPK_KL-OumCjE4wPVI8U_eFyRonn9XjGXJw,7196
|
3
3
|
tokenator/anthropic/stream_interceptors.py,sha256=4VHC_-WkG3Pa10YizmFLrHcbz0Tm2MR_YB5-uohKp5A,5221
|
4
|
-
tokenator/base_wrapper.py,sha256=
|
4
|
+
tokenator/base_wrapper.py,sha256=Qhd7efdasNHyatR95uxIzbKKVgouT3OZ72DJ7ZrHcrQ,5015
|
5
5
|
tokenator/create_migrations.py,sha256=k9IHiGK21dLTA8MYNsuhO0-kUVIcMSViMFYtY4WU2Rw,730
|
6
|
+
tokenator/gemini/__init__.py,sha256=XphFSP33w0j3j7oNn2PSHTwdjnqvAxitQs6qe6URDOY,132
|
7
|
+
tokenator/gemini/client_gemini.py,sha256=9dJxOG2HczLzDvqepekrU2hY89oHVMRtOQNrvrgV6sQ,8090
|
8
|
+
tokenator/gemini/stream_interceptors.py,sha256=i8DEWAsp1MlZ1xVJZGPzkiCjpX9o6oCJgdkL2k2dung,2266
|
6
9
|
tokenator/migrations/env.py,sha256=JoF5MJ4ae0wJW5kdBHuFlG3ZqeCCDvbMcU8fNA_a6hM,1396
|
7
10
|
tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
|
8
11
|
tokenator/migrations/versions/f028b8155fed_adding_detailed_input_and_output_token_.py,sha256=WIZN5HdNRXlRdfpUJpJFaPD4G1s-SgRdTMQl4WDB-hA,2189
|
@@ -13,9 +16,9 @@ tokenator/openai/client_openai.py,sha256=pbdJ-aZPuJs-7OT1VEv0DW36cCYbRAVKhSQEprx
|
|
13
16
|
tokenator/openai/stream_interceptors.py,sha256=ez1MnjRZW_rEalv2SIPAvrU9oMD6OJoD9vht-057fDM,5243
|
14
17
|
tokenator/schemas.py,sha256=kBmShqgpQ3W-ILAP1NuCaFgqFplQM4OH0MmJteLqrwI,2371
|
15
18
|
tokenator/state.py,sha256=xdqDC-rlEA88-VgqQqHnAOXQ5pNTpnHcgOtohDIImPY,262
|
16
|
-
tokenator/usage.py,sha256=
|
17
|
-
tokenator/utils.py,sha256=
|
18
|
-
tokenator-0.1.
|
19
|
-
tokenator-0.1.
|
20
|
-
tokenator-0.1.
|
21
|
-
tokenator-0.1.
|
19
|
+
tokenator/usage.py,sha256=ZxTRBpmQZysC9x6WWNiQ4aJvCMywF7LnRKGglxLDF_U,25237
|
20
|
+
tokenator/utils.py,sha256=sLC3UxnPWxTFoxuQjGROQHT_POcOKJ-32p8-E0B7hwo,2438
|
21
|
+
tokenator-0.2.1.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
|
22
|
+
tokenator-0.2.1.dist-info/METADATA,sha256=eIRt34a286YLbSMu6iNbWznPIUZw3XxamRTBvOG-QHw,7554
|
23
|
+
tokenator-0.2.1.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
|
24
|
+
tokenator-0.2.1.dist-info/RECORD,,
|
File without changes
|