diff --git a/src/examples/evaluate_agents_as_tools.py b/src/examples/evaluate_agents_as_tools.py
index 9a3e1a07..395ddd0a 100644
--- a/src/examples/evaluate_agents_as_tools.py
+++ b/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())
diff --git a/src/examples/evaluate_graph.py b/src/examples/evaluate_graph.py
index 1c18f6d4..94e8c3a5 100644
--- a/src/examples/evaluate_graph.py
+++ b/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 ###
diff --git a/src/examples/evaluate_swarm.py b/src/examples/evaluate_swarm.py
index 9bd601df..e8050786 100644
--- a/src/examples/evaluate_swarm.py
+++ b/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."
)
diff --git a/src/examples/judge_output.py b/src/examples/judge_output.py
index 02640ec8..c00e222d 100644
--- a/src/examples/judge_output.py
+++ b/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,
)
diff --git a/src/examples/try_dataset_generator.py b/src/examples/try_dataset_generator.py
new file mode 100644
index 00000000..e413c3a9
--- /dev/null
+++ b/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())
diff --git a/src/strands_evaluation/case.py b/src/strands_evaluation/case.py
index 27619c42..4adcdde4 100644
--- a/src/strands_evaluation/case.py
+++ b/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
diff --git a/src/strands_evaluation/dataset.py b/src/strands_evaluation/dataset.py
index c53b4888..191cc7da 100644
--- a/src/strands_evaluation/dataset.py
+++ b/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.
diff --git a/src/strands_evaluation/evaluators/interactions_evaluator.py b/src/strands_evaluation/evaluators/interactions_evaluator.py
index fdbc9134..8a337ed1 100644
--- a/src/strands_evaluation/evaluators/interactions_evaluator.py
+++ b/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:
diff --git a/src/strands_evaluation/evaluators/utils/case_prompt_template.py b/src/strands_evaluation/evaluators/utils/case_prompt_template.py
index 0611dff8..96f4ff8c 100644
--- a/src/strands_evaluation/evaluators/utils/case_prompt_template.py
+++ b/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"\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"{evaluation_case.expected_output}\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"{evaluation_case.actual_trajectory}\n"
diff --git a/src/strands_evaluation/evaluators/utils/helper_funcs.py b/src/strands_evaluation/evaluators/utils/helper_funcs.py
index 8c32eecd..25340bde 100644
--- a/src/strands_evaluation/evaluators/utils/helper_funcs.py
+++ b/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.
diff --git a/src/strands_evaluation/generators/dataset_generator.py b/src/strands_evaluation/generators/dataset_generator.py
new file mode 100644
index 00000000..9b4be377
--- /dev/null
+++ b/src/strands_evaluation/generators/dataset_generator.py
@@ -0,0 +1,379 @@
+import asyncio
+
+from pydantic import create_model
+from strands import Agent
+from strands_evaluation.evaluators.evaluator import Evaluator
+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 typing_extensions import Any, Generic, TypeVar
+
+from ..case import Case
+from ..dataset import Dataset
+from ..types.evaluation import Interaction
+from .utils.prompt_templates import generate_case_template as CASE_SYSTEM_PROMPT
+from .utils.prompt_templates import generate_rubric_template as RUBRIC_SYSTEM_PROMPT
+
+InputT = TypeVar("InputT")
+OutputT = TypeVar("OutputT")
+
+
+class DatasetGenerator(Generic[InputT, OutputT]):
+ """
+ Generates evaluation datasets with test cases and rubrics for LLM-based evaluators for agent assessment.
+
+ This class creates structured test cases and evaluation rubrics tailored to specific tasks
+ and domains, enabling comprehensive evaluation of agents' performance.
+ """
+
+ _default_evaluators = {
+ OutputEvaluator: "evaluates only the output response, don't include information about trajectory nor interactions even if provided",
+ TrajectoryEvaluator: "evaluates the trajectory and output if provided, don't include info about interactions even if provided",
+ InteractionsEvaluator: "evaluates the interactions and output if provided, don't include info about trajectory even if provided",
+ }
+
+ def __init__(
+ self,
+ input_type: type,
+ output_type: type,
+ trajectory_type: type = None,
+ include_expected_output: bool = True,
+ include_expected_trajectory: bool = False,
+ include_expected_interactions: bool = False,
+ include_metadata: bool = False,
+ model: str | None = None,
+ max_parallel_num_cases: int = 10,
+ rubric_system_prompt: str = RUBRIC_SYSTEM_PROMPT,
+ case_system_prompt: str = CASE_SYSTEM_PROMPT,
+ ):
+ """
+ Initialize the dataset generator with configuration for test case structure.
+
+ Args:
+ input_type: Type of input data for test cases (e.g., str, dict)
+ output_type: Type of expected output data (e.g., str, int)
+ trajectory_type: Type for trajectory elements, defaults to Any if None
+ include_expected_output: Whether to include expected outputs in test cases
+ include_expected_trajectory: Whether to include expected tool/action trajectories
+ include_expected_interactions: Whether to include expected interaction sequences
+ include_metadata: Whether to include metadata fields in test cases
+ model: Model identifier for the generation agent, defaults to strands' default model.
+ max_parallel_num_cases: Maximum number of test cases to generate in parallel asynchronously
+ rubric_system_prompt: System prompt for rubric generation, defaults to one of the available templates.
+ case_system_prompt: System prompt for test case generation, defaults to one of the available templates.
+ """
+ self.model = model
+ self.input_type = input_type
+ self.output_type = output_type
+ self.include_expected_output = include_expected_output
+ self.include_expected_trajectory = include_expected_trajectory
+ self.include_expected_interactions = include_expected_interactions
+ self.include_metadata = include_metadata
+ self.max_parallel_num_cases = max_parallel_num_cases
+
+ self.rubric_system_prompt = rubric_system_prompt
+ self.case_system_prompt = case_system_prompt
+
+ # Create class structure for Case with stricter/literal types, excluding any fields not needed
+ fields = {"name": (str, ...), "input": (self.input_type, ...)}
+ if self.include_expected_output:
+ fields["expected_output"] = (self.output_type, ...)
+ if self.include_expected_trajectory:
+ fields["expected_trajectory"] = (list[trajectory_type], ...) if trajectory_type else (list[Any], ...)
+ if self.include_expected_interactions:
+ fields["expected_interactions"] = (list[Interaction], ...)
+ if self.include_metadata:
+ fields["metadata"] = (dict[str, Any], ...)
+ self._Case = create_model("_Case", **fields)
+
+ async def _case_worker(self, queue: asyncio.Queue, prompt: str, message_history: list, results: list):
+ """
+ Worker that generates cases from the queue.
+
+ Args:
+ queue: Queue containing cases to process
+ prompt: Generation prompt describing the test case requirements
+ message_history: Optional conversation history to provide context to the generation agent
+ results: List to store results
+
+ """
+ case_generator = Agent(
+ model=self.model,
+ system_prompt=self.case_system_prompt,
+ callback_handler=None,
+ messages=message_history if message_history else [],
+ )
+
+ while True:
+ try:
+ difficulty = queue.get_nowait()
+ except asyncio.QueueEmpty:
+ break
+
+ try:
+ gen_case = await case_generator.structured_output_async(self._Case, prompt + f"Ensure that the test case has a difficulty level of {difficulty}.")
+ results.append(Case(**gen_case.model_dump()))
+ except Exception as e:
+ print(f"Error generating case: {e}")
+ finally:
+ queue.task_done()
+
+ async def generate_cases_async(self, prompt: str, num_cases: int = 5, message_history: list = None) -> list[Case]:
+ """
+ Generate test cases asynchronously using parallel workers.
+
+ Args:
+ prompt: Generation prompt describing the test case requirements
+ num_cases: Number of test cases to generate
+ message_history: Optional conversation history to provide context to the generation agent
+
+ Returns:
+ List of generated Case objects matching the configured schema
+ """
+ queue = asyncio.Queue()
+ generated_cases = []
+
+ # Fill queue with tasks
+ for i in range(num_cases):
+ difficulty = "medium"
+ if i < num_cases*0.3:
+ difficulty = "easy"
+ elif i > num_cases*0.8:
+ difficulty = "hard"
+ queue.put_nowait(difficulty)
+
+ num_workers = min(self.max_parallel_num_cases, num_cases)
+
+ workers = [
+ asyncio.create_task(self._case_worker(queue, prompt, message_history, generated_cases))
+ for _ in range(num_workers)
+ ]
+
+ await queue.join()
+ for worker in workers:
+ worker.cancel()
+ await asyncio.gather(*workers, return_exceptions=True)
+
+ return generated_cases
+
+ async def construct_evaluator_async(
+ self, prompt: str, evaluator: Evaluator, message_history: list = None
+ ) -> Evaluator:
+ """
+ Create an evaluator instance with a generated rubric.
+
+ Currently supports default evaluators: OutputEvaluator, TrajectoryEvaluator,
+ and InteractionsEvaluator. Generates task-specific rubrics for evaluation.
+
+ Args:
+ prompt: Prompt describing the evaluation context and requirements
+ evaluator: Evaluator class to instantiate (must be a default evaluator)
+ message_history: Optional conversation history to provide context to the rubric generation agent
+
+ Returns:
+ Configured evaluator instance with generated rubric
+
+ Raises:
+ ValueError: If evaluator is not one of the supported default evaluators
+ """
+ if evaluator not in self._default_evaluators:
+ raise ValueError(
+ f"{evaluator} is not a default evaluator that needs a rubric. Please use one of the default evaluators: {list(self._default_evaluators.keys())}."
+ )
+
+ rubric_generator_agent = Agent(
+ model=self.model,
+ system_prompt=self.rubric_system_prompt,
+ callback_handler=None,
+ messages=message_history if message_history else [],
+ )
+ final_prompt = (
+ prompt
+ + f"""The evaluator selected is {evaluator.get_type_name()}. This evaluator {self._default_evaluators[evaluator]}.
+ IMPORTANT: Your response must be ONLY a few sentences describing how to evaluate the test cases."""
+ )
+
+ rubric = await rubric_generator_agent.invoke_async(final_prompt)
+ return evaluator(rubric=str(rubric))
+
+ async def from_scratch_async(
+ self, topics: list[str], task_description: str, num_cases: int = 5, evaluator: Evaluator = None
+ ) -> Dataset:
+ """
+ Generate a dataset from scratch based on specified topics and task description.
+
+ Creates diverse test cases covering the given topics for the specified task,
+ with optional evaluator and rubric generation.
+
+ Args:
+ topics: List of topics/domains to cover in test cases
+ task_description: Description of the task the AI system will perform
+ num_cases: Number of test cases to generate
+ evaluator: Optional evaluator class for assessment (generates rubric if provided).
+
+ Returns:
+ Dataset containing generated test cases and evaluator. Use the generic Evaluator as placeholder if no evaluator is passed in.
+ """
+ cases = await self.generate_cases_async(
+ f"""Create test cases for the following topics: {' '.join(topics)} for this task:
+ {task_description}.""",
+ num_cases,
+ )
+ if evaluator:
+ _evaluator = await self.construct_evaluator_async(
+ prompt=f"""Create a rubric for the following topics: {' '.join(topics)} for this task:
+ {task_description}.""",
+ evaluator=evaluator,
+ )
+ return Dataset(cases=cases, evaluator=_evaluator)
+ else:
+ return Dataset(cases=cases)
+
+ async def from_context_async(
+ self, context: str, task_description: str, num_cases: int = 5, evaluator: Evaluator = None
+ ) -> Dataset:
+ """
+ Generate a dataset based on specific context that test cases should reference.
+
+ Creates test cases that can be answered using the provided context,
+ useful for testing knowledge retrieval, context understanding, or domain-specific tasks.
+
+ Args:
+ context: Specific context/information that test cases should reference. If there's any tools they need to use, specify them here too.
+ Be sure to include as much information as you can about tools or sub-agents for generating interaction and/or trajectory.
+ task_description: Description of the task the AI system will perform
+ num_cases: Number of test cases to generate
+ evaluator: Optional evaluator class for assessment (generates rubric if provided), use Evaluator() as a placeholder.
+
+ Returns:
+ Dataset containing context-based test cases and evaluator. Use the generic Evaluator as placeholder if no evaluator is passed in.
+ """
+ cases = await self.generate_cases_async(
+ f"""Create test cases with the following context: {context}. Ensure that the questions can be answer using the provided context for this task: {task_description} """,
+ num_cases=num_cases,
+ )
+ if evaluator:
+ _evaluator = await self.construct_evaluator_async(
+ prompt=f"""Create a rubric with the following context: {context} for this task: {task_description} """,
+ evaluator=evaluator,
+ )
+ return Dataset(cases=cases, evaluator=_evaluator)
+ else:
+ return Dataset(cases=cases)
+
+ async def from_dataset_async(
+ self, source_dataset: Dataset, task_description: str, num_cases: int = 5, extra_information: str = None
+ ) -> Dataset:
+ """
+ Generate a new dataset using an existing dataset as reference.
+
+ Creates new test cases that are similar in style and structure to the source dataset,
+ while adapting them for the specified task. If the source dataset uses a default
+ evaluator with a rubric, generates a new rubric based on the original.
+
+ Args:
+ source_dataset: Original dataset to use as reference for generating new test cases
+ task_description: Description of the task the AI system will perform
+ num_cases: Number of test cases to generate
+ extra_information: Optional additional context or requirements for the new test cases and rubric,
+ be sure to include as much information as you can about tools or sub-agents
+ for generating interaction and/or trajectory.
+
+ Returns:
+ A new Dataset containing test cases inspired by the source dataset but adapted
+ for the new task. Uses an updated evaluator with new rubric if the source
+ evaluator is a default type, otherwise uses generic Evaluator.
+ """
+ source_cases = source_dataset.cases
+ source_evaluator = source_dataset.evaluator
+
+ # construct messages to initialize the agent with context about the previous test cases
+ messages = [{"role": "user", "content": [{"text": "Here are the reference test cases: "}]}]
+ cases_string_list = []
+ for i, case in enumerate(source_cases):
+ cases_string_list.append({"text": f"{i}. {case.model_dump()}"})
+ messages.append({"role": "user", "content": cases_string_list})
+ new_cases = await self.generate_cases_async(
+ prompt=f"Create new test cases similar to the reference cases. Ensure that the input and output are relevant for this task: {task_description}. Here are some extra information: {extra_information}.",
+ num_cases=num_cases,
+ message_history=messages,
+ )
+ new_evaluator = Evaluator()
+ if type(source_evaluator) in self._default_evaluators:
+ source_rubric = source_evaluator.rubric
+ new_evaluator = await self.construct_evaluator_async(
+ prompt=f"Create a new rubric based on the reference rubric. Ensure that the rubric is relevant for this task: {task_description}. Here are some extra information: {extra_information}.",
+ evaluator=type(source_evaluator),
+ message_history=[{"role": "user", "content": [{"text": source_rubric}]}],
+ )
+
+ return Dataset(cases=new_cases, evaluator=new_evaluator)
+
+ async def update_current_dataset_async(
+ self,
+ source_dataset: Dataset,
+ task_description: str,
+ num_cases: int = 5,
+ context: str = None,
+ add_new_cases: bool = True,
+ add_new_rubric: bool = True,
+ new_evaluator_type: type = None,
+ ) -> Dataset:
+ """
+ Update an existing dataset by adding new test cases and/or updating the evaluator.
+
+ Extends the source dataset with additional test cases that complement the existing ones,
+ and optionally updates the evaluation rubric. Useful for iteratively improving datasets
+ or adapting them to new requirements while preserving the original test cases.
+
+ Args:
+ source_dataset: Original dataset to extend and update
+ task_description: Description of the task the AI system will perform
+ num_cases: Number of new test cases to add (if add_new_cases is True)
+ context: Additional context or requirements for new test cases and rubric,
+ be sure to include as much information as you can about tools or sub-agents
+ for generating interaction and/or trajectory.
+ add_new_cases: Whether to generate and add new test cases to the dataset
+ add_new_rubric: Whether to generate a new evaluation rubric
+ new_evaluator_type: Optional new evaluator type to use instead of the source evaluator type
+
+ Returns:
+ Updated Dataset containing original cases plus new cases (if requested) and
+ updated evaluator with new rubric (if requested and evaluator supports it).
+ """
+ source_cases = source_dataset.cases
+ source_evaluator = source_dataset.evaluator
+
+ if add_new_cases:
+ # construct messages to initialize the agent with context about the previous test cases
+ messages = [{"role": "user", "content": [{"text": "Here are the current test cases: "}]}]
+ cases_string_list = []
+ for i, case in enumerate(source_cases):
+ cases_string_list.append({"text": f"{i}. {case.model_dump()}"})
+ messages.append({"role": "user", "content": cases_string_list})
+ new_cases = await self.generate_cases_async(
+ prompt=f"Create new test cases, expanding on previous cases for the following context: {context}. Ensure that the input and output are relevant for this task: {task_description}.",
+ num_cases=num_cases,
+ message_history=messages,
+ )
+
+ if add_new_rubric:
+ if new_evaluator_type:
+ new_evaluator = new_evaluator_type
+ else:
+ new_evaluator = type(source_evaluator) # use the previous evaluator if no new evaluator is passed in
+
+ if new_evaluator in self._default_evaluators:
+ source_rubric = source_evaluator.rubric if type(source_evaluator) in self._default_evaluators else None
+ new_evaluator = await self.construct_evaluator_async(
+ prompt=f"Create a new rubric based on the reference rubric if provided for the following context: {context}. Ensure that the rubric is relevant for this task: {task_description}.",
+ evaluator=new_evaluator,
+ message_history=[{"role": "user", "content": [{"text": source_rubric}]}],
+ )
+ else: # use the original if it's not supported
+ new_evaluator = source_evaluator
+
+ return Dataset(
+ cases=source_cases + new_cases if add_new_cases else source_cases,
+ evaluator=new_evaluator if add_new_rubric else source_evaluator,
+ )
diff --git a/src/strands_evaluation/generators/utils/prompt_templates.py b/src/strands_evaluation/generators/utils/prompt_templates.py
new file mode 100644
index 00000000..9753dfdd
--- /dev/null
+++ b/src/strands_evaluation/generators/utils/prompt_templates.py
@@ -0,0 +1,64 @@
+generate_case_template = """
+You are an expert test case generator for AI evaluation datasets. Your role is to create high-quality, diverse test cases that thoroughly evaluate AI systems across different domains and capabilities.
+
+When given a task description, you will generate test cases specifically designed to evaluate how well an AI system can perform that task.
+
+CORE PRINCIPLES:
+- Generate realistic, practical test cases that reflect real-world usage patterns for the given task
+- Ensure comprehensive coverage of the task requirements and potential challenges
+- Create test cases that are specific, unambiguous, and measurable within the task context
+- Balance difficulty levels to assess different capability thresholds for the task
+- Include edge cases, corner scenarios, and potential failure modes relevant to the task
+
+TEST CASE DESIGN:
+- Easy Level (30%): Basic task functionality, straightforward scenarios, common use cases
+- Medium Level (50%): Multi-step reasoning, moderate complexity, realistic task challenges
+- Hard Level (20%): Complex task scenarios, edge cases, advanced reasoning, error handling
+
+QUALITY STANDARDS:
+- Each test case should have a clear, well-defined input relevant to the task
+- Expected outputs should be accurate, complete, and verifiable for the task
+- Test cases should be independent and not rely on previous context
+- Avoid repetitive or overly similar scenarios within the task scope
+- Ensure cultural sensitivity and avoid biased content
+
+TASK-SPECIFIC CONSIDERATIONS:
+When creating test cases, consider:
+- What inputs will the AI system receive for this task?
+- What outputs should it produce?
+- What tools or capabilities might it need to use?
+- What are the success criteria for this task?
+- What could go wrong or be challenging about this task?
+
+Remember: You are creating evaluation data to measure AI performance on specific tasks. Quality and diversity are paramount for meaningful assessment.
+"""
+
+generate_rubric_template = """
+You are an expert evaluation specialist focused on creating concise, actionable rubrics for AI agent system assessment.
+
+When given a task description, you will create a rubric that captures the essential criteria for evaluating
+how well an AI agent system performs that specific task for a particular information type (eg. output, trajectory, and/or interactions).
+
+RUBRIC REQUIREMENTS:
+- Should be clear, comprehensive, and easy to understand for the specific task
+- Focus on what makes a response high-quality when performing the given task
+- Include key evaluation dimensions relevant to the task (accuracy, completeness, clarity, tool usage, etc.)
+- Be specific enough to guide evaluation but general enough to apply across test cases for the task
+- Consider the task's success criteria and potential failure modes
+- Avoid mentioning specific test case details or examples
+
+TASK-AWARE EVALUATION:
+When creating rubrics, consider:
+- What does successful task completion look like?
+- What are the key quality indicators for this task?
+- What tools, reasoning, or capabilities should be demonstrated?
+- What are common failure modes or errors for this task?
+- How should edge cases or complex scenarios be handled?
+
+FORMAT:
+- Use active, measurable criteria specific to the task
+- Keep concise but comprehensive
+- Focus on observable, evaluable qualities
+
+Focus on creating a rubric that evaluators can consistently apply to measure how well AI systems perform the given task. Starts with "Scoring should ..."
+"""
diff --git a/src/strands_evaluation/types/evaluation.py b/src/strands_evaluation/types/evaluation.py
index 348da61e..4483f0f8 100644
--- a/src/strands_evaluation/types/evaluation.py
+++ b/src/strands_evaluation/types/evaluation.py
@@ -1,10 +1,36 @@
from pydantic import BaseModel
-from typing_extensions import Any, Generic, TypeVar
+from typing_extensions import Any, Generic, TypedDict, TypeVar
InputT = TypeVar("InputT")
OutputT = TypeVar("OutputT")
+class Interaction(TypedDict, total=False):
+ """
+ Represents a single interaction in a multi-agent or multi-step system.
+
+ Used to capture the communication flow and dependencies between different
+ components (agents, tools, or processing nodes) during task execution.
+ All fields are optional to accommodate different interaction patterns.
+
+ Attributes:
+ node_name: Identifier for the agent, tool, or component involved in this interaction
+ dependencies: List of other nodes/components this interaction depends on or references
+ messages: Sequence of messages, responses, or communication exchanged during this interaction
+
+ Example:
+ interaction = {
+ "node_name": "calculator_agent",
+ "dependencies": ["input_parser", "math_validator"],
+ "messages": ["Calculate 2+2"]
+ }
+ """
+
+ node_name: str
+ dependencies: list[str] | None
+ messages: list[str] | None
+
+
class EvaluationData(BaseModel, Generic[InputT, OutputT]):
"""
A record of all of the context for the evaluator to evaluate a test case.
@@ -27,9 +53,9 @@ class EvaluationData(BaseModel, Generic[InputT, OutputT]):
expected_output: OutputT | None = None
expected_trajectory: list[Any] | None = None
actual_trajectory: list[Any] | None = None
- metadata: dict = {}
- actual_interactions: list[dict] | None = None
- expected_interactions: list[dict] | None = None
+ metadata: dict[str, Any] | None = None
+ actual_interactions: list[Interaction] | None = None
+ expected_interactions: list[Interaction] | None = None
class EvaluationOutput(BaseModel):
diff --git a/src/strands_evaluation/types/evaluation_report.py b/src/strands_evaluation/types/evaluation_report.py
index 3b029bbd..faba9408 100644
--- a/src/strands_evaluation/types/evaluation_report.py
+++ b/src/strands_evaluation/types/evaluation_report.py
@@ -179,7 +179,7 @@ def from_dict(cls, data: dict):
"""
return cls.model_validate(data)
- def to_file(self, file_name: str, format: str, directory: str = "report_files"):
+ def to_file(self, file_name: str, format: str = "json", directory: str = "report_files"):
"""
Write the report to a file.
@@ -198,7 +198,7 @@ def to_file(self, file_name: str, format: str, directory: str = "report_files"):
raise ValueError(f"Unsupported format: {format}")
@classmethod
- def from_file(cls, file_path: str, format: str):
+ def from_file(cls, file_path: str, format: str = "json"):
"""
Create an EvaluationReport instance from a file.
diff --git a/tests/test_cases.py b/tests/test_cases.py
index 03471584..6d3dc783 100644
--- a/tests/test_cases.py
+++ b/tests/test_cases.py
@@ -12,7 +12,7 @@ def test_create_simple_case(self):
assert case.expected_output is None
assert case.expected_trajectory is None
assert case.expected_interactions is None
- assert case.metadata == {}
+ assert case.metadata is None
def test_create_full_case(self):
"""Test creating a Case with all fields"""
@@ -21,7 +21,7 @@ def test_create_full_case(self):
input="What is 2+2?",
expected_output="4",
expected_trajectory=["calculator"],
- expected_interactions=[{"node_name": "math_agent", "message": "2x2 is 4.", "dependencies": []}],
+ expected_interactions=[{"node_name": "math_agent", "messages": ["2x2 is 4."], "dependencies": []}],
metadata={"category": "math", "difficulty": "easy"},
)
@@ -29,7 +29,9 @@ def test_create_full_case(self):
assert case.input == "What is 2+2?"
assert case.expected_output == "4"
assert case.expected_trajectory == ["calculator"]
- assert case.expected_interactions == [{"node_name": "math_agent", "message": "2x2 is 4.", "dependencies": []}]
+ assert case.expected_interactions == [
+ {"node_name": "math_agent", "messages": ["2x2 is 4."], "dependencies": []}
+ ]
assert case.metadata == {"category": "math", "difficulty": "easy"}
def test_case_with_different_types(self):
@@ -42,8 +44,8 @@ def test_case_with_different_types(self):
def test_case_with_interactions(self):
"""Test Case with expected_interactions"""
interactions = [
- {"agent": "planner", "message": "plan", "dependencies": []},
- {"agent": "executor", "message": "execute", "dependencies": []},
+ {"node_name": "planner", "messages": ["plan"], "dependencies": []},
+ {"node_name": "executor", "messages": ["execute"]},
]
case = Case[str, str](input="Complex task", expected_interactions=interactions)
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 41233425..1824d116 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -55,7 +55,7 @@ def simple_task(input_val):
assert result.name == "test"
assert result.expected_trajectory is None
assert result.actual_trajectory is None
- assert result.metadata == {}
+ assert result.metadata is None
assert result.actual_interactions is None
assert result.expected_interactions is None
@@ -74,7 +74,7 @@ def dict_task(input_val):
def test_run_task_dict_output_with_interactions(self, mock_evaluator):
"""Test _run_task with dictionary output containing interactions"""
- interactions = [{"node_name": "agent1", "dependencies": [], "message": "hello"}]
+ interactions = [{"node_name": "agent1", "dependencies": [], "messages": ["hello"]}]
case = Case(name="test", input="hello", expected_output="world", expected_interactions=interactions)
dataset = Dataset(cases=[case], evaluator=mock_evaluator)
@@ -89,7 +89,7 @@ def dict_task(input_val):
assert result.actual_output == "response to hello"
assert result.actual_trajectory == ["step1", "step2"]
- assert result.actual_interactions == [{"node_name": "agent1", "dependencies": [], "message": "hello"}]
+ assert result.actual_interactions == interactions
assert result.expected_output == "world"
assert result.expected_trajectory is None
assert result.expected_interactions == interactions
@@ -146,7 +146,7 @@ def test_to_dict_non_empty(self, mock_evaluator):
"expected_output": "world",
"expected_trajectory": None,
"expected_interactions": None,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {"evaluator_type": "MockEvaluator"},
@@ -167,7 +167,7 @@ def test_to_dict_LLM_evaluator_full(self):
"expected_output": "world",
"expected_trajectory": None,
"expected_interactions": None,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {
@@ -195,7 +195,7 @@ def test_to_dict_LLM_evaluator_default(self):
"expected_output": "world",
"expected_trajectory": None,
"expected_interactions": None,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {"evaluator_type": "OutputEvaluator", "rubric": "rubric"},
@@ -217,7 +217,7 @@ def test_to_dict_Trajectory_evaluator_default(self):
"expected_output": "world",
"expected_trajectory": ["step1", "step2"],
"expected_interactions": None,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {"evaluator_type": "TrajectoryEvaluator", "rubric": "rubric"},
@@ -240,7 +240,7 @@ def test_to_dict_Trajectory_evaluator_full(self):
"expected_output": "world",
"expected_trajectory": ["step1", "step2"],
"expected_interactions": None,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {
@@ -256,7 +256,7 @@ def test_to_dict_Interactions_evaluator_default(self):
"""Test converting dataset with Interactions evaluator to dictionary with defaults."""
from src.strands_evaluation.evaluators.interactions_evaluator import InteractionsEvaluator
- interactions = [{"node_name": "agent1", "dependencies": [], "message": "hello"}]
+ interactions = [{"node_name": "agent1", "dependencies": [], "messages": ["hello"]}]
cases = [Case(name="test", input="hello", expected_output="world", expected_interactions=interactions)]
evaluator = InteractionsEvaluator(rubric="rubric")
dataset = Dataset(cases=cases, evaluator=evaluator)
@@ -268,7 +268,7 @@ def test_to_dict_Interactions_evaluator_default(self):
"expected_output": "world",
"expected_trajectory": None,
"expected_interactions": interactions,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {"evaluator_type": "InteractionsEvaluator", "rubric": "rubric"},
@@ -278,7 +278,7 @@ def test_to_dict_Interactions_evaluator_full(self):
"""Test converting dataset with Interactions evaluator to dictionary with no defaults."""
from src.strands_evaluation.evaluators.interactions_evaluator import InteractionsEvaluator
- interactions = [{"node_name": "agent1", "dependencies": [], "message": "hello"}]
+ interactions = [{"node_name": "agent1", "dependencies": [], "messages": ["hello"]}]
cases = [Case(name="test", input="hello", expected_output="world", expected_interactions=interactions)]
evaluator = InteractionsEvaluator(
rubric="rubric", model="model", include_inputs=False, system_prompt="system prompt"
@@ -292,7 +292,7 @@ def test_to_dict_Interactions_evaluator_full(self):
"expected_output": "world",
"expected_trajectory": None,
"expected_interactions": interactions,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {
@@ -306,7 +306,7 @@ def test_to_dict_Interactions_evaluator_full(self):
def test_to_dict_case_dict(self):
"""Test converting dataset with Case with dictionaries as types."""
- case = Case(name="test", input={"field1": "hello"}, expected_output={"field2": "world"})
+ case = Case(name="test", input={"field1": "hello"}, expected_output={"field2": "world"}, metadata={})
evaluator = MockEvaluator()
dataset = Dataset(cases=[case], evaluator=evaluator)
assert dataset.to_dict() == {
@@ -340,7 +340,7 @@ def simple_echo(query):
"expected_output": None,
"expected_trajectory": None,
"expected_interactions": None,
- "metadata": {},
+ "metadata": None,
}
],
"evaluator": {"evaluator_type": "MockEvaluator"},
@@ -548,7 +548,7 @@ def failing_task(input_str):
def test_run_evaluations_with_interactions(self):
"""Test evaluation run with interactions data"""
- interactions = [{"node_name": "agent1", "dependencies": [], "message": "test message"}]
+ interactions = [{"node_name": "agent1", "dependencies": [], "messages": ["test message"]}]
case = Case(name="test", input="hello", expected_output="world", expected_interactions=interactions)
dataset = Dataset(cases=[case], evaluator=MockEvaluator())
@@ -560,3 +560,39 @@ def task_with_interactions(input_val):
assert len(report.cases) == 1
assert report.cases[0]["actual_interactions"] == interactions
assert report.cases[0]["expected_interactions"] == interactions
+
+ def test_cases_getter_deep_copy(self):
+ """Test cases getter should return deep copies"""
+ case = Case(name="test", input="hello", expected_output="world")
+ dataset = Dataset(cases=[case], evaluator=MockEvaluator())
+
+ retrieved = dataset.cases
+ retrieved[0].name = "modified"
+
+ assert dataset.cases == [case]
+
+ def test_cases_setter(self):
+ """Test cases setter updates dataset"""
+ case1 = Case(name="test1", input="hello", expected_output="world")
+ case2 = Case(name="test2", input="hi", expected_output="there")
+ dataset = Dataset(cases=[case1], evaluator=MockEvaluator())
+
+ dataset.cases = [case2]
+ assert dataset.cases == [case2]
+
+ def test_evaluator_getter(self):
+ """Test evaluator getter returns evaluator"""
+ evaluator = MockEvaluator()
+ dataset = Dataset(cases=[], evaluator=evaluator)
+
+ retrieved = dataset.evaluator
+ assert retrieved == evaluator
+
+ def test_evaluator_setter(self):
+ """Test evaluator setter updates dataset"""
+ eval1 = Evaluator()
+ eval2 = MockEvaluator()
+ dataset = Dataset(cases=[], evaluator=eval1)
+
+ dataset.evaluator = eval2
+ assert dataset.evaluator == eval2
diff --git a/tests/test_dataset_generator.py b/tests/test_dataset_generator.py
new file mode 100644
index 00000000..91146fc7
--- /dev/null
+++ b/tests/test_dataset_generator.py
@@ -0,0 +1,312 @@
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from strands_evaluation.case import Case
+from strands_evaluation.dataset import Dataset
+from strands_evaluation.evaluators.evaluator import Evaluator
+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.generators.dataset_generator import DatasetGenerator
+
+
+class TestDatasetGenerator:
+ def test_init(self):
+ """Test initialization"""
+ generator = DatasetGenerator(
+ str,
+ int,
+ trajectory_type=str,
+ include_expected_output=False,
+ include_expected_trajectory=True,
+ include_expected_interactions=True,
+ include_metadata=True,
+ model="test-model",
+ max_parallel_num_cases=5,
+ )
+ assert generator.input_type == str
+ assert generator.output_type == int
+ assert generator.include_expected_output is False
+ assert generator.include_expected_trajectory is True
+ assert generator.include_expected_interactions is True
+ assert generator.include_metadata is True
+ assert generator.model == "test-model"
+ assert generator.max_parallel_num_cases == 5
+
+ @pytest.mark.asyncio
+ async def test_case_worker(self):
+ """Test case worker functionality"""
+ generator = DatasetGenerator(str, str)
+ queue = asyncio.Queue()
+ queue.put_nowait(None)
+ results = []
+
+ mock_agent = AsyncMock()
+ mock_case_data = MagicMock()
+ mock_case_data.model_dump.return_value = {"name": "test", "input": "hello"}
+ mock_agent.structured_output_async.return_value = mock_case_data
+
+ with patch("strands_evaluation.generators.dataset_generator.Agent", return_value=mock_agent):
+ await generator._case_worker(queue, "test prompt", [], results)
+
+ assert len(results) == 1
+ assert isinstance(results[0], Case)
+
+ @pytest.mark.asyncio
+ async def test_generate_cases_async(self):
+ """Test async case generation"""
+ generator = DatasetGenerator(str, str, max_parallel_num_cases=2)
+
+ mock_agent = AsyncMock()
+ mock_case_data = MagicMock()
+ mock_case_data.model_dump.return_value = {"name": "test", "input": "hello"}
+ mock_agent.structured_output_async.return_value = mock_case_data
+
+ with patch("strands_evaluation.generators.dataset_generator.Agent", return_value=mock_agent):
+ cases = await generator.generate_cases_async("test prompt", num_cases=3)
+
+ assert len(cases) == 3
+ assert all(isinstance(case, Case) for case in cases)
+
+ @pytest.mark.asyncio
+ async def test_construct_evaluator_async_output(self):
+ """Test constructing OutputEvaluator"""
+ generator = DatasetGenerator(str, str)
+
+ mock_agent = AsyncMock()
+ mock_agent.invoke_async.return_value = "Generated rubric"
+
+ with patch("strands_evaluation.generators.dataset_generator.Agent", return_value=mock_agent):
+ evaluator = await generator.construct_evaluator_async("test prompt", OutputEvaluator)
+
+ assert isinstance(evaluator, OutputEvaluator)
+ assert evaluator.rubric == "Generated rubric"
+
+ @pytest.mark.asyncio
+ async def test_construct_evaluator_async_trajectory(self):
+ """Test constructing TrajectoryEvaluator"""
+ generator = DatasetGenerator(str, str)
+
+ mock_agent = AsyncMock()
+ mock_agent.invoke_async.return_value = "Generated rubric"
+
+ with patch("strands_evaluation.generators.dataset_generator.Agent", return_value=mock_agent):
+ evaluator = await generator.construct_evaluator_async("test prompt", TrajectoryEvaluator)
+
+ assert isinstance(evaluator, TrajectoryEvaluator)
+ assert evaluator.rubric == "Generated rubric"
+
+ @pytest.mark.asyncio
+ async def test_construct_evaluator_async_interactions(self):
+ """Test constructing InteractionsEvaluator"""
+ generator = DatasetGenerator(str, str)
+
+ mock_agent = AsyncMock()
+ mock_agent.invoke_async.return_value = "Generated rubric"
+
+ with patch("strands_evaluation.generators.dataset_generator.Agent", return_value=mock_agent):
+ evaluator = await generator.construct_evaluator_async("test prompt", InteractionsEvaluator)
+
+ assert isinstance(evaluator, InteractionsEvaluator)
+ assert evaluator.rubric == "Generated rubric"
+
+ @pytest.mark.asyncio
+ async def test_construct_evaluator_async_invalid(self):
+ """Test constructing evaluator with invalid type"""
+ generator = DatasetGenerator(str, str)
+
+ class CustomEvaluator(Evaluator):
+ pass
+
+ with pytest.raises(ValueError, match="is not a default evaluator"):
+ await generator.construct_evaluator_async("test prompt", CustomEvaluator)
+
+ @pytest.mark.asyncio
+ async def test_from_scratch_async_no_evaluator(self):
+ """Test generating dataset from scratch without evaluator"""
+ generator = DatasetGenerator(str, str)
+
+ mock_cases = [Case(name="test", input="hello")]
+
+ with patch.object(generator, "generate_cases_async", return_value=mock_cases):
+ dataset = await generator.from_scratch_async(["topic1"], "test task", num_cases=1)
+
+ assert isinstance(dataset, Dataset)
+ assert dataset.cases == mock_cases
+ assert isinstance(dataset.evaluator, Evaluator)
+
+ @pytest.mark.asyncio
+ async def test_from_scratch_async_with_evaluator(self):
+ """Test generating dataset from scratch with evaluator"""
+ generator = DatasetGenerator(str, str)
+
+ mock_cases = [Case(name="test", input="hello")]
+ mock_evaluator = OutputEvaluator(rubric="test rubric")
+
+ with (
+ patch.object(generator, "generate_cases_async", return_value=mock_cases),
+ patch.object(generator, "construct_evaluator_async", return_value=mock_evaluator),
+ ):
+ dataset = await generator.from_scratch_async(["topic1"], "test task", evaluator=OutputEvaluator)
+
+ assert isinstance(dataset, Dataset)
+ assert dataset.cases == mock_cases
+ assert dataset.evaluator == mock_evaluator
+
+ @pytest.mark.asyncio
+ async def test_from_context_async_no_evaluator(self):
+ """Test generating dataset from context without evaluator"""
+ generator = DatasetGenerator(str, str)
+
+ mock_cases = [Case(name="test", input="hello")]
+
+ with patch.object(generator, "generate_cases_async", return_value=mock_cases):
+ dataset = await generator.from_context_async("test context", "test task", num_cases=1)
+
+ assert isinstance(dataset, Dataset)
+ assert dataset.cases == mock_cases
+ assert isinstance(dataset.evaluator, Evaluator)
+
+ @pytest.mark.asyncio
+ async def test_from_context_async_with_evaluator(self):
+ """Test generating dataset from context with evaluator"""
+ generator = DatasetGenerator(str, str)
+
+ mock_cases = [Case(name="test", input="hello")]
+ mock_evaluator = OutputEvaluator(rubric="test rubric")
+
+ with (
+ patch.object(generator, "generate_cases_async", return_value=mock_cases),
+ patch.object(generator, "construct_evaluator_async", return_value=mock_evaluator),
+ ):
+ dataset = await generator.from_context_async("test context", "test task", evaluator=OutputEvaluator)
+
+ assert isinstance(dataset, Dataset)
+ assert dataset.cases == mock_cases
+ assert dataset.evaluator == mock_evaluator
+
+ @pytest.mark.asyncio
+ async def test_from_dataset_async_generic_evaluator(self):
+ """Test generating dataset from existing dataset with generic evaluator"""
+ generator = DatasetGenerator(str, str)
+
+ source_cases = [Case(name="source", input="source_input")]
+ source_dataset = Dataset(cases=source_cases, evaluator=Evaluator())
+ mock_new_cases = [Case(name="new", input="new_input")]
+
+ with patch.object(generator, "generate_cases_async", return_value=mock_new_cases):
+ dataset = await generator.from_dataset_async(source_dataset, "test task")
+
+ assert isinstance(dataset, Dataset)
+ assert dataset.cases == mock_new_cases
+ assert isinstance(dataset.evaluator, Evaluator)
+
+ @pytest.mark.asyncio
+ async def test_from_dataset_async_default_evaluator(self):
+ """Test generating dataset from existing dataset with default evaluator"""
+ generator = DatasetGenerator(str, str)
+
+ source_cases = [Case(name="source", input="source_input")]
+ source_evaluator = OutputEvaluator(rubric="source rubric")
+ source_dataset = Dataset(cases=source_cases, evaluator=source_evaluator)
+ mock_new_cases = [Case(name="new", input="new_input")]
+ mock_new_evaluator = OutputEvaluator(rubric="new rubric")
+
+ with (
+ patch.object(generator, "generate_cases_async", return_value=mock_new_cases),
+ patch.object(generator, "construct_evaluator_async", return_value=mock_new_evaluator),
+ ):
+ dataset = await generator.from_dataset_async(source_dataset, "test task")
+
+ assert isinstance(dataset, Dataset)
+ assert dataset.cases == mock_new_cases
+ assert dataset.evaluator == mock_new_evaluator
+
+ @pytest.mark.asyncio
+ async def test_update_current_dataset_async_add_cases_only(self):
+ """Test updating dataset by adding new cases only"""
+ generator = DatasetGenerator(str, str)
+
+ source_cases = [Case(name="source", input="source_input")]
+ source_evaluator = Evaluator()
+ source_dataset = Dataset(cases=source_cases, evaluator=source_evaluator)
+ mock_new_cases = [Case(name="new", input="new_input")]
+
+ with patch.object(generator, "generate_cases_async", return_value=mock_new_cases):
+ dataset = await generator.update_current_dataset_async(
+ source_dataset, "test task", add_new_cases=True, add_new_rubric=False
+ )
+
+ assert len(dataset.cases) == 2
+ assert dataset.cases == source_cases + mock_new_cases
+ assert dataset.evaluator == source_evaluator
+
+ @pytest.mark.asyncio
+ async def test_update_current_dataset_async_add_rubric_only(self):
+ """Test updating dataset by adding new rubric only"""
+ generator = DatasetGenerator(str, str)
+
+ source_cases = [Case(name="source", input="source_input")]
+ source_evaluator = OutputEvaluator(rubric="source rubric")
+ source_dataset = Dataset(cases=source_cases, evaluator=source_evaluator)
+ mock_new_evaluator = OutputEvaluator(rubric="new rubric")
+
+ with patch.object(generator, "construct_evaluator_async", return_value=mock_new_evaluator):
+ dataset = await generator.update_current_dataset_async(
+ source_dataset, "test task", add_new_cases=False, add_new_rubric=True
+ )
+
+ assert dataset.cases == source_cases
+ print("here", dataset.evaluator)
+ assert dataset.evaluator == mock_new_evaluator
+
+ @pytest.mark.asyncio
+ async def test_update_current_dataset_async_new_evaluator_type(self):
+ """Test updating dataset with new evaluator type"""
+ generator = DatasetGenerator(str, str)
+
+ source_cases = [Case(name="source", input="source_input")]
+ source_evaluator = OutputEvaluator(rubric="source rubric")
+ source_dataset = Dataset(cases=source_cases, evaluator=source_evaluator)
+ mock_new_evaluator = TrajectoryEvaluator(rubric="new rubric")
+
+ with patch.object(generator, "construct_evaluator_async", return_value=mock_new_evaluator):
+ dataset = await generator.update_current_dataset_async(
+ source_dataset,
+ "test task",
+ add_new_cases=False,
+ add_new_rubric=True,
+ new_evaluator_type=TrajectoryEvaluator,
+ )
+
+ assert dataset.cases == source_cases
+ assert isinstance(dataset.evaluator, TrajectoryEvaluator)
+
+ @pytest.mark.asyncio
+ async def test_update_current_dataset_async_unsupported_evaluator_type(self):
+ """Test updating dataset with unsupported evaluator type falls back to original"""
+ generator = DatasetGenerator(str, str)
+
+ source_cases = [Case(name="source", input="source_input")]
+ source_evaluator = OutputEvaluator(rubric="source rubric")
+ source_dataset = Dataset(cases=source_cases, evaluator=source_evaluator)
+
+ class UnsupportedEvaluator(Evaluator):
+ pass
+
+ dataset = await generator.update_current_dataset_async(
+ source_dataset, "test task", add_new_cases=False, add_new_rubric=True,
+ new_evaluator_type=UnsupportedEvaluator
+ )
+
+ assert dataset.cases == source_cases
+ assert dataset.evaluator == source_evaluator
+
+ def test_default_evaluators_mapping(self):
+ """Test that default evaluators are properly mapped"""
+ generator = DatasetGenerator(str, str)
+
+ assert OutputEvaluator in generator._default_evaluators
+ assert TrajectoryEvaluator in generator._default_evaluators
+ assert InteractionsEvaluator in generator._default_evaluators
diff --git a/tests/test_evaluator.py b/tests/test_evaluator.py
index b9cf2714..2e7fac2f 100644
--- a/tests/test_evaluator.py
+++ b/tests/test_evaluator.py
@@ -32,7 +32,7 @@ async def test_async_base_evaluator_not_implemented(self, evaluation_data):
with pytest.raises(
NotImplementedError,
- match="This method should be implemented in subclasses," \
+ match="This method should be implemented in subclasses,"
" especially if you want to run evaluations asynchronously.",
):
await evaluator.evaluate_async(evaluation_data)
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 23eb49ed..a9374c26 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -44,8 +44,8 @@ def interaction_case():
input="hello",
expected_output="world",
expected_interactions=[
- {"node_name": "agent1", "dependencies": [], "messages": "processing hello"},
- {"node_name": "agent2", "dependencies": ["agent1"], "messages": "final result"},
+ {"node_name": "agent1", "dependencies": [], "messages": ["processing hello"]},
+ {"node_name": "agent2", "dependencies": ["agent1"], "messages": ["final result"]},
],
)
]
@@ -280,8 +280,8 @@ def task_with_interactions(input_val):
return {
"output": "world",
"interactions": [
- {"node_name": "agent1", "dependencies": [], "messages": "processing hello"},
- {"node_name": "agent2", "dependencies": ["agent1"], "messages": "final result"},
+ {"node_name": "agent1", "dependencies": [], "messages": ["processing hello"]},
+ {"node_name": "agent2", "dependencies": ["agent1"], "messages": ["final result"]},
],
}
@@ -303,8 +303,8 @@ async def async_interactions_task(input_val):
return {
"output": "world",
"interactions": [
- {"node_name": "agent1", "dependencies": [], "message": "processing hello"},
- {"node_name": "agent2", "dependencies": ["agent1"], "message": "final result"},
+ {"node_name": "agent1", "dependencies": [], "messages": ["processing hello"]},
+ {"node_name": "agent2", "dependencies": ["agent1"], "messages": ["final result"]},
],
}
diff --git a/tests/test_interactions_evaluator.py b/tests/test_interactions_evaluator.py
index 0f8189c3..36ad563f 100644
--- a/tests/test_interactions_evaluator.py
+++ b/tests/test_interactions_evaluator.py
@@ -36,12 +36,12 @@ def evaluation_data():
actual_output="Climate change affects agriculture through drought, temperature, and pests.",
expected_output="Climate change impacts agriculture via multiple factors.",
actual_interactions=[
- {"node_name": "planner", "dependencies": [], "messages": "Breaking down the analysis task"},
- {"node_name": "researcher", "dependencies": ["planner"], "messages": "Found key climate impacts"},
+ {"node_name": "planner", "dependencies": [], "messages": ["Breaking down the analysis task"]},
+ {"node_name": "researcher", "dependencies": ["planner"], "messages": ["Found key climate impacts"]},
],
expected_interactions=[
- {"node_name": "planner", "dependencies": [], "messages": "Plan the analysis approach"},
- {"node_name": "researcher", "dependencies": ["planner"], "messages": "Research climate data"},
+ {"node_name": "planner", "dependencies": [], "messages": ["Plan the analysis approach"]},
+ {"node_name": "researcher", "dependencies": ["planner"], "messages": ["Research climate data"]},
],
name="climate_analysis_test",
)
@@ -132,12 +132,11 @@ def test_evaluate_without_inputs(self, mock_agent_class, evaluation_data, mock_a
def test_evaluate_missing_interactions(self):
"""Test evaluation raises exception when interactions are missing"""
evaluator = InteractionsEvaluator(rubric="Test rubric")
- evaluation_data = EvaluationData(
- input="test",
- actual_output="result"
- )
-
- with pytest.raises(Exception, match="Please make sure the task function returns a dictionary with the key 'interactions'"):
+ evaluation_data = EvaluationData(input="test", actual_output="result")
+
+ with pytest.raises(
+ Exception, match="Please make sure the task function returns a dictionary with the key 'interactions'"
+ ):
evaluator.evaluate(evaluation_data)
@patch("src.strands_evaluation.evaluators.interactions_evaluator.Agent")
@@ -149,11 +148,12 @@ def test_evaluate_missing_interaction_fields(self, mock_agent_class, mock_agent)
input="test",
actual_interactions=[{"node_name": "test"}], # Missing dependencies and message
)
-
+
with pytest.raises(
Exception,
- match="Please make sure the task function returns a dictionary" \
- " with the key 'interactions' that contains 'node_name', 'dependencies', and 'messages'"):
+ match="Please make sure the task function returns a dictionary"
+ " with the key 'interactions' that contains 'node_name', 'dependencies', and 'messages'",
+ ):
evaluator.evaluate(evaluation_data)
@patch("src.strands_evaluation.evaluators.interactions_evaluator.Agent")
@@ -169,8 +169,8 @@ def test_evaluate_with_dict_rubric(self, mock_agent_class, mock_agent):
evaluation_data = EvaluationData(
input="Analyze climate change",
actual_interactions=[
- {"node_name": "planner", "dependencies": [], "messages": "Breaking down analysis"},
- {"node_name": "researcher", "dependencies": ["planner"], "messages": "Research findings"},
+ {"node_name": "planner", "dependencies": [], "messages": ["Breaking down analysis"]},
+ {"node_name": "researcher", "dependencies": ["planner"], "messages": ["Research findings"]},
],
)
@@ -199,7 +199,7 @@ def test_evaluate_dict_rubric_missing_node(self, mock_agent_class, mock_agent):
evaluation_data = EvaluationData(
input="test",
- actual_interactions=[{"node_name": "missing_node", "dependencies": [], "messages": "test message"}],
+ actual_interactions=[{"node_name": "missing_node", "dependencies": [], "messages": ["test message"]}],
)
with pytest.raises(Exception, match="Please make sure the rubric dictionary contains the key 'missing_node'"):
@@ -215,8 +215,8 @@ async def test_evaluate_async_with_dict_rubric(self, mock_agent_class, mock_asyn
evaluation_data = EvaluationData(
input="Climate analysis",
actual_interactions=[
- {"node_name": "planner", "dependencies": [], "messages": "Plan analysis"},
- {"node_name": "researcher", "dependencies": ["planner"], "messages": "Research data"},
+ {"node_name": "planner", "dependencies": [], "messages": ["Plan analysis"]},
+ {"node_name": "researcher", "dependencies": ["planner"], "messages": ["Research data"]},
],
)