diff --git a/src/chatbot_plugin/llm/base_provider.py b/src/chatbot_plugin/llm/base_provider.py index 6add29b..573928e 100644 --- a/src/chatbot_plugin/llm/base_provider.py +++ b/src/chatbot_plugin/llm/base_provider.py @@ -79,3 +79,7 @@ async def generate(self, system_prompt: str, human_prompt: str) -> str | None: except Exception as e: logger.warning("provider_generate_failed", model=self._model, error=str(e)) return None + + async def aclose(self) -> None: + """Clean up async resources. Override in subclasses that create clients.""" + pass diff --git a/src/chatbot_plugin/llm/claude_provider.py b/src/chatbot_plugin/llm/claude_provider.py index f572809..ec1398a 100644 --- a/src/chatbot_plugin/llm/claude_provider.py +++ b/src/chatbot_plugin/llm/claude_provider.py @@ -4,6 +4,7 @@ import structlog from chatbot_plugin.llm.base_provider import BaseProvider +from chatbot_plugin.llm.rate_limit.quota_strategy import RateLimitExhausted logger = structlog.get_logger() @@ -16,12 +17,16 @@ def __init__(self, api_key: str, model: str) -> None: self._client = anthropic.AsyncAnthropic(api_key=api_key) async def _call_api(self, system_prompt: str, human_prompt: str) -> str: - response = await self._client.messages.create( - model=self._model, - max_tokens=2048, - system=system_prompt, - messages=[{"role": "user", "content": human_prompt}], - ) + try: + response = await self._client.messages.create( + model=self._model, + max_tokens=2048, + system=system_prompt, + messages=[{"role": "user", "content": human_prompt}], + ) + except anthropic.RateLimitError as e: + raise RateLimitExhausted(f"Claude rate limit: {e}") from e + logger.info( "claude_api_called", model=self._model, diff --git a/src/chatbot_plugin/llm/gemini_provider.py b/src/chatbot_plugin/llm/gemini_provider.py index 00491ac..ca5dea3 100644 --- a/src/chatbot_plugin/llm/gemini_provider.py +++ b/src/chatbot_plugin/llm/gemini_provider.py @@ -47,7 +47,7 @@ def _sync_generate(self, system_prompt: str, human_prompt: str) -> str: # Check for blocked/safety-filtered responses if not response.candidates: logger.warning("gemini_no_candidates", model=self._model) - return "" + return None candidate = response.candidates[0] if hasattr(candidate, "finish_reason") and candidate.finish_reason not in (1, "STOP"): @@ -56,7 +56,7 @@ def _sync_generate(self, system_prompt: str, human_prompt: str) -> str: model=self._model, finish_reason=str(candidate.finish_reason), ) - return "" + return None token_counts = {} if hasattr(response, "usage_metadata") and response.usage_metadata: @@ -66,4 +66,4 @@ def _sync_generate(self, system_prompt: str, human_prompt: str) -> str: } logger.info("gemini_api_called", model=self._model, **token_counts) - return response.text or "" + return response.text or None diff --git a/src/chatbot_plugin/llm/openrouter_provider.py b/src/chatbot_plugin/llm/openrouter_provider.py index 42c533a..f6a94bc 100644 --- a/src/chatbot_plugin/llm/openrouter_provider.py +++ b/src/chatbot_plugin/llm/openrouter_provider.py @@ -4,6 +4,7 @@ import structlog from chatbot_plugin.llm.base_provider import BaseProvider +from chatbot_plugin.llm.rate_limit.quota_strategy import RateLimitExhausted logger = structlog.get_logger() @@ -30,6 +31,8 @@ async def _call_api(self, system_prompt: str, human_prompt: str) -> str: ], }, ) + if response.status_code == 429: + raise RateLimitExhausted(f"OpenRouter rate limit: HTTP 429") response.raise_for_status() data = response.json() @@ -43,3 +46,7 @@ async def _call_api(self, system_prompt: str, human_prompt: str) -> str: output_tokens=usage.get("completion_tokens", 0), ) return content + + async def aclose(self) -> None: + """Close the httpx async client.""" + await self._client.aclose() diff --git a/src/chatbot_plugin/llm/rate_limit/sliding_window_strategy.py b/src/chatbot_plugin/llm/rate_limit/sliding_window_strategy.py index 4a9302e..32bfe9a 100644 --- a/src/chatbot_plugin/llm/rate_limit/sliding_window_strategy.py +++ b/src/chatbot_plugin/llm/rate_limit/sliding_window_strategy.py @@ -3,6 +3,7 @@ import asyncio import time from collections import deque +from datetime import date from chatbot_plugin.llm.rate_limit.quota_strategy import QuotaStrategy, RateLimitExhausted @@ -12,7 +13,7 @@ class SlidingWindowStrategy(QuotaStrategy): Uses asyncio.Lock for thread safety in async contexts. Two 60-second rolling windows track requests and tokens. - A daily counter tracks total requests. + A daily counter tracks total requests, resetting at calendar day boundary. """ def __init__(self, rpm: int, tpm: int, rpd: int) -> None: @@ -23,6 +24,7 @@ def __init__(self, rpm: int, tpm: int, rpd: int) -> None: self._rpm_window: deque[float] = deque() self._tpm_window: deque[tuple[float, int]] = deque() self._daily_count = 0 + self._daily_date: date = date.today() async def acquire(self, estimated_tokens: int = 0) -> None: """Wait until a request slot is available. @@ -49,7 +51,12 @@ async def _compute_wait(self, estimated_tokens: int) -> float: async with self._lock: now = time.monotonic() - # Check daily quota + # Check daily quota (reset counter if new day) + today = date.today() + if today != self._daily_date: + self._daily_count = 0 + self._daily_date = today + if self._rpd > 0 and self._daily_count >= self._rpd: raise RateLimitExhausted( f"Daily quota reached: {self._daily_count}/{self._rpd}" diff --git a/src/chatbot_plugin/service.py b/src/chatbot_plugin/service.py index 068bff0..0a7304e 100644 --- a/src/chatbot_plugin/service.py +++ b/src/chatbot_plugin/service.py @@ -55,7 +55,7 @@ async def chat(self, message: str, user_id: str | None = None) -> ChatMessageRes reply = await rag_generate(message, articles, self.llm_service) except RuntimeError: raise HTTPException(status_code=503, detail="LLM provider unavailable") - except Exception as e: + except (ConnectionError, TimeoutError) as e: raise HTTPException(status_code=503, detail="LLM provider unavailable") from e # 3. Build response with article references diff --git a/src/tests/conftest.py b/src/tests/conftest.py index c8cf470..8386059 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -37,7 +37,9 @@ def app(mock_llm_service: AsyncMock) -> FastAPI: set_llm_service(mock_llm_service) app = FastAPI() app.include_router(chat_router, prefix="/chat", tags=["chat"]) - return app + yield app + # Teardown: reset LLM service to avoid leaking between tests + set_llm_service(None) @pytest.fixture diff --git a/src/tests/llm/test_provider_implementations.py b/src/tests/llm/test_provider_implementations.py index 77ce497..edacdb9 100644 --- a/src/tests/llm/test_provider_implementations.py +++ b/src/tests/llm/test_provider_implementations.py @@ -45,6 +45,22 @@ async def test_call_api_uses_correct_params(self): assert call_kwargs.kwargs["system"] == "system-instr" assert call_kwargs.kwargs["messages"] == [{"role": "user", "content": "user-msg"}] + @pytest.mark.asyncio + async def test_call_api_rate_limit_raises(self): + with patch("chatbot_plugin.llm.claude_provider.anthropic") as mock_anthropic: + import anthropic as real_anthropic + mock_client = AsyncMock() + mock_anthropic.AsyncAnthropic.return_value = mock_client + mock_anthropic.RateLimitError = real_anthropic.RateLimitError + from chatbot_plugin.llm.claude_provider import ClaudeProvider + provider = ClaudeProvider(api_key="sk-test", model="claude-sonnet-4-6-20250514") + mock_client.messages.create.side_effect = real_anthropic.RateLimitError( + message="rate limit", response=MagicMock(), body=None + ) + + with pytest.raises(RateLimitExhausted): + await provider._call_api("sys", "human") + # ── GeminiProvider ── @@ -84,11 +100,11 @@ async def test_call_api_no_candidates_returns_empty(self, mock_genai): mock_client.models.generate_content.return_value = mock_response result = await provider._call_api("sys", "human") - assert result == "" + assert result is None @pytest.mark.asyncio @patch("chatbot_plugin.llm.gemini_provider.genai") - async def test_call_api_blocked_finish_reason_returns_empty(self, mock_genai): + async def test_call_api_blocked_finish_reason_returns_none(self, mock_genai): mock_client = MagicMock() mock_genai.Client.return_value = mock_client mock_genai.GenerateContentConfig = MagicMock() @@ -101,7 +117,7 @@ async def test_call_api_blocked_finish_reason_returns_empty(self, mock_genai): mock_client.models.generate_content.return_value = mock_response result = await provider._call_api("sys", "human") - assert result == "" + assert result is None @pytest.mark.asyncio @patch("chatbot_plugin.llm.gemini_provider.genai") @@ -208,3 +224,29 @@ async def test_call_api_missing_usage_defaults_to_zero(self): result = await provider._call_api("sys", "human") assert result == "ok" + + @pytest.mark.asyncio + async def test_call_api_429_raises_rate_limit_exhausted(self): + with patch("chatbot_plugin.llm.openrouter_provider.httpx") as mock_httpx: + mock_client = AsyncMock() + mock_httpx.AsyncClient.return_value = mock_client + from chatbot_plugin.llm.openrouter_provider import OpenRouterProvider + provider = OpenRouterProvider(api_key="sk-test", model="test-model") + + mock_response = MagicMock() + mock_response.status_code = 429 + mock_client.post.return_value = mock_response + + with pytest.raises(RateLimitExhausted): + await provider._call_api("sys", "human") + + @pytest.mark.asyncio + async def test_aclose_closes_client(self): + with patch("chatbot_plugin.llm.openrouter_provider.httpx") as mock_httpx: + mock_client = AsyncMock() + mock_httpx.AsyncClient.return_value = mock_client + from chatbot_plugin.llm.openrouter_provider import OpenRouterProvider + provider = OpenRouterProvider(api_key="sk-test", model="test-model") + + await provider.aclose() + mock_client.aclose.assert_called_once() diff --git a/src/tests/llm/test_rate_limit.py b/src/tests/llm/test_rate_limit.py index 727db5c..d9c7108 100644 --- a/src/tests/llm/test_rate_limit.py +++ b/src/tests/llm/test_rate_limit.py @@ -100,3 +100,18 @@ async def test_tpm_wait_returns_positive_when_full(self): # TPM is full — compute_wait should return > 0 wait = await strategy._compute_wait(20) assert wait > 0 + + @pytest.mark.asyncio + async def test_daily_counter_resets_on_new_day(self): + """Daily counter resets when the calendar day changes.""" + from datetime import date + strategy = SlidingWindowStrategy(rpm=100, tpm=100000, rpd=2) + await strategy.acquire(10) + await strategy.acquire(10) + # Daily quota exhausted + with pytest.raises(RateLimitExhausted): + await strategy.acquire(10) + # Simulate day change + strategy._daily_date = date(2000, 1, 1) + # Should succeed again after day resets + await strategy.acquire(10) diff --git a/src/tests/test_service.py b/src/tests/test_service.py index 357fefc..df036e9 100644 --- a/src/tests/test_service.py +++ b/src/tests/test_service.py @@ -56,10 +56,10 @@ async def test_chat_llm_failure_raises_503(service, mock_db, mock_llm_service): @pytest.mark.asyncio -async def test_chat_generic_exception_raises_503(service, mock_db, mock_llm_service): - """Non-RuntimeError exceptions from rag_generate also produce 503.""" +async def test_chat_connection_error_raises_503(service, mock_db, mock_llm_service): + """Connection errors from LLM provider produce 503.""" service.retriever.search = AsyncMock(return_value=[]) - mock_llm_service.generate.side_effect = Exception("unexpected error") + mock_llm_service.generate.side_effect = ConnectionError("connection refused") with pytest.raises(HTTPException) as exc_info: await service.chat("hello")