From 3c778081ee128694f8c76728fb62c2e9926c2fd9 Mon Sep 17 00:00:00 2001
From: Dmitry Paramonov <asmatic075@gmail.com>
Date: Mon, 14 Oct 2024 14:45:03 +0300
Subject: [PATCH] feat: Add token usage counter

---
 agents-api/agents_api/metrics/counters.py         | 2 +-
 agents-api/agents_api/routers/sessions/chat.py    | 5 +++++
 agents-api/agents_api/routers/sessions/metrics.py | 7 +++++++
 3 files changed, 13 insertions(+), 1 deletion(-)
 create mode 100644 agents-api/agents_api/routers/sessions/metrics.py

diff --git a/agents-api/agents_api/metrics/counters.py b/agents-api/agents_api/metrics/counters.py
index 6c2a53272..0c8b1c180 100644
--- 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
index 1963db316..25550a517 100644
--- 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
diff --git a/agents-api/agents_api/routers/sessions/metrics.py b/agents-api/agents_api/routers/sessions/metrics.py
new file mode 100644
index 000000000..b38937313
--- /dev/null
+++ b/agents-api/agents_api/routers/sessions/metrics.py
@@ -0,0 +1,7 @@
+from prometheus_client import Counter
+
+total_tokens_per_user = Counter(
+    "total_tokens_per_user",
+    "Total token count per user",
+    labels=("developer_id",),
+)