# -*- coding: utf-8 -*-
"""The judge module for tuner."""
from typing import Any, Callable, Dict, Awaitable
from pydantic import BaseModel, Field
from ..model import ChatModelBase


class JudgeOutput(BaseModel):
    """The output of a judge function."""

    reward: float = Field(
        description="The reward value assigned by the judge function.",
    )

    metrics: Dict[str, float] | None = Field(
        description="Metrics from the judge function.",
        default=None,
    )


JudgeType = Callable[
    [Dict, Any, Dict[str, ChatModelBase]],
    Awaitable[JudgeOutput],
]
# A judge function type for tuning.

# Args:
#     task (`Dict`):
#         The task information for the corresponding workflow.
#     response (`Any`):
#         The response field of the WorkflowOutput generated by the
#         corresponding workflow.
#     auxiliary_models (`Dict[str, ChatModelBase] | None`):
#         A dictionary of additional chat models available for LLM-as-a-Judge
#         usage. The keys are model names, and the values are the corresponding
#         `ChatModelBase` instances.
# Returns:
#     `JudgeOutput`:
#         The reward value assigned by the judge function along with optional
#         metrics.
