Track token in/out per prompt on PromptMetric #18
@@ -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):
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
),
|
||||
]
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user