fix: harden public inference endpoints

.env.example

@@ -47,7 +47,7 @@ LLM_MODEL=llama2
 # ---------------------------------------------------------------------------
 # Inference API authentication
 # ---------------------------------------------------------------------------
-# Required to enable POST /infer. Use a unique local secret; do not commit it.
+# Required to enable POST /infer and POST /predict. Use a unique local secret; do not commit it.
 API_BEARER_TOKEN=
 
 # ---------------------------------------------------------------------------

README.md

@@ -213,7 +213,7 @@ All configuration is via environment variables. Copy `.env.example` to `.env` an
 | `AWS_REGION` | `us-east-1` | AWS region used by the Firehose client |
 | `OLLAMA_HOST` | `http://ollama:11434` | Ollama endpoint (optional) |
 | `LLM_MODEL` | `llama2` | LLM model name |
-| `API_BEARER_TOKEN` | *(unset)* | Required shared bearer token for `POST /infer`; the endpoint returns 503 until configured |
+| `API_BEARER_TOKEN` | *(unset)* | Required shared bearer token for `POST /infer` and `POST /predict`; both endpoints return 503 until configured |
 
 **Snowflake (optional):** set `WAREHOUSE_MODE=snowflake` and fill in `SNOWFLAKE_ACCOUNT`, `SNOWFLAKE_USER`, `SNOWFLAKE_PASSWORD`, `SNOWFLAKE_DATABASE`, `SNOWFLAKE_SCHEMA`, `SNOWFLAKE_WAREHOUSE`.
 

api/main.py

@@ -1,12 +1,12 @@
 import hmac
 import os
 
-from fastapi import Depends, FastAPI, HTTPException, Header, Request
+from fastapi import Depends, FastAPI, Header, HTTPException, Request
 from slowapi import Limiter, _rate_limit_exceeded_handler
 from slowapi.util import get_remote_address
 from slowapi.errors import RateLimitExceeded
 from api.inference import run_inference
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 from typing import List, Optional
 
 limiter = Limiter(key_func=get_remote_address)
@@ -19,7 +19,9 @@
 FEATURE_DIMENSIONS = 16
 
 class RequestModel(BaseModel):
-    """Fixed-width feature vector required by the shipped SentinelModel."""
+    """Fixed-width, finite feature vector required by the shipped SentinelModel."""
+
+    model_config = ConfigDict(allow_inf_nan=False)
 
     features: List[float] = Field(
         min_length=FEATURE_DIMENSIONS,
@@ -29,7 +31,9 @@ class RequestModel(BaseModel):
 
 
 class PromptRequest(BaseModel):
-    prompt: str
+    """Bounded text request for the inference stub."""
+
+    prompt: str = Field(min_length=1, max_length=4_000)
 
 
 def _require_bearer(authorization: Optional[str] = Header(default=None)) -> str:
@@ -61,8 +65,15 @@ def health():
 
 
 @app.post("/predict")
-def predict(request: RequestModel):
-    result = run_inference(request.features)
+@limiter.limit("30/minute")
+def predict(
+    request: Request,
+    body: RequestModel,
+    _token: str = Depends(_require_bearer),
+):
+    """Run the model only for authenticated, bounded requests."""
+
+    result = run_inference(body.features)
     return {"prediction": result}
 
 

tests/test_auth.py

@@ -2,8 +2,10 @@ def test_infer_fails_closed_when_bearer_token_is_unconfigured(client, monkeypatc
     monkeypatch.delenv("API_BEARER_TOKEN", raising=False)
 
     response = client.post("/infer", json={"prompt": "Hello"})
+    predict_response = client.post("/predict", json={"features": [0.0] * 16})
 
     assert response.status_code == 503
+    assert predict_response.status_code == 503
 
 
 def test_infer_rejects_missing_and_invalid_bearer_tokens(client, monkeypatch):
@@ -15,3 +17,9 @@ def test_infer_rejects_missing_and_invalid_bearer_tokens(client, monkeypatch):
         json={"prompt": "Hello"},
         headers={"Authorization": "Bearer wrong-token"},
     ).status_code == 401
+    assert client.post("/predict", json={"features": [0.0] * 16}).status_code == 401
+    assert client.post(
+        "/predict",
+        json={"features": [0.0] * 16},
+        headers={"Authorization": "Bearer wrong-token"},
+    ).status_code == 401

tests/test_inference.py

@@ -7,3 +7,15 @@ def test_inference_success(client, monkeypatch):
 
     assert response.status_code == 200
     assert "response" in response.json()
+
+
+def test_inference_rejects_empty_or_oversized_prompts(client, monkeypatch):
+    monkeypatch.setenv("API_BEARER_TOKEN", "test-token")
+    headers = {"Authorization": "Bearer test-token"}
+
+    assert client.post("/infer", json={"prompt": ""}, headers=headers).status_code == 422
+    assert client.post(
+        "/infer",
+        json={"prompt": "x" * 4_001},
+        headers=headers,
+    ).status_code == 422

tests/test_predict.py

@@ -3,15 +3,37 @@
 
 def test_predict_accepts_exact_model_feature_dimension(client, monkeypatch):
     monkeypatch.setattr(main, "run_inference", lambda features: [0.42])
+    monkeypatch.setenv("API_BEARER_TOKEN", "test-token")
 
-    response = client.post("/predict", json={"features": [0.0] * 16})
+    response = client.post(
+        "/predict",
+        json={"features": [0.0] * 16},
+        headers={"Authorization": "Bearer test-token"},
+    )
 
     assert response.status_code == 200
     assert response.json() == {"prediction": [0.42]}
 
 
-def test_predict_rejects_feature_vectors_with_the_wrong_dimension(client):
-    response = client.post("/predict", json={"features": [0.0] * 15})
+def test_predict_rejects_feature_vectors_with_the_wrong_dimension(client, monkeypatch):
+    monkeypatch.setenv("API_BEARER_TOKEN", "test-token")
+    response = client.post(
+        "/predict",
+        json={"features": [0.0] * 15},
+        headers={"Authorization": "Bearer test-token"},
+    )
 
     assert response.status_code == 422
     assert any(error["loc"] == ["body", "features"] for error in response.json()["detail"])
+
+
+def test_predict_rejects_non_finite_features(client, monkeypatch):
+    monkeypatch.setenv("API_BEARER_TOKEN", "test-token")
+
+    response = client.post(
+        "/predict",
+        json={"features": ["NaN"] * 16},
+        headers={"Authorization": "Bearer test-token"},
+    )
+
+    assert response.status_code == 422