

## Automagik Codebase Guide for AI Agents

**Objective:** Your primary goal is to understand, enhance, and maintain the Automagik codebase. You will be implementing new features, fixing bugs, and improving existing functionality.

**Context:** Automagik is an automated workflow management system designed to integrate with LangFlow for flow synchronization and execution. It allows users to define, schedule, and run complex workflows. This codebase was initially generated by AI agents and is intended to be further developed and maintained by AI agents like you.

**Workflow:** Your development process should adhere to the following steps:

1. **Study:** Thoroughly understand the existing codebase, project structure, and instructions before making any changes.
2. **Implement:** Write clean, efficient, and well-documented code. Avoid code duplication and follow established patterns.
3. **Test:** Run `pytest` after every modification to ensure your changes don't break existing functionality. Write new tests for new features or bug fixes. Use the command `pytest`
4. **Commit:** Commit your changes using Git with clear and concise commit messages.
5. **Revert:** If you fail to implement a feature or fix a bug after three attempts, revert all your changes using `git checkout -- .` and `git clean -fd`. Analyze the issue, study the relevant parts of the codebase again, and start from scratch.

**Important Considerations:**

*   **Error Handling:** Implement robust error handling in your code. Log errors appropriately using the `logging` module.
*   **Asynchronous Programming:** Automagik heavily relies on asynchronous programming (using `asyncio`). Ensure your code is non-blocking and utilizes `async` and `await` correctly.
*   **Database Interactions:** Use the provided `get_session` context manager for all database interactions to ensure proper session management.

---

## Automagik Project Overview

### 1. Project Structure

The Automagik project is organized as follows:

\`\`\`
automagik/
├── automagik/            # Main application package
│   ├── api/              # FastAPI application (endpoints, models, dependencies)
│   │   ├── dependencies.py   # API dependencies (e.g., API key verification, database session)
│   │   ├── __init__.py
│   │   ├── models.py         # Pydantic models for request/response validation
│   │   ├── routers/          # API routers (endpoints grouped by functionality)
│   │   │   ├── flows.py      # Flow-related endpoints
│   │   │   ├── __init__.py
│   │   │   ├── schedules.py  # Schedule-related endpoints
│   │   │   ├── tasks.py      # Task-related endpoints
│   │   │   └── workers.py    # Worker-related endpoints
│   │   ├── app.py            # Main FastAPI app initialization
│   │   └── config.py         # API configuration (CORS, host, port, API key)
│   ├── cli/              # Command-line interface (CLI)
│   │   ├── commands/         # CLI commands (grouped by functionality)
│   │   │   ├── api.py        # Start automagik api server
│   │   │   ├── db.py         # Database management commands
│   │   │   ├── flow.py       # Flow management commands
│   │   │   ├── __init__.py
│   │   │   ├── schedule.py   # Schedule management commands
│   │   │   ├── task.py       # Task management commands
│   │   │   └── worker.py     # Worker process management commands
│   │   ├── __init__.py
│   │   └── cli.py            # Main CLI entry point and setup
│   ├── core/             # Core application logic
│   │   ├── database/         # Database models and session management
│   │   │   ├── base.py       # SQLAlchemy base model
│   │   │   ├── __init__.py
│   │   │   ├── models.py     # Database models (Flow, Task, Schedule, etc.)
│   │   │   └── session.py    # Database session creation and management
│   │   ├── flows/            # Flow management (sync, analysis, execution)
│   │   │   ├── analyzer.py   # Flow component analyzer (for input/output detection)
│   │   │   ├── __init__.py
│   │   │   ├── local.py      # Local flow management (database interactions)
│   │   │   ├── manager.py    # Main flow manager (handles both local and remote flows)
│   │   │   ├── remote.py     # Remote flow management (LangFlow API interactions)
│   │   │   ├── sync.py       # Flow synchronization logic
│   │   │   └── task.py       # Task execution logic
│   │   ├── scheduler/        # Scheduling and task running
│   │   │   ├── __init__.py
│   │   │   ├── manager.py    # Scheduler management (create, update, delete schedules)
│   │   │   ├── scheduler.py  # Flow scheduling and execution based on cron/interval
│   │   │   └── task_runner.py # Task execution, including retries
│   │   └── config.py         # Core configuration settings
│   ├── tests/            # Unit and integration tests
│   │   ├── core/
│   │   │   └── flows/
│   │   │       ├── __init__.py
│   │   │       ├── test_flow_execution.py
│   │   │       └── test_task_manager.py
│   │   └── __init__.py
│   ├── __init__.py
│   ├── __main__.py           # Main entry point for the CLI
│   └── __version__.py        # Version information
├── migrations/       # Alembic database migrations
│   ├── env.py
│   ├── README
│   ├── script.py.mako
│   └── versions/
│       ├── 20250127_2156_a23006a86de3_.py
│       ├── 4ff8fb67c367_initial_schema.py
│       └── a323736b0635_add_workers_table.py
├── scripts/          # Helper scripts (setup, testing, etc.)
│   ├── reset_db.sh
│   ├── run_tests.sh
│   └── setup.sh
├── tests/            # Test suite
│   ├── api/              # API tests
│   │   ├── conftest.py
│   │   ├── __init__.py
│   │   ├── test_api.py
│   │   ├── test_auth.py
│   │   ├── test_config.py
│   │   ├── test_flows.py.skip
│   │   ├── test_schedules.py.skip
│   │   ├── test_tasks.py.skip
│   │   └── test_workers.py.skip
│   ├── cli/              # CLI tests
│   │   └── commands/
│   │       ├── __init__.py
│   │       ├── test_task.py
│   │       ├── test_worker.py
│   │       └── test_worker_commands.py
│   ├── conftest.py       # Pytest fixtures and configuration
│   ├── core/             # Core logic tests
│   │   ├── flows/
│   │   │   ├── components/
│   │   │   │   ├── __init__.py
│   │   │   │   └── test_components.py
│   │   │   ├── listing/
│   │   │   │   ├── __init__.py
│   │   │   │   └── test_flows.py
│   │   │   ├── scheduling/
│   │   │   │   ├── __init__.py
│   │   │   │   └── test_scheduling.py
│   │   │   ├── sync/
│   │   │   │   ├── __init__.py
│   │   │   │   └── test_sync.py
│   │   │   ├── __init__.py
│   │   │   ├── test_delete_flow.py
│   │   │   ├── test_sync_flow.py
│   │   │   └── test_task_manager.py
│   │   └── scheduler/
│   │       ├── __init__.py
│   │       └── test_scheduler_manager.py
│   ├── __init__.py
│   └── mock_data/        # Mock data for testing
│       └── flows/
│           ├── flows.json
│           └── folders.json
├── .env.example      # Example environment variables
├── alembic.ini       # Alembic configuration
├── docker/
│   ├── docker-compose.yml
│   └── Dockerfile
├── pytest.ini        # Pytest configuration
├── README.md         # Project README
├── redis.conf        # Redis configuration (if applicable)
├── requirements.txt  # Project dependencies
└── setup.py          # Project setup and metadata

\`\`\`

### 2. Core Features and Modules

#### 2.1. API (`automagik/api/`)

*   **Purpose:** Provides a RESTful API for interacting with Automagik.
*   **Key Files:**
    *   `app.py`: Initializes the FastAPI application, includes routers, and sets up middleware (e.g., CORS).
    *   `models.py`: Defines Pydantic models for request and response validation. These models ensure data integrity and provide automatic API documentation.
    *   `dependencies.py`: Contains dependency injection functions, such as `get_session` for database access and `verify_api_key` for API key authentication.
    *   `routers/`: Each file in this directory defines API endpoints for a specific resource (e.g., `flows.py`, `schedules.py`, `tasks.py`, `workers.py`).
*   **Dependencies:**
    *   `fastapi`: Web framework for building the API.
    *   `uvicorn`: ASGI server for running the FastAPI application.
    *   `pydantic`: Data validation and settings management.

#### 2.2. CLI (`automagik/cli/`)

*   **Purpose:** Provides a command-line interface for managing Automagik (e.g., creating schedules, listing tasks, starting the worker).
*   **Key Files:**
    *   `cli.py`: Main entry point for the CLI. Defines the main command group and handles overall CLI setup.
    *   `commands/`: Each file in this directory implements a set of related CLI commands (e.g., `flow.py`, `schedule.py`, `task.py`, `db.py`, `worker.py`).
*   **Dependencies:**
    *   `click`: Library for creating command-line interfaces.
    *   `tabulate`: For pretty-printing tabular data in the CLI.

#### 2.3. Core Logic (`automagik/core/`)

*   **Purpose:** Contains the core business logic of Automagik, including flow management, scheduling, and task execution.
*   **Key Files:**
    *   `flows/`:
        *   `manager.py`: Main `FlowManager` class, which orchestrates flow-related operations (local and remote).
        *   `local.py`: `LocalFlowManager` handles interactions with the local database for flows.
        *   `remote.py`: `RemoteFlowManager` interacts with the LangFlow API for flow synchronization.
        *   `sync.py`: `FlowSync` handles the details of executing a flow and updating task status.
        *   `analyzer.py`: `FlowAnalyzer` provides utilities for analyzing flow components (input/output detection, parameter extraction).
        *   `task.py`: `TaskManager` manages the creation, listing, retrieval, and retrying of tasks.
    *   `scheduler/`:
        *   `manager.py`: `SchedulerManager` is responsible for creating, updating, deleting, and managing schedules.
        *   `scheduler.py`: `FlowScheduler` handles the actual scheduling logic, determining when to run flows based on their schedules.
        *   `task_runner.py`: `TaskRunner` executes individual tasks, including retries.
    *   `database/`:
        *   `models.py`: SQLAlchemy models representing database tables (e.g., `Flow`, `Task`, `Schedule`, `Worker`).
        *   `session.py`: Handles the creation and management of database sessions using `async_session` and the `get_session` context manager.
*   **Dependencies:**
    *   `sqlalchemy`: ORM for database interactions.
    *   `asyncpg`: Asynchronous PostgreSQL driver.
    *   `httpx`: HTTP client for interacting with the LangFlow API.
    *   `croniter`: Library for working with cron expressions.

#### 2.4. Database (`automagik/core/database/`)

*   **Purpose:** Defines the database schema and provides utilities for database interaction.
*   **Key Files:**
    *   `models.py`: Defines the SQLAlchemy models for the database tables (Flow, Task, Schedule, FlowComponent, TaskLog).
    *   `session.py`: Provides the `get_session` dependency for accessing the database within the application and `DATABASE_URL` configuration.
    *   `base.py`: Defines the SQLAlchemy declarative base.
*   **Dependencies:**
    *   `SQLAlchemy`: Database toolkit and ORM.

#### 2.5 Migrations (`automagik/migrations/`)

*   **Purpose:** Alembic database migrations for managing database schema changes.
*   **Key Files:**
    *   `env.py`: Alembic environment configuration.
    *   `script.py.mako`: Alembic migration script template.
    *   `versions/`: Contains the actual migration scripts (each file represents a schema change).

#### 2.6. Tests (`automagik/tests/`)

*   **Purpose:** Contains unit and integration tests for the application.
*   **Key Files:**
    *   `conftest.py`: Pytest fixtures for setting up test environments (e.g., in-memory database).
    *   `api/`: Tests for the API endpoints.
    *   `cli/`: Tests for the CLI commands.
    *   `core/`: Tests for the core logic (flows, scheduler, etc.).
    *   `mock_data/`: Contains mock data (e.g., JSON files) used in tests.
*   **Dependencies:**
    *   `pytest`: Testing framework.
    *   `pytest-asyncio`: Pytest plugin for testing asynchronous code.
    *   `fastapi.testclient`: For testing FastAPI applications.

### 3. Dependencies

*   **Core Dependencies:**
    *   `click`: CLI creation.
    *   `SQLAlchemy`: Database ORM.
    *   `asyncpg`: Asynchronous PostgreSQL driver.
    *   `python-dotenv`: Loading environment variables.
    *   `tabulate`: Pretty-printing tables in CLI.
    *   `croniter`: Cron expression parsing.
    *   `httpx`: HTTP client (for LangFlow API interactions).
    *   `alembic`: Database migrations.
    *   `fastapi`: Web framework.
    *   `uvicorn`: ASGI server.
    *   `pydantic`: Data validation.
    *   `psutil`: Process and system monitoring
    *   `rich`: Rich text and beautiful formatting in the terminal

*   **Development Dependencies:**
    *   `pytest`: Testing framework.
    *   `pytest-asyncio`: Asyncio support for pytest.
    *   `black`: Code formatter.
    *   `isort`: Import sorter.
    *   `flake8`: Style guide enforcement.
    *   `mypy`: Static type checker.

### 4. Key Concepts

#### 4.1. Flows

*   A **Flow** represents a workflow definition.
*   Flows are defined in the `Flow` database model.
*   Flows can be synced from LangFlow (remote) or created locally.
*   Flows have properties like `name`, `description`, `source`, `source_id`, `input_component`, `output_component`, and `data`.
*   The `FlowManager` handles flow-related operations.
*   The `FlowSync` class handles the execution of the flow via Langflow API.

#### 4.2. Tasks

*   A **Task** represents a single execution of a flow.
*   Tasks are defined in the `Task` database model.
*   Tasks have properties like `flow_id`, `status` (pending, running, completed, failed), `input_data`, `output_data`, `error`, `tries`, and `max_retries`.
*   The `TaskManager` handles task-related operations.

#### 4.3. Schedules

*   A **Schedule** defines when a flow should be run (e.g., using cron expressions or intervals).
*   Schedules are defined in the `Schedule` database model.
*   Schedules have properties like `flow_id`, `schedule_type`, `schedule_expr`, `flow_params`, `status`, and `next_run_at`.
*   The `SchedulerManager` handles schedule-related operations.
*   The `FlowScheduler` monitors schedules and triggers task creation when a flow is due to run.

#### 4.4. Worker

*   The **Worker** is a background process that monitors schedules and executes tasks.
*   The worker is managed using the `worker_group` CLI commands (`start`, `stop`, `status`).
*   The worker uses the `SchedulerManager` and `FlowManager` to process schedules and run flows.
*   A worker is registered in the database table `Worker` with its hostname, pid and status.

#### 4.5. API

*   The **API** provides endpoints for interacting with flows, tasks, and schedules.
*   API endpoints are defined in the `automagik/api/routers/` directory.
*   API models (request/response schemas) are defined in `automagik/api/models.py`.
*   API authentication is handled using an API key (passed in the `X-API-Key` header).

#### 4.6. CLI

*   The **CLI** provides commands for managing the application (e.g., `automagik flow list`, `automagik schedule create`, `automagik worker start`).
*   CLI commands are defined in the `automagik/cli/commands/` directory.

### 5. Development Guidelines

*   **Code Style:** Follow PEP 8 style guidelines. Use `black` for code formatting and `isort` for import sorting.
*   **Error Handling:** Implement proper error handling and logging.
*   **Database Interactions:** Use the provided `get_session` context manager for all database interactions.
*   **Testing:** Write unit and integration tests for all new features and bug fixes. Aim for high test coverage.
*   **Documentation:** Document your code using docstrings and comments. Update this system prompt as needed.

### 6. Example Use Cases

#### 6.1. Implementing a New Feature

1. **Identify where the code should go:**
    *   If it's a new API endpoint, add it to the relevant file in `automagik/api/routers/`.
    *   If it's a new CLI command, add it to the relevant file in `automagik/cli/commands/`.
    *   If it's a core logic change, modify the appropriate module in `automagik/core/`.
2. **Write tests:** Create new test files or update existing ones in the `tests/` directory.
3. **Implement the feature:** Write the code, following the established patterns and guidelines.
4. **Run tests:** Use `pytest` to run the entire test suite.
5. **Commit changes:** Use `git commit` with a clear commit message.

#### 6.2. Fixing a Bug

1. **Reproduce the bug:** Write a test that reproduces the bug.
2. **Identify the cause:** Use debugging tools and logs to find the root cause.
3. **Fix the bug:** Modify the relevant code.
4. **Run tests:** Ensure the bug is fixed and no regressions are introduced.
5. **Commit changes:** Use `git commit` with a descriptive commit message.

#### 6.3. Refactoring Code

1. **Identify areas for improvement:** Look for code that is duplicated, complex, or inefficient.
2. **Plan the refactoring:** Determine how to improve the code while maintaining existing functionality.
3. **Refactor:** Make the changes, ensuring that tests still pass.
4. **Run tests:** Verify that the refactoring didn't break anything.
5. **Commit changes:** Use `git commit` with a message explaining the refactoring.

#### 6.4 Adding a new CLI command

1. **Create a new `.py` file** for your command within the appropriate subgroup (e.g., `automagik/cli/commands/flow.py` for a flow-related command).
2. **Define your command function** using the `@click.command()` decorator. Use `@click.option()` and `@click.argument()` to define command options and arguments.
3. **Add your command function** to the appropriate command group using `@<command_group_name>.command()`.
4. **Implement the command logic** within your function. Use the `get_session` context manager for database interactions.
5. **Add tests** for your new command in the corresponding test file (e.g., `tests/cli/commands/test_flow.py`).
6. **Run tests** using `pytest` to ensure your command works as expected.

---

