trpc-agent-python/examples/evaluation/callbacks/test_callbacks.py at main · trpc-group/trpc-agent-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
# -*- coding: utf-8 -*-
#
# Copyright @ 2025 Tencent.com
"""回调示例：在评测中注册 8 个生命周期 Callbacks，打日志并演示 context 传递。"""

import os
import pytest
from trpc_agent_sdk.evaluation import (
    AgentEvaluator,
    Callbacks,
    Callback,
    CallbackResult,
    BeforeInferenceSetArgs,
    AfterInferenceSetArgs,
    BeforeInferenceCaseArgs,
    AfterInferenceCaseArgs,
    BeforeEvaluateSetArgs,
    AfterEvaluateSetArgs,
    BeforeEvaluateCaseArgs,
    AfterEvaluateCaseArgs,
)

triggered: list[str] = []


def before_inference_set(ctx, args: BeforeInferenceSetArgs):
    triggered.append("before_inference_set")
    print("[callback] 推理集开始", args.request.eval_set_id, flush=True)
    return None


def after_inference_set(ctx, args: AfterInferenceSetArgs):
    triggered.append("after_inference_set")
    n = len(args.results) if args.results else 0
    print("[callback] 推理集结束，共", n, "个用例", flush=True)
    return None


def before_inference_case(ctx, args: BeforeInferenceCaseArgs):
    triggered.append("before_inference_case")
    print("[callback] 用例推理开始", args.eval_case_id, flush=True)
    return None


def after_inference_case(ctx, args: AfterInferenceCaseArgs):
    triggered.append("after_inference_case")
    print("[callback] 用例推理结束", args.result.eval_case_id, flush=True)
    return None


def before_evaluate_set(ctx, args: BeforeEvaluateSetArgs):
    triggered.append("before_evaluate_set")
    n = len(args.request.inference_results)
    print("[callback] 打分集开始 cases=", n, flush=True)
    return CallbackResult(context={"phase": "evaluate"})


def after_evaluate_set(ctx, args: AfterEvaluateSetArgs):
    triggered.append("after_evaluate_set")
    n = len(args.result.eval_case_results) if args.result else 0
    phase = (ctx.get("context") or {}).get("phase", "?")
    print("[callback] 打分集结束，共", n, "个用例，ctx.phase=", phase, flush=True)
    return None


def before_evaluate_case(ctx, args: BeforeEvaluateCaseArgs):
    triggered.append("before_evaluate_case")
    print("[callback] 用例打分开始", args.eval_case_id, flush=True)
    return None


def after_evaluate_case(ctx, args: AfterEvaluateCaseArgs):
    triggered.append("after_evaluate_case")
    print("[callback] 用例打分结束", args.result.eval_id, flush=True)
    return None


@pytest.mark.asyncio
async def test_with_callbacks():
    triggered.clear()
    test_dir = os.path.dirname(os.path.abspath(__file__))
    eval_set_path = os.path.join(test_dir, "agent", "callbacks_example.evalset.json")

    callbacks = Callbacks()
    callbacks.register(
        "demo",
        Callback(
            before_inference_set=before_inference_set,
            after_inference_set=after_inference_set,
            before_inference_case=before_inference_case,
            after_inference_case=after_inference_case,
            before_evaluate_set=before_evaluate_set,
            after_evaluate_set=after_evaluate_set,
            before_evaluate_case=before_evaluate_case,
            after_evaluate_case=after_evaluate_case,
        ),
    )

    await AgentEvaluator.evaluate(
        agent_module="agent",
        agent_name="weather_agent",
        eval_dataset_file_path_or_dir=eval_set_path,
        callbacks=callbacks,
    )

    expected = [
        "before_inference_set",
        "before_inference_case",
        "after_inference_case",
        "after_inference_set",
        "before_evaluate_set",
        "before_evaluate_case",
        "after_evaluate_case",
        "after_evaluate_set",
    ]
    assert triggered == expected, triggered