112 lines
4.4 KiB
Python
112 lines
4.4 KiB
Python
"""Metrics: usage forwarding from the LLM stream and first-audio-chunk (TTS TTFB) timing."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
os.environ.setdefault("CONFIG_DIR", tempfile.mkdtemp(prefix="ttvc-metrics-test-config-"))
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
import routes.conversation as conversation
|
|
|
|
|
|
class _Resp:
|
|
def __init__(self, status=200, lines=()):
|
|
self.status_code, self._lines = status, lines
|
|
|
|
def raise_for_status(self):
|
|
pass
|
|
|
|
def iter_lines(self):
|
|
return iter(self._lines)
|
|
|
|
|
|
def _chunk(content=None, usage=None):
|
|
o = {"choices": [{"delta": {"content": content}}] if content is not None else []}
|
|
if usage is not None:
|
|
o["usage"] = usage
|
|
return b"data: " + json.dumps(o).encode()
|
|
|
|
|
|
class UsageParsing(unittest.TestCase):
|
|
def test_numbers_only(self):
|
|
self.assertEqual(conversation._usage_from_chunk({"usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}}),
|
|
{"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12})
|
|
|
|
def test_absent_or_junk_is_none(self):
|
|
for o in ({}, {"usage": None}, {"usage": {}}, {"usage": {"prompt_tokens": "x", "total_tokens": True}}):
|
|
self.assertIsNone(conversation._usage_from_chunk(o))
|
|
|
|
|
|
class RetryWithoutOptionalKeys(unittest.TestCase):
|
|
def test_400_drops_stream_options_and_thinking(self):
|
|
sent = []
|
|
resps = [_Resp(400), _Resp(200)]
|
|
|
|
def fake(u, json=None, **kw):
|
|
sent.append(dict(json))
|
|
return resps.pop(0)
|
|
|
|
payload = {"stream": True, "stream_options": {"include_usage": True}, "chat_template_kwargs": {}, "think": False}
|
|
with mock.patch.object(conversation.requests, "post", side_effect=fake):
|
|
r = conversation._post_llm_stream("http://h/v1/chat/completions", payload, {})
|
|
self.assertEqual(r.status_code, 200)
|
|
self.assertIn("stream_options", sent[0])
|
|
self.assertEqual(sent[1], {"stream": True})
|
|
|
|
def test_no_retry_when_nothing_optional(self):
|
|
with mock.patch.object(conversation.requests, "post", return_value=_Resp(400)) as p:
|
|
self.assertEqual(conversation._post_llm_stream("u", {"stream": True}, {}).status_code, 400)
|
|
self.assertEqual(p.call_count, 1)
|
|
|
|
|
|
class TurnStream(unittest.TestCase):
|
|
def _turn(self, lines_fn, tts_delay=0.0):
|
|
app = FastAPI()
|
|
app.include_router(conversation.router)
|
|
sent = []
|
|
|
|
def fake_post(u, json=None, **kw):
|
|
sent.append(json)
|
|
return _Resp(lines=lines_fn())
|
|
|
|
def fake_tts(*a, **k):
|
|
time.sleep(tts_delay)
|
|
return b"RIFF", "audio/wav"
|
|
|
|
with mock.patch.object(conversation, "_load_settings", return_value={"llm_url": "http://h:17861/v1"}), \
|
|
mock.patch.object(conversation.requests, "post", side_effect=fake_post), \
|
|
mock.patch.object(conversation, "_preview_request_audio", side_effect=fake_tts):
|
|
r = TestClient(app).post("/api/conversation/turn", data={"text": "Tell me about the sea please"})
|
|
evts = [json.loads(l[5:]) for l in r.text.splitlines() if l.startswith("data:")]
|
|
return sent, next(e for e in evts if e["type"] == "stats")
|
|
|
|
def test_usage_forwarded_and_requested(self):
|
|
sent, stats = self._turn(lambda: [_chunk("The sea is wide and deep. "), _chunk(usage={"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13}), b"data: [DONE]"])
|
|
self.assertEqual(sent[0]["stream_options"], {"include_usage": True})
|
|
self.assertEqual(stats["usage"], {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13})
|
|
|
|
def test_no_usage_is_null_not_zero(self):
|
|
_, stats = self._turn(lambda: [_chunk("The sea is wide and deep. "), b"data: [DONE]"])
|
|
self.assertIsNone(stats["usage"])
|
|
|
|
def test_tts_ms_is_first_chunk_not_llm_duration(self):
|
|
def lines():
|
|
yield _chunk("The sea is wide and deep. ")
|
|
time.sleep(0.6) # LLM keeps streaming long after the first sentence's audio is ready
|
|
yield _chunk("It is also very blue today. ")
|
|
yield b"data: [DONE]"
|
|
_, stats = self._turn(lines)
|
|
self.assertGreaterEqual(stats["llm_total_ms"], 600)
|
|
self.assertLess(stats["tts_ms"], 400)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|