smart-org/apps/api/tests/test_credit_chat_gate.py

"""The manager chat does not call a model for an organization with no credit.

What these tests defend:

* a zero-credit organization gets the note as the manager's reply: the model is
  never called and the tags in the order are not executed (prepaid, decision D-B)
* a missing wallet is no credit either
* once the owner credits the organization, the same chat reaches the model
* a manager built without a wallet (the direct construction the other manager
  tests use) keeps its old behaviour

The organization fixture and the manager's model profile come from
``tests.test_engine_manager``.
"""

from __future__ import annotations

from app.application.container_org import build_memory
from app.application.engine.manager import CHAT_NO_CREDIT_MESSAGE, ManagerTurns
from app.domain.workforce.models import MessageRole, RoleKey
from app.infrastructure.db.models import CreditWalletRow, ThreadRow
from app.infrastructure.repositories.ai_control import (
    AiModelRepository,
    ModelCredentialRepository,
    RateCardRepository,
)
from app.infrastructure.repositories.audit import UsageLedgerRepository
from app.infrastructure.repositories.billing import CreditWalletRepository
from app.infrastructure.repositories.engine import (
    EmployeeProfileRepository,
    PipelineStateRepository,
)
from app.infrastructure.repositories.workforce import (
    EmployeeRepository,
    MessageRepository,
    TaskRepository,
    ThreadRepository,
)
from tests.billing_helpers import fund_org
from tests.engine_fakes import FakeGrants, ScriptedModel, ScriptedReply
from tests.test_engine_manager import _with_model_profile
from tests.test_engine_manager import world as manager_world

#: The engine manager tests' organization fixture. Binding it under its own name
#: registers it for this module's tests (a bare import would be flagged as unused).
world = manager_world

_ORDER = '[[CREATE_TASK title="ورود" kind="feature"]]ساخت فرم ورود[[/CREATE_TASK]]'
_PLAIN_REPLY = "آماده‌ام."


async def _thread_with_manager(world, session) -> str:
    """A manager with a strong model profile, and a thread to talk in."""
    manager = await world.employee(session, "مدیر", RoleKey.MANAGER)
    await _with_model_profile(world, session, manager.id, world.strong_id)
    thread = await ThreadRepository(session).create(
        ThreadRow(organization_id=world.organization_id, title="گفتگو")
    )
    await session.commit()
    return thread.id


async def _manager(session, *, scripted, grants, wired: bool) -> ManagerTurns:
    """The engine manager on real repositories; ``wired`` adds the prepaid gate."""
    return ManagerTurns(
        provider=scripted,
        employees=EmployeeRepository(session),
        messages=MessageRepository(session),
        tasks=TaskRepository(session),
        threads=ThreadRepository(session),
        models=AiModelRepository(session),
        credentials=ModelCredentialRepository(session),
        rate_cards=RateCardRepository(session),
        ledger=UsageLedgerRepository(session),
        grants=grants,
        memory=build_memory(session),
        profiles=EmployeeProfileRepository(session),
        states=PipelineStateRepository(session),
        runner=None,
        wallets=CreditWalletRepository(session) if wired else None,
    )


async def test_a_zero_credit_org_gets_the_note_and_no_model_call(world, session_factory):
    async with session_factory() as session:
        session.add(
            CreditWalletRow(organization_id=world.organization_id, balance=0.0, currency="USD")
        )
        thread_id = await _thread_with_manager(world, session)
    scripted = ScriptedModel([ScriptedReply(content=_ORDER)])
    grants = FakeGrants({world.strong_id: "strong"})

    async with session_factory() as session:
        manager = await _manager(session, scripted=scripted, grants=grants, wired=True)
        user_message, assistant = await manager.chat(
            organization_id=world.organization_id,
            thread_id=thread_id,
            content="یک فرم ورود بساز.",
            user_id=world.user_id,
        )
        await session.commit()
        tasks = await TaskRepository(session).list_for_org(world.organization_id)

    assert user_message.content == "یک فرم ورود بساز."
    assert assistant.role is MessageRole.ASSISTANT
    assert assistant.content == CHAT_NO_CREDIT_MESSAGE
    assert scripted.calls == 0
    # The order's CREATE_TASK tag was not executed.
    assert tasks == []


async def test_a_missing_wallet_counts_as_no_credit(world, session_factory):
    async with session_factory() as session:
        thread_id = await _thread_with_manager(world, session)
    scripted = ScriptedModel([ScriptedReply(content=_PLAIN_REPLY)])
    grants = FakeGrants({world.strong_id: "strong"})

    async with session_factory() as session:
        manager = await _manager(session, scripted=scripted, grants=grants, wired=True)
        _, assistant = await manager.chat(
            organization_id=world.organization_id,
            thread_id=thread_id,
            content="سلام",
            user_id=world.user_id,
        )
        await session.commit()

    assert assistant.content == CHAT_NO_CREDIT_MESSAGE
    assert scripted.calls == 0


async def test_after_the_owner_credits_the_org_the_model_answers(world, session_factory):
    async with session_factory() as session:
        thread_id = await _thread_with_manager(world, session)
    await fund_org(session_factory, world.organization_id, 50.0)
    scripted = ScriptedModel([ScriptedReply(content=_PLAIN_REPLY)])
    grants = FakeGrants({world.strong_id: "strong"})

    async with session_factory() as session:
        manager = await _manager(session, scripted=scripted, grants=grants, wired=True)
        _, assistant = await manager.chat(
            organization_id=world.organization_id,
            thread_id=thread_id,
            content="سلام",
            user_id=world.user_id,
        )
        await session.commit()

    assert assistant.content == _PLAIN_REPLY
    assert scripted.calls == 1


async def test_a_manager_built_without_a_wallet_keeps_the_old_behaviour(world, session_factory):
    """No ``wallets`` means no gate: the direct constructions in other tests rely on it."""
    async with session_factory() as session:
        thread_id = await _thread_with_manager(world, session)
    scripted = ScriptedModel([ScriptedReply(content=_PLAIN_REPLY)])
    grants = FakeGrants({world.strong_id: "strong"})

    async with session_factory() as session:
        manager = await _manager(session, scripted=scripted, grants=grants, wired=False)
        _, assistant = await manager.chat(
            organization_id=world.organization_id,
            thread_id=thread_id,
            content="سلام",
            user_id=world.user_id,
        )
        await session.commit()

    assert assistant.content == _PLAIN_REPLY
    assert scripted.calls == 1