LLM hardening: async wrapper, rate limiting, timeouts, history caps, …

app/config.py

@@ -25,6 +25,16 @@ class Settings(BaseSettings):
     # Public webhook URL override (useful when behind proxy / ngrok)
     WEBHOOK_PUBLIC_URL: str = ""
 
+    # LLM / OpenAI settings
+    LLM_MODEL: str = "gpt-3.5-turbo"
+    LLM_RATE_LIMIT_RPS: float = 1.0
+    LLM_MAX_HISTORY_MESSAGES: int = 8
+    LLM_TIMEOUT_SECONDS: int = 10
+    LLM_MAX_TOKENS: int = 200
+
+    # Observability / Sentry
+    SENTRY_DSN: str = ""
+
     class Config:
         env_file = ".env"
 

app/reply_generator.py

@@ -1,8 +1,17 @@
 import os
 import openai
 import asyncio
-from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
+import logging
+from typing import List
 from app.config import settings
+from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
+
+try:
+    from aiolimiter import AsyncLimiter
+except Exception:
+    AsyncLimiter = None
+
+logger = logging.getLogger(__name__)
 
 openai.api_key = settings.OPENAI_API_KEY
 
@@ -12,37 +21,55 @@
     "Keep replies short (1-3 sentences) and actionable."
 )
 
+# Rate limiter (simple per-process limiter). If aiolimiter isn't installed, we'll fall back to no rate-limiting.
+_limiter = AsyncLimiter(max_rate=settings.LLM_RATE_LIMIT_RPS, time_period=1) if AsyncLimiter else None
 
-@retry(wait=wait_exponential(min=1, max=10), stop=stop_after_attempt(3), retry=retry_if_exception_type(Exception))
+
+@retry(wait=wait_exponential(min=1, max=8), stop=stop_after_attempt(3), retry=retry_if_exception_type(Exception))
 def _call_openai_sync(messages, model="gpt-3.5-turbo"):
-    resp = openai.ChatCompletion.create(model=model, messages=messages, max_tokens=200, temperature=0.3)
-    return resp.choices[0].message.content.strip()
+    # Synchronous call to OpenAI ChatCompletion (run in thread) -- wrapped with retries
+    resp = openai.ChatCompletion.create(model=model, messages=messages, max_tokens=settings.LLM_MAX_TOKENS, temperature=0.3)
+    # Compatibility: some SDKs return choices with message.content; adjust as needed
+    try:
+        return resp.choices[0].message.content.strip()
+    except Exception:
+        # fallback for older SDK response shape
+        return resp.choices[0].text.strip()
 
 
-async def generate_reply(incoming_text: str, history):
-    # If OPENAI_API_KEY not set, return fallback
+async def generate_reply(incoming_text: str, history: List[dict]):
+    # Short-circuit if no OpenAI key
     if not settings.OPENAI_API_KEY:
+        logger.info("OPENAI_API_KEY not set; using fallback reply")
         return "Thanks — I received your message. Can you tell me more?"
 
-    # If history is long, compress/summarize the earlier part (simple approach: keep last 8 messages)
-    recent = history[-8:] if history and len(history) > 8 else history
-
-    messages = [
-        {"role": "system", "content": SYSTEM_PROMPT}
-    ]
+    # Cap history length
+    if history and len(history) > settings.LLM_MAX_HISTORY_MESSAGES:
+        history = history[-settings.LLM_MAX_HISTORY_MESSAGES:]
 
-    # Add history as user messages
-    if recent:
-        for item in recent:
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    if history:
+        for item in history:
             text = item.get("text", "")
             messages.append({"role": "user", "content": text})
 
     messages.append({"role": "user", "content": incoming_text})
 
-    # Call OpenAI synchronously inside thread to keep compatibility with sync SDK
-    loop = asyncio.get_event_loop()
+    # Acquire rate limiter token if available
     try:
-        reply = await loop.run_in_executor(None, _call_openai_sync, messages)
-        return reply
+        if _limiter:
+            await _limiter.acquire()
+        # run blocking OpenAI call in a thread with timeout
+        loop = asyncio.get_event_loop()
+        try:
+            raw = await asyncio.wait_for(loop.run_in_executor(None, _call_openai_sync, messages, settings.LLM_MODEL), timeout=settings.LLM_TIMEOUT_SECONDS)
+            return raw
+        except asyncio.TimeoutError:
+            logger.exception("OpenAI call timed out")
+            return "Thanks — I received your message. Can you tell me more?"
     except Exception:
+        logger.exception("LLM generation failed")
         return "Thanks — I received your message. Can you tell me more?"
+    finally:
+        # No explicit release needed for aiolimiter
+        pass

docs/llm_hardening.md

@@ -0,0 +1,20 @@
+LLM hardening notes
+
+This change implements safer OpenAI/LLM usage:
+
+- Rate limiting: a simple per-process rate limiter using aiolimiter (optional). Configure LLM_RATE_LIMIT_RPS in the .env.
+- Timeouts: OpenAI calls are wrapped with asyncio.wait_for and will timeout after LLM_TIMEOUT_SECONDS (default 10s).
+- History caps: only the most recent LLM_MAX_HISTORY_MESSAGES are sent to the model to avoid unbounded prompts.
+- Retries: network/SDK errors are retried with exponential backoff via tenacity.
+- Fallback: if OpenAI is unavailable or the request times out, a short canned reply is used.
+
+Environment variables (defaults set in app/config.py):
+- LLM_MODEL (default: gpt-3.5-turbo)
+- LLM_RATE_LIMIT_RPS (default: 1.0)
+- LLM_MAX_HISTORY_MESSAGES (default: 8)
+- LLM_TIMEOUT_SECONDS (default: 10)
+- LLM_MAX_TOKENS (default: 200)
+
+Install extra requirements:
+  pip install aiolimiter tenacity
+

requirements-llm.txt

@@ -0,0 +1 @@
+aiolimiter

tests/test_reply_generator.py

@@ -0,0 +1,20 @@
+import pytest
+import asyncio
+from unittest.mock import patch, MagicMock
+from app.reply_generator import generate_reply
+
+
+@pytest.mark.asyncio
+async def test_generate_reply_with_mocked_openai(monkeypatch):
+    # Mock the sync OpenAI call used inside reply_generator
+    mocked = MagicMock()
+    class Choice:
+        def __init__(self, text):
+            self.text = text
+            self.message = type('M', (), {'content': text})
+    mocked.return_value = type('R', (), {'choices': [Choice('Mocked reply')]})
+
+    monkeypatch.setattr('app.reply_generator._call_openai_sync', mocked)
+
+    reply = await generate_reply('Hello', [{'text': 'previous message'}])
+    assert 'Mocked reply' in reply