Skip to content

Commit de7f929

Browse files
📝 Add docstrings to fix/type-annotations-llm-generator
Docstrings generation was requested by @aviralgarg05. * #74 (comment) The following files were modified: * `src/agentunit/generators/llm_generator.py`
1 parent 947296b commit de7f929

1 file changed

Lines changed: 28 additions & 18 deletions

File tree

src/agentunit/generators/llm_generator.py

Lines changed: 28 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -102,15 +102,19 @@ def _create_generation_prompt(self, domain: str, task_description: str) -> str:
102102
async def generate(
103103
self, domain: str, task_description: str, constraints: list[str] | None = None
104104
) -> DatasetSource:
105-
"""Generate synthetic dataset.
106-
107-
Args:
108-
domain: Domain of the task (e.g., "customer service", "code review")
109-
task_description: Detailed description of what the agent should do
110-
constraints: Optional list of constraints or requirements
111-
105+
"""
106+
Generate a dataset of synthetic test cases for a given domain and task.
107+
108+
Parameters:
109+
domain (str): Domain name for the generated cases (e.g., "customer service").
110+
task_description (str): Description of the task the cases should exercise.
111+
constraints (list[str] | None): Optional list of additional constraints to apply when generating cases.
112+
112113
Returns:
113-
DatasetSource with generated cases
114+
DatasetSource: A dataset containing generated DatasetCase objects. Each case includes an `id`, `query`, optional `expected_output`, and `metadata` containing at least `difficulty`, `generated` (True), and `domain`.
115+
116+
Raises:
117+
json.JSONDecodeError: If the model response cannot be parsed as JSON.
114118
"""
115119
prompt = self._create_generation_prompt(domain, task_description)
116120

@@ -238,16 +242,22 @@ async def generate(
238242
constraints: list[str] | None = None,
239243
seed_examples: list[dict[str, Any]] | None = None,
240244
) -> DatasetSource:
241-
"""Generate synthetic dataset.
242-
243-
Args:
244-
domain: Domain of the task
245-
task_description: Detailed description of what the agent should do
246-
constraints: Optional list of constraints
247-
seed_examples: Optional seed examples to guide generation
248-
245+
"""
246+
Generate a synthetic dataset of test cases for a given domain and task.
247+
248+
Constructs a chat prompt from the domain, task description, optional constraints, and optional seed examples, sends it to the configured OpenAI chat model, parses the model's JSON array response into DatasetCase objects, and returns a DatasetSource containing those cases.
249+
250+
Parameters:
251+
domain (str): Domain or subject area for which to generate test cases.
252+
task_description (str): Description of the task or problem the generated cases should exercise.
253+
constraints (list[str] | None): Optional list of additional constraints to include in the prompt.
254+
seed_examples (list[dict[str, Any]] | None): Optional list of example cases to guide generation; included verbatim in the prompt.
255+
249256
Returns:
250-
DatasetSource with generated cases
257+
DatasetSource: A dataset source containing generated DatasetCase entries with metadata (including difficulty, generated=True, and domain).
258+
259+
Raises:
260+
json.JSONDecodeError: If the model's response cannot be parsed as a JSON array of cases.
251261
"""
252262
messages = [
253263
{"role": "system", "content": "You are an expert test case generator."},
@@ -330,4 +340,4 @@ def generate_sync(
330340
"GeneratorConfig",
331341
"LlamaDatasetGenerator",
332342
"OpenAIDatasetGenerator",
333-
]
343+
]

0 commit comments

Comments
 (0)