Add API key authentication for fine-tuning service

.gitignore

@@ -14,4 +14,6 @@ postgres_data/
 .DS_Store
 Thumbs.db
 
-.idea/
+.idea/
+
+whitelist.txt

stacks/finetuning-service/app/auth.py

@@ -0,0 +1,62 @@
+"""Authentication dependency for the fine-tuning service."""
+
+import os
+
+import httpx
+from fastapi import Depends, HTTPException
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+
+_bearer = HTTPBearer()
+
+LITELLM_URL = os.environ["LITELLM_URL"]
+LITELLM_MASTER_KEY = os.environ["LITELLM_MASTER_KEY"]
+
+
+def _load_allowed_users() -> list[str]:
+    """Load allowed user IDs from whitelist.txt, one ID per line.
+
+    Returns an empty list if the file does not exist, which permits
+    any valid LiteLLM key.
+    """
+    try:
+        with open("/app/whitelist.txt") as f:
+            return [line.strip() for line in f if line.strip()]
+    except FileNotFoundError:
+        return []
+
+
+async def verify_litellm_key(
+    credentials: HTTPAuthorizationCredentials = Depends(_bearer),
+) -> None:
+    """Verify a Bearer token is a valid LiteLLM API key.
+
+    Calls LiteLLM's /key/info endpoint using the service master key.
+    Returns normally if the key is valid and the associated user is
+    on the allowlist (if one is configured); raises 401 otherwise.
+
+    Args:
+        credentials: The Bearer token extracted from the Authorization
+            header by FastAPI's HTTPBearer scheme.
+
+    Raises:
+        HTTPException: 401 if the token is invalid or the user is not
+            on the allowlist.
+    """
+    async with httpx.AsyncClient() as client:
+        response = await client.get(
+            f"{LITELLM_URL}/key/info",
+            params={"key": credentials.credentials},
+            headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"},
+        )
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=401, detail="Invalid or inactive API key."
+        )
+
+    allowed_users = _load_allowed_users()
+    if allowed_users:
+        user_id = response.json().get("info", {}).get("user_id")
+        if user_id not in allowed_users:
+            raise HTTPException(
+                status_code=401, detail="User not authorised for this service."
+            )

stacks/finetuning-service/app/main.py

@@ -3,8 +3,9 @@
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator
 
-from fastapi import FastAPI
+from fastapi import Depends, FastAPI
 
+from .auth import verify_litellm_key
 from .database import init_db, recover_running_jobs
 from .routes import router
 
@@ -31,7 +32,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     title="Splinter Fine-Tuning Service",
     lifespan=lifespan,
 )
-app.include_router(router, prefix="/v1/fine_tuning")
+app.include_router(
+    router,
+    prefix="/v1/fine_tuning",
+    dependencies=[Depends(verify_litellm_key)],
+)
 
 
 @app.get("/health")

stacks/finetuning-service/docker-compose.yml

@@ -22,15 +22,17 @@ services:
     restart: unless-stopped
     environment:
       FINETUNING_PORT: ${FINETUNING_PORT}
-      LITELLM_URL: "http://host.docker.internal:${LITELLM_PORT}"
+      LITELLM_URL: "http://litellm-proxy:${LITELLM_PORT}"
       LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY}
     volumes:
       - ./config.yaml:/app/config.yaml:ro
+      - ./whitelist.txt:/app/whitelist.txt:ro
       - finetuning_jobs:/data
     ports:
       - "127.0.0.1:${FINETUNING_PORT}:${FINETUNING_PORT}"
-    extra_hosts:
-      - "host.docker.internal:host-gateway"
+    networks:
+      - default
+      - finetuning
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:${FINETUNING_PORT}/health"]
       interval: 30s
@@ -73,3 +75,9 @@ services:
 
 volumes:
   finetuning_jobs:
+
+networks:
+  default:
+  finetuning:
+    external: true
+    name: finetuning_default

stacks/finetuning-service/requirements.api.txt

@@ -1,4 +1,5 @@
 uvicorn
 fastapi
+httpx
 pyyaml
 typing_extensions

stacks/finetuning-service/whitelist.txt.example

@@ -0,0 +1,6 @@
+# Allowed LiteLLM user IDs (one per line).
+# Copy this file to whitelist.txt and add real user IDs.
+# If whitelist.txt is absent, any valid LiteLLM key is accepted.
+
+user_abc123
+user_xyz456

stacks/llm-service/docker-compose.yml

@@ -299,6 +299,7 @@ services:
     networks:
       - default
       - monitoring
+      - finetuning
 
 volumes:
   postgres_data:
@@ -309,4 +310,7 @@ networks:
   default:
   monitoring:
     external: true
-    name: monitoring_default
+    name: monitoring_default
+  finetuning:
+    external: true
+    name: finetuning_default