feat: Add token usage counter

julep-ai · Oct 14, 2024 · 464d538 · 464d538
1 parent ccac99c
commit 464d538
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/agents-api/agents_api/metrics/counters.py b/agents-api/agents_api/metrics/counters.py
@@ -16,7 +16,7 @@ def decor(func: Callable[P, T] | None = None):
         )
 
         @wraps(func)
-        def wrapper(*args: P.args, **kwargs: P.kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             metric.labels(kwargs.get(id_field_name, "not_set")).inc()
             return func(*args, **kwargs)
 

diff --git a/agents-api/agents_api/routers/sessions/chat.py b/agents-api/agents_api/routers/sessions/chat.py
@@ -20,6 +20,7 @@
 from ...models.chat.gather_messages import gather_messages
 from ...models.chat.prepare_chat_context import prepare_chat_context
 from ...models.entry.create_entries import create_entries
+from .metrics import total_tokens_per_user
 from .router import router
 
 
@@ -166,4 +167,8 @@ async def chat(
         choices=[choice.model_dump() for choice in model_response.choices],
     )
 
+    total_tokens_per_user.labels(str(developer.id)).inc(
+        amount=chat_response.usage.total_tokens or 0
+    )
+
     return chat_response