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": "dashscope", "model": "qwen-plus", "api_env": "DASHSCOPE_API_KEY"},
                    "high": {"provider": "openai", "model": "GLM-5", "api_env": "OPENAI_API_KEY"},
                    "safety": {"provider": "dashscope", "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("CONTROL_PLANE_ROUTING_ENABLED", "true")

    runtime_rule = llm_service._match_control_plane_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_model_routing_rejects_low_override_for_protected_adaptive_rag_functions(db_env) -> None:
    db_session, _maker = db_env

    with pytest.raises(ValueError, match="保护场景不能降级"):
        await ModelRoutingService.upsert_override(
            db_session,
            tenant_id=1,
            function_name="adaptive_rag_answer_deep",
            tier_alias="low",
            updated_by=9,
        )


@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": "dashscope",
        "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"] == "dashscope"
    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_and_model_metadata(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",
        provider_request_payload={"model": "GLM-5", "messages": [{"role": "user", "content": "planner prompt"}]},
        provider_response_payload={"choices": [{"message": {"content": "planner response"}}]},
        provider_attempts=[{"status": "success", "request_payload": {"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 row.provider == "openai"
    assert row.model == "GLM-5"
    assert row.provider_request_payload["model"] == "GLM-5"
    assert row.provider_response_payload["choices"][0]["message"]["content"] == "planner response"
    assert row.provider_attempts[0]["status"] == "success"


@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


@pytest.mark.asyncio
async def test_get_embedding_async_falls_back_to_local_ollama_embed_when_active_engine_is_not_embedding_capable(
    db_env,
    monkeypatch,
) -> None:
    db_session, _maker = db_env
    captured: dict[str, object] = {}

    class _FakeResponse:
        def raise_for_status(self) -> None:
            return None

        def json(self) -> dict[str, object]:
            return {"embeddings": [[0.11, 0.22, 0.33]]}

    class _FakeAsyncClient:
        def __init__(self, *args, **kwargs) -> None:
            _ = args
            captured["init_kwargs"] = dict(kwargs)

        async def __aenter__(self):
            return self

        async def __aexit__(self, exc_type, exc, tb) -> None:
            _ = (exc_type, exc, tb)
            return None

        async def post(self, url: str, headers: dict[str, str], json: dict[str, object]):
            captured["url"] = url
            captured["headers"] = dict(headers)
            captured["payload"] = dict(json)
            return _FakeResponse()

    async def _fake_resolve_active_engine(db, *, tenant_id: int):
        _ = (db, tenant_id)
        return {
            "provider": "dashscope",
            "model": "qwen-plus",
            "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "api_key": "dashscope-key",
        }

    monkeypatch.setattr(llm_service, "_resolve_active_engine", _fake_resolve_active_engine)
    monkeypatch.setattr("app.services.llm_service.mixin_task_helpers.httpx.AsyncClient", _FakeAsyncClient)
    monkeypatch.setenv("OLLAMA_HOST", "http://127.0.0.1:11434")
    monkeypatch.setenv("OLLAMA_EMBED_MODEL", "nomic-embed-text:latest")

    embedding = await llm_service.get_embedding_async(
        "广州今天天气如何",
        db=db_session,
        tenant_id=1,
    )

    assert embedding == [0.11, 0.22, 0.33]
    assert captured["url"] == "http://127.0.0.1:11434/api/embed"
    assert captured["payload"] == {"model": "nomic-embed-text:latest", "input": "广州今天天气如何"}
    assert captured["init_kwargs"]["trust_env"] is False


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

    class _FailingAsyncClient:
        def __init__(self, *args, **kwargs) -> None:
            _ = (args, kwargs)

        async def __aenter__(self):
            return self

        async def __aexit__(self, exc_type, exc, tb) -> None:
            _ = (exc_type, exc, tb)
            return None

        async def post(self, url: str, headers: dict[str, str], json: dict[str, object]):
            _ = (url, headers, json)
            raise RuntimeError("embedding backend unavailable")

    async def _fake_resolve_active_engine(db, *, tenant_id: int):
        _ = (db, tenant_id)
        return {
            "provider": "dashscope",
            "model": "qwen-plus",
            "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "api_key": "dashscope-key",
        }

    monkeypatch.setattr(llm_service, "_resolve_active_engine", _fake_resolve_active_engine)
    monkeypatch.setattr("app.services.llm_service.mixin_task_helpers.httpx.AsyncClient", _FailingAsyncClient)
    monkeypatch.setenv("OLLAMA_HOST", "http://127.0.0.1:11434")
    monkeypatch.setenv("OLLAMA_EMBED_MODEL", "nomic-embed-text:latest")

    embedding = await llm_service.get_embedding_async(
        "今天荔湾区天气如何",
        db=db_session,
        tenant_id=1,
    )

    assert embedding == []