Source code for agentscope.tuner._workflow

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




class WorkflowOutput(BaseModel):
    """The output of a workflow function."""

    reward: float | None = Field(
        description=(
            "The reward obtained from the workflow function. "
            "Used for direct reward output."
        ),
        default=None,
    )
    response: Any | None = Field(
        description=(
            "The response generated by the workflow function. "
            "Used as judge input."
        ),
        default=None,
    )

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



WorkflowType = Callable[
    [Dict, ChatModelBase, Dict[str, ChatModelBase]],
    Awaitable[WorkflowOutput],
]
# An agent workflow function type for tuning.

# Args:
#     task (`Dict`):
#         The task information for the workflow run.
#     model (`ChatModelBase`):
#         The primary chat model used in the workflow, this is the main model
#         being tuned.
#     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. Note that these auxiliary models are not
#         tuned during the workflow.

# Returns:
#     `WorkflowOutput`:
#         The workflow execution results, including optional reward, raw
#         response and metrics.