diff --git a/.gitignore b/.gitignore index 7af86f2..d135cdb 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,9 @@ npm-debug.log* yarn-debug.log* yarn-error.log* +**/logs/** +*.log +*.log.* # Environment variables .env diff --git a/app/api/auth.py b/app/api/auth.py index 301831c..bc0632f 100644 --- a/app/api/auth.py +++ b/app/api/auth.py @@ -4,11 +4,38 @@ from email_validator import validate_email, EmailNotValidError from sqlalchemy.orm import Session import logging +from logging.handlers import TimedRotatingFileHandler import secrets import os from urllib.parse import urlparse +from pathlib import Path + +# Configure auth logger +auth_logger = logging.getLogger("auth") +auth_logger.setLevel(logging.INFO) + +# Create logs directory if it doesn't exist +log_dir = Path("logs") +log_dir.mkdir(exist_ok=True) + +# Configure file handler with daily rotation +file_handler = TimedRotatingFileHandler( + filename=log_dir / "auth.log", + when="midnight", + interval=1, + backupCount=30, # Keep logs for 30 days + encoding="utf-8" +) + +# Configure formatter +formatter = logging.Formatter( + "%(asctime)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S" +) +file_handler.setFormatter(formatter) -logger = logging.getLogger(__name__) +# Add handler to logger +auth_logger.addHandler(file_handler) from app.db.database import get_db from app.schemas.models import ( @@ -156,13 +183,16 @@ async def login( detail="Invalid login data. Please provide username and password in either form data or JSON format." ) + auth_logger.info(f"Login attempt for user: {login_data.username}") user = db.query(DBUser).filter(DBUser.email == login_data.username).first() if not user or not verify_password(login_data.password, user.hashed_password): + auth_logger.warning(f"Failed login attempt for user: {login_data.username}") raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect email or password" ) + auth_logger.info(f"Successful login for user: {login_data.username}") return create_and_set_access_token(response, user.email) @router.post("/logout") @@ -243,7 +273,11 @@ async def update_user_me( return current_user @router.post("/register", response_model=User) -async def register(user: UserCreate, db: Session = Depends(get_db)): +async def register( + request: Request, + user: UserCreate, + db: Session = Depends(get_db) +): """ Register a new user account. @@ -252,9 +286,11 @@ async def register(user: UserCreate, db: Session = Depends(get_db)): After registration, you'll need to login to get an access token. """ + auth_logger.info(f"Registration attempt for user: {user.email}") # Check if user with this email exists db_user = db.query(DBUser).filter(DBUser.email == user.email).first() if db_user: + auth_logger.warning(f"Registration failed - email already exists: {user.email}") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Email already registered" @@ -270,6 +306,7 @@ async def register(user: UserCreate, db: Session = Depends(get_db)): db.add(db_user) db.commit() db.refresh(db_user) + auth_logger.info(f"Successfully registered new user: {user.email}") return db_user @router.post("/validate-email") @@ -304,14 +341,17 @@ async def validate_email( pass if not email: + auth_logger.warning("Email validation attempt with no email provided") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Email is required" ) + auth_logger.info(f"Email validation attempt for: {email}") try: email_validator.validate_email(email, check_deliverability=False) except EmailNotValidError as e: + auth_logger.warning(f"Invalid email format for {email}: {e}") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid email format: {e}" @@ -322,8 +362,10 @@ async def validate_email( user = db.query(DBUser).filter(DBUser.email == email).first() if user: email_template = 'returning-user-code' + auth_logger.info(f"Sending validation code to existing user: {email}") else: email_template = 'new-user-code' + auth_logger.info(f"Sending validation code to new user: {email}") # Send the validation code via email ses_service = SESService() @@ -336,12 +378,13 @@ async def validate_email( ) if not email_sent: - logger.error(f"Failed to send validation code email to {email}") + auth_logger.error(f"Failed to send validation code email to {email}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to send validation code email" ) + auth_logger.info(f"Successfully sent validation code to: {email}") return { "message": "Validation code has been generated and sent" } @@ -437,16 +480,19 @@ async def sign_in( with them as the admin. """ if not sign_in_data: + auth_logger.warning("Sign-in attempt with invalid data format") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid sign in data. Please provide username and verification code in either form data or JSON format." ) + auth_logger.info(f"Sign-in attempt for user: {sign_in_data.username}") # Verify the code using DynamoDB first dynamodb_service = DynamoDBService() stored_code = dynamodb_service.read_validation_code(sign_in_data.username) if not stored_code or stored_code.get('code').upper() != sign_in_data.verification_code.upper(): + auth_logger.warning(f"Invalid verification code for user: {sign_in_data.username}") raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect email or verification code" @@ -457,6 +503,7 @@ async def sign_in( # If user doesn't exist, create a new user and team if not user: + auth_logger.info(f"Creating new user and team for: {sign_in_data.username}") # First create the team team_data = TeamCreate( name=f"Team {sign_in_data.username}", @@ -476,5 +523,7 @@ async def sign_in( db.add(user) db.commit() db.refresh(user) + auth_logger.info(f"Successfully created new user and team for: {sign_in_data.username}") + auth_logger.info(f"Successful sign-in for user: {sign_in_data.username}") return create_and_set_access_token(response, user.email) \ No newline at end of file diff --git a/app/api/private_ai_keys.py b/app/api/private_ai_keys.py index 1a19706..9264da0 100644 --- a/app/api/private_ai_keys.py +++ b/app/api/private_ai_keys.py @@ -556,14 +556,17 @@ async def update_budget_period( # Get updated spend information spend_data = await litellm_service.get_key_info(private_ai_key.litellm_token) + info = spend_data.get("info", {}) + # Only set default for spend field spend_info = { - "spend": spend_data.get("spend", 0.0), - **spend_data + "spend": info.get("spend", 0.0), + **info } return PrivateAIKeySpend.model_validate(spend_info) except Exception as e: + logger.error(f"Failed to update budget period: {str(e)}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to update budget period: {str(e)}" diff --git a/app/core/config.py b/app/core/config.py index b8cc50f..260d849 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -19,6 +19,7 @@ class Settings(BaseSettings): "http://localhost:8800" ] ALLOWED_HOSTS: list[str] = ["*"] # In production, restrict this + PUBLIC_PATHS: list[str] = ["/health", "/docs", "/openapi.json", "/metrics"] model_config = ConfigDict(env_file=".env") diff --git a/app/core/security.py b/app/core/security.py index 3db38e0..952d40d 100644 --- a/app/core/security.py +++ b/app/core/security.py @@ -2,7 +2,7 @@ from typing import Optional, Literal, Dict from jose import JWTError, jwt from passlib.context import CryptContext -from fastapi import Depends, HTTPException, status, Cookie, Header +from fastapi import Depends, HTTPException, status, Cookie, Header, Request from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials import logging from app.core.config import settings @@ -83,9 +83,20 @@ async def get_current_user( async def get_current_user_from_auth( access_token: Optional[str] = Cookie(None, alias="access_token"), authorization: Optional[str] = Header(None), - db: Session = Depends(get_db) + db: Session = Depends(get_db), + request: Request = None ) -> DBUser: """Get current user from either JWT token (in cookie or Authorization header) or API token.""" + # First check if user is already in request state (set by AuthMiddleware) + if request and hasattr(request.state, 'user') and request.state.user is not None: + # If we have a dict from middleware, load the full user object + if isinstance(request.state.user, dict): + user = db.query(DBUser).filter(DBUser.id == request.state.user["id"]).first() + if user: + return user + else: + return request.state.user + if not access_token and not authorization: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, diff --git a/app/main.py b/app/main.py index 85c47d0..c514898 100644 --- a/app/main.py +++ b/app/main.py @@ -1,11 +1,10 @@ -from fastapi import FastAPI, Depends, HTTPException -from fastapi.security import OAuth2PasswordBearer -from sqlalchemy.orm import Session +from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.trustedhost import TrustedHostMiddleware from starlette.middleware.base import BaseHTTPMiddleware from fastapi.openapi.docs import get_swagger_ui_html from fastapi.openapi.utils import get_openapi +from prometheus_fastapi_instrumentator import Instrumentator, metrics import os import logging @@ -27,6 +26,8 @@ async def dispatch(self, request, call_next): from app.core.config import settings from app.db.database import get_db from app.middleware.audit import AuditLogMiddleware +from app.middleware.prometheus import PrometheusMiddleware +from app.middleware.auth import AuthMiddleware app = FastAPI( title="Private AI Keys as a Service", @@ -79,6 +80,12 @@ async def dispatch(self, request, call_next): # Add HTTPS redirect middleware first app.add_middleware(HTTPSRedirectMiddleware) +# Add Auth middleware (must be before Prometheus and Audit middleware) +app.add_middleware(AuthMiddleware) + +# Add Prometheus middleware +app.add_middleware(PrometheusMiddleware) + # Configure CORS app.add_middleware( CORSMiddleware, @@ -96,6 +103,24 @@ async def dispatch(self, request, call_next): app.add_middleware(AuditLogMiddleware, db=next(get_db())) +# Setup Prometheus instrumentation +instrumentator = Instrumentator( + should_group_status_codes=False, + should_ignore_untemplated=True, + should_respect_env_var=True, + should_instrument_requests_inprogress=True, + excluded_handlers=["/metrics"], + env_var_name="ENABLE_METRICS", + inprogress_name="fastapi_inprogress", + inprogress_labels=True, +) + +# Add default metrics +instrumentator.add(metrics.default()) + +# Instrument the app +instrumentator.instrument(app).expose(app) + @app.get("/health") async def health_check(): return {"status": "healthy"} diff --git a/app/middleware/audit.py b/app/middleware/audit.py index 95173a9..8d4d7ac 100644 --- a/app/middleware/audit.py +++ b/app/middleware/audit.py @@ -2,12 +2,11 @@ from starlette.middleware.base import BaseHTTPMiddleware from sqlalchemy.orm import Session from app.db.models import DBAuditLog -from app.api.auth import get_current_user_from_auth from app.db.database import get_db -import json +from app.middleware.prometheus import audit_events_total, audit_event_duration_seconds import logging -from fastapi import Cookie, Header -from typing import Optional +import time +from app.core.config import settings logger = logging.getLogger(__name__) @@ -18,9 +17,11 @@ def __init__(self, app, db: Session): async def dispatch(self, request: Request, call_next): # Skip audit logging for certain paths - if request.url.path in ["/health", "/docs", "/openapi.json", "/audit/logs", "/auth/me"]: + if request.url.path in {*settings.PUBLIC_PATHS, "/audit/logs", "/auth/me"}: return await call_next(request) + start_time = time.time() + # Get the response response = await call_next(request) @@ -28,30 +29,13 @@ async def dispatch(self, request: Request, call_next): # Get a fresh database session for each request db = next(get_db()) - # Try to get the current user from cookies or authorization header + # Get user_id from request state (set by AuthMiddleware) user_id = None - try: - # Get access token from cookie or authorization header - cookies = request.cookies - headers = request.headers - access_token = cookies.get("access_token") - auth_header = headers.get("authorization") - - if auth_header: - parts = auth_header.split() - if len(parts) == 2 and parts[0].lower() == "bearer": - access_token = parts[1] - - if access_token: - user = await get_current_user_from_auth( - access_token=access_token if access_token else None, - authorization=auth_header if auth_header else None, - db=db - ) - user_id = user.id if user else None - except Exception as e: - logger.debug(f"Could not get user for audit log: {str(e)}") - user_id = None + if hasattr(request.state, 'user') and request.state.user: + if isinstance(request.state.user, dict): + user_id = request.state.user.get('id') + else: + user_id = request.state.user.id # Extract path parameters path_params = request.path_params @@ -67,13 +51,16 @@ async def dispatch(self, request: Request, call_next): request_source = "frontend" else: # If no origin/referer and has auth header, likely direct API call - request_source = "api" if auth_header else None + request_source = "api" if request.headers.get("authorization") else None + + # Get resource type from path + resource_type = request.url.path.split("/")[1] # First path segment # Create audit log entry audit_log = DBAuditLog( user_id=user_id, event_type=request.method, - resource_type=request.url.path.split("/")[1], # First path segment + resource_type=resource_type, resource_id=str(resource_id) if resource_id else None, action=f"{request.method} {request.url.path}", details={ @@ -89,6 +76,21 @@ async def dispatch(self, request: Request, call_next): db.add(audit_log) db.commit() + # Record audit metrics + audit_events_total.labels( + event_type=request.method, + resource_type=resource_type, + request_source=request_source or "unknown", + status_code=response.status_code + ).inc() + + # Record audit event duration + duration = time.time() - start_time + audit_event_duration_seconds.labels( + event_type=request.method, + resource_type=resource_type + ).observe(duration) + except Exception as e: logger.error(f"Failed to create audit log: {str(e)}", exc_info=True) # Don't re-raise the exception - we don't want to break the request if audit logging fails diff --git a/app/middleware/auth.py b/app/middleware/auth.py new file mode 100644 index 0000000..21f3046 --- /dev/null +++ b/app/middleware/auth.py @@ -0,0 +1,55 @@ +from fastapi import Request +from starlette.middleware.base import BaseHTTPMiddleware +from app.core.security import get_current_user_from_auth +from app.db.database import get_db +import logging +from app.core.config import settings + +logger = logging.getLogger(__name__) + +class AuthMiddleware(BaseHTTPMiddleware): + async def dispatch(self, request: Request, call_next): + # Skip auth for certain paths + if request.url.path in settings.PUBLIC_PATHS: + return await call_next(request) + + # Initialize user as None + request.state.user = None + + try: + # Get access token from cookie or authorization header + cookies = request.cookies + headers = request.headers + access_token = cookies.get("access_token") + auth_header = headers.get("authorization") + + if auth_header: + parts = auth_header.split() + if len(parts) == 2 and parts[0].lower() == "bearer": + access_token = parts[1] + + if access_token: + # Get a fresh database session + db = next(get_db()) + try: + user = await get_current_user_from_auth( + access_token=access_token if access_token else None, + authorization=auth_header if auth_header else None, + db=db + ) + # Store essential user data instead of the full SQLAlchemy object + request.state.user = { + "id": user.id, + "email": user.email, + "is_admin": user.is_admin, + "role": user.role, + "team_id": user.team_id + } + except Exception as e: + logger.debug(f"Could not get user for request: {str(e)}") + finally: + db.close() + except Exception as e: + logger.debug(f"Error in auth middleware: {str(e)}") + + return await call_next(request) \ No newline at end of file diff --git a/app/middleware/prometheus.py b/app/middleware/prometheus.py new file mode 100644 index 0000000..4d555f6 --- /dev/null +++ b/app/middleware/prometheus.py @@ -0,0 +1,77 @@ +from prometheus_client import Counter, Histogram +from fastapi import Request +from starlette.middleware.base import BaseHTTPMiddleware +import logging +from app.core.config import settings + +logger = logging.getLogger(__name__) + +def normalize_path(path: str) -> str: + """Replace numeric segments in path with {id}.""" + return '/'.join('{id}' if segment.isdigit() else segment for segment in path.split('/')) + +# Audit Metrics +audit_events_total = Counter( + "audit_events_total", + "Total number of audit events", + ["event_type", "resource_type", "request_source", "status_code"] +) + +audit_event_duration_seconds = Histogram( + "audit_event_duration_seconds", + "Audit event processing duration in seconds", + ["event_type", "resource_type"], + buckets=(0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0, float("inf")) +) + +# User Metrics - grouped by user type and endpoint +requests_by_user_type = Counter( + "requests_by_user_type", + "Number of requests grouped by user type", + ["user_type", "endpoint", "method"] +) + +# Auth Metrics - simplified to track success/failure +auth_requests_total = Counter( + "auth_requests_total", + "Total number of authentication requests", + ["endpoint", "status"] # status will be "success" or "failure" +) + +class PrometheusMiddleware(BaseHTTPMiddleware): + async def dispatch(self, request: Request, call_next): + # Skip metrics for certain paths + if request.url.path in settings.PUBLIC_PATHS: + return await call_next(request) + + # Track auth requests for specific endpoints + is_auth_endpoint = request.url.path in [ + "/auth/login", + "/auth/register", + "/auth/validate-email", + "/auth/sign-in" + ] + + response = await call_next(request) + + if is_auth_endpoint: + auth_requests_total.labels( + endpoint=request.url.path, + status="success" + ).inc() + + # Get user type from request state (set by AuthMiddleware) + user_type = "anonymous" + if hasattr(request.state, 'user') and request.state.user: + # Group users by their role or type + user_type = request.state.user.role if hasattr(request.state.user, 'role') else "authenticated" + + # Record requests by user type with normalized path + normalized_path = normalize_path(request.url.path) + requests_by_user_type.labels( + user_type=user_type, + endpoint=normalized_path, + method=request.method + ).inc() + + return response diff --git a/docker-compose.yml b/docker-compose.yml index 9d7be83..69cb3e4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -24,7 +24,7 @@ services: environment: DATABASE_URL: postgresql://postgres:postgres@postgres/postgres_service SECRET_KEY: "dKq2BK3pqGQfNqC7SK8ZxNCdqJnGV4F9" # More secure key for development - ENV_SUFFIX: "local" + ENABLE_METRICS: "true" # Enable Prometheus metrics ports: - "8800:8800" volumes: @@ -88,8 +88,42 @@ services: labels: lagoon.type: none + prometheus: + image: prom/prometheus:latest + ports: + - "9090:9090" + volumes: + - ./prometheus:/etc/prometheus + - prometheus_data:/prometheus + command: + - '--config.file=/etc/prometheus/prometheus.yml' + - '--storage.tsdb.path=/prometheus' + - '--web.console.libraries=/usr/share/prometheus/console_libraries' + - '--web.console.templates=/usr/share/prometheus/consoles' + depends_on: + - backend + labels: + lagoon.type: none + + grafana: + image: grafana/grafana:latest + ports: + - "3001:3000" + volumes: + - ./grafana/provisioning:/etc/grafana/provisioning + - grafana_data:/var/lib/grafana + environment: + - GF_SECURITY_ADMIN_USER=admin + - GF_SECURITY_ADMIN_PASSWORD=admin + - GF_USERS_ALLOW_SIGN_UP=false + depends_on: + - prometheus + labels: + lagoon.type: none volumes: postgres_data: litellm_postgres_data: - name: litellm_postgres_data # Named volume for Postgres data persistence \ No newline at end of file + name: litellm_postgres_data # Named volume for Postgres data persistence + prometheus_data: + grafana_data: \ No newline at end of file diff --git a/grafana/provisioning/dashboards/fastapi.json b/grafana/provisioning/dashboards/fastapi.json new file mode 100644 index 0000000..43919ba --- /dev/null +++ b/grafana/provisioning/dashboards/fastapi.json @@ -0,0 +1,625 @@ +{ + "annotations": { + "list": [ + { + "builtIn": 1, + "datasource": { + "type": "grafana", + "uid": "-- Grafana --" + }, + "enable": true, + "hide": true, + "iconColor": "rgba(0, 211, 255, 1)", + "name": "Annotations & Alerts", + "type": "dashboard" + } + ] + }, + "editable": true, + "fiscalYearStartMonth": 0, + "graphTooltip": 0, + "id": 4, + "links": [], + "panels": [ + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 10, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + } + ] + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 0 + }, + "id": 1, + "options": { + "legend": { + "calcs": [ + "max" + ], + "displayMode": "table", + "placement": "right", + "showLegend": true, + "sortBy": "Max", + "sortDesc": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "12.0.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "editorMode": "builder", + "expr": "rate(http_requests_total[5m])", + "legendFormat": "{{method}} - {{handler}}", + "range": true, + "refId": "A" + } + ], + "title": "Request Rate", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 10, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + } + ] + }, + "unit": "s" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 0 + }, + "id": 2, + "options": { + "legend": { + "calcs": [ + "max" + ], + "displayMode": "table", + "placement": "right", + "showLegend": true, + "sortBy": "Max", + "sortDesc": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "12.0.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "editorMode": "code", + "expr": "rate(http_request_duration_seconds_sum[5m]) / rate(http_request_duration_seconds_count[5m])", + "legendFormat": "{{method}} - {{handler}}", + "range": true, + "refId": "A" + } + ], + "title": "Request Duration", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 10, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 8 + }, + "id": 3, + "options": { + "legend": { + "calcs": [ + "max" + ], + "displayMode": "table", + "placement": "right", + "showLegend": true, + "sortBy": "Max", + "sortDesc": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "12.0.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "expr": "rate(requests_by_user_type_total[5m])", + "legendFormat": "{{user_type}} - {{endpoint}}", + "refId": "A" + } + ], + "title": "Requests by User Type", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 10, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 8 + }, + "id": 4, + "options": { + "legend": { + "calcs": [ + "max" + ], + "displayMode": "table", + "placement": "right", + "showLegend": true, + "sortBy": "Max", + "sortDesc": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "12.0.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "expr": "rate(auth_requests_total[5m])", + "legendFormat": "{{endpoint}} - {{status}}", + "refId": "A" + } + ], + "title": "Auth Request Rate", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 10, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + } + ] + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 16 + }, + "id": 5, + "options": { + "legend": { + "calcs": [ + "max" + ], + "displayMode": "table", + "placement": "right", + "showLegend": true, + "sortBy": "Max", + "sortDesc": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "12.0.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "rate(http_requests_total{status=~\"4..\"}[5m])", + "fullMetaSearch": false, + "includeNullMetadata": true, + "legendFormat": "{{method}} {{handler}} - {{status}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "4xx Client Error Rate", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 10, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + } + ] + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 16 + }, + "id": 6, + "options": { + "legend": { + "calcs": [ + "max" + ], + "displayMode": "table", + "placement": "right", + "showLegend": true, + "sortBy": "Max", + "sortDesc": true + }, + "tooltip": { + "hideZeros": false, + "mode": "single", + "sort": "none" + } + }, + "pluginVersion": "12.0.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "prometheus" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "rate(http_requests_total{status=~\"5..\"}[5m])", + "fullMetaSearch": false, + "includeNullMetadata": true, + "legendFormat": "{{method}} {{handler}} - {{status}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "5xx Server Error Rate", + "type": "timeseries" + } + ], + "preload": false, + "refresh": "30s", + "schemaVersion": 41, + "tags": [], + "templating": { + "list": [] + }, + "time": { + "from": "now-1h", + "to": "now" + }, + "timepicker": {}, + "timezone": "", + "title": "amazee.ai Backend", + "uid": "amazeeai-backend", + "version": 6 +} \ No newline at end of file diff --git a/grafana/provisioning/dashboards/fastapi.yml b/grafana/provisioning/dashboards/fastapi.yml new file mode 100644 index 0000000..1db346a --- /dev/null +++ b/grafana/provisioning/dashboards/fastapi.yml @@ -0,0 +1,14 @@ +apiVersion: 1 + +providers: + - name: 'amazee.ai Backend' + orgId: 1 + folder: '' + type: file + disableDeletion: false + editable: true + updateIntervalSeconds: 10 + allowUiUpdates: true + options: + path: /etc/grafana/provisioning/dashboards + foldersFromFilesStructure: true \ No newline at end of file diff --git a/grafana/provisioning/datasources/prometheus.yml b/grafana/provisioning/datasources/prometheus.yml new file mode 100644 index 0000000..ea52740 --- /dev/null +++ b/grafana/provisioning/datasources/prometheus.yml @@ -0,0 +1,14 @@ +apiVersion: 1 + +datasources: + - name: Prometheus + type: prometheus + access: proxy + url: http://prometheus:9090 + isDefault: true + editable: false + uid: prometheus + jsonData: + timeInterval: 15s + queryTimeout: 30s + httpMethod: GET \ No newline at end of file diff --git a/prometheus/prometheus.yml b/prometheus/prometheus.yml new file mode 100644 index 0000000..0c0021a --- /dev/null +++ b/prometheus/prometheus.yml @@ -0,0 +1,10 @@ +global: + scrape_interval: 15s + evaluation_interval: 15s + +scrape_configs: + - job_name: 'fastapi' + static_configs: + - targets: ['backend:8800'] + metrics_path: '/metrics' + scheme: 'http' \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 4101a22..6a85b14 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,3 +14,5 @@ alembic==1.15.2 boto3==1.35.1 markdown==3.5.2 email-validator==2.1.2 +prometheus-client==0.19.0 +prometheus-fastapi-instrumentator==6.1.0 diff --git a/tests/test_private_ai.py b/tests/test_private_ai.py index 4b72c0a..6a7ba4d 100644 --- a/tests/test_private_ai.py +++ b/tests/test_private_ai.py @@ -979,6 +979,79 @@ def test_update_budget_period_as_key_creator(mock_post, client, team_key_creator db.delete(test_key) db.commit() +@patch("app.services.litellm.requests.post") +@patch("app.services.litellm.requests.get") +def test_update_budget_duration_as_team_admin(mock_get, mock_post, client, team_admin_token, test_region, mock_litellm_response, db, test_team): + """Test that a team admin can update the budget duration for a team-owned key""" + # Mock the LiteLLM API responses + mock_post.return_value.status_code = 200 + mock_post.return_value.json.return_value = mock_litellm_response + mock_post.return_value.raise_for_status.return_value = None + + # Mock the key info response + mock_get.return_value.status_code = 200 + mock_get.return_value.json.return_value = { + "info": { + "spend": 0.0, + "expires": "2024-12-31T23:59:59Z", + "created_at": "2024-01-01T00:00:00Z", + "updated_at": "2024-01-02T00:00:00Z", + "max_budget": 100.0, + "budget_duration": "monthly", + "budget_reset_at": "2024-02-01T00:00:00Z" + } + } + mock_get.return_value.raise_for_status.return_value = None + + # Create a test key owned by the team + test_key = DBPrivateAIKey( + database_name="test-db-team", + name="Test Team Key", + database_host="test-host", + database_username="test-user", + database_password="test-pass", + litellm_token="test-token-team", + litellm_api_url="https://test-litellm.com", + team_id=test_team.id, + region_id=test_region.id + ) + db.add(test_key) + db.commit() + db.refresh(test_key) + + # Update the budget duration as team admin + response = client.put( + f"/private-ai-keys/{test_key.id}/budget-period", + headers={"Authorization": f"Bearer {team_admin_token}"}, + json={"budget_duration": "monthly"} + ) + + # Verify the response + assert response.status_code == 200 + data = response.json() + assert data["budget_duration"] == "monthly" + + # Verify that the LiteLLM API was called with the correct parameters + mock_post.assert_called_with( + f"{test_region.litellm_api_url}/key/update", + headers={"Authorization": f"Bearer {test_region.litellm_api_key}"}, + json={ + "key": test_key.litellm_token, + "budget_duration": "monthly" + } + ) + + # Verify that the key info was checked + mock_get.assert_called_with( + f"{test_region.litellm_api_url}/key/info", + headers={"Authorization": f"Bearer {test_region.litellm_api_key}"}, + params={"key": test_key.litellm_token} + ) + + # Clean up the test key + db.delete(test_key) + db.commit() + @patch("app.services.litellm.requests.post") def test_create_llm_token_as_system_admin(mock_post, client, admin_token, test_region, mock_litellm_response): """Test that a system admin can create an LLM token for themselves"""