PMFlow/backend/tests/test_model_routing_service.py

from __future__ import annotations

import json
import os

import pytest
import pytest_asyncio
from sqlalchemy import JSON, and_, select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

os.environ["DATABASE_URL"] = "postgresql+asyncpg://pmflow:pmflow_pass@127.0.0.1:5435/pmflow"

from app.models.ai_call_log import AICallLog
from app.models.base import Base
from app.models.integration import Integration
from app.models.user import User
from app.services.llm_service import llm_service
from app.services.llm_service import mixin_prompt_profile_runtime
from app.services.model.model_config import load_model_config
from app.services.model.model_context_bundle import ModelContextBundleService
from app.services.model.model_routing import ModelRoutingService


def _use_sqlite_json_types() -> None:
    for table in (Integration.__table__, User.__table__):
        for column in table.columns:
            if column.name in {"config", "assigned_domains"}:
                column.type = JSON()


@pytest_asyncio.fixture
async def db_env():
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    _use_sqlite_json_types()
    async with engine.begin() as conn:
        await conn.run_sync(
            lambda sync_conn: Base.metadata.create_all(
                sync_conn,
                tables=[
                    User.__table__,
                    Integration.__table__,
                    AICallLog.__table__,
                ],
            )
        )
    maker = async_sessionmaker(engine, expire_on_commit=False)
    async with maker() as session:
        yield session, maker
    await engine.dispose()


def test_load_model_config_rejects_safety_tier_configuration(tmp_path) -> None:
    config_path = tmp_path / "pmflow-routing-rules.json"
    config_path.write_text(
        json.dumps(
            {
                "version": "v-test",
                "updated_at": "2026-03-26T22:00:00+08:00",
                "tier_defaults": {
                    "low": {"provider": "glm", "model": "qwen-plus", "api_env": "DASHSCOPE_API_KEY"},
                    "high": {"provider": "openai", "model": "GLM-5", "api_env": "OPENAI_API_KEY"},
                    "safety": {"provider": "glm", "model": "qwen-plus", "api_env": "DASHSCOPE_API_KEY"},
                },
                "rules": [
                    {
                        "name": "bad_safety_rule",
                        "enabled": True,
                        "priority": 1,
                        "tier": "safety",
                        "match_all": True,
                    }
                ],
            },
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )

    with pytest.raises(ValueError, match="low/high"):
        load_model_config(refresh=True, path=config_path)


@pytest.mark.asyncio
async def test_model_routing_matches_runtime_rule_and_returns_logical_view(db_env, monkeypatch) -> None:
    db_session, _maker = db_env
    monkeypatch.setenv("OPENCLAW_ROUTING_ENABLED", "true")

    runtime_rule = llm_service._match_openclaw_rule("classify_intent_complex")
    route = await ModelRoutingService.get_effective_routing(
        db_session,
        tenant_id=1,
        function_name="classify_intent_complex",
    )

    assert runtime_rule is not None
    assert route["function_name"] == "classify_intent_complex"
    assert route["tier_alias"] == "high"
    assert route["source"] == "default"
    assert route["rule_name"] == runtime_rule["name"]
    assert route["rule_priority"] == runtime_rule["priority"]
    assert route["last_execution"] is None
    assert "provider" not in route
    assert "model" not in route


@pytest.mark.asyncio
async def test_model_routing_override_is_tier_only_and_reads_last_execution(db_env) -> None:
    db_session, _maker = db_env

    route = await ModelRoutingService.upsert_override(
        db_session,
        tenant_id=1,
        function_name="signal_planner_analyze",
        tier_alias="low",
        updated_by=9,
    )
    db_session.add(
        AICallLog(
            tenant_id=1,
            user_id=None,
            function_name="signal_planner_analyze",
            prompt="prompt",
            response="response",
            tokens_used=12,
            cost=0.0001,
            status="success_failover",
            planned_tier="low",
            executed_tier="high",
            routing_decision="low_to_high_escalated",
            route_source="override",
            route_rule_name="engine_high_reasoning_glm5",
            routing_config_version="v3",
            routing_config_updated_at="2026-03-26T21:30:00+08:00",
        )
    )
    await db_session.commit()

    refreshed = await ModelRoutingService.get_effective_routing(
        db_session,
        tenant_id=1,
        function_name="signal_planner_analyze",
    )
    row = (
        await db_session.execute(
            select(Integration).where(
                and_(
                    Integration.tenant_id == 1,
                    Integration.type == "model_routing",
                )
            )
        )
    ).scalar_one()
    mapping = row.config["mappings"]["signal_planner_analyze"]

    assert route["tier_alias"] == "low"
    assert refreshed["source"] == "override"
    assert refreshed["last_execution"] == {
        "planned_tier": "low",
        "executed_tier": "high",
        "routing_decision": "low_to_high_escalated",
        "created_at": refreshed["last_execution"]["created_at"],
    }
    assert mapping["tier_alias"] == "low"
    assert "provider" not in mapping
    assert "model" not in mapping
    assert "api_env" not in mapping


@pytest.mark.asyncio
async def test_runtime_low_override_does_not_downgrade_protected_adaptive_rag_functions(db_env) -> None:
    db_session, _maker = db_env

    await ModelRoutingService.upsert_override(
        db_session,
        tenant_id=1,
        function_name="adaptive_rag_answer_deep",
        tier_alias="low",
        updated_by=9,
    )

    low_engine = {
        "id": "low-1",
        "provider": "glm",
        "model": "qwen-plus",
        "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
        "api_key": "low-key",
    }
    high_engine = {
        "id": "high-1",
        "provider": "openai",
        "model": "GLM-5",
        "base_url": "https://api.openai.com/v1",
        "api_key": "high-key",
    }

    routed = await llm_service._resolve_function_override_engine(
        db_session,
        tenant_id=1,
        function_name="adaptive_rag_answer_deep",
        engines=[low_engine, high_engine],
        active_engine=low_engine,
    )

    assert routed is None


@pytest.mark.asyncio
async def test_runtime_low_override_allows_adaptive_rag_answer_normal_to_use_low_tier(db_env) -> None:
    db_session, _maker = db_env

    await ModelRoutingService.upsert_override(
        db_session,
        tenant_id=1,
        function_name="adaptive_rag_answer_normal",
        tier_alias="low",
        updated_by=9,
    )

    low_engine = {
        "id": "low-1",
        "provider": "glm",
        "model": "qwen-plus",
        "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
        "api_key": "low-key",
    }
    high_engine = {
        "id": "high-1",
        "provider": "openai",
        "model": "GLM-5",
        "base_url": "https://api.openai.com/v1",
        "api_key": "high-key",
    }

    routed = await llm_service._resolve_function_override_engine(
        db_session,
        tenant_id=1,
        function_name="adaptive_rag_answer_normal",
        engines=[low_engine, high_engine],
        active_engine=low_engine,
    )

    assert routed is not None
    assert routed["provider"] == "glm"
    assert routed["model"] == "qwen-plus"


@pytest.mark.asyncio
async def test_model_routing_list_includes_current_business_entry_points(db_env) -> None:
    db_session, _maker = db_env

    items = await ModelRoutingService.list_effective_routing(db_session, tenant_id=1)
    names = {item["function_name"] for item in items}

    assert "ai_pmo_signal_title_refine" in names
    assert "ai_pmo_signal_batch_analysis" in names
    assert "ai_service_stream_chat" in names
    assert "draft_generator_v1" in names
    assert "signal_source_onboarding_plan" in names
    assert "signal_source_ai_worker_execution" in names
    assert "project_draft_analysis_glm47_fallback" not in names


@pytest.mark.asyncio
async def test_log_ai_call_persists_routing_metadata_without_model_columns(db_env, monkeypatch) -> None:
    db_session, maker = db_env
    bundle = load_model_config(refresh=True)
    monkeypatch.setattr(mixin_prompt_profile_runtime, "AsyncSessionLocal", maker)

    await llm_service._log_ai_call(
        tenant_id=1,
        function_name="signal_planner_analyze",
        prompt="planner prompt",
        response="planner response",
        status="success_failover",
        provider="openai",
        model="GLM-5",
        routing_meta={
            "planned_tier": "low",
            "executed_tier": "high",
            "routing_decision": "low_to_high_escalated",
            "route_source": "default",
            "route_rule_name": "engine_high_reasoning_glm5",
            "routing_config_version": str(bundle.get("version") or ""),
            "routing_config_updated_at": str(bundle.get("updated_at") or ""),
        },
    )

    row = (
        await db_session.execute(
            select(AICallLog).where(AICallLog.function_name == "signal_planner_analyze")
        )
    ).scalar_one()

    assert row.planned_tier == "low"
    assert row.executed_tier == "high"
    assert row.routing_decision == "low_to_high_escalated"
    assert row.route_source == "default"
    assert row.route_rule_name == "engine_high_reasoning_glm5"
    assert row.routing_config_version == bundle["version"]
    assert row.routing_config_updated_at == bundle["updated_at"]
    assert not hasattr(row, "model")


@pytest.mark.asyncio
async def test_model_context_bundle_summary_uses_routing_names(db_env, monkeypatch) -> None:
    db_session, _maker = db_env

    async def _fake_get_overview(self):
        _ = self
        return {
            "entity_total": 3,
            "relation_total": 2,
            "evidence_total": 5,
            "pending_events": 0,
            "latest_sync_run": {"id": 1, "run_type": "manual", "status": "success"},
        }

    async def _fake_list_templates(db, *, tenant_id: int):
        _ = (db, tenant_id)
        return [
            {
                "id": "tpl-low",
                "name": "Low Template",
                "description": "low",
                "updated_at": "2026-03-27T10:00:00+08:00",
            }
        ]

    async def _fake_contract_sync(db, *, tenant_id: int):
        _ = (db, tenant_id)
        return {"all_synced": True, "files": []}

    async def _fake_get_effective_routing(db, *, tenant_id: int, function_name: str):
        _ = (db, tenant_id, function_name)
        return {
            "tier_alias": "high" if function_name == "classify_intent_complex" else "low",
            "source": "default",
            "rule_name": f"{function_name}_rule",
        }

    monkeypatch.setattr(
        "app.services.model.model_context_bundle.mixin_build.ProjectKnowledgeGraphService.get_overview",
        _fake_get_overview,
    )
    monkeypatch.setattr(
        "app.services.model.model_context_bundle.mixin_build.TaskBriefService.list_templates",
        _fake_list_templates,
    )
    monkeypatch.setattr(
        "app.services.model.model_context_bundle.mixin_build.TaskBriefService.get_contract_sync_status",
        _fake_contract_sync,
    )
    monkeypatch.setattr(
        "app.services.model.model_context_bundle.mixin_build.ModelRoutingService.get_effective_routing",
        _fake_get_effective_routing,
    )

    summary = await ModelContextBundleService.build_capability_summary(db_session, tenant_id=1)

    assert "function_routing_template_matrix" in summary
    assert "function_model_template_matrix" not in summary
    first_row = summary["function_routing_template_matrix"][0]
    assert "routing_snapshot" in first_row
    assert "model_routing" not in first_row