AgentEvalTool/backend/agenteval/evaluation/engine.py

"""Evaluation execution engine.

Async-first implementation: channels and LLM calls are awaited cooperatively,
so multiple cases can run concurrently and a run can be cancelled mid-flight
via an ``asyncio.Event`` cancel token.
"""

import asyncio
import uuid
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Callable, Optional

from agenteval.channels.base import EvalChannel, ExchangeOutcome, ExchangeStatus, SendResult
from agenteval.channels.factory import ChannelFactory
from agenteval.config import get_settings
from agenteval.evaluation.implicit_rules import derive_implicit_rules
from agenteval.evaluation.judgement import CaseOutcome, RuleOutcome, combine_case_outcome
from agenteval.evaluation.rules import RuleResult, get_rule
from agenteval.evaluation.run_summary import build_run_summary
from agenteval.model_gateway import ModelGateway
from agenteval.models import (
    Case,
    CaseType,
    EvalResult,
    EvalRun,
    EvalTarget,
    ModelCapability,
    ModelPurpose,
    RunStatus,
    RunTrigger,
    Scenario,
    Turn,
)
from agenteval.services.model_configs import ModelConfigService, ModelRuntimeConfig
from agenteval.storage.db import get_session, utc_now
from agenteval.storage.repository import ResultRepository, RunRepository
from agenteval.utils.llm import parse_json_from_llm_text

# Progress callbacks may be sync or async; the engine awaits the result if
# it is a coroutine, otherwise treats it as a plain function.
ProgressCallback = Callable[[str, dict[str, Any]], Any]


class CancelledError(RuntimeError):
    """Raised inside the engine when the cancel token fires."""


@dataclass
class TimeoutConfig:
    """Per-operation timeouts (all in seconds)."""

    poll_reply: float = 30.0
    llm_generate: float = 60.0


def _build_send_message(content: str) -> dict[str, Any]:
    return {
        "msgType": "text",
        "msgBody": {"content": content},
    }


def _reply_payload(outcome: ExchangeOutcome) -> Optional[dict[str, Any]]:
    """Recover the adapter payload kept in the opaque exchange diagnostic."""
    if not outcome.ok:
        return None
    diagnostic = outcome.diagnostic
    if isinstance(diagnostic, dict):
        reply = diagnostic.get("reply")
        if isinstance(reply, dict):
            return reply
    return {"msgBody": {"content": outcome.reply_text or ""}}


# 需要模型资源的规则类型 → 评测岗位（ModelPurpose）；其余规则无需模型
RULE_PURPOSE = {
    "llm_score": ModelPurpose.JUDGE,
    "semantic_similarity": ModelPurpose.EMBEDDING,
    "safety": ModelPurpose.MODERATION,
}


class EvalEngine:
    """Execute evaluation scenarios against targets.

    The engine is async so the (slow, network-bound) channel and LLM calls
    can be awaited cooperatively. Database writes remain synchronous for now
    (SQLite + StaticPool); they are fast enough not to block the event loop
    in practice.
    """

    def __init__(
        self,
        target: EvalTarget,
        scenario: Scenario,
        session=None,
        run_repo: Optional[RunRepository] = None,
        result_repo: Optional[ResultRepository] = None,
        cancel_token: Optional[asyncio.Event] = None,
        timeout_config: Optional[TimeoutConfig] = None,
        max_concurrent_cases: int = 1,
        triggered_by: RunTrigger = RunTrigger.MANUAL,
    ):
        self.target = target
        self.scenario = scenario
        self.channel: EvalChannel = ChannelFactory.create(target)
        self.session = session or get_session()
        self.run_repo = run_repo or RunRepository(self.session)
        self.result_repo = result_repo or ResultRepository(self.session)
        self.cancel_token = cancel_token or asyncio.Event()
        self.timeout_config = timeout_config or TimeoutConfig(
            poll_reply=get_settings().poll_reply_timeout,
        )
        self.triggered_by = triggered_by
        self._case_semaphore = asyncio.Semaphore(max(1, max_concurrent_cases))
        self._rule_semaphore = asyncio.Semaphore(
            max(1, get_settings().max_concurrent_rules),
        )
        self._state_lock = asyncio.Lock()
        # 分岗位的评测侧 token 用量（ModelPurpose.value → 用量），由各规则/生成器
        # 的按次用量归集而来，避免并发下用网关总量做差产生竞态。
        self._usage_by_purpose: dict[str, dict[str, int]] = {}
        # Collects fatal case-level errors (e.g. dynamic message generation
        # failures) so their cause is persisted into run.summary — not just
        # emitted transiently over WebSocket.
        self._case_errors: list[dict[str, str]] = []
        self.model_gateway = ModelGateway(timeout=self.timeout_config.llm_generate)
        self.model_service = ModelConfigService(self.session)
        self._resolved_models: dict[ModelPurpose, ModelRuntimeConfig] = {}

    # ── public entry point ────────────────────────────────────────────

    async def run(
        self,
        progress_callback: Optional[ProgressCallback] = None,
        existing_run: Optional[EvalRun] = None,
    ) -> EvalRun:
        """Run the evaluation and return the completed run record.

        Cancellation is cooperative: set ``cancel_token`` and the engine will
        mark the run as FAILED with ``cancelled_by_user`` at the next checkpoint.
        """
        if existing_run:
            run = existing_run
            run.status = RunStatus.RUNNING
            run.started_at = utc_now()
            run = self.run_repo.update(run) or run
        else:
            run = EvalRun(
                id=str(uuid.uuid4()),
                target_id=self.target.id or "",
                scenario_id=self.scenario.id or "",
                scenario_version=self.scenario.version or 1,
                status=RunStatus.RUNNING,
                triggered_by=self.triggered_by,
                started_at=utc_now(),
            )
            run = self.run_repo.create(run)

        try:
            total_cases = len(self.scenario.cases)
            case_outcomes: dict[str, CaseOutcome] = {}

            async def _case_worker(idx: int, case: Case) -> None:
                self._check_cancel()
                await self._emit(
                    progress_callback,
                    "case_start",
                    {
                        "index": idx,
                        "total": total_cases,
                        "case_id": case.id,
                    },
                )
                async with self._case_semaphore:
                    outcome, rule_pass, rule_total = await self._run_case(
                        run,
                        case,
                        progress_callback,
                    )
                async with self._state_lock:
                    case_outcomes[case.id] = outcome
                await self._emit(
                    progress_callback,
                    "case_end",
                    {
                        "index": idx,
                        "total": total_cases,
                        "case_id": case.id,
                        "passed": outcome.passed,
                        "rule_pass_count": rule_pass,
                        "rule_total": rule_total,
                    },
                )

            workers = [
                _case_worker(idx, case)
                for idx, case in enumerate(self.scenario.cases, start=1)
            ]
            try:
                await asyncio.gather(*workers)
            except Exception:
                raise

            async with self._state_lock:
                resolved_snapshot = {
                    purpose.value: config.snapshot()
                    for purpose, config in self._resolved_models.items()
                }

            results = self.run_repo.get_results(run.id)
            turns = self.run_repo.get_turns(run.id)
            usage = self.model_gateway.total_usage
            summary = build_run_summary(
                case_outcomes=case_outcomes,
                latencies=[t.latency_ms for t in turns if t.latency_ms is not None],
                rule_passes=[r.passed for r in results],
                case_errors=self._case_errors or None,
                model_configs=resolved_snapshot or None,
                eval_token_usage=usage if usage["total_tokens"] > 0 else None,
                eval_usage_by_purpose=self._usage_by_purpose or None,
            )
            run.status = RunStatus.COMPLETED
            run.completed_at = utc_now()
            run.summary = summary
            await self._emit(
                progress_callback,
                "run_completed",
                {
                    "status": "completed",
                    "summary": summary.model_dump(),
                },
            )
        except CancelledError:
            run.status = RunStatus.FAILED
            run.completed_at = utc_now()
            run.summary = {
                "error": {"code": "cancelled_by_user", "message": "评测已手动停止"},
            }
            await self._emit(
                progress_callback,
                "run_completed",
                {
                    "status": "failed",
                    "reason": "cancelled",
                    "error": {"code": "cancelled_by_user", "message": "评测已手动停止"},
                },
            )
        except Exception as exc:
            run.status = RunStatus.FAILED
            run.completed_at = utc_now()
            run.summary = {"error": str(exc)}
            await self._emit(progress_callback, "error", {"error": str(exc)})
            await self._emit(
                progress_callback,
                "run_completed",
                {
                    "status": "failed",
                    "error": str(exc),
                },
            )
            raise
        finally:
            run = self.run_repo.update(run) or run
            # Best-effort cleanup of the channel's HTTP client.
            for resource in (self.channel, self.model_gateway):
                close = getattr(resource, "close", None)
                if callable(close):
                    try:
                        result = close()
                        if asyncio.iscoroutine(result):
                            await result
                    except Exception:
                        pass
            try:
                self.session.close()
            except Exception:
                pass

        return run

    # ── case / turn execution ─────────────────────────────────────────

    def _persist_turn(
        self,
        run: EvalRun,
        case: Case,
        round_index: int,
        message: str,
        sent_at: datetime,
        *,
        question_msg_id: Optional[str] = None,
        reply: Optional[dict] = None,
        received_at: Optional[datetime] = None,
        latency_ms: Optional[int] = None,
    ) -> Turn:
        """Build, persist, and return a Turn. The three call sites (send-fail /
        poll-except / happy path) differ only in which optional fields are set."""
        turn = Turn(
            id=str(uuid.uuid4()),
            run_id=run.id,
            case_id=case.id,
            round_index=round_index,
            sent_message=_build_send_message(message),
            sent_at=sent_at,
            question_msg_id=question_msg_id,
            reply=reply,
            received_at=received_at,
            latency_ms=latency_ms,
        )
        self.result_repo.save_turn(turn)
        return turn

    def _complete_turn(self, turn: Turn, outcome: ExchangeOutcome, received_at: datetime) -> Turn:
        """Attach exchange facts without replacing the already-sent ledger row."""
        reply = _reply_payload(outcome)
        if (
            self.result_repo.update_turn_exchange(
                turn.id or "",
                question_msg_id=outcome.correlation_id,
                reply=reply,
                received_at=received_at,
                latency_ms=outcome.latency_ms,
            )
            is None
        ):
            raise RuntimeError(f"persisted turn disappeared: {turn.id}")
        turn.question_msg_id = outcome.correlation_id
        turn.reply = reply
        turn.received_at = received_at
        turn.latency_ms = outcome.latency_ms
        return turn

    async def _run_case(
        self,
        run: EvalRun,
        case: Case,
        progress_callback: Optional[ProgressCallback],
    ) -> tuple[CaseOutcome, int, int]:
        """Run a single case; returns (outcome, passed_rules, total_rules)."""
        failed = CaseOutcome(passed=False, connectivity=False)
        if case.type == CaseType.DYNAMIC:
            generated = await self._generate_messages(case, progress_callback)
            if not generated:
                await self._emit(
                    progress_callback,
                    "error",
                    {
                        "error": "LLM 未能生成测试消息",
                        "case_id": case.id,
                    },
                )
                return failed, 0, 0
            case = case.model_copy(update={"messages": generated})

        dialog: list[Turn] = []

        for round_index, message in enumerate(case.messages, start=1):
            self._check_cancel()
            await self._emit(
                progress_callback,
                "turn_start",
                {
                    "run_id": run.id,
                    "case_id": case.id,
                    "round": round_index,
                    "message": message,
                },
            )

            sent_at = utc_now()
            turn: Optional[Turn] = None

            async def record_sent(send_result: SendResult) -> None:
                nonlocal turn
                turn = self._persist_turn(
                    run,
                    case,
                    round_index,
                    message,
                    sent_at,
                    question_msg_id=send_result.question_msg_id,
                )

            outcome = await self.channel.exchange(
                message,
                timeout=self.timeout_config.poll_reply,
                on_sent=record_sent,
            )
            if outcome.status is ExchangeStatus.SEND_FAILED:
                turn = self._persist_turn(run, case, round_index, message, sent_at)
                await self._save_rule_results(run, case, turn, [], progress_callback)
                await self._emit(
                    progress_callback,
                    "turn_error",
                    {
                        "case_id": case.id,
                        "round": round_index,
                        "error": outcome.reason,
                    },
                )
                return CaseOutcome(passed=False, connectivity=False, abandoned=bool(dialog)), 0, 0

            if turn is None:
                raise RuntimeError("channel exchange succeeded without invoking the sent hook")

            received_at = utc_now()
            turn = self._complete_turn(turn, outcome, received_at)
            if outcome.status is ExchangeStatus.POLL_FAILED:
                await self._emit(
                    progress_callback,
                    "turn_error",
                    {
                        "case_id": case.id,
                        "round": round_index,
                        "error": f"poll_reply 异常: {outcome.reason}",
                    },
                )
                return CaseOutcome(passed=False, connectivity=False, abandoned=bool(dialog)), 0, 0

            dialog.append(turn)

            await self._emit(
                progress_callback,
                "turn_end",
                {
                    "run_id": run.id,
                    "case_id": case.id,
                    "round": round_index,
                    "latency_ms": outcome.latency_ms,
                    "reply_text": outcome.reply_text or "",
                },
            )

        if not dialog:
            return failed, 0, 0

        return await self._save_rule_results(run, case, dialog[-1], dialog, progress_callback)

    async def _save_rule_results(
        self,
        run: EvalRun,
        case: Case,
        turn: Turn,
        dialog: list[Turn],
        progress_callback: Optional[ProgressCallback],
    ) -> tuple[CaseOutcome, int, int]:
        """Apply rules and save results; returns (outcome, passed_count, total_count).

        判定组合本身在 judgement.combine_case_outcome（单一权威）——
        本方法只负责执行规则、持久化结果并把规则输出规范化为 RuleOutcome。
        """
        from agenteval.models import EvalRuleConfig

        rules_config: list[EvalRuleConfig] = list(case.eval_rules)
        implicit_config = derive_implicit_rules(case.expectations)

        all_replied = bool(dialog) and all(t.reply is not None for t in dialog)

        if not rules_config and not implicit_config:
            # 连通用例：收到全部回复才通过（无回复=故障，ADR-0002）
            return combine_case_outcome(all_replied=all_replied), 0, 0

        passed_count = 0
        total_count = 0
        explicit_outcomes: list[RuleOutcome] = []
        implicit_outcomes: list[RuleOutcome] = []

        all_rules = [(cfg, False) for cfg in rules_config] + [(cfg, True) for cfg in implicit_config]

        async def _eval_one(
            rule_config: EvalRuleConfig,
            is_implicit: bool,
        ) -> tuple[EvalRuleConfig, bool, RuleResult]:
            async with self._rule_semaphore:
                purpose = RULE_PURPOSE.get(rule_config.type)
                try:
                    model_config = await self._resolve_model(purpose) if purpose else None
                    rule = get_rule(
                        rule_config.type,
                        rule_config.params,
                        model_config=model_config,
                        gateway=self.model_gateway if model_config else None,
                    )
                    result = await rule.evaluate(case, dialog)
                    if purpose:
                        await self._accumulate_usage(purpose.value, rule.llm_usage)
                except Exception as exc:
                    result = RuleResult(passed=False, reason=f"模型配置解析失败: {exc}")
                return rule_config, is_implicit, result

        rule_results = await asyncio.gather(
            *[_eval_one(cfg, is_impl) for cfg, is_impl in all_rules]
        )

        passed_count = 0
        total_count = 0
        explicit_outcomes: list[RuleOutcome] = []
        implicit_outcomes: list[RuleOutcome] = []

        for rule_config, is_implicit, result in rule_results:
            reason = f"[期望] {result.reason}" if is_implicit else result.reason
            eval_result = EvalResult(
                id=str(uuid.uuid4()),
                run_id=run.id,
                case_id=case.id,
                turn_id=turn.id or "",
                rule_type=rule_config.type,
                passed=result.passed,
                score=result.score,
                reason=reason,
            )
            self.result_repo.save_result(eval_result)
            total_count += 1
            if result.passed:
                passed_count += 1

            rule_outcome = RuleOutcome(passed=result.passed, score=result.score, weight=rule_config.weight)
            if is_implicit:
                implicit_outcomes.append(rule_outcome)
            else:
                explicit_outcomes.append(rule_outcome)

            await self._emit(
                progress_callback,
                "rule_result",
                {
                    "run_id": run.id,
                    "case_id": case.id,
                    "rule_type": rule_config.type,
                    "passed": result.passed,
                    "score": result.score,
                    "reason": reason,
                    "weight": rule_config.weight,
                },
            )

        outcome = combine_case_outcome(
            all_replied=all_replied,
            explicit=explicit_outcomes,
            implicit=implicit_outcomes,
            rule_logic=case.rule_logic,
            threshold=case.rule_pass_threshold,
        )
        return outcome, passed_count, total_count

    async def _generate_messages(
        self,
        case: Case,
        progress_callback: Optional[ProgressCallback],
    ) -> list[str]:
        """Use LLM to generate test messages for dynamic cases."""

        async def _fail(msg: str) -> list[str]:
            # Persist the reason into run-level case_errors (surfaced in summary),
            # not just a transient WebSocket emit that's lost after the run.
            async with self._state_lock:
                self._case_errors.append({"case_id": case.id, "stage": "generate_messages", "error": msg})
            await self._emit(progress_callback, "error", {"error": msg, "case_id": case.id})
            return []

        turns = case.turns or 3
        prompt = case.prompt or "请生成一些测试问题"

        system_prompt = (
            f"你需要扮演一个真实的用户/患者，根据以下要求生成 {turns} 条独立的测试问题。\n\n"
            f"要求：{prompt}\n\n"
            "输出格式要求：只输出一个 JSON 数组，包含 " + str(turns) + " 个字符串，每个字符串是一条消息。"
            "不要输出任何解释、markdown 或其他内容。"
        )

        messages_payload = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"请生成 {turns} 条测试消息"},
        ]

        try:
            model_config = await self._resolve_model(ModelPurpose.GENERATOR)
            if model_config:
                content, usage = await self.model_gateway.chat_with_usage(
                    model_config, messages_payload, temperature=0.7
                )
                await self._accumulate_usage(ModelPurpose.GENERATOR.value, usage)
            else:
                content = await self._generate_messages_legacy(messages_payload)

            try:
                parsed = parse_json_from_llm_text(content)
            except (ValueError, Exception) as parse_exc:
                return await _fail(f"LLM 返回无法解析为数组: {parse_exc}")

            if not isinstance(parsed, list):
                return await _fail("LLM 返回的不是数组")

            messages = [str(m) for m in parsed if isinstance(m, str) and m.strip()]
            if not messages:
                return await _fail("LLM 返回的消息为空")

            await self._emit(
                progress_callback,
                "messages_generated",
                {
                    "case_id": case.id,
                    "messages": messages,
                },
            )
            return messages

        except Exception as exc:
            return await _fail(f"LLM 生成消息失败: {exc}")

    async def _generate_messages_legacy(self, messages: list[dict[str, str]]) -> str:
        """Temporary fallback for scenarios not yet migrated to model bindings."""
        import httpx

        from agenteval.utils.llm import extract_content_from_llm_response

        llm_config = self.scenario.llm_config
        if not llm_config or not llm_config.get("api_url"):
            raise ValueError("动态用例未绑定生成模型，且兼容 llm_config 缺少 api_url")
        headers = {"Content-Type": "application/json"}
        if llm_config.get("api_key"):
            headers["Authorization"] = f"Bearer {llm_config['api_key']}"
        payload = {
            "model": llm_config.get("model", "doubao-seed-2.0-lite"),
            "messages": messages,
            "temperature": 0.7,
        }
        async with httpx.AsyncClient(timeout=self.timeout_config.llm_generate) as client:
            response = await client.post(llm_config["api_url"], headers=headers, json=payload)
        response.raise_for_status()
        content = extract_content_from_llm_response(response.json())
        if not content:
            raise ValueError("LLM 返回内容为空或无法解析")
        return content

    async def _accumulate_usage(self, purpose_key: str, usage: dict[str, int] | None) -> None:
        if not usage or not usage.get("total_tokens"):
            return
        async with self._state_lock:
            bucket = self._usage_by_purpose.setdefault(
                purpose_key, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
            )
            for key in bucket:
                bucket[key] += int(usage.get(key) or 0)

    async def _resolve_model(self, purpose: ModelPurpose | None) -> ModelRuntimeConfig | None:
        if purpose is None:
            return None
        if purpose in self._resolved_models:
            return self._resolved_models[purpose]
        async with self._state_lock:
            if purpose in self._resolved_models:
                return self._resolved_models[purpose]
            config_id = self.scenario.model_bindings.get(purpose)
            if not config_id:
                return None
            expected = {
                ModelPurpose.GENERATOR: ModelCapability.CHAT,
                ModelPurpose.JUDGE: ModelCapability.CHAT,
                ModelPurpose.EMBEDDING: ModelCapability.EMBEDDING,
                ModelPurpose.MODERATION: ModelCapability.MODERATION,
            }[purpose]
            runtime = self.model_service.resolve(config_id, expected)
            self._resolved_models[purpose] = runtime
            return runtime

    # ── helpers ────────────────────────────────────────────────────────

    def _check_cancel(self) -> None:
        if self.cancel_token.is_set():
            raise CancelledError("run cancelled")

    async def _emit(
        self,
        callback: Optional[ProgressCallback],
        event: str,
        data: dict[str, Any],
    ) -> None:
        if not callback:
            return
        try:
            result = callback(event, data)
            if asyncio.iscoroutine(result):
                await result
        except Exception:
            pass