aiqtoolkit 1.1.0__py3-none-any.whl

aiq/eval/dataset_handler/dataset_handler.py

@@ -0,0 +1,169 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import pandas as pd
+
+from aiq.data_models.dataset_handler import EvalDatasetConfig
+from aiq.data_models.dataset_handler import EvalDatasetJsonConfig
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.data_models.intermediate_step import IntermediateStepType
+from aiq.eval.dataset_handler.dataset_downloader import DatasetDownloader
+from aiq.eval.dataset_handler.dataset_filter import DatasetFilter
+from aiq.eval.evaluator.evaluator_model import EvalInput
+from aiq.eval.evaluator.evaluator_model import EvalInputItem
+
+
+class DatasetHandler:
+    """
+    Read the datasets and pre-process (apply filters, deduplicate etc.) before turning them into EvalInput objects.
+    One DatasetHandler object is needed for each dataset to be evaluated.
+    """
+
+    def __init__(self, dataset_config: EvalDatasetConfig, reps: int):
+        from aiq.eval.intermediate_step_adapter import IntermediateStepAdapter
+
+        self.dataset_config = dataset_config
+        self.dataset_filter = DatasetFilter(dataset_config.filter)
+        self.reps = reps
+        # Helpers
+        self.intermediate_step_adapter = IntermediateStepAdapter()
+
+    def is_structured_input(self) -> bool:
+        '''Check if the input is structured or unstructured'''
+        return not self.dataset_config.structure.disable
+
+    @property
+    def id_key(self) -> str:
+        return self.dataset_config.id_key
+
+    @property
+    def question_key(self) -> str:
+        return self.dataset_config.structure.question_key
+
+    @property
+    def answer_key(self) -> str:
+        return self.dataset_config.structure.answer_key
+
+    @property
+    def generated_answer_key(self) -> str:
+        return self.dataset_config.structure.generated_answer_key
+
+    @property
+    def trajectory_key(self) -> str:
+        return self.dataset_config.structure.trajectory_key
+
+    @property
+    def expected_trajectory_key(self) -> str:
+        return self.dataset_config.structure.expected_trajectory_key
+
+    def get_eval_input_from_df(self, input_df: pd.DataFrame) -> EvalInput:
+
+        def create_eval_item(row: pd.Series, structured: bool) -> EvalInputItem:
+            """Helper function to create EvalInputItem."""
+            return EvalInputItem(
+                id=row.get(self.id_key, ""),
+                input_obj=row.to_json() if not structured else row.get(self.question_key, ""),
+                expected_output_obj=row.get(self.answer_key, "") if structured else "",
+                output_obj=row.get(self.generated_answer_key, "") if structured else "",
+                trajectory=row.get(self.trajectory_key, []) if structured else [],
+                expected_trajectory=row.get(self.expected_trajectory_key, []) if structured else [],
+            )
+
+        # if input dataframe is empty return an empty list
+        if input_df.empty:
+            return EvalInput(eval_input_items=[])
+
+        structured = self.is_structured_input()
+        if structured:
+            # For structured input, question is mandatory. Ignore rows with missing or empty questions
+            input_df = input_df[input_df[self.question_key].notnull() & input_df[self.question_key].str.strip().ne("")]
+        eval_input_items = [create_eval_item(row, structured) for _, row in input_df.iterrows()]
+
+        return EvalInput(eval_input_items=eval_input_items)
+
+    def setup_reps(self, input_df: pd.DataFrame) -> pd.DataFrame:
+        """replicate the rows and update the id to id_key + "_rep" + rep_number"""
+        # Replicate the rows
+        input_df = pd.concat([input_df] * self.reps, ignore_index=True)
+        # Compute repetition index
+        rep_index = input_df.groupby(self.dataset_config.id_key).cumcount().astype(str)
+        # Convert id_key to string (id can be integer) if needed and update IDs
+        input_df[self.dataset_config.id_key] = input_df[self.dataset_config.id_key].astype(str) + "_rep" + rep_index
+        # Ensure unique ID values after modification
+        input_df.drop_duplicates(subset=[self.dataset_config.id_key], inplace=True)
+
+        return input_df
+
+    def get_eval_input_from_dataset(self, dataset: str) -> EvalInput:
+        # read the dataset and convert it to EvalInput
+
+        # if a dataset file has been provided in the command line, use that
+        dataset_config = EvalDatasetJsonConfig(file_path=dataset) if dataset else self.dataset_config
+
+        # Download the dataset if it is remote
+        downloader = DatasetDownloader(dataset_config=dataset_config)
+        downloader.download_dataset()
+
+        parser, kwargs = dataset_config.parser()
+        # Parse the dataset into a DataFrame
+        input_df = parser(dataset_config.file_path, **kwargs)
+
+        # Apply filters and deduplicate
+        input_df = self.dataset_filter.apply_filters(input_df)
+        input_df.drop_duplicates(subset=[self.dataset_config.id_key], inplace=True)
+
+        # If more than one repetition is needed, replicate the rows
+        if self.reps > 1:
+            input_df = self.setup_reps(input_df)
+
+        # Convert the DataFrame to a list of EvalInput objects
+        return self.get_eval_input_from_df(input_df)
+
+    def filter_intermediate_steps(self,
+                                  intermediate_steps: list[IntermediateStep],
+                                  event_filter: list[IntermediateStepType] = None) -> list[dict]:
+        """
+        Filter out the intermediate steps that are not relevant for evaluation.
+        The output is written with with the intention of re-running the evaluation using the original config file.
+        """
+        if event_filter is None:
+            event_filter = self.intermediate_step_adapter.DEFAULT_EVENT_FILTER
+        filtered_steps = self.intermediate_step_adapter.filter_intermediate_steps(intermediate_steps, event_filter)
+        return self.intermediate_step_adapter.serialize_intermediate_steps(filtered_steps)
+
+    def publish_eval_input(self, eval_input, workflow_output_step_filter: list[IntermediateStepType] = None) -> str:
+        """
+        Convert the EvalInput object to a JSON output for storing in a file. Use the orginal keys to
+        allow re-running evaluation using the orignal config file and '--skip_workflow' option.
+        """
+
+        indent = 2
+        if self.is_structured_input():
+            # Extract structured data from EvalInputItems
+            data = [{
+                self.id_key: item.id,
+                self.question_key: item.input_obj,
+                self.answer_key: item.expected_output_obj,
+                self.generated_answer_key: item.output_obj,
+                self.trajectory_key: self.filter_intermediate_steps(item.trajectory, workflow_output_step_filter),
+                self.expected_trajectory_key: self.filter_intermediate_steps(item.expected_trajectory),
+            } for item in eval_input.eval_input_items]
+        else:
+            # Unstructured case: return only raw output objects as a JSON array
+            data = [json.loads(item.output_obj) for item in eval_input.eval_input_items]
+
+        return json.dumps(data, indent=indent, ensure_ascii=False, default=str)

aiq/eval/evaluate.py

@@ -0,0 +1,325 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+import shutil
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+from tqdm import tqdm
+
+from aiq.data_models.evaluate import EvalConfig
+from aiq.eval.config import EvaluationRunConfig
+from aiq.eval.config import EvaluationRunOutput
+from aiq.eval.dataset_handler.dataset_handler import DatasetHandler
+from aiq.eval.evaluator.evaluator_model import EvalInput
+from aiq.eval.evaluator.evaluator_model import EvalInputItem
+from aiq.eval.evaluator.evaluator_model import EvalOutput
+from aiq.eval.utils.output_uploader import OutputUploader
+from aiq.runtime.session import AIQSessionManager
+
+logger = logging.getLogger(__name__)
+
+
+class EvaluationRun:  # pylint: disable=too-many-public-methods
+    """
+    Instantiated for each evaluation run and used to store data for that single run.
+    """
+
+    def __init__(self, config: EvaluationRunConfig):
+        """
+        Initialize an EvaluationRun with configuration.
+        """
+        from aiq.eval.intermediate_step_adapter import IntermediateStepAdapter
+
+        # Run-specific configuration
+        self.config: EvaluationRunConfig = config
+        self.eval_config: EvalConfig | None = None
+
+        # Helpers
+        self.intermediate_step_adapter: IntermediateStepAdapter = IntermediateStepAdapter()
+
+        # Metadata
+        self.eval_input: EvalInput | None = None
+        self.workflow_interrupted: bool = False
+
+        # evaluation_results is list of tuples (evaluator_name, EvalOutput)
+        self.evaluation_results: list[tuple[str, EvalOutput]] = []
+
+        # workflow output file
+        self.workflow_output_file: Path | None = None
+
+        # evaluation output files
+        self.evaluator_output_files: list[Path] = []
+
+    async def run_workflow_local(self, session_manager: AIQSessionManager):
+        '''
+        Launch the workflow with the specified questions and extract the output using the jsonpath
+        '''
+        # import function level dependencies
+        from jsonpath_ng import parse
+
+        from aiq.eval.runtime_event_subscriber import pull_intermediate
+
+        # Run the workflow
+        jsonpath_expr = parse(self.config.result_json_path)
+        stop_event = asyncio.Event()
+
+        async def run_one(item: EvalInputItem):
+            if stop_event.is_set():
+                return "", []
+
+            async with session_manager.run(item.input_obj) as runner:
+                try:
+                    # Start usage stats and intermediate steps collection in parallel
+                    intermediate_future = pull_intermediate()
+
+                    if session_manager.workflow.has_single_output:
+                        base_output = await runner.result()
+                    else:
+                        # raise an error if the workflow has multiple outputs
+                        raise NotImplementedError("Multiple outputs are not supported")
+                    intermediate_steps = await intermediate_future
+                except NotImplementedError as e:
+                    # raise original error
+                    raise e
+                except Exception as e:
+                    logger.exception("Failed to run the workflow: %s", e, exc_info=True)
+                    # stop processing if a workflow error occurs
+                    self.workflow_interrupted = True
+                    stop_event.set()
+                    return
+
+                try:
+                    base_output = runner.convert(base_output, to_type=str)
+                except ValueError:
+                    pass
+
+                # if base_output is a pydantic model dump it to json
+                if isinstance(base_output, BaseModel):
+                    output = base_output.model_dump_json(indent=2)
+                else:
+                    m = jsonpath_expr.find(base_output)
+                    if (not m):
+                        raise RuntimeError(f"Failed to extract output using jsonpath: {self.config.result_json_path}")
+                    if (len(m) > 1):
+                        logger.warning("Multiple matches found for jsonpath at row '%s'. Matches: %s. Using the first",
+                                       base_output,
+                                       m)
+                    output = m[0].value
+
+                item.output_obj = output
+                item.trajectory = self.intermediate_step_adapter.validate_intermediate_steps(intermediate_steps)
+
+        async def wrapped_run(item: EvalInputItem) -> None:
+            await run_one(item)
+            pbar.update(1)
+
+        # if self.config.skip_complete is set skip eval_input_items with a non-empty output_obj
+        if self.config.skip_completed_entries:
+            eval_input_items = [item for item in self.eval_input.eval_input_items if not item.output_obj]
+            if not eval_input_items:
+                logger.warning("All items have a non-empty output. Skipping workflow pass altogether.")
+                return
+        else:
+            eval_input_items = self.eval_input.eval_input_items
+        pbar = tqdm(total=len(eval_input_items), desc="Running workflow")
+        await asyncio.gather(*[wrapped_run(item) for item in eval_input_items])
+        pbar.close()
+
+    async def run_workflow_remote(self):
+        from aiq.eval.remote_workflow import EvaluationRemoteWorkflowHandler
+        handler = EvaluationRemoteWorkflowHandler(self.config, self.eval_config.general.max_concurrency)
+        await handler.run_workflow_remote(self.eval_input)
+
+    async def profile_workflow(self):
+        """
+        Profile a dataset
+        """
+
+        if not self.eval_config.general.profiler:
+            logger.info("Profiler is not enabled. Skipping profiling.")
+            return
+
+        from aiq.profiler.profile_runner import ProfilerRunner
+
+        all_stats = []
+        for input_item in self.eval_input.eval_input_items:
+            all_stats.append(input_item.trajectory)
+
+        profiler_runner = ProfilerRunner(self.eval_config.general.profiler, self.eval_config.general.output_dir)
+
+        await profiler_runner.run(all_stats)
+
+    def cleanup_output_directory(self):
+        '''Remove contents of the output directory if it exists'''
+        if self.eval_config.general.output and self.eval_config.general.output.dir and \
+                self.eval_config.general.output.dir.exists():
+            logger.info("Cleaning up output directory %s", self.eval_config.general.output.dir)
+            shutil.rmtree(self.eval_config.general.output.dir)
+
+    def write_output(self, dataset_handler: DatasetHandler):
+        workflow_output_file = self.eval_config.general.output_dir / "workflow_output.json"
+        workflow_output_file.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write the workflow output to a file (this can be used for re-running the evaluation)
+
+        step_filter = self.eval_config.general.output.workflow_output_step_filter \
+            if self.eval_config.general.output else None
+        workflow_output = dataset_handler.publish_eval_input(self.eval_input, step_filter)
+        with open(workflow_output_file, "w", encoding="utf-8") as f:
+            # set indent to 2 for pretty printing
+            f.write(workflow_output)
+        self.workflow_output_file = workflow_output_file
+        logger.info("Workflow output written to %s", workflow_output_file)
+
+        # Write the output of each evaluator to a separate json file
+        for evaluator_name, eval_output in self.evaluation_results:
+            output_file = self.eval_config.general.output_dir / f"{evaluator_name}_output.json"
+            output_file.parent.mkdir(parents=True, exist_ok=True)
+            # create json content using the evaluation results
+            output = eval_output.model_dump_json(indent=2)
+            with open(output_file, "w", encoding="utf-8") as f:
+                f.write(output)
+            self.evaluator_output_files.append(output_file)
+            logger.info("Evaluation results written to %s", output_file)
+
+        if self.workflow_interrupted:
+            # Issue a warning if the workflow was not completed on all datasets
+            msg = ("Workflow execution was interrupted due to an error. The results may be incomplete. "
+                   "You can re-execute evaluation for incomplete results by running "
+                   "`eval` with the --skip_completed_entries flag.")
+            logger.warning(msg)
+
+    async def run_single_evaluator(self, evaluator_name: str, evaluator: Any):
+        """Run a single evaluator and store its results."""
+        try:
+            eval_output = await evaluator.evaluate_fn(self.eval_input)
+            self.evaluation_results.append((evaluator_name, eval_output))
+        except Exception as e:
+            logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e, exc_info=True)
+
+    async def run_evaluators(self, evaluators: dict[str, Any]):
+        """Run all configured evaluators asynchronously."""
+        tasks = [self.run_single_evaluator(name, evaluator) for name, evaluator in evaluators.items() if evaluator]
+
+        if not tasks:
+            logger.warning("All evaluators were empty or invalid.")
+            return
+
+        try:
+            await asyncio.gather(*tasks)
+        except Exception as e:
+            logger.exception("An error occurred while running evaluators: %s", e, exc_info=True)
+            raise
+
+    def apply_overrides(self):
+        from aiq.cli.cli_utils.config_override import load_and_override_config
+        from aiq.data_models.config import AIQConfig
+        from aiq.runtime.loader import PluginTypes
+        from aiq.runtime.loader import discover_and_register_plugins
+        from aiq.utils.data_models.schema_validator import validate_schema
+
+        # Register plugins before validation
+        discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
+
+        config_dict = load_and_override_config(self.config.config_file, self.config.override)
+        config = validate_schema(config_dict, AIQConfig)
+        return config
+
+    async def run_and_evaluate(self,
+                               session_manager: AIQSessionManager | None = None,
+                               job_id: str | None = None) -> EvaluationRunOutput:
+        """
+        Run the workflow with the specified config file and evaluate the dataset
+        """
+        logger.info("Starting evaluation run with config file: %s", self.config.config_file)
+
+        from aiq.builder.eval_builder import WorkflowEvalBuilder
+        from aiq.runtime.loader import load_config
+
+        # Load and override the config
+        if self.config.override:
+            config = self.apply_overrides()
+        else:
+            config = load_config(self.config.config_file)
+        self.eval_config = config.eval
+        logger.debug("Loaded evaluation configuration: %s", self.eval_config)
+
+        # Cleanup the output directory
+        if self.eval_config.general.output and self.eval_config.general.output.cleanup:
+            self.cleanup_output_directory()
+
+        # If a job id is provided keep the data per-job
+        if job_id:
+            self.eval_config.general.output_dir = self.eval_config.general.output_dir / f"jobs/{job_id}"
+            if self.eval_config.general.output:
+                self.eval_config.general.output.dir = self.eval_config.general.output_dir
+
+        # Load the input dataset
+        # For multiple datasets, one handler per dataset can be created
+        dataset_config = self.eval_config.general.dataset  # Currently only one dataset is supported
+        if not dataset_config:
+            logger.info("No dataset found, nothing to evaluate")
+            return EvaluationRunOutput(
+                workflow_output_file=self.workflow_output_file,
+                evaluator_output_files=self.evaluator_output_files,
+                workflow_interrupted=self.workflow_interrupted,
+            )
+
+        dataset_handler = DatasetHandler(dataset_config=dataset_config, reps=self.config.reps)
+        self.eval_input = dataset_handler.get_eval_input_from_dataset(self.config.dataset)
+        if not self.eval_input.eval_input_items:
+            logger.info("Dataset is empty. Nothing to evaluate.")
+            return EvaluationRunOutput(
+                workflow_output_file=self.workflow_output_file,
+                evaluator_output_files=self.evaluator_output_files,
+                workflow_interrupted=self.workflow_interrupted,
+            )
+
+        # Run workflow and evaluate
+        async with WorkflowEvalBuilder.from_config(config=config) as eval_workflow:
+            if self.config.endpoint:
+                await self.run_workflow_remote()
+            else:
+                if not self.config.skip_workflow:
+                    if session_manager is None:
+                        session_manager = AIQSessionManager(eval_workflow.build(),
+                                                            max_concurrency=self.eval_config.general.max_concurrency)
+                    await self.run_workflow_local(session_manager)
+
+            # Evaluate
+            evaluators = {name: eval_workflow.get_evaluator(name) for name in self.eval_config.evaluators}
+            await self.run_evaluators(evaluators)
+
+        # Profile the workflow
+        await self.profile_workflow()
+
+        # Write the results to the output directory
+        self.write_output(dataset_handler)
+
+        # Run custom scripts and upload evaluation outputs to S3
+        if self.eval_config.general.output:
+            output_uploader = OutputUploader(self.eval_config.general.output, job_id=job_id)
+            output_uploader.run_custom_scripts()
+            await output_uploader.upload_directory()
+
+        return EvaluationRunOutput(
+            workflow_output_file=self.workflow_output_file,
+            evaluator_output_files=self.evaluator_output_files,
+            workflow_interrupted=self.workflow_interrupted,
+        )

aiq/eval/evaluator/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

aiq/eval/evaluator/evaluator_model.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from pydantic import BaseModel
+
+from aiq.data_models.intermediate_step import IntermediateStep
+
+
+class EvalInputItem(BaseModel):
+    id: typing.Any
+    input_obj: typing.Any
+    expected_output_obj: typing.Any
+    output_obj: typing.Any
+    expected_trajectory: list[IntermediateStep]
+    trajectory: list[IntermediateStep]
+
+
+class EvalInput(BaseModel):
+    eval_input_items: list[EvalInputItem]
+
+
+class EvalOutputItem(BaseModel):
+    id: typing.Any  # id or input_obj from EvalInputItem
+    score: typing.Any  # float or any serializable type
+    reasoning: typing.Any
+
+
+class EvalOutput(BaseModel):
+    average_score: typing.Any  # float or any serializable type
+    eval_output_items: list[EvalOutputItem]

aiq/eval/intermediate_step_adapter.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from langchain_core.agents import AgentAction
+
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.data_models.intermediate_step import IntermediateStepType
+
+logger = logging.getLogger(__name__)
+
+
+class IntermediateStepAdapter:
+    DEFAULT_EVENT_FILTER = [IntermediateStepType.LLM_END, IntermediateStepType.TOOL_END]
+
+    def filter_intermediate_steps(self,
+                                  intermediate_steps: list[IntermediateStep],
+                                  event_filter: list[IntermediateStepType]) -> list[IntermediateStep]:
+        """ Filters intermediate steps"""
+        if not event_filter:
+            return intermediate_steps
+        return [step for step in intermediate_steps if step.event_type in event_filter]
+
+    def validate_intermediate_steps(self, intermediate_steps: list[dict]) -> list[IntermediateStep]:
+        validated_steps = []
+        for step_data in intermediate_steps:
+            try:
+                validated_steps.append(IntermediateStep.model_validate(step_data))
+            except Exception as e:
+                logger.exception("Validation failed for step: %r, Error: %s", step_data, e, exc_info=True)
+        return validated_steps
+
+    def serialize_intermediate_steps(self, intermediate_steps: list[IntermediateStep]) -> list[dict]:
+        """Converts a list of IntermediateStep objects to a list of dictionaries."""
+        return [step.model_dump() for step in intermediate_steps]
+
+    @staticmethod
+    def agent_action_to_dict(action) -> dict:
+        """Convert AgentAction to a JSON-serializable dictionary."""
+        return {
+            "tool": action.tool,
+            "tool_input": action.tool_input,
+            "log": action.log,
+            "type": action.type,
+        }
+
+    def get_agent_action_single(self, step: IntermediateStep,
+                                last_llm_end_step: IntermediateStep | None) -> tuple[AgentAction, str]:
+        """Converts a single intermediate step to Tuple[AgentAction, str]."""
+        # use the previous llm output as log
+        log = getattr(last_llm_end_step.data, "output", "") if last_llm_end_step else ""
+        tool_name = step.name or ""
+        tool_input = getattr(step.data, "input", "") if step.data else ""
+        tool_output = getattr(step.data, "output", "") if step.data else ""
+
+        action = AgentAction(tool=tool_name, tool_input=tool_input, log=log)
+
+        return action, tool_output
+
+    def get_agent_actions(self, intermediate_steps: list[IntermediateStep],
+                          event_filter: list[IntermediateStepType]) -> list[tuple[AgentAction, str]]:
+        """Converts a list of intermediate steps to a list of (AgentAction, output)."""
+        steps = self.filter_intermediate_steps(intermediate_steps, event_filter)
+        last_llm_end_step = None
+        agent_actions = []
+        for step in steps:
+            if step.event_type == IntermediateStepType.LLM_END:
+                last_llm_end_step = step
+            else:
+                action = self.get_agent_action_single(step, last_llm_end_step)
+                agent_actions.append(action)
+
+        return agent_actions
+
+    def get_context(self, intermediate_steps: list[IntermediateStep]) -> list[str]:
+        """Grab the output of all the tools and return them as retrieved context."""
+        return [
+            str(step.data.output) for step in intermediate_steps
+            if step.event_type == IntermediateStepType.TOOL_END and step.data and step.data.output
+        ]