aws-ai-capacity 0.1.2
AWS GPU Capacity Management Agent with Pydantic AI
Requires Python: >=3.12

# AWS AI Capacity - GPU Capacity Management Agent

A Pydantic AI-powered agent for managing and querying AWS GPU capacity, with a Chainlit chat interface and CLI for automated reporting.

## Features

- **SageMaker Training Plans**: Search offerings, list plans, check utilization
- **EC2 Capacity**: Query reservations, check instance availability
- **GPU Specs**: Compare instance types and their capabilities
- **Chat Interface**: Interactive Chainlit UI for ad-hoc queries
- **CLI Reports**: Automated report generation for cron jobs

## Quick Start (uvx)

Run directly without installing:

```bash
# With inline environment variables
AWS_PROFILE=myprofile uvx aws-ai-capacity chat "What GPU capacity is available?"

# Or export first
export AWS_PROFILE=myprofile
uvx aws-ai-capacity chat "What p5 training plans are available?"
uvx aws-ai-capacity report daily
```

## Installation (Development)

```bash
# Clone and install
git clone https://github.com/drewdresser/aws-ai-capacity.git
cd aws-ai-capacity
uv sync

# Copy environment template
cp .env.example .env
```

## Configuration

The agent uses **AWS Bedrock** with **Claude Opus 4.5** (`anthropic.claude-opus-4-5-20251101-v1:0`) by default.

Required AWS permissions:
- `bedrock:InvokeModel` - For Claude model access
- `sagemaker:SearchTrainingPlanOfferings` - Query available plans
- `sagemaker:ListTrainingPlans` - List existing plans
- `sagemaker:DescribeTrainingPlan` - Get plan details
- `ec2:DescribeCapacityReservations` - Query reservations
- `ec2:DescribeInstanceTypeOfferings` - Check availability

### Environment Variables

For development or to override defaults, create a `.env` file:

```bash
# AWS credentials
AWS_PROFILE=myprofile
AWS_REGION=us-east-1

# Override the default model
BEDROCK_MODEL_ID=us.anthropic.claude-opus-4-5-20251101-v1:0
BEDROCK_REGION=us-east-1

# Other options
AGENT_MAX_RETRIES=3
LOG_LEVEL=DEBUG
```

See `.env.example` for all available options.

## Usage

### Chat Interface

Start the Chainlit UI:
```bash
uv run aws-ai-capacity serve
```

Then open http://localhost:8000

### CLI Commands

```bash
# Single query
uv run aws-ai-capacity chat "What p5 training plans are available?"

# Debug mode - shows all tool calls made by the agent
uv run aws-ai-capacity chat "What p5 training plans are available?" --debug

# Generate reports
uv run aws-ai-capacity report daily
uv run aws-ai-capacity report availability
uv run aws-ai-capacity report training-plans

# Save report to file
uv run aws-ai-capacity report daily -o capacity-report.md

# Generate all reports (for cron)
uv run aws-ai-capacity cron-report -d ./reports

# List GPU instance types
uv run aws-ai-capacity list-instance-types
```

### Cron Job Example

```bash
# Daily capacity report at 8am
0 8 * * * cd /path/to/ai-capacity && uv run aws-ai-capacity cron-report -d ./reports
```

## Example Questions

- "What p5.48xlarge training plans are available for the next week?"
- "Show me all active capacity reservations"
- "Which regions have p4d.24xlarge available?"
- "Compare the specs of p5 vs p4d instances"
- "Generate a daily capacity report"
- "Search for available H100 capacity"

## Project Structure

```
ai-capacity/
├── src/ai_capacity/
│   ├── agent/          # Pydantic AI agent definition
│   ├── tools/          # AWS API tools (SageMaker, EC2)
│   ├── cli/            # Typer CLI commands
│   ├── ui/             # Chainlit chat interface
│   └── config.py       # Settings management
├── chainlit.md         # Chat welcome message
└── .env.example        # Environment template
```

## Development

```bash
# Install with dev dependencies
uv sync --all-extras

# Run tests
uv run pytest

# Lint
uv run ruff check .
```