diff --git a/llm_be/chat_backend/admin.py b/llm_be/chat_backend/admin.py index 28fa9c2..8c4a631 100644 --- a/llm_be/chat_backend/admin.py +++ b/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): diff --git a/llm_be/chat_backend/consumers.py b/llm_be/chat_backend/consumers.py index 89d6720..b1d3d59 100644 --- a/llm_be/chat_backend/consumers.py +++ b/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") diff --git a/llm_be/chat_backend/consumers_graph.py b/llm_be/chat_backend/consumers_graph.py index 21148a9..2dafff4 100644 --- a/llm_be/chat_backend/consumers_graph.py +++ b/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: diff --git a/llm_be/chat_backend/migrations/0023_promptmetric_tokens_in_promptmetric_tokens_out.py b/llm_be/chat_backend/migrations/0023_promptmetric_tokens_in_promptmetric_tokens_out.py new file mode 100644 index 0000000..304e5a0 --- /dev/null +++ b/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, + ), + ), + ] diff --git a/llm_be/chat_backend/models.py b/llm_be/chat_backend/models.py index 95ce118..6f7b854 100644 --- a/llm_be/chat_backend/models.py +++ b/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 diff --git a/llm_be/chat_backend/tests/test_admin.py b/llm_be/chat_backend/tests/test_admin.py new file mode 100644 index 0000000..c49e689 --- /dev/null +++ b/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) diff --git a/llm_be/chat_backend/tests/test_consumers.py b/llm_be/chat_backend/tests/test_consumers.py index 5693cc8..f440e8e 100644 --- a/llm_be/chat_backend/tests/test_consumers.py +++ b/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): diff --git a/llm_be/chat_backend/tests/test_models.py b/llm_be/chat_backend/tests/test_models.py index 00ee191..198ffb0 100644 --- a/llm_be/chat_backend/tests/test_models.py +++ b/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): diff --git a/llm_be/chat_backend/tests/test_utils.py b/llm_be/chat_backend/tests/test_utils.py index 9bb91f9..6b3f612 100644 --- a/llm_be/chat_backend/tests/test_utils.py +++ b/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): diff --git a/llm_be/chat_backend/utils.py b/llm_be/chat_backend/utils.py index 2d33454..9972693 100644 --- a/llm_be/chat_backend/utils.py +++ b/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)