| # AegisLM Red Team API |
|
|
| Production-ready SaaS backend for AI red team evaluations and security assessments. |
|
|
| ## 🚀 Overview |
|
|
| This backend system converts the existing AI red team pipeline into a scalable, API-driven SaaS platform with: |
|
|
| - **Authentication System**: JWT-based user auth with signup/login |
| - **API Key Management**: Secure API key generation and validation |
| - **Evaluation API**: Async AI evaluation processing with Celery |
| - **Results API**: Comprehensive result storage and analysis |
| - **Benchmark API**: Multi-model comparison and leaderboard |
| - **Production Features**: Rate limiting, security headers, monitoring |
|
|
| ## 🏗️ Architecture |
|
|
| ``` |
| backend/ |
| ├── api/ |
| │ ├── routes/ # API endpoints |
| │ └── dependencies/ # Dependency injection |
| ├── core/ # Core configuration |
| ├── models/ # Database models |
| ├── schemas/ # Pydantic schemas |
| ├── services/ # Business logic layer |
| ├── workers/ # Celery workers |
| ├── tasks/ # Background tasks |
| ├── middleware/ # Custom middleware |
| ├── main.py # FastAPI application |
| └── requirements.txt # Dependencies |
| ``` |
|
|
| ## 🛠️ Tech Stack |
|
|
| - **Framework**: FastAPI with async support |
| - **Database**: PostgreSQL with SQLAlchemy |
| - **Cache/Queue**: Redis + Celery |
| - **Authentication**: JWT tokens + API keys |
| - **Validation**: Pydantic models |
| - **Security**: Rate limiting, CORS, security headers |
|
|
| ## 📋 Features |
|
|
| ### 🔐 Authentication |
| - User registration and login |
| - JWT token generation and validation |
| - Password hashing with bcrypt |
| - API key generation per user |
|
|
| ### 🚀 Evaluation System |
| - Async evaluation processing |
| - Job tracking with status updates |
| - Configurable attack types and parameters |
| - Integration with existing AI pipeline |
|
|
| ### 📊 Results Management |
| - Comprehensive result storage |
| - Export functionality (JSON, CSV) |
| - Result comparison and analytics |
| - Performance metrics |
|
|
| ### 🏆 Benchmark System |
| - Multi-model comparison |
| - Leaderboard generation |
| - Risk profiling |
| - Statistical analysis |
|
|
| ### ⚡ Production Features |
| - Rate limiting with Redis |
| - Security headers middleware |
| - Request logging and monitoring |
| - Health check endpoints |
| - Error handling and validation |
|
|
| ## 🚀 Getting Started |
|
|
| ### Prerequisites |
|
|
| - Python 3.8+ |
| - PostgreSQL |
| - Redis |
| - Docker (optional) |
|
|
| ### Installation |
|
|
| 1. **Clone and install dependencies**: |
| ```bash |
| cd backend |
| pip install -r requirements.txt |
| ``` |
|
|
| 2. **Set up environment variables**: |
| ```bash |
| cp .env.example .env |
| # Edit .env with your configuration |
| ``` |
|
|
| 3. **Database setup**: |
| ```bash |
| # Create database |
| createdb aegislm |
| |
| # Run migrations (if using Alembic) |
| alembic upgrade head |
| ``` |
|
|
| 4. **Start Redis**: |
| ```bash |
| redis-server |
| ``` |
|
|
| 5. **Start Celery worker**: |
| ```bash |
| celery -A workers.celery_worker worker --loglevel=info |
| ``` |
|
|
| 6. **Start the API server**: |
| ```bash |
| uvicorn main:app --reload --host 0.0.0.0 --port 8000 |
| ``` |
|
|
| ### Docker Setup |
|
|
| ```bash |
| # Build and run with Docker Compose |
| docker-compose up -d |
| ``` |
|
|
| ## 📚 API Documentation |
|
|
| Once running, visit: |
| - **Swagger UI**: `http://localhost:8000/docs` |
| - **ReDoc**: `http://localhost:8000/redoc` |
|
|
| ### Core Endpoints |
|
|
| #### Authentication |
| - `POST /api/v1/auth/signup` - User registration |
| - `POST /api/v1/auth/login` - User login |
| - `GET /api/v1/auth/me` - Get current user |
|
|
| #### Evaluations |
| - `POST /api/v1/evaluations/` - Create evaluation |
| - `GET /api/v1/evaluations/` - List evaluations |
| - `GET /api/v1/evaluations/{id}` - Get evaluation |
| - `POST /api/v1/evaluations/{id}/cancel` - Cancel evaluation |
|
|
| #### Results |
| - `GET /api/v1/results/job/{job_id}` - Get result |
| - `GET /api/v1/results/` - List results |
| - `POST /api/v1/results/job/{job_id}/export` - Export result |
|
|
| #### Benchmarks |
| - `POST /api/v1/benchmarks/` - Create benchmark |
| - `GET /api/v1/benchmarks/{id}/status` - Get benchmark status |
| - `GET /api/v1/benchmarks/{id}/result` - Get benchmark result |
|
|
| ## 🔧 Configuration |
|
|
| ### Environment Variables |
|
|
| ```bash |
| # Application |
| APP_NAME="AegisLM Red Team API" |
| DEBUG=false |
| API_V1_STR="/api/v1" |
| |
| # Security |
| SECRET_KEY="your-secret-key-here" |
| ACCESS_TOKEN_EXPIRE_MINUTES=30 |
| |
| # Database |
| DATABASE_URL="postgresql://user:pass@localhost/aegislm" |
| |
| # Redis |
| REDIS_URL="redis://localhost:6379/0" |
| |
| # Rate Limiting |
| RATE_LIMIT_PER_MINUTE=60 |
| RATE_LIMIT_BURST=10 |
| ``` |
|
|
| ## 🧪 Testing |
|
|
| ```bash |
| # Run tests |
| pytest |
| |
| # Run with coverage |
| pytest --cov=. |
| |
| # Run specific test file |
| pytest tests/test_auth.py |
| ``` |
|
|
| ## 📊 Monitoring |
|
|
| ### Health Checks |
| - `GET /health` - Overall system health |
| - `GET /metrics` - Basic metrics |
| - `GET /api/v1/status` - API status |
|
|
| ### Logging |
| - Request/response logging |
| - Error tracking |
| - Performance metrics |
|
|
| ## 🔒 Security Features |
|
|
| - **Rate Limiting**: Redis-based sliding window |
| - **Authentication**: JWT + API key support |
| - **Security Headers**: XSS, CSRF protection |
| - **Input Validation**: Pydantic schemas |
| - **Password Security**: Bcrypt hashing |
|
|
| ## 🚀 Deployment |
|
|
| ### Production Setup |
|
|
| 1. **Environment Configuration**: |
| ```bash |
| export DEBUG=false |
| export SECRET_KEY="production-secret-key" |
| export DATABASE_URL="postgresql://prod_user:pass@db_host/aegislm" |
| ``` |
|
|
| 2. **Database Migrations**: |
| ```bash |
| alembic upgrade head |
| ``` |
|
|
| 3. **Start Services**: |
| ```bash |
| # Start Celery workers |
| celery -A workers.celery_worker worker --loglevel=info |
| |
| # Start API server |
| gunicorn main:app -w 4 -k uvicorn.workers.UvicornWorker |
| ``` |
|
|
| ### Docker Deployment |
|
|
| ```bash |
| # Build image |
| docker build -t aegislm-api . |
| |
| # Run container |
| docker run -d \ |
| --name aegislm-api \ |
| -p 8000:8000 \ |
| -e DATABASE_URL=$DATABASE_URL \ |
| -e REDIS_URL=$REDIS_URL \ |
| aegislm-api |
| ``` |
|
|
| ## 📈 Performance |
|
|
| - **Async Processing**: FastAPI + async/await |
| - **Connection Pooling**: SQLAlchemy connection pools |
| - **Caching**: Redis for rate limiting and results |
| - **Background Tasks**: Celery for evaluation processing |
|
|
| ## 🔄 Integration with AI Engine |
|
|
| The backend integrates seamlessly with the existing red team pipeline: |
|
|
| ```python |
| # Service layer integration |
| from ai.pipelines.redteam_pipeline import run_redteam_pipeline |
| |
| # Celery task execution |
| pipeline_result = run_redteam_pipeline(model, config, job_id) |
| ``` |
|
|
| ## 🛠️ Development |
|
|
| ### Code Quality |
|
|
| ```bash |
| # Code formatting |
| black . |
| isort . |
| |
| # Linting |
| flake8 . |
| |
| # Type checking |
| mypy . |
| ``` |
|
|
| ### Adding New Features |
|
|
| 1. **Add Models**: Define in `models/` |
| 2. **Add Schemas**: Define in `schemas/` |
| 3. **Add Services**: Implement in `services/` |
| 4. **Add Routes**: Define in `api/routes/` |
| 5. **Add Tests**: Write in `tests/` |
|
|
| ## 📝 License |
|
|
| This project is part of the AegisLM Red Team Engine. |
|
|
| ## 🤝 Contributing |
|
|
| 1. Fork the repository |
| 2. Create a feature branch |
| 3. Make your changes |
| 4. Add tests |
| 5. Submit a pull request |
|
|
| ## 📞 Support |
|
|
| For support and questions: |
| - Create an issue in the repository |
| - Check the API documentation |
| - Review the code comments |
|
|
| --- |
|
|
| **Built with ❤️ for secure AI evaluation** |
|
|