diff --git a/.env.example b/.env.example index dd6da70..d4a30f5 100644 --- a/.env.example +++ b/.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= # --------------------------------------------------------------------------- diff --git a/README.md b/README.md index e9de58b..25a371e 100644 --- a/README.md +++ b/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`. diff --git a/api/main.py b/api/main.py index affa456..c3c2352 100644 --- a/api/main.py +++ b/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} diff --git a/tests/test_auth.py b/tests/test_auth.py index 2850bbe..f8e8bf8 100644 --- a/tests/test_auth.py +++ b/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 diff --git a/tests/test_inference.py b/tests/test_inference.py index 85ca306..fa68243 100644 --- a/tests/test_inference.py +++ b/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 diff --git a/tests/test_predict.py b/tests/test_predict.py index c60e8e6..168c78e 100644 --- a/tests/test_predict.py +++ b/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