Show native prefill and generation timing in chat tests

test_chat_test.py

@@ -41,6 +41,15 @@ class ChatTestsTests(unittest.TestCase):
    with self.assertRaises(ValueError):self.manager.unload()
  def test_truncated_stream_is_not_success(self):
   self.worker.conn.getresponse.return_value=Reply([b'data: {"choices":[]}\n\n']);self.manager.start(self.request());self.assertEqual(self.finish()['state'],'failed')
+ def test_native_timings_and_usage_from_final_empty_choices(self):
+  event=dict(choices=[],timings=dict(prompt_n=100,prompt_ms=200,prompt_per_second=500,predicted_n=20,predicted_ms=1000,predicted_per_second=20,secret='excluded'),usage=dict(prompt_tokens=150,completion_tokens=20,total_tokens=170))
+  self.worker.conn.getresponse.return_value=Reply([('data: '+json.dumps(event)+'\n\n').encode(),b'data: [DONE]\n\n'])
+  self.manager.start(self.request());job=self.finish()
+  self.assertEqual(job['timings']['prompt_per_second'],500);self.assertEqual(job['usage']['prompt_tokens'],150);self.assertNotIn('secret',job['timings'])
+  body=json.loads(self.worker.conn.request.call_args.args[2]);self.assertTrue(body['stream_options']['include_usage'])
+ def test_missing_and_invalid_timings_are_not_invented(self):
+  self.worker.conn.getresponse.return_value=Reply([b'data: {"choices":[],"timings":{"prompt_n":true,"prompt_ms":-1,"predicted_per_second":"42"},"usage":null}\n\n',b'data: [DONE]\n\n'])
+  self.manager.start(self.request());job=self.finish();self.assertEqual(job['timings'],{});self.assertEqual(job['usage'],{})
  def test_validation(self):
   for data in [dict(self.request(),max_tokens=True),dict(self.request(),messages=[dict(role='user',content='')]),dict(self.request(),profile_id='missing')]:
    with self.assertRaises(ValueError):self.manager.start(data)