From 5b0c2434bdb00d45b805cf39e4ba21db695da9d6 Mon Sep 17 00:00:00 2001
From: umsangdon <umsangdon@gmail.com>
Date: Tue, 6 Oct 2026 15:43:34 +0900
Subject: [PATCH] fix: report Gemini model overload

---
 B01_Dashboard/B01_Dashboard_Knowledge_Chat.py | 32 +++++++++++++++----
 resources/tester/test_b01_knowledge_chat.py   | 32 +++++++++++++++++++
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/B01_Dashboard/B01_Dashboard_Knowledge_Chat.py b/B01_Dashboard/B01_Dashboard_Knowledge_Chat.py
index ee0ce55bc..653e64637 100644
--- a/B01_Dashboard/B01_Dashboard_Knowledge_Chat.py
+++ b/B01_Dashboard/B01_Dashboard_Knowledge_Chat.py
@@ -1,5 +1,6 @@
 """지식DB Markdown 검색과 Gemini 설명. 제외 폴더는 읽지 않는다."""
 
+import asyncio
 import os
 import re
 from functools import lru_cache
@@ -150,21 +151,38 @@ async def ask_knowledge_chat(payload: ChatRequest, _session=Depends(verify_sessi
     )
     try:
         async with httpx.AsyncClient(timeout=30) as client:
-            response = await client.post(
-                f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent",
-                headers={"x-goog-api-key": key},
-                json={"contents": [{"parts": [{"text": prompt}]}]},
-            )
+            for attempt in range(3):
+                response = await client.post(
+                    f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent",
+                    headers={"x-goog-api-key": key},
+                    json={"contents": [{"parts": [{"text": prompt}]}]},
+                )
+                if response.status_code != 503 or attempt == 2:
+                    break
+                await asyncio.sleep(attempt + 1)
+            if response.status_code == 503:
+                raise HTTPException(
+                    status_code=503,
+                    detail="Gemini 모델에 요청이 몰려 있습니다. 잠시 후 다시 질문해 주세요.",
+                )
             response.raise_for_status()
             answer = "".join(
                 part.get("text", "")
                 for candidate in response.json().get("candidates", [])
                 for part in candidate.get("content", {}).get("parts", [])
             ).strip()
-    except (httpx.HTTPError, ValueError) as exc:
+    except httpx.HTTPStatusError as exc:
+        detail = {
+            400: "Gemini 요청 형식이 올바르지 않습니다.",
+            401: "Gemini API 키를 확인해 주세요.",
+            403: "Gemini API 키의 사용 권한을 확인해 주세요.",
+            429: "Gemini API 사용 한도에 도달했습니다. 잠시 후 다시 시도해 주세요.",
+        }.get(exc.response.status_code, "Gemini 요청에 실패했습니다. 잠시 후 다시 시도해 주세요.")
+        raise HTTPException(status_code=502, detail=detail) from exc
+    except (httpx.RequestError, ValueError) as exc:
         raise HTTPException(
             status_code=502,
-            detail="Gemini 응답을 받지 못했습니다. API 키와 사용 한도를 확인해 주세요.",
+            detail="Gemini 서버에 연결하지 못했습니다. 잠시 후 다시 시도해 주세요.",
         ) from exc
     if not answer:
         raise HTTPException(status_code=502, detail="Gemini가 답변을 생성하지 못했습니다.")
diff --git a/resources/tester/test_b01_knowledge_chat.py b/resources/tester/test_b01_knowledge_chat.py
index 1860ba666..81ef1c8be 100644
--- a/resources/tester/test_b01_knowledge_chat.py
+++ b/resources/tester/test_b01_knowledge_chat.py
@@ -35,6 +35,36 @@ def test_missing_api_key_is_reported_before_search(monkeypatch):
     assert error.value.status_code == 503
 
 
+def test_busy_gemini_is_retried_and_reported(monkeypatch):
+    monkeypatch.setenv("GEMINI_API_KEY", "test-key")
+    monkeypatch.setattr(chat, "search", lambda _question: [])
+    calls = []
+
+    class Client:
+        def __init__(self, timeout):
+            pass
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *_args):
+            pass
+
+        async def post(self, *_args, **_kwargs):
+            calls.append(1)
+            return type("Response", (), {"status_code": 503})()
+
+    async def no_wait(_seconds):
+        pass
+
+    monkeypatch.setattr(chat.httpx, "AsyncClient", Client)
+    monkeypatch.setattr(chat.asyncio, "sleep", no_wait)
+    with pytest.raises(HTTPException) as error:
+        asyncio.run(chat.ask_knowledge_chat(chat.ChatRequest(question="일위대가가 뭐야"), {}))
+    assert len(calls) == 3
+    assert "요청이 몰려" in error.value.detail
+
+
 def test_chat_searches_all_knowledge_markdown_except_excluded_folders(monkeypatch, tmp_path):
     public = tmp_path / "original" / "법률"
     guide = tmp_path / "technical_info"
@@ -57,6 +87,8 @@ def test_chat_searches_all_knowledge_markdown_except_excluded_folders(monkeypatc
     sent = {}
 
     class Response:
+        status_code = 200
+
         def raise_for_status(self):
             pass