smart-org/apps/api/app/seed.py

"""Idempotent development seed.

Run with ``pnpm db:seed``. Creates a platform owner, a demo user, a sample
organization that both of them own, the system agent templates, routing
policies, rate cards and a mock provider with two models.

Safe to run repeatedly: every step checks for existing rows first.

The two accounts it prints at the end are meant to be usable immediately —
``tests/test_seed.py`` asserts exactly that, because a seed whose credentials the
login endpoint rejects is worse than no seed at all.
"""

from __future__ import annotations

import asyncio
import logging

from sqlalchemy import select

from app.core.database import Base, SessionFactory, engine
from app.core.security import hash_password
from app.domain.ai_control.models import (
    AiModel,
    Capability,
    Provider,
    ProviderType,
    RateCard,
    RateUnit,
    RoutingMode,
    RoutingPolicy,
)
from app.domain.identity.models import Membership, MembershipStatus, Organization, Role, User
from app.domain.sandbox.models import SandboxProviderKind, SandboxTemplate
from app.domain.workforce.models import AgentTemplate, RoleKey
from app.infrastructure.db.models import (
    AgentTemplateRow,
    AiModelRow,
    AiProviderRow,
    CreditWalletRow,
    RateCardRow,
    RoutingPolicyRow,
    SandboxTemplateRow,
    UserRow,
)

logging.basicConfig(level=logging.INFO, format="%(levelname)-7s %(message)s")
logger = logging.getLogger("agy.seed")

#: ``example.com`` is a documentation domain and passes email validation.
#: Do NOT use ``.local``/``.test``/``.invalid``: those are special-use TLDs and
#: ``email-validator`` rejects them, so the seed would create an account that the
#: API then refuses to log in — the user exists but is unreachable.
PLATFORM_OWNER_EMAIL = "owner@example.com"
PLATFORM_OWNER_PASSWORD = "owner12345"
DEMO_EMAIL = "demo@example.com"
DEMO_PASSWORD = "demo12345"
#: Stable identifiers for the sample organization. Both must stay constant for
#: the seed to be idempotent — see ``_ensure_org``.
DEMO_ORG_NAME = "سازمان نمونه"
DEMO_ORG_SLUG = "sample-org"

AGENT_TEMPLATES: list[tuple[str, RoleKey, str, list[str]]] = [
    (
        "مدیر سازمان", RoleKey.MANAGER,
        "کارها را دسته‌بندی می‌کند و بین کارمندها تقسیم می‌کند.",
        ["knowledge.search", "artifact.create"],
    ),
    (
        "تحلیل‌گر و معمار", RoleKey.PLANNER,
        "نیازمندی‌ها را تحلیل و پلن اجرایی می‌سازد.",
        ["knowledge.search"],
    ),
    (
        "برنامه‌نویس", RoleKey.DEVELOPER,
        "کد می‌نویسد و در sandbox اجرا و تست می‌کند.",
        ["git.ensure_repo", "git.create_branch", "sandbox.run_command",
         "sandbox.read_file", "sandbox.write_file"],
    ),
    (
        "بازبین کد", RoleKey.REVIEWER,
        "PRها را بازبینی می‌کند و بازخورد می‌دهد.",
        ["git.open_pr"],
    ),
    (
        "تستر", RoleKey.TESTER,
        "تست می‌نویسد و اجرا می‌کند.",
        ["sandbox.run_command", "sandbox.write_file"],
    ),
    (
        "محقق", RoleKey.RESEARCHER,
        "منابع را جمع و خلاصه می‌کند.",
        ["knowledge.search", "artifact.create"],
    ),
    (
        "تولیدکننده محتوا", RoleKey.CONTENT_CREATOR,
        "متن و محتوای سازمانی تولید می‌کند.",
        ["artifact.create"],
    ),
    (
        "پشتیبان مشتری", RoleKey.SUPPORT,
        "به پرسش‌های مشتری پاسخ می‌دهد.",
        ["knowledge.search"],
    ),
]

ROUTING_POLICIES: list[tuple[str, RoutingMode, bool, str]] = [
    ("سریع و ارزان", RoutingMode.CHEAP, True, "برای کارهای روزمره و پرتکرار."),
    ("متعادل", RoutingMode.BALANCED, False, "پیش‌فرض متعادل بین هزینه و کیفیت."),
    ("استدلال عمیق", RoutingMode.REASONING, False, "برای تحلیل و کارهای پیچیده."),
    ("کدنویسی", RoutingMode.CODING, False, "بهینه برای تولید و اصلاح کد."),
    ("چندرسانه‌ای", RoutingMode.VISION, False, "برای ورودی تصویری."),
]

RATE_CARDS: list[tuple[RateUnit, float, float]] = [
    (RateUnit.MODEL_INPUT_TOKENS, 0.0000004, 0.0000002),
    (RateUnit.MODEL_OUTPUT_TOKENS, 0.0000016, 0.0000008),
    (RateUnit.SANDBOX_CPU_SECOND, 0.00002, 0.00001),
    (RateUnit.SANDBOX_MEMORY_GB_SECOND, 0.000004, 0.000002),
    (RateUnit.STORAGE_GB_DAY, 0.0004, 0.0002),
    (RateUnit.TOOL_CALL, 0.0005, 0.0001),
    (RateUnit.PREMIUM_EMPLOYEE_RUN, 0.02, 0.008),
]

#: (name, image, description, cpu, memory_mb, disk_mb)
SANDBOX_TEMPLATES: list[tuple[str, str, str, float, int, int]] = [
    (
        "python-3.13", "python:3.13-slim",
        "محیط پایتون برای اجرای تست و اسکریپت.", 1.0, 1024, 2048,
    ),
    (
        "node-22", "node:22-slim",
        "محیط Node برای build و اجرای پروژه‌های فرانت‌اند.", 1.0, 2048, 4096,
    ),
    (
        "ubuntu-base", "ubuntu:24.04",
        "محیط عمومی لینوکس برای ابزارهای خط فرمان.", 1.0, 1024, 2048,
    ),
]


async def _ensure_user(
    session, *, email: str, password: str, name: str, platform_owner: bool
) -> User:
    existing = (
        await session.execute(select(UserRow).where(UserRow.email == email))
    ).scalar_one_or_none()
    if existing:
        return User(
            id=existing.id, email=existing.email, password_hash=existing.password_hash,
            name=existing.name, is_platform_owner=existing.is_platform_owner,
        )
    from app.infrastructure.repositories.identity import UserRepository

    user = await UserRepository(session).create(
        User(
            email=email,
            password_hash=hash_password(password),
            name=name,
            is_platform_owner=platform_owner,
        )
    )
    logger.info("created user %s", email)
    return user


async def _ensure_org(session, *, members: list[User]) -> Organization:
    """Get or create the sample organization, with every given user in it.

    Idempotent in all three parts: the organization, its memberships and its
    credit wallet. A developer who runs ``pnpm db:seed`` twice must end up with
    exactly one organization and no duplicate rows.
    """
    from app.infrastructure.repositories.identity import (
        MembershipRepository,
        OrganizationRepository,
    )

    orgs = OrganizationRepository(session)
    memberships = MembershipRepository(session)

    # A *fixed* slug, not `slugify(name)`. The name is Persian, and `slugify`
    # falls back to a random `x<hex>` when it cannot transliterate — so deriving
    # the slug here would produce a different value on every run, `get_by_slug`
    # would always miss, and re-seeding would silently create a duplicate
    # organization (and a duplicate wallet) each time.
    org = await orgs.get_by_slug(DEMO_ORG_SLUG)
    if org is None:
        org = await orgs.create(
            Organization(name=DEMO_ORG_NAME, slug=DEMO_ORG_SLUG, owner_user_id=members[0].id)
        )
        logger.info("created organization %s", org.slug)

    # Every listed user becomes an owner of the sample org. Without this the
    # demo account can log in but belongs to no organization, so every screen
    # renders empty and the account looks broken.
    for user in members:
        if await memberships.get(org.id, user.id) is None:
            await memberships.add(
                Membership(
                    organization_id=org.id, user_id=user.id,
                    role=Role.OWNER, status=MembershipStatus.ACTIVE,
                )
            )
            logger.info("added %s to %s as owner", user.email, org.slug)

    wallet = (
        await session.execute(
            select(CreditWalletRow).where(CreditWalletRow.organization_id == org.id)
        )
    ).scalars().first()
    if wallet is None:
        session.add(CreditWalletRow(organization_id=org.id, balance=25.0, currency="USD"))
    return org


async def seed() -> None:
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

    async with SessionFactory() as session:
        owner = await _ensure_user(
            session, email=PLATFORM_OWNER_EMAIL, password=PLATFORM_OWNER_PASSWORD,
            name="مالک سامانه", platform_owner=True,
        )
        demo = await _ensure_user(
            session, email=DEMO_EMAIL, password=DEMO_PASSWORD,
            name="کاربر نمونه", platform_owner=False,
        )
        # Both accounts join the sample organization, so a developer who logs in
        # as either one lands on a populated dashboard rather than an empty one.
        await _ensure_org(session, members=[owner, demo])

        if not (await session.execute(select(AiProviderRow))).scalars().first():
            provider = Provider(name="mock", type=ProviderType.MOCK, base_url="")
            session.add(
                AiProviderRow(
                    id=provider.id, name=provider.name, type=provider.type.value,
                    base_url=provider.base_url, is_active=True,
                )
            )
            await session.flush()
            for internal, display, caps in (
                ("mock-fast", "Mock Fast", [Capability.TEXT]),
                (
                    "mock-reasoning",
                    "Mock Reasoning",
                    [Capability.TEXT, Capability.REASONING, Capability.CODE],
                ),
            ):
                model = AiModel(
                    provider_id=provider.id,
                    internal_name=internal,
                    display_name_for_admin=display,
                    capabilities=[str(c) for c in caps],
                    input_price_per_1m=0.3,
                    output_price_per_1m=1.2,
                )
                session.add(
                    AiModelRow(
                        id=model.id, provider_id=model.provider_id,
                        internal_name=model.internal_name,
                        display_name_for_admin=model.display_name_for_admin,
                        capabilities=model.capabilities,
                        input_price_per_1m=model.input_price_per_1m,
                        output_price_per_1m=model.output_price_per_1m,
                        is_active=True,
                    )
                )
            logger.info("created mock provider + 2 models")

        if not (await session.execute(select(AgentTemplateRow))).scalars().first():
            for name, role_key, description, tools in AGENT_TEMPLATES:
                template = AgentTemplate(
                    name=name, role_key=role_key, description=description,
                    default_prompt=f"تو {name} هستی. وظیفه‌ات: {description}",
                    default_tools_json=tools,
                )
                session.add(
                    AgentTemplateRow(
                        id=template.id, name=template.name, role_key=template.role_key.value,
                        description=template.description, default_prompt=template.default_prompt,
                        default_tools_json=template.default_tools_json, is_system=True,
                    )
                )
            logger.info("created %d agent templates", len(AGENT_TEMPLATES))

        if not (await session.execute(select(RoutingPolicyRow))).scalars().first():
            for name, mode, is_default, description in ROUTING_POLICIES:
                policy = RoutingPolicy(
                    name=name, mode=mode, description=description, is_default=is_default
                )
                session.add(
                    RoutingPolicyRow(
                        id=policy.id, name=policy.name, description=policy.description,
                        mode=policy.mode.value, is_default=policy.is_default, is_active=True,
                    )
                )
            logger.info("created %d routing policies", len(ROUTING_POLICIES))

        if not (await session.execute(select(RateCardRow))).scalars().first():
            for unit, customer_price, internal_cost in RATE_CARDS:
                card = RateCard(
                    name=unit.value.replace("_", " "),
                    unit=unit,
                    customer_price=customer_price,
                    internal_cost_estimate=internal_cost,
                )
                session.add(
                    RateCardRow(
                        id=card.id, name=card.name, unit=card.unit.value,
                        customer_price=card.customer_price,
                        internal_cost_estimate=card.internal_cost_estimate, is_active=True,
                    )
                )
            logger.info("created %d rate cards", len(RATE_CARDS))

        if not (await session.execute(select(SandboxTemplateRow))).scalars().first():
            for name, image, description, cpu, memory_mb, disk_mb in SANDBOX_TEMPLATES:
                # Named distinctly from the agent-template loop above so the two
                # entities cannot be confused.
                sandbox_template = SandboxTemplate(
                    name=name, provider=SandboxProviderKind.MOCK, image=image,
                    description=description, default_cpu=cpu,
                    default_memory_mb=memory_mb, default_disk_mb=disk_mb,
                )
                session.add(
                    SandboxTemplateRow(
                        id=sandbox_template.id, name=sandbox_template.name,
                        provider=str(sandbox_template.provider),
                        image=sandbox_template.image,
                        description=sandbox_template.description,
                        default_cpu=sandbox_template.default_cpu,
                        default_memory_mb=sandbox_template.default_memory_mb,
                        default_disk_mb=sandbox_template.default_disk_mb,
                        is_system=True,
                    )
                )
            logger.info("created %d sandbox templates", len(SANDBOX_TEMPLATES))

        await session.commit()

    logger.info("seed complete")
    logger.info("platform owner: %s / %s", PLATFORM_OWNER_EMAIL, PLATFORM_OWNER_PASSWORD)
    logger.info("demo user:      %s / %s", DEMO_EMAIL, DEMO_PASSWORD)


if __name__ == "__main__":
    try:
        asyncio.run(seed())
    except Exception as exc:
        # A failed seed is almost always a first-run problem — the very moment
        # there is nobody watching a terminal. Record it before exiting non-zero.
        from app.domain.audit.models import ErrorSeverity, ErrorSource
        from app.infrastructure.error_log import record_error

        asyncio.run(
            record_error(
                kind="seed_failed",
                message=str(exc) or type(exc).__name__,
                exc=exc,
                source=ErrorSource.STARTUP,
                severity=ErrorSeverity.CRITICAL,
            )
        )
        raise