/mlflow/demo/generators/evaluation.py

from __future__ import annotations

import contextlib
import hashlib
import io
import logging
import os
from collections.abc import Callable
from typing import TYPE_CHECKING, Literal

import mlflow

if TYPE_CHECKING:
    from mlflow.genai.datasets import EvaluationDataset

from mlflow.demo.base import (
    DEMO_EXPERIMENT_NAME,
    BaseDemoGenerator,
    DemoFeature,
    DemoResult,
)
from mlflow.demo.data import EXPECTED_ANSWERS
from mlflow.demo.generators.traces import DEMO_TRACE_TYPE_TAG, DEMO_VERSION_TAG, TracesDemoGenerator
from mlflow.entities.assessment import AssessmentSource, Expectation, Feedback
from mlflow.entities.trace import Trace
from mlflow.entities.view_type import ViewType
from mlflow.genai.datasets import create_dataset, delete_dataset, search_datasets
from mlflow.genai.scorers import scorer

_logger = logging.getLogger(__name__)


@contextlib.contextmanager
def _suppress_evaluation_output():
    """Suppress tqdm progress bars and evaluation completion messages."""
    original_tqdm_disable = os.environ.get("TQDM_DISABLE")
    os.environ["TQDM_DISABLE"] = "1"
    try:
        # Suppress both stdout (evaluation messages) and stderr (tqdm progress bars)
        with (
            contextlib.redirect_stdout(io.StringIO()),
            contextlib.redirect_stderr(io.StringIO()),
        ):
            yield
    finally:
        if original_tqdm_disable is None:
            os.environ.pop("TQDM_DISABLE", None)
        else:
            os.environ["TQDM_DISABLE"] = original_tqdm_disable


DEMO_DATASET_TRACE_LEVEL_NAME = "demo-trace-level-dataset"
DEMO_DATASET_BASELINE_SESSION_NAME = "demo-baseline-session-dataset"
DEMO_DATASET_IMPROVED_SESSION_NAME = "demo-improved-session-dataset"


def _get_relevance_rationale(is_relevant: bool) -> str:
    if is_relevant:
        return "The response directly addresses the question with relevant information."
    return "The response is not sufficiently relevant to the question asked."


def _get_correctness_rationale(is_correct: bool) -> str:
    if is_correct:
        return "The response accurately captures the key information from the expected answer."
    return (
        "The response contains relevant information but differs "
        "significantly from the expected answer."
    )


def _get_groundedness_rationale(is_grounded: bool) -> str:
    if is_grounded:
        return "The response is well-grounded in the provided context with clear references."
    return "The response includes claims not supported by the provided context."


def _get_safety_rationale(is_safe: bool) -> str:
    if is_safe:
        return "The response contains no harmful, offensive, or inappropriate content."
    return "The response may contain potentially harmful or inappropriate content."


def _create_quality_aware_scorer(
    name: str,
    baseline_pass_rate: float,
    improved_pass_rate: float,
    rationale_fn: Callable[[bool], str],
):
    """Create a deterministic scorer that simulates quality-aware evaluation.

    The scorer detects response quality based on content characteristics:
    - Longer, more detailed responses get evaluated with higher pass rates
    - Shorter, less detailed responses get evaluated with lower pass rates

    This simulates the real-world scenario where improved model outputs
    naturally score better when evaluated by the same scorers.
    """
    quality_threshold = 400

    @scorer(name=name)
    def quality_aware_scorer(inputs, outputs, trace) -> Feedback:
        content = str(inputs) + str(outputs)
        output_str = str(outputs)

        if len(output_str) > quality_threshold:
            effective_pass_rate = improved_pass_rate
        else:
            effective_pass_rate = baseline_pass_rate

        # Use content hash for deterministic but varied results
        hash_input = f"{content}:{name}"
        hash_val = int(hashlib.md5(hash_input.encode(), usedforsecurity=False).hexdigest()[:8], 16)
        normalized = hash_val / 0xFFFFFFFF
        is_passing = normalized < effective_pass_rate

        # Use the trace timestamp so the quality overview chart shows a trend
        # across days instead of a single dot at the current time.
        trace_timestamp_ms = trace.info.timestamp_ms if trace else None

        return Feedback(
            value="yes" if is_passing else "no",
            rationale=rationale_fn(is_passing),
            source=AssessmentSource(
                source_type="LLM_JUDGE",
                source_id=f"judges/{name}",
            ),
            create_time_ms=trace_timestamp_ms,
            last_update_time_ms=trace_timestamp_ms,
        )

    return quality_aware_scorer


SCORER_PASS_RATES = {
    "relevance": {"baseline": 0.65, "improved": 0.92},
    "correctness": {"baseline": 0.58, "improved": 0.88},
    "groundedness": {"baseline": 0.52, "improved": 0.85},
    "safety": {"baseline": 0.95, "improved": 1.0},
}


class EvaluationDemoGenerator(BaseDemoGenerator):
    """Generates demo evaluation data.

    Creates:
    - Ground truth expectations on all demo traces
    - Three datasets and evaluation runs, each in a single mode:
      - trace-level-evaluation: non-session traces (v1 + v2 combined)
      - baseline-session-evaluation: v1 session traces
      - improved-session-evaluation: v2 session traces

    Assessment timestamps are spread to match trace timestamps so the
    quality overview chart shows a trend across days.
    """

    name = DemoFeature.EVALUATION
    version = 2

    def generate(self) -> DemoResult:
        traces_generator = TracesDemoGenerator()
        if not traces_generator.is_generated():
            traces_generator.generate()
            traces_generator.store_version()

        experiment = mlflow.get_experiment_by_name(DEMO_EXPERIMENT_NAME)
        experiment_id = experiment.experiment_id

        # Fetch traces split by session vs non-session
        v1_non_session = self._fetch_demo_traces(experiment_id, "v1", session=False)
        v2_non_session = self._fetch_demo_traces(experiment_id, "v2", session=False)
        v1_session = self._fetch_demo_traces(experiment_id, "v1", session=True)
        v2_session = self._fetch_demo_traces(experiment_id, "v2", session=True)

        all_traces = v1_non_session + v2_non_session + v1_session + v2_session
        self._add_expectations_to_traces(all_traces)

        # Re-fetch to include expectations
        v1_non_session = self._fetch_demo_traces(experiment_id, "v1", session=False)
        v2_non_session = self._fetch_demo_traces(experiment_id, "v2", session=False)
        v1_session = self._fetch_demo_traces(experiment_id, "v1", session=True)
        v2_session = self._fetch_demo_traces(experiment_id, "v2", session=True)

        trace_level_traces = v1_non_session + v2_non_session

        # Create datasets
        self._create_evaluation_dataset(
            trace_level_traces, experiment_id, DEMO_DATASET_TRACE_LEVEL_NAME
        )
        self._create_evaluation_dataset(
            v1_session, experiment_id, DEMO_DATASET_BASELINE_SESSION_NAME
        )
        self._create_evaluation_dataset(
            v2_session, experiment_id, DEMO_DATASET_IMPROVED_SESSION_NAME
        )

        # Create evaluation runs
        trace_level_run_id = self._create_evaluation_run(
            traces=trace_level_traces,
            experiment_id=experiment_id,
            run_name="trace-level-evaluation",
        )

        baseline_session_run_id = self._create_evaluation_run(
            traces=v1_session,
            experiment_id=experiment_id,
            run_name="baseline-session-evaluation",
        )

        improved_session_run_id = self._create_evaluation_run(
            traces=v2_session,
            experiment_id=experiment_id,
            run_name="improved-session-evaluation",
        )

        return DemoResult(
            feature=self.name,
            entity_ids=[trace_level_run_id, baseline_session_run_id, improved_session_run_id],
            navigation_url=f"#/experiments/{experiment_id}/evaluation-runs",
        )

    def _data_exists(self) -> bool:
        experiment = mlflow.get_experiment_by_name(DEMO_EXPERIMENT_NAME)
        if experiment is None or experiment.lifecycle_stage != "active":
            return False

        try:
            client = mlflow.MlflowClient()
            runs = client.search_runs(
                experiment_ids=[experiment.experiment_id],
                filter_string="params.demo = 'true'",
                max_results=1,
            )
            return len(runs) > 0
        except Exception:
            _logger.debug("Failed to check if evaluation demo exists", exc_info=True)
            return False

    def delete_demo(self) -> None:
        experiment = mlflow.get_experiment_by_name(DEMO_EXPERIMENT_NAME)
        if experiment is None:
            return

        try:
            client = mlflow.MlflowClient()
            runs = client.search_runs(
                experiment_ids=[experiment.experiment_id],
                filter_string="params.demo = 'true'",
                run_view_type=ViewType.ALL,
                max_results=100,
            )
            for run in runs:
                try:
                    if run.info.lifecycle_stage == "deleted":
                        client.restore_run(run.info.run_id)
                    client.delete_run(run.info.run_id)
                except Exception:
                    _logger.debug("Failed to delete run %s", run.info.run_id, exc_info=True)
        except Exception:
            _logger.debug("Failed to delete evaluation demo runs", exc_info=True)

        for name in [
            DEMO_DATASET_TRACE_LEVEL_NAME,
            DEMO_DATASET_BASELINE_SESSION_NAME,
            DEMO_DATASET_IMPROVED_SESSION_NAME,
        ]:
            self._delete_demo_dataset(experiment.experiment_id, name)

    def _fetch_demo_traces(
        self,
        experiment_id: str,
        version: Literal["v1", "v2"],
        session: bool | None = None,
    ) -> list[Trace]:
        filter_parts = [f"metadata.`{DEMO_VERSION_TAG}` = '{version}'"]
        operator = "=" if session else "!="
        filter_parts.append(f"metadata.`{DEMO_TRACE_TYPE_TAG}` {operator} 'session'")
        return mlflow.search_traces(
            locations=[experiment_id],
            filter_string=" AND ".join(filter_parts),
            max_results=100,
            return_type="list",
            flush=True,
        )

    def _add_expectations_to_traces(self, traces: list[Trace]) -> int:
        expectation_count = 0

        for trace in traces:
            trace_id = trace.info.trace_id
            trace_timestamp_ms = trace.info.timestamp_ms

            root_span = next((span for span in trace.data.spans if span.parent_id is None), None)
            if root_span is None:
                continue

            inputs = root_span.inputs or {}
            query = inputs.get("query") or inputs.get("message")

            if expected_answer := self._find_expected_answer(query):
                try:
                    expectation = Expectation(
                        name="expected_response",
                        value=expected_answer,
                        source=AssessmentSource(
                            source_type="HUMAN",
                            source_id="demo_annotator",
                        ),
                        metadata={"demo": "true"},
                        trace_id=trace_id,
                        create_time_ms=trace_timestamp_ms,
                        last_update_time_ms=trace_timestamp_ms,
                    )
                    mlflow.log_assessment(trace_id=trace_id, assessment=expectation)
                    expectation_count += 1
                except Exception:
                    _logger.debug("Failed to log expectation for trace %s", trace_id, exc_info=True)

        return expectation_count

    def _find_expected_answer(self, query: str | None) -> str | None:
        if not query:
            return None
        query_lower = query.lower().strip()
        if query_lower in EXPECTED_ANSWERS:
            return EXPECTED_ANSWERS[query_lower]
        for q, answer in EXPECTED_ANSWERS.items():
            if q in query_lower or query_lower in q:
                return answer
        return None

    def _create_evaluation_dataset(
        self, traces: list[Trace], experiment_id: str, dataset_name: str
    ) -> "EvaluationDataset":
        from mlflow.genai.datasets import get_dataset

        dataset = create_dataset(
            name=dataset_name,
            experiment_id=experiment_id,
            tags={"demo": "true", "description": f"Demo evaluation dataset: {dataset_name}"},
        )

        dataset.merge_records(traces)
        return get_dataset(dataset_id=dataset.dataset_id)

    def _delete_demo_dataset(self, experiment_id: str, dataset_name: str) -> None:
        datasets = search_datasets(
            experiment_ids=[experiment_id],
            filter_string=f"name = '{dataset_name}'",
            max_results=10,
        )
        for ds in datasets:
            try:
                delete_dataset(dataset_id=ds.dataset_id)
            except Exception:
                _logger.debug("Failed to delete dataset %s", ds.dataset_id, exc_info=True)

    def _create_evaluation_run(
        self,
        traces: list[Trace],
        experiment_id: str,
        run_name: str,
    ) -> str:
        demo_scorers = [
            _create_quality_aware_scorer(
                name="relevance",
                baseline_pass_rate=SCORER_PASS_RATES["relevance"]["baseline"],
                improved_pass_rate=SCORER_PASS_RATES["relevance"]["improved"],
                rationale_fn=_get_relevance_rationale,
            ),
            _create_quality_aware_scorer(
                name="correctness",
                baseline_pass_rate=SCORER_PASS_RATES["correctness"]["baseline"],
                improved_pass_rate=SCORER_PASS_RATES["correctness"]["improved"],
                rationale_fn=_get_correctness_rationale,
            ),
            _create_quality_aware_scorer(
                name="groundedness",
                baseline_pass_rate=SCORER_PASS_RATES["groundedness"]["baseline"],
                improved_pass_rate=SCORER_PASS_RATES["groundedness"]["improved"],
                rationale_fn=_get_groundedness_rationale,
            ),
            _create_quality_aware_scorer(
                name="safety",
                baseline_pass_rate=SCORER_PASS_RATES["safety"]["baseline"],
                improved_pass_rate=SCORER_PASS_RATES["safety"]["improved"],
                rationale_fn=_get_safety_rationale,
            ),
        ]

        mlflow.set_experiment(experiment_id=experiment_id)

        with _suppress_evaluation_output():
            result = mlflow.genai.evaluate(
                data=traces,
                scorers=demo_scorers,
            )

        client = mlflow.MlflowClient()
        client.set_tag(result.run_id, "mlflow.runName", run_name)
        client.log_param(result.run_id, "demo", "true")

        return result.run_id