ALM-2 / backend /README.md
ACA050's picture
Upload 520 files
2ed8996 verified
|
Raw
History Blame Contribute Delete
6.98 kB
# 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**