diff --git a/src/aituner/http_client.py b/src/aituner/http_client.py index 8f8f5b6..e840731 100644 --- a/src/aituner/http_client.py +++ b/src/aituner/http_client.py @@ -310,7 +310,7 @@ def stream_chat_completion( return StreamMetrics( ttft_ms=ttft_ms, tpot_ms=tpot_ms, - completion_tokens=used_tokens if used_tokens > 0 else None, + completion_tokens=used_tokens if used_tokens is not None and used_tokens > 0 else None, completion_tokens_source=completion_tokens_source, streamed_chunk_count=chunk_token_count, ) diff --git a/src/aituner/worker.py b/src/aituner/worker.py index ae9f8bf..4de67ec 100644 --- a/src/aituner/worker.py +++ b/src/aituner/worker.py @@ -8,6 +8,7 @@ import statistics import subprocess import threading import time +import traceback from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait from dataclasses import dataclass from pathlib import Path @@ -596,6 +597,7 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]: "best_request_count": None, "failure_stage": failure_stage, "failure_reason": str(exc), + "failure_traceback": traceback.format_exc(), "probes": probe_history, } StudyStore.write_json(Path(trial.result_path), result) diff --git a/tests/test_core_flow.py b/tests/test_core_flow.py index 813a361..528eea8 100644 --- a/tests/test_core_flow.py +++ b/tests/test_core_flow.py @@ -11,7 +11,13 @@ from unittest import mock from aituner.cli import main as cli_main from aituner.compare import _aggregate_summary, load_compare_spec, run_compare from aituner.engine import build_launch_recipe -from aituner.http_client import StreamMetrics, _auth_headers, _openai_url, _should_bypass_proxy +from aituner.http_client import ( + StreamMetrics, + _auth_headers, + _openai_url, + _should_bypass_proxy, + stream_chat_completion, +) from aituner.job import append_job, build_trial_job from aituner.harness import ( build_harness_context, @@ -3767,6 +3773,29 @@ class CoreFlowTests(unittest.TestCase): "http://example.com/v1/chat/completions", ) + def test_stream_chat_completion_handles_missing_usage_and_chunks(self) -> None: + class FakeResponse: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback): + return False + + def __iter__(self): + return iter([b"data: {\"choices\": []}\n", b"data: [DONE]\n"]) + + with mock.patch("aituner.http_client._urlopen", return_value=FakeResponse()): + metrics = stream_chat_completion( + base_url="http://127.0.0.1:8000", + body={"model": "m", "messages": [{"role": "user", "content": "x"}]}, + timeout_s=1.0, + ) + + self.assertIsNone(metrics.ttft_ms) + self.assertIsNone(metrics.tpot_ms) + self.assertIsNone(metrics.completion_tokens) + self.assertEqual(metrics.completion_tokens_source, "none") + def test_loopback_urls_bypass_proxy(self) -> None: self.assertTrue(_should_bypass_proxy("http://127.0.0.1:8000/v1/models")) self.assertTrue(_should_bypass_proxy("http://localhost:8000/health"))