# Auror

A distributed task processing framework for AI inference with MQTT messaging and OAuth2 authentication.

## Overview

Auror is a Python framework designed for distributed AI inference workloads. It provides a robust infrastructure for processing tasks through MQTT messaging with built-in OAuth2 authentication via Auth0. The framework features an extensible worker pattern that makes it easy to add support for different types of AI inference tasks.

### Key Features

- **MQTT-based messaging**: Reliable task distribution using MQTT protocol
- **OAuth2 authentication**: Secure authentication via Auth0
- **Extensible worker pattern**: Easy to add new inference types
- **Production-ready**: Comprehensive logging, error handling, and cleanup
- **CLI interface**: Simple command-line interface for running workers
- **Type hints**: Full type annotations for better IDE support

## Architecture

Auror separates concerns into distinct components:

- **Core messaging** (`auror.core.messaging`): Handles MQTT connections and authentication
- **Worker base** (`auror.core.worker`): Abstract base class for all workers
- **Worker implementations** (`auror.workers`): Specific inference implementations (e.g., LTX Video)
- **Configuration** (`auror.config`): Centralized configuration management

This separation makes it easy to add new inference types without modifying the core infrastructure.

## Installation

### From PyPI (when published)

```bash
pip install auror
```

### From source

```bash
git clone https://github.com/auror-ai/auror
cd auror
pip install -e .
```

### With LTXV support

```bash
pip install auror[ltxv]
```

### Development installation

```bash
pip install -e ".[dev]"
```

## Quick Start

### Configuration with .env file

The easiest way to configure Auror is using a `.env` file:

```bash
# Copy the example .env file
cp .env.example .env

# Edit .env with your credentials
# Then run the worker - it will automatically load the .env file
auror ltxv
```

Your `.env` file should contain:

```bash
AUROR_MQTT_HOST=your-mqtt-broker.com
AUROR_MQTT_PORT=1883
AUROR_MQTT_TOPIC=video/ltx
AUROR_CLIENT_ID=your-client-id
AUROR_CLIENT_SECRET=your-client-secret
AUROR_AUDIENCE=https://your-audience.com
```

### Using the CLI

Alternatively, you can set environment variables manually:

```bash
# Set environment variables
export AUROR_MQTT_HOST="your-mqtt-broker.com"
export AUROR_MQTT_PORT="1883"
export AUROR_CLIENT_ID="your-client-id"
export AUROR_CLIENT_SECRET="your-client-secret"
export AUROR_AUDIENCE="https://your-audience.com"

# Run the LTXV worker
auror ltxv
```

### Using Python

```python
from auror import MessagingCore
from auror.workers import LTXVWorker
from auror.config import AurorConfig

# Create configuration
config = AurorConfig(
    mqtt_host="your-mqtt-broker.com",
    mqtt_port=1883,
    mqtt_topic="video/ltx",
    client_id="your-client-id",
    client_secret="your-client-secret",
    audience="https://your-audience.com",
)

# Create worker
worker = LTXVWorker()

# Create messaging core
core = MessagingCore(
    host=config.mqtt_host,
    port=config.mqtt_port,
    client_id=config.client_id,
    client_secret=config.client_secret,
    audience=config.audience,
    topic=config.mqtt_topic,
)

# Start processing
core.start(worker)
```

## Configuration

Auror can be configured through environment variables or programmatically.

### Environment Variables

| Variable | Required | Default | Description |
|----------|----------|---------|-------------|
| `AUROR_MQTT_HOST` | Yes | - | MQTT broker hostname |
| `AUROR_MQTT_PORT` | No | `1883` | MQTT broker port |
| `AUROR_MQTT_TOPIC` | No | `video/ltx` | MQTT topic to subscribe to |
| `AUROR_CLIENT_ID` | Yes | - | OAuth2 client ID |
| `AUROR_CLIENT_SECRET` | Yes | - | OAuth2 client secret |
| `AUROR_AUDIENCE` | Yes | - | OAuth2 audience |
| `AUROR_AUTH0_TOKEN_URL` | No | Auth0 default | Auth0 token endpoint |
| `AUROR_PIPELINE_CONFIG` | No | Default config | Pipeline configuration path |

### Programmatic Configuration

```python
from auror.config import AurorConfig

# From environment variables
config = AurorConfig.from_env()

# From dictionary
config = AurorConfig.from_dict({
    "mqtt_host": "broker.example.com",
    "mqtt_port": 1883,
    "client_id": "...",
    "client_secret": "...",
    "audience": "https://api.example.com",
})

# Direct instantiation
config = AurorConfig(
    mqtt_host="broker.example.com",
    mqtt_port=1883,
    mqtt_topic="my/topic",
    client_id="...",
    client_secret="...",
    audience="https://api.example.com",
)
```

## CLI Usage

### Running the LTXV Worker

```bash
# Using environment variables
auror ltxv

# With command-line overrides
auror ltxv --host broker.example.com --port 1883 --topic custom/topic

# Using default configuration (development only)
auror ltxv --use-defaults

# With verbose logging
auror ltxv -v
```

### CLI Options

```
auror ltxv [OPTIONS]

Options:
  --host TEXT              MQTT broker hostname
  --port INTEGER           MQTT broker port
  --topic TEXT             MQTT topic to subscribe to
  --client-id TEXT         OAuth2 client ID
  --client-secret TEXT     OAuth2 client secret
  --audience TEXT          OAuth2 audience
  --pipeline-config TEXT   Path to LTX pipeline config
  --use-defaults          Use default configuration (dev only)
  -v, --verbose           Enable verbose logging
```

## Creating Custom Workers

To add support for a new inference type, create a worker class that inherits from `Worker`:

```python
from typing import Any, Dict
from auror.core.worker import Worker
import logging

logger = logging.getLogger(__name__)

class MyCustomWorker(Worker):
    """Worker for my custom inference type."""

    def __init__(self, model_path: str):
        super().__init__()
        self.model_path = model_path
        # Initialize your model here

    def process(self, job: Dict[str, Any]) -> None:
        """
        Process a job.

        Args:
            job: Job dictionary with 'id' and other parameters
        """
        job_id = job["id"]

        try:
            logger.info(f"[{job_id}] Starting custom inference")

            # Your inference logic here
            input_data = job["input"]
            result = self.run_inference(input_data)

            # Upload result
            self.upload_result(result, job["input"]["uploadUrl"])

            logger.info(f"[{job_id}] Job completed")

        except Exception as e:
            logger.error(f"[{job_id}] Error: {e}", exc_info=True)
            raise

    def run_inference(self, input_data):
        # Your inference implementation
        pass

    def upload_result(self, result, url):
        # Your upload implementation
        pass
```

Then use it with the messaging core:

```python
from auror import MessagingCore
from my_module import MyCustomWorker

worker = MyCustomWorker(model_path="/path/to/model")
core = MessagingCore(...)
core.start(worker)
```

## Job Format

Jobs are received as JSON messages with the following structure:

```json
{
  "id": "unique-job-id",
  "input": {
    "uploadUrl": "https://presigned-url.com/upload",
    // ... other job-specific parameters
  }
}
```

### LTXV Job Format

For LTX Video jobs:

```json
{
  "id": "job-123",
  "input": {
    "prompt": "A beautiful sunset over the ocean",
    "conditioningMedias": [
      "https://example.com/image1.jpg"
    ],
    "conditioningStartFrames": [0],
    "uploadUrl": "https://storage.com/output.mp4",
    "seed": 42,
    "height": 720,
    "width": 1280,
    "numFrames": 120
  }
}
```

## Logging

Auror uses Python's standard logging module. Configure logging in your application:

```python
import logging

# Basic configuration
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

# Or use the CLI with -v flag for verbose output
```

## Error Handling

Workers automatically handle errors and log them without crashing the worker process. Failed jobs are logged with full stack traces for debugging.

## Development

### Setup Development Environment

```bash
# Clone the repository
git clone https://github.com/auror-ai/auror
cd auror

# Install in development mode with dev dependencies
pip install -e ".[dev]"
```

### Running Tests

```bash
pytest
```

### Code Formatting

```bash
black auror/
```

### Type Checking

```bash
mypy auror/
```

## Publishing to PyPI

### Build the package

```bash
python -m pip install build twine
python -m build
```

### Upload to PyPI

```bash
# Test PyPI first
python -m twine upload --repository testpypi dist/*

# Production PyPI
python -m twine upload dist/*
```

## License

MIT License - see LICENSE file for details

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

## Support

For issues and questions, please use the [GitHub issue tracker](https://github.com/auror-ai/auror/issues).
