fix: report Gemini model overload
This commit is contained in:
@@ -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가 답변을 생성하지 못했습니다.")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user