Skip to content

Trace

Trace

A context manager which captures steps taken during evaluation.

Once the context manager exits, the trace is dumped to a JSON file.

Attributes:

Name Type Description
test_name str

Name of the test.
trace_dir str

Directory to store the trace.
start_time datetime

Start time of the trace.
end_time datetime

End time of the trace.
steps list

List of steps in the trace.
Source code in src/agenteval/trace.py 
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
class Trace:
    """A context manager which captures steps taken during evaluation.

    Once the context manager exits, the trace is dumped to a JSON file.

    Attributes:
        test_name (str): Name of the test.
        trace_dir (str): Directory to store the trace.
        start_time (datetime): Start time of the trace.
        end_time (datetime): End time of the trace.
        steps (list): List of steps in the trace.

    """

    def __init__(self, test_name: str, work_dir: str):
        """
        Initialize the trace handler.

        Args:
            test_name (str): Name of the test.
            work_dir (str): Directory to store the trace.
        """
        self.test_name = test_name
        self.trace_dir = os.path.join(work_dir, _TRACE_DIR)
        self.start_time = None
        self.end_time = None
        self.steps = []

    def __enter__(self):
        self.start_time = datetime.now(timezone.utc)
        return self

    def __exit__(self, *exc):
        self.end_time = datetime.now(timezone.utc)
        self._dump_trace()

    def _dump_trace(self):
        os.makedirs(self.trace_dir, exist_ok=True)

        with open(os.path.join(self.trace_dir, f"{self.test_name}.json"), "w") as f:
            json.dump(self._get_trace(), f, default=str)

    def _get_trace(self) -> str:
        return {
            "test_name": self.test_name,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "steps": self.steps,
        }

    def add_step(self, step_name: Optional[str] = None, **kwargs):
        """Add a step to the trace.

        Args:
            step_name (Optional[str]): The name of the step. Defaults to
                the name of the caller function
        """
        step_name = step_name or inspect.stack()[1].function
        step = {"timestamp": datetime.now(timezone.utc), "step_name": step_name}
        step.update(kwargs)
        self.steps.append(step)

__init__(test_name, work_dir)

Initialize the trace handler.

Parameters:

Name Type Description Default
test_name str

Name of the test.

 required
work_dir str

Directory to store the trace.

 required
Source code in src/agenteval/trace.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
def __init__(self, test_name: str, work_dir: str):
    """
    Initialize the trace handler.

    Args:
        test_name (str): Name of the test.
        work_dir (str): Directory to store the trace.
    """
    self.test_name = test_name
    self.trace_dir = os.path.join(work_dir, _TRACE_DIR)
    self.start_time = None
    self.end_time = None
    self.steps = []

add_step(step_name=None, **kwargs)

Add a step to the trace.

Parameters:

Name Type Description Default
step_name Optional[str]

The name of the step. Defaults to the name of the caller function

 None
Source code in src/agenteval/trace.py 
63
64
65
66
67
68
69
70
71
72
73
def add_step(self, step_name: Optional[str] = None, **kwargs):
    """Add a step to the trace.

    Args:
        step_name (Optional[str]): The name of the step. Defaults to
            the name of the caller function
    """
    step_name = step_name or inspect.stack()[1].function
    step = {"timestamp": datetime.now(timezone.utc), "step_name": step_name}
    step.update(kwargs)
    self.steps.append(step)