Add support for dataset generation with strands agents

src/examples/evaluate_agents_as_tools.py

@@ -7,6 +7,7 @@
 from strands_evaluation.evaluators.interactions_evaluator import InteractionsEvaluator
 from strands_evaluation.evaluators.trajectory_evaluator import TrajectoryEvaluator
 from strands_evaluation.evaluators.utils import helper_funcs
+from strands_evaluation.types.evaluation import Interaction
 from strands_tools import http_request, retrieve
 
 
@@ -207,7 +208,9 @@ def trip_planning_assistant(query: str) -> str:
         interactions = []
         for tool_used in tools_used:
             interactions.append(
-                {"node_name": tool_used.get("name"), "dependencies": [], "messages": tool_used.get("message")}
+                Interaction(
+                    **{"node_name": tool_used.get("name"), "dependencies": [], "messages": tool_used.get("message")}
+                )
             )
 
         # This helper function does not include message
@@ -219,12 +222,12 @@ def trip_planning_assistant(query: str) -> str:
 
 if __name__ == "__main__":
     # run the file as a module: eg. python -m examples.evaluate_agents_as_tools
-    # start = datetime.datetime.now()
-    # report = asyncio.run(async_evaluate_tools_as_agents_example())
-    # end = datetime.datetime.now()
-    # print("Time: ", end - start)
-    # report.to_file("async_tools_as_agents_report_w_output", "json")
-    # report.run_display(include_actual_trajectory=True)
+    start = datetime.datetime.now()
+    report = asyncio.run(async_evaluate_tools_as_agents_example())
+    end = datetime.datetime.now()
+    print("Time: ", end - start)
+    report.to_file("async_tools_as_agents_report_w_output", "json")
+    report.run_display(include_actual_trajectory=True)
 
     start = datetime.datetime.now()
     report = asyncio.run(async_agents_as_tools_interaction_example())

src/examples/evaluate_graph.py

@@ -98,8 +98,10 @@ async def async_graph_interaction_history_example():
         " and fact check and synthesize the information into a coherent report.",
     }
     # if want to use the same rubric
-    basic_rubric = "The graph system should ultilized the agents as expected with relevant information." \
-    " The actual interactions should include more information than expected."
+    basic_rubric = (
+        "The graph system should ultilized the agents as expected with relevant information."
+        " The actual interactions should include more information than expected."
+    )
     evaluator = InteractionsEvaluator(rubric=rubric)
 
     ### Step 3: Create dataset ###

src/examples/evaluate_swarm.py

@@ -68,7 +68,7 @@ async def async_swarm_interactions_example():
 
     ### Step 2: Create evaluator ###
     evaluator = InteractionsEvaluator(
-        rubric="The interaction sequence should represent a logical and optimal handoff of tasks" \
+        rubric="The interaction sequence should represent a logical and optimal handoff of tasks"
         " from one agent to another."
     )
 

src/examples/judge_output.py

@@ -42,16 +42,16 @@ def output_judge_example():
 
     ### Step 2: Create evaluator ###
     LLM_judge = OutputEvaluator(
-        rubric="The output should represent a reasonable answer to the input. 1 if the output is concise and correct." \
+        rubric="The output should represent a reasonable answer to the input. 1 if the output is concise and correct."
         " 0 if the output is wrong and full of unnecessary text.",
         include_inputs=True,
     )
     ## or
     LLM_judge_w_prompt = OutputEvaluator(
-        rubric="The output should represent a reasonable answer to the input. 1 if the output is concise and correct." \
+        rubric="The output should represent a reasonable answer to the input. 1 if the output is concise and correct."
         " 0 if the output is wrong and full of unnecessary text.",
-        system_prompt="You are an expert AI evaluator. Your job is to assess the quality of the response" \
-        " based according to a user-specified rubric." \
+        system_prompt="You are an expert AI evaluator. Your job is to assess the quality of the response"
+        " based according to a user-specified rubric."
         " You respond with a JSON object with this structure: {reason: string, pass: boolean, score: number}",
         include_inputs=True,
     )
@@ -104,16 +104,16 @@ async def async_output_judge_example():
 
     ### Step 2: Create evaluator ###
     LLM_judge = OutputEvaluator(
-        rubric="The output should represent a reasonable answer to the input." \
+        rubric="The output should represent a reasonable answer to the input."
         " 1 if the output is concise and correct. 0 if the output is wrong and full of unnecessary text.",
         include_inputs=True,
     )
     ## or
     LLM_judge_w_prompt = OutputEvaluator(
-        rubric="The output should represent a reasonable answer to the input." \
+        rubric="The output should represent a reasonable answer to the input."
         " 1 if the output is concise and correct. 0 if the output is wrong and full of unnecessary text.",
-        system_prompt="You are an expert AI evaluator. Your job is to assess the quality of the response based" \
-        " according to a user-specified rubric." \
+        system_prompt="You are an expert AI evaluator. Your job is to assess the quality of the response based"
+        " according to a user-specified rubric."
         " You respond with a JSON object with this structure: {reason: string, pass: boolean, score: number}",
         include_inputs=True,
     )

src/examples/try_dataset_generator.py

@@ -0,0 +1,92 @@
+import asyncio
+
+from strands import Agent
+from strands_evaluation.dataset import Dataset
+from strands_evaluation.evaluators.interactions_evaluator import InteractionsEvaluator
+from strands_evaluation.evaluators.output_evaluator import OutputEvaluator
+from strands_evaluation.evaluators.trajectory_evaluator import TrajectoryEvaluator
+from strands_evaluation.evaluators.utils import helper_funcs
+from strands_evaluation.generators.dataset_generator import DatasetGenerator
+from typing_extensions import TypedDict
+
+
+async def test_dataset_generator():
+    class TrajectoryType(TypedDict):
+        tool: str
+        input: dict
+
+    generator = DatasetGenerator[str, float](
+        str, float, trajectory_type=TrajectoryType, include_expected_trajectory=True
+    )
+    dataset = await generator.from_context_async(
+        "Create test cases about math given that you have access to these tools: calculator, python.",
+        task_description="Getting the response from an AI agent with access to tools.",
+        num_cases=10,
+        evaluator=TrajectoryEvaluator,
+    )
+    print(len(dataset.cases))
+    dataset.to_file("generated_traj_dataset_context")
+
+    generator = DatasetGenerator[str, str](str, str, trajectory_type=TrajectoryType, include_expected_interactions=True)
+    dataset = await generator.from_context_async(
+        "Create test cases about research for a multi-agent system with the following agents: researcher, analyst, fact_checker, and report_writer.",
+        task_description="Getting the response from an AI agent with access to tools.",
+        num_cases=5,
+        evaluator=InteractionsEvaluator,
+    )
+    print(len(dataset.cases))
+    dataset.to_file("generated_interaction_dataset_context")
+
+    generator = DatasetGenerator[str, str](str, str, trajectory_type=TrajectoryType, include_expected_interactions=True)
+    dataset = await generator.from_scratch_async(
+        ["math", "science"],
+        task_description="Getting the response from an AI agent.",
+        num_cases=5,
+        evaluator=OutputEvaluator,
+    )
+    dataset.to_file("generated_output_dataset_scratch")
+    print(len(dataset.cases))
+
+    # try from_dataset
+    generator = DatasetGenerator[str, str](str, str, trajectory_type=TrajectoryType, include_expected_interactions=True)
+    dataset = Dataset.from_file("dataset_files/generated_interaction_dataset_context.json")
+    new_dataset = await generator.update_current_dataset_async(
+        dataset,
+        "Getting the response from an AI agent with access to tools.",
+        context="Create test cases about research for a multi-agent system with the following agents: researcher, analyst, fact_checker, and report_writer.",
+        num_cases=5,
+        new_evaluator_type=OutputEvaluator,
+    )
+    new_dataset.to_file("generated_output_dataset_from_dataset")
+    print(len(new_dataset.cases))
+
+    # new from dataset
+    generator = DatasetGenerator[str, str](str, str, trajectory_type=TrajectoryType, include_expected_trajectory=True)
+    dataset = Dataset.from_file("dataset_files/generated_traj_dataset_context.json")
+    new_dataset = await generator.from_dataset_async(
+        dataset,
+        "Getting the response from an AI agent with access to tools.",
+        extra_information="Create test cases about math given that you have access to these tools: calculator, python.",
+    )
+
+    new_dataset.to_file("generated_traj_dataset_from_dataset")
+    print(len(new_dataset.cases))
+
+    # try evaluating them
+    dataset = Dataset.from_file("dataset_files/generated_traj_dataset_context.json")
+
+    def task_func(input: str) -> dict:
+        agent = Agent(system_prompt="You are a helpful assistant that can do math.", callback_handler=None)
+        output = agent(input)
+        return {
+            "output": str(output),
+            "trajectory": helper_funcs.extract_agent_tools_used_from_messages(agent.messages),
+        }
+
+    report = await dataset.run_evaluations_async(task_func)
+    report.run_display()
+
+
+if __name__ == "__main__":
+    # python -m examples.try_dataset_generator
+    asyncio.run(test_dataset_generator())

src/strands_evaluation/case.py

@@ -1,6 +1,8 @@
 from pydantic import BaseModel
 from typing_extensions import Any, Generic, TypeVar
 
+from .types.evaluation import Interaction
+
 InputT = TypeVar("InputT")
 OutputT = TypeVar("OutputT")
 
@@ -43,5 +45,5 @@ class Case(BaseModel, Generic[InputT, OutputT]):
     input: InputT
     expected_output: OutputT | None = None
     expected_trajectory: list[Any] | None = None
-    expected_interactions: list[dict] | None = None
-    metadata: dict[str, Any] = {}
+    expected_interactions: list[Interaction] | None = None
+    metadata: dict[str, Any] | None = None

src/strands_evaluation/dataset.py

@@ -3,7 +3,6 @@
 import os
 from collections.abc import Callable
 
-from pydantic import BaseModel
 from typing_extensions import Any, Generic, TypeVar
 
 from .case import Case
@@ -18,7 +17,7 @@
 OutputT = TypeVar("OutputT")
 
 
-class Dataset(BaseModel, Generic[InputT, OutputT]):
+class Dataset(Generic[InputT, OutputT]):
     """
     A collection of test cases, representing a dataset.
 
@@ -47,8 +46,52 @@ class Dataset(BaseModel, Generic[InputT, OutputT]):
         )
     """
 
-    cases: list[Case[InputT, OutputT]]
-    evaluator: Evaluator[InputT, OutputT]
+    def __init__(self, cases: list[Case[InputT, OutputT]] = None, evaluator: Evaluator[InputT, OutputT] = Evaluator()):
+        self._cases = cases if cases else []
+        self._evaluator = evaluator
+
+    @property
+    def cases(self) -> list[Case[InputT, OutputT]]:
+        """
+        Get a deep copy of all test cases in the dataset.
+
+        Returns deep copies to prevent accidental mutation of the original test cases.
+        Users can safely modify the returned cases without affecting the dataset.
+
+        Returns:
+            List of Case objects (deep copies) containing all test cases in the dataset
+        """
+        return [case.model_copy(deep=True) for case in self._cases]
+
+    @property
+    def evaluator(self) -> Evaluator[InputT, OutputT]:
+        """
+        Get the evaluator used for assessing test case performance.
+
+        Returns:
+            The evaluator instance configured for this dataset
+        """
+        return self._evaluator
+
+    @cases.setter
+    def cases(self, new_cases: list[Case[InputT, OutputT]]):
+        """
+        Set the test cases for this dataset.
+
+        Args:
+            new_cases: List of Case objects to use as the dataset's test cases
+        """
+        self._cases = new_cases
+
+    @evaluator.setter
+    def evaluator(self, new_evaluator: Evaluator[InputT, OutputT]):
+        """
+        Set the evaluator for assessing test case performance.
+
+        Args:
+            new_evaluator: Evaluator instance to use for evaluating test cases
+        """
+        self._evaluator = new_evaluator
 
     def _run_task(
         self, task: Callable[[InputT], OutputT | dict[str, Any]], case: Case[InputT, OutputT]
@@ -175,7 +218,7 @@ def run_evaluations(self, task: Callable[[InputT], OutputT | dict[str, Any]]) ->
         test_passes = []
         cases = []
         reasons = []
-        for case in self.cases:
+        for case in self._cases:
             try:
                 evaluation_context = self._run_task(task, case)
                 evaluation_output = self.evaluator.evaluate(evaluation_context)
@@ -217,10 +260,10 @@ async def run_evaluations_async(self, task: Callable, max_workers: int = 10) ->
         queue = asyncio.Queue()
         results = []
 
-        for case in self.cases:
+        for case in self._cases:
             queue.put_nowait(case)
 
-        num_workers = min(max_workers, len(self.cases))
+        num_workers = min(max_workers, len(self._cases))
 
         # Create and start workers
         workers = [asyncio.create_task(self._worker(queue, task, results)) for _ in range(num_workers)]
@@ -252,9 +295,9 @@ def to_dict(self) -> dict:
         Return:
             A dictionary representation of the dataset.
         """
-        return {"cases": [case.model_dump() for case in self.cases], "evaluator": self.evaluator.to_dict()}
+        return {"cases": [case.model_dump() for case in self._cases], "evaluator": self.evaluator.to_dict()}
 
-    def to_file(self, file_name: str, format: str, directory: str = "dataset_files"):
+    def to_file(self, file_name: str, format: str = "json", directory: str = "dataset_files"):
         """
         Write the dataset to a file.
 
@@ -304,7 +347,7 @@ def from_dict(cls, data: dict, custom_evaluators: list[Evaluator] = []):
         return cls(cases=cases, evaluator=evaluator)
 
     @classmethod
-    def from_file(cls, file_path: str, format: str, custom_evaluators: list[Evaluator] = []):
+    def from_file(cls, file_path: str, format: str = "json", custom_evaluators: list[Evaluator] = []):
         """
         Create a dataset from a file.
 

src/strands_evaluation/evaluators/interactions_evaluator.py

@@ -146,8 +146,10 @@ def evaluate(self, evaluation_case: EvaluationData[InputT, OutputT]) -> Evaluati
         Returns:
             The results of the evaluation as EvaluationOutput.
         """
-        if not evaluation_case.actual_interactions:
-            raise Exception("Please make sure the task function returns a dictionary with the key 'interactions'.")
+        if evaluation_case.actual_interactions is None:
+            raise Exception(
+                "Please make sure the task function returns a dictionary with the key 'interactions' of type Interaction."
+            )
         num_interactions = len(evaluation_case.actual_interactions)
 
         # keep all of the context
@@ -168,7 +170,7 @@ def evaluate(self, evaluation_case: EvaluationData[InputT, OutputT]) -> Evaluati
 
             ## Evaluate ##
             result = evaluator_agent.structured_output(EvaluationOutput, evaluation_prompt)
-            
+
         return result
 
     async def evaluate_async(self, evaluation_case: EvaluationData[InputT, OutputT]) -> EvaluationOutput:

src/strands_evaluation/evaluators/utils/case_prompt_template.py

@@ -38,7 +38,7 @@ def compose_test_prompt(
         if evaluation_case.actual_output:
             evaluation_prompt += f"<Output>{evaluation_case.actual_output}</Output>\n"
     else:
-        if not evaluation_case.actual_output:
+        if evaluation_case.actual_output is None:
             raise Exception(
                 "Please make sure the task function return the output or a dictionary with the key 'output'."
             )
@@ -48,7 +48,7 @@ def compose_test_prompt(
         evaluation_prompt += f"<ExpectedOutput>{evaluation_case.expected_output}</ExpectedOutput>\n"
 
     if uses_trajectory:  # trajectory evaluations require actual_trajectory
-        if not evaluation_case.actual_trajectory:
+        if evaluation_case.actual_trajectory is None:
             raise Exception("Please make sure the task function return a dictionary with the key 'trajectory'.")
         evaluation_prompt += f"<Trajectory>{evaluation_case.actual_trajectory}</Trajectory>\n"
 

src/strands_evaluation/evaluators/utils/helper_funcs.py

@@ -71,6 +71,7 @@ def extract_swarm_interactions(swarm_result: SwarmResult) -> list[dict]:
     handoff_info = extract_swarm_handoffs(swarm_result)
     return extract_swarm_interactions_from_handoffs(handoff_info)
 
+
 def extract_graph_interactions(graph_result: GraphResult):
     """
     Extract interaction information from graph execution results.
@@ -90,6 +91,7 @@ def extract_graph_interactions(graph_result: GraphResult):
         message_info.append({"node_name": node_name, "dependencies": dependencies, "messages": node_messages})
     return message_info
 
+
 def extract_agent_tools_used_from_messages(agent_messages):
     """
     Extract tool usage information from agent message history.
@@ -146,6 +148,7 @@ def extract_agent_tools_used_from_metrics(agent_result):
         )
     return tools_used
 
+
 def extract_tools_description(agent: Agent, is_short: bool = True):
     """
     Extract a dictionary of all tools used in a given agent.