Metadata-Version: 2.4
Name: aws-ai-capacity
Version: 0.1.3
Summary: AWS GPU Capacity Management Agent with Pydantic AI
License-File: LICENSE
Requires-Python: >=3.12
Requires-Dist: aioboto3>=13.0
Requires-Dist: boto3>=1.35
Requires-Dist: chainlit>=1.3
Requires-Dist: pydantic-ai>=0.1.0
Requires-Dist: pydantic-settings>=2.6
Requires-Dist: pydantic>=2.10
Requires-Dist: python-dotenv>=1.0
Requires-Dist: rich>=13.9
Requires-Dist: structlog>=24.4
Requires-Dist: typer>=0.12
Provides-Extra: dev
Requires-Dist: moto[ec2,sagemaker]>=5.0; extra == 'dev'
Requires-Dist: pytest-asyncio>=0.24; extra == 'dev'
Requires-Dist: pytest>=8.3; extra == 'dev'
Requires-Dist: ruff>=0.8; extra == 'dev'
Description-Content-Type: text/markdown

# 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
- **Spot Capacity**: Spot placement scores (1-10) and price history across regions
- **On-Demand Capacity**: Launch-and-terminate availability check (opt-in, costs money)
- **GPU Utilization**: View running GPU instances with spot vs on-demand breakdown
- **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
- `ec2:DescribeInstances` - List running GPU instances
- `ec2:DescribeSpotPriceHistory` - Spot price trends
- `ec2:GetSpotPlacementScores` - Spot capacity scores
- `ec2:RunInstances` / `ec2:TerminateInstances` - On-demand capacity check (opt-in)
- `ssm:GetParameter` - AMI lookup for on-demand check

### 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
uv run aws-ai-capacity report spot

# 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"
- "What are the spot placement scores for p5.48xlarge across all regions?"
- "Show me spot price history for p4d.24xlarge in us-east-1"
- "Show me all running GPU instances in my account"
- "Check on-demand capacity for g5.xlarge in us-east-1"

## Project Structure

```
ai-capacity/
├── src/ai_capacity/
│   ├── agent/          # Pydantic AI agent definition
│   ├── tools/          # AWS API tools (SageMaker, EC2, Spot/On-Demand)
│   ├── 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 .
```
