feat: enhance query logic to handle empty content scenarios and add error handling for judge responses

Author: sanjeevafk

api/routers/query.py

@@ -420,7 +420,8 @@ async def close_stream(stream):
                 record_chunk()
                 yield emit("chunk", {"chunk": chunk})
 
-            if (start_timeout or timed_out) and not full_content.strip():
+            no_chunks = chunk_count == 0 and not full_content.strip()
+            if (start_timeout or timed_out or no_chunks) and not full_content.strip():
                 fallback_used = True
                 try:
                     fallback_content = await asyncio.wait_for(

api/services/ensemble.py

@@ -67,13 +67,22 @@ async def judge_responses(topic: str, responses: list[str], mode: str = LEARNING
     telemetry_sink = kwargs.get("telemetry_sink") if isinstance(kwargs.get("telemetry_sink"), dict) else None
 
     model_start = time.perf_counter()
-    result = await create_chat_completion(
-        model="judge",
-        messages=[{"role": "user", "content": prompt}],
-        temperature=0,
-        max_tokens=512,
-        request_id=request_id,
-    )
+    try:
+        result = await create_chat_completion(
+            model="judge",
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0,
+            max_tokens=512,
+            request_id=request_id,
+        )
+    except Exception as exc:
+        logger.warning(
+            "judge_completion_failed",
+            request_id=request_id,
+            user_id_hash=anonymized_user_id,
+            error=str(exc),
+        )
+        return responses[0]
     usage = None
     usage_obj = getattr(result, "usage", None)
     if usage_obj is not None:

api/tests/test_ensemble.py

@@ -44,3 +44,18 @@ async def fake_generate_explanation(_topic, _level, _model, **_kwargs):
 
     with pytest.raises(RuntimeError):
         await ensemble_module.ensemble_generate("topic", "eli5", use_premium=False, mode="learning")
+
+
+@pytest.mark.asyncio
+async def test_ensemble_falls_back_when_judge_fails(monkeypatch):
+    async def fake_generate_explanation(_topic, _level, _model, **_kwargs):
+        return "candidate"
+
+    async def fake_create_chat_completion(*_args, **_kwargs):
+        raise RuntimeError("judge down")
+
+    monkeypatch.setattr(ensemble_module, "generate_explanation", fake_generate_explanation)
+    monkeypatch.setattr(ensemble_module, "create_chat_completion", fake_create_chat_completion)
+
+    result = await ensemble_module.ensemble_generate("topic", "eli5", use_premium=False, mode="learning")
+    assert result == "candidate"