Track token in/out per prompt on PromptMetric

Closes #15

Add nullable tokens_in / tokens_out IntegerFields to PromptMetric so we
record real prompt/completion token counts per turn. Values come from the
LLM provider usage payload via a new extract_token_usage helper; when a
provider does not report usage the fields stay null (never estimated).

- consumers/consumers_graph create_prompt_metric + finish_prompt_metric
  now accept and persist optional tokens_in / tokens_out.
- Admin: PromptMetric lists tokens_in/tokens_out (+ event/model filters);
  Conversation admin shows summed tokens_in/out/total per conversation.
- Tests cover the model fields, extract_token_usage parsing, metric
  lifecycle token persistence, and admin aggregation.

Follow-ups: #16 (frontend token display) and #17 (billing/quotas).
Co-authored-by: Cursor <cursoragent@cursor.com>

llm_be/chat_backend/admin.py

@@ -1,4 +1,5 @@
 from django.contrib import admin
+from django.db.models import Sum
 from .models import (
     CustomUser,
     Announcement,
@@ -60,10 +61,35 @@ class PromptInline(admin.TabularInline):
 
 class ConversationAdmin(admin.ModelAdmin):
     model = Conversation
-    list_display = ("title", "get_user_email", "deleted")
+    list_display = (
+        "title",
+        "get_user_email",
+        "deleted",
+        "tokens_in",
+        "tokens_out",
+        "tokens_total",
+    )
     search_fields = ("title",)
     inlines = [PromptInline,]
 
+    def _token_sum(self, conversation, field):
+        total = PromptMetric.objects.filter(
+            conversation_id=conversation.id
+        ).aggregate(total=Sum(field))["total"]
+        return total or 0
+
+    @admin.display(description="Tokens in")
+    def tokens_in(self, conversation):
+        return self._token_sum(conversation, "tokens_in")
+
+    @admin.display(description="Tokens out")
+    def tokens_out(self, conversation):
+        return self._token_sum(conversation, "tokens_out")
+
+    @admin.display(description="Tokens total")
+    def tokens_total(self, conversation):
+        return self.tokens_in(conversation) + self.tokens_out(conversation)
+
 
 class PromptAdmin(admin.ModelAdmin):
     model = Prompt
@@ -79,11 +105,14 @@ class PromptMetricAdmin(admin.ModelAdmin):
         "model_name",
         "prompt_length",
         "reponse_length",
+        "tokens_in",
+        "tokens_out",
         "has_file",
         "file_type",
         "get_duration",
         "created"
     )
+    list_filter = ("event", "model_name", "has_file")
 
 
 class DocumentWorkspaceAdmin(admin.ModelAdmin):

llm_be/chat_backend/consumers.py

@@ -141,12 +141,13 @@ def save_generated_message(conversation_id, message):
 
 @database_sync_to_async
 def create_prompt_metric(
-    prompt_id, prompt, has_file, file_type, model_name, conversation_id
+    prompt_id, prompt, has_file, file_type, model_name, conversation_id, tokens_in=None
 ):
     prompt_metric = PromptMetric.objects.create(
         prompt_id=prompt_id,
         start_time=timezone.now(),
         prompt_length=len(prompt),
+        tokens_in=tokens_in,
         has_file=has_file,
         file_type=file_type,
         model_name=model_name,
@@ -163,12 +164,19 @@ def update_prompt_metric(prompt_metric, status):
 
 
 @database_sync_to_async
-def finish_prompt_metric(prompt_metric, response_length):
+def finish_prompt_metric(prompt_metric, response_length, tokens_in=None, tokens_out=None):
     logger.info(f"finish_prompt_metric: {response_length}")
     prompt_metric.end_time = timezone.now()
     prompt_metric.reponse_length = response_length
     prompt_metric.event = "FINISHED"
-    prompt_metric.save(update_fields=["end_time", "reponse_length", "event"])
+    update_fields = ["end_time", "reponse_length", "event"]
+    if tokens_in is not None:
+        prompt_metric.tokens_in = tokens_in
+        update_fields.append("tokens_in")
+    if tokens_out is not None:
+        prompt_metric.tokens_out = tokens_out
+        update_fields.append("tokens_out")
+    prompt_metric.save(update_fields=update_fields)
     logger.info("finish_prompt_metric saved")
 
 

llm_be/chat_backend/consumers_graph.py

@@ -114,11 +114,12 @@ def save_generated_message(conversation_id, message):
         print(serializer.errors)
 
 @database_sync_to_async
-def create_prompt_metric(prompt_id, prompt, has_file, file_type, model_name, conversation_id):
+def create_prompt_metric(prompt_id, prompt, has_file, file_type, model_name, conversation_id, tokens_in=None):
     prompt_metric = PromptMetric.objects.create(
         prompt_id=prompt_id,
         start_time=timezone.now(),
         prompt_length=len(prompt),
+        tokens_in=tokens_in,
         has_file=has_file,
         file_type=file_type,
         model_name=model_name,
@@ -127,11 +128,18 @@ def create_prompt_metric(prompt_id, prompt, has_file, file_type, model_name, con
     return prompt_metric
 
 @database_sync_to_async
-def finish_prompt_metric(prompt_metric, response_length):
+def finish_prompt_metric(prompt_metric, response_length, tokens_in=None, tokens_out=None):
     prompt_metric.end_time = timezone.now()
     prompt_metric.reponse_length = response_length
     prompt_metric.event = "FINISHED"
-    prompt_metric.save(update_fields=["end_time", "reponse_length", "event"])
+    update_fields = ["end_time", "reponse_length", "event"]
+    if tokens_in is not None:
+        prompt_metric.tokens_in = tokens_in
+        update_fields.append("tokens_in")
+    if tokens_out is not None:
+        prompt_metric.tokens_out = tokens_out
+        update_fields.append("tokens_out")
+    prompt_metric.save(update_fields=update_fields)
 
 async def get_conversation_file_async(conversation_id):
     try:

llm_be/chat_backend/migrations/0023_promptmetric_tokens_in_promptmetric_tokens_out.py

@@ -0,0 +1,31 @@
+# Generated by Django 6.0 on 2026-07-26 12:58
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("chat_backend", "0022_db_file_storage"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="promptmetric",
+            name="tokens_in",
+            field=models.IntegerField(
+                blank=True,
+                help_text="Prompt/input tokens reported by the LLM provider usage payload. Null when the provider did not report usage (never estimated).",
+                null=True,
+            ),
+        ),
+        migrations.AddField(
+            model_name="promptmetric",
+            name="tokens_out",
+            field=models.IntegerField(
+                blank=True,
+                help_text="Completion/output tokens reported by the LLM provider usage payload. Null when the provider did not report usage (never estimated).",
+                null=True,
+            ),
+        ),
+    ]

llm_be/chat_backend/models.py

@@ -216,6 +216,22 @@ class PromptMetric(TimeInfoBase):
     reponse_length = models.IntegerField(
         blank=True, null=True, help_text="How many characters are in the response"
     )
+    tokens_in = models.IntegerField(
+        blank=True,
+        null=True,
+        help_text=(
+            "Prompt/input tokens reported by the LLM provider usage payload. "
+            "Null when the provider did not report usage (never estimated)."
+        ),
+    )
+    tokens_out = models.IntegerField(
+        blank=True,
+        null=True,
+        help_text=(
+            "Completion/output tokens reported by the LLM provider usage payload. "
+            "Null when the provider did not report usage (never estimated)."
+        ),
+    )
     has_file = models.BooleanField(help_text="Is there a file")
     file_type = models.CharField(
         max_length=16, help_text="The file type, if any", blank=True, null=True

llm_be/chat_backend/tests/test_admin.py

@@ -0,0 +1,59 @@
+from django.contrib.admin.sites import AdminSite
+from django.test import TestCase
+from django.utils import timezone
+
+from chat_backend.admin import ConversationAdmin
+from chat_backend.models import Conversation, PromptMetric
+
+from .factories import make_conversation
+
+
+class ConversationAdminTokenTotalsTestCase(TestCase):
+    def setUp(self):
+        self.site = AdminSite()
+        self.admin = ConversationAdmin(Conversation, self.site)
+        self.conversation = make_conversation(title="Token usage")
+
+    def _metric(self, **kwargs):
+        defaults = {
+            "prompt_id": 1,
+            "conversation_id": self.conversation.id,
+            "model_name": "llama3.2",
+            "start_time": timezone.now(),
+            "prompt_length": 10,
+            "has_file": False,
+        }
+        defaults.update(kwargs)
+        return PromptMetric.objects.create(**defaults)
+
+    def test_totals_sum_across_metrics(self):
+        self._metric(tokens_in=10, tokens_out=20)
+        self._metric(tokens_in=5, tokens_out=7)
+
+        self.assertEqual(self.admin.tokens_in(self.conversation), 15)
+        self.assertEqual(self.admin.tokens_out(self.conversation), 27)
+        self.assertEqual(self.admin.tokens_total(self.conversation), 42)
+
+    def test_totals_default_to_zero_without_usage(self):
+        self._metric()
+
+        self.assertEqual(self.admin.tokens_in(self.conversation), 0)
+        self.assertEqual(self.admin.tokens_out(self.conversation), 0)
+        self.assertEqual(self.admin.tokens_total(self.conversation), 0)
+
+    def test_totals_ignore_other_conversations(self):
+        other = make_conversation(title="Other")
+        self._metric(tokens_in=100, tokens_out=200)
+        PromptMetric.objects.create(
+            prompt_id=2,
+            conversation_id=other.id,
+            model_name="llama3.2",
+            start_time=timezone.now(),
+            prompt_length=10,
+            has_file=False,
+            tokens_in=999,
+            tokens_out=999,
+        )
+
+        self.assertEqual(self.admin.tokens_in(self.conversation), 100)
+        self.assertEqual(self.admin.tokens_out(self.conversation), 200)

llm_be/chat_backend/tests/test_consumers.py

@@ -132,6 +132,7 @@ class DatabaseHelperTestCase(TransactionTestCase):
 
         self.assertEqual(metric.prompt_length, len("what is the weather"))
         self.assertEqual(metric.event, "CREATED")
+        self.assertIsNone(metric.tokens_in)
 
         await module.finish_prompt_metric(metric, 120)
 
@@ -139,6 +140,46 @@ class DatabaseHelperTestCase(TransactionTestCase):
         self.assertEqual(refreshed.event, "FINISHED")
         self.assertEqual(refreshed.reponse_length, 120)
         self.assertIsNotNone(refreshed.end_time)
+        self.assertIsNone(refreshed.tokens_out)
+
+    @parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
+    async def test_prompt_metric_records_token_usage(self, _name, module):
+        metric = await module.create_prompt_metric(
+            prompt_id=8,
+            prompt="what is the weather",
+            has_file=False,
+            file_type="",
+            model_name="llama3.2",
+            conversation_id=self.conversation.id,
+            tokens_in=15,
+        )
+
+        self.assertEqual(metric.tokens_in, 15)
+
+        await module.finish_prompt_metric(metric, 120, tokens_out=42)
+
+        refreshed = await sync_to_async(PromptMetric.objects.get)(id=metric.id)
+        self.assertEqual(refreshed.tokens_in, 15)
+        self.assertEqual(refreshed.tokens_out, 42)
+
+    @parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
+    async def test_finish_prompt_metric_leaves_tokens_null_when_absent(
+        self, _name, module
+    ):
+        metric = await module.create_prompt_metric(
+            prompt_id=9,
+            prompt="what is the weather",
+            has_file=False,
+            file_type="",
+            model_name="llama3.2",
+            conversation_id=self.conversation.id,
+        )
+
+        await module.finish_prompt_metric(metric, 120)
+
+        refreshed = await sync_to_async(PromptMetric.objects.get)(id=metric.id)
+        self.assertIsNone(refreshed.tokens_in)
+        self.assertIsNone(refreshed.tokens_out)
 
 
 class GraphNodeTestCase(TransactionTestCase):

llm_be/chat_backend/tests/test_models.py

@@ -144,6 +144,17 @@ class PromptMetricTestCase(TestCase):
     def test_default_event_is_created(self):
         self.assertEqual(self._metric().event, "CREATED")
 
+    def test_tokens_default_to_none(self):
+        metric = self._metric()
+        self.assertIsNone(metric.tokens_in)
+        self.assertIsNone(metric.tokens_out)
+
+    def test_tokens_persist_when_provided(self):
+        metric = self._metric(tokens_in=42, tokens_out=108)
+        metric.refresh_from_db()
+        self.assertEqual(metric.tokens_in, 42)
+        self.assertEqual(metric.tokens_out, 108)
+
 
 class FeedbackTestCase(TestCase):
     def test_get_user_email(self):

llm_be/chat_backend/tests/test_utils.py

@@ -11,7 +11,53 @@ from chat_backend.ollama_config import (
     ollama_llm_kwargs,
     ollama_model,
 )
-from chat_backend.utils import last_day_of_month
+from chat_backend.utils import extract_token_usage, last_day_of_month
+
+
+class ExtractTokenUsageTestCase(SimpleTestCase):
+    def test_returns_none_none_when_no_source(self):
+        self.assertEqual(extract_token_usage(None), (None, None))
+
+    def test_returns_none_none_when_usage_absent(self):
+        self.assertEqual(extract_token_usage({"content": "hi"}), (None, None))
+
+    def test_langchain_usage_metadata_keys(self):
+        self.assertEqual(
+            extract_token_usage({"input_tokens": 12, "output_tokens": 34}),
+            (12, 34),
+        )
+
+    def test_openai_style_keys(self):
+        usage = {"usage": {"prompt_tokens": 5, "completion_tokens": 7}}
+        self.assertEqual(extract_token_usage(usage), (5, 7))
+
+    def test_ollama_style_keys(self):
+        usage = {"prompt_eval_count": 40, "eval_count": 100}
+        self.assertEqual(extract_token_usage(usage), (40, 100))
+
+    def test_reads_usage_metadata_attribute(self):
+        class Message:
+            usage_metadata = {"input_tokens": 3, "output_tokens": 9}
+
+        self.assertEqual(extract_token_usage(Message()), (3, 9))
+
+    def test_reads_nested_response_metadata_attribute(self):
+        class Message:
+            response_metadata = {"token_usage": {"prompt_tokens": 8}}
+
+        self.assertEqual(extract_token_usage(Message()), (8, None))
+
+    def test_ignores_boolean_values(self):
+        self.assertEqual(
+            extract_token_usage({"input_tokens": True, "output_tokens": 2}),
+            (None, 2),
+        )
+
+    def test_coerces_integral_floats(self):
+        self.assertEqual(
+            extract_token_usage({"input_tokens": 10.0, "output_tokens": 20.0}),
+            (10, 20),
+        )
 
 
 class LastDayOfMonthTestCase(SimpleTestCase):

llm_be/chat_backend/utils.py

@@ -6,3 +6,56 @@ def last_day_of_month(any_day):
     next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
     # subtracting the number of the current day brings us back one month
     return next_month - datetime.timedelta(days=next_month.day)
+
+
+# Keys different providers use for input/output token counts. We only ever read
+# real usage the provider reports; we never estimate, so absence maps to None.
+_TOKENS_IN_KEYS = ("input_tokens", "prompt_tokens", "prompt_eval_count")
+_TOKENS_OUT_KEYS = ("output_tokens", "completion_tokens", "eval_count")
+
+
+def _first_int(mapping, keys):
+    for key in keys:
+        value = mapping.get(key)
+        if isinstance(value, bool):
+            continue
+        if isinstance(value, int):
+            return value
+        if isinstance(value, float) and value.is_integer():
+            return int(value)
+    return None
+
+
+def _as_usage_mapping(source):
+    """Best-effort pull of a usage dict out of a provider response.
+
+    Accepts a raw dict, a LangChain message (``usage_metadata`` /
+    ``response_metadata``), or any object exposing those attributes.
+    """
+    if source is None:
+        return None
+    if isinstance(source, dict):
+        for nested_key in ("usage_metadata", "usage", "token_usage"):
+            nested = source.get(nested_key)
+            if isinstance(nested, dict):
+                return nested
+        return source
+    for attr in ("usage_metadata", "response_metadata"):
+        nested = getattr(source, attr, None)
+        if isinstance(nested, dict):
+            mapping = _as_usage_mapping(nested)
+            if mapping:
+                return mapping
+    return None
+
+
+def extract_token_usage(source):
+    """Return ``(tokens_in, tokens_out)`` from a provider usage payload.
+
+    Values are only returned when the provider actually reports them; anything
+    missing comes back as ``None`` so callers never persist estimated counts.
+    """
+    mapping = _as_usage_mapping(source)
+    if not mapping:
+        return None, None
+    return _first_int(mapping, _TOKENS_IN_KEYS), _first_int(mapping, _TOKENS_OUT_KEYS)