diff --git a/.gitignore b/.gitignore index 7bdb2c4..040754b 100755 --- a/.gitignore +++ b/.gitignore @@ -14,4 +14,6 @@ postgres_data/ .DS_Store Thumbs.db -.idea/ \ No newline at end of file +.idea/ + +whitelist.txt \ No newline at end of file diff --git a/stacks/finetuning-service/app/auth.py b/stacks/finetuning-service/app/auth.py new file mode 100644 index 0000000..190c9f0 --- /dev/null +++ b/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." + ) diff --git a/stacks/finetuning-service/app/main.py b/stacks/finetuning-service/app/main.py index 902072c..ad05ef5 100644 --- a/stacks/finetuning-service/app/main.py +++ b/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") diff --git a/stacks/finetuning-service/docker-compose.yml b/stacks/finetuning-service/docker-compose.yml index 01e23e6..0a567cc 100644 --- a/stacks/finetuning-service/docker-compose.yml +++ b/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 diff --git a/stacks/finetuning-service/requirements.api.txt b/stacks/finetuning-service/requirements.api.txt index 36b7287..d0f2079 100644 --- a/stacks/finetuning-service/requirements.api.txt +++ b/stacks/finetuning-service/requirements.api.txt @@ -1,4 +1,5 @@ uvicorn fastapi +httpx pyyaml typing_extensions diff --git a/stacks/finetuning-service/whitelist.txt.example b/stacks/finetuning-service/whitelist.txt.example new file mode 100644 index 0000000..70a47ca --- /dev/null +++ b/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 diff --git a/stacks/llm-service/docker-compose.yml b/stacks/llm-service/docker-compose.yml index 5adb749..ebb6a7e 100755 --- a/stacks/llm-service/docker-compose.yml +++ b/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 \ No newline at end of file + name: monitoring_default + finetuning: + external: true + name: finetuning_default \ No newline at end of file