class CaseGenerator:
def __init__(
self,
url: str | None = None,
model_id: str | None = None,
api_key: str | None = None,
) -> None:
resolved_url = url or os.environ.get("DATAROBOT_ENDPOINT")
resolved_model = model_id or os.environ.get("LLM_DEFAULT_MODEL")
if not resolved_url:
raise ValueError("url is required. Pass it explicitly or set DATAROBOT_ENDPOINT.")
if not resolved_model:
raise ValueError("model_id is required. Pass it explicitly or set LLM_DEFAULT_MODEL.")
# Strip /api/v2 suffix so litellm receives the gateway base URL
self._api_base = resolved_url.removesuffix("/api/v2")
self._model = (
resolved_model
if resolved_model.startswith("datarobot/")
else f"datarobot/{resolved_model}"
)
self._api_key = api_key or os.environ.get("DATAROBOT_API_TOKEN")
def generate(
self,
agent_description: str,
n_good: int,
n_bad: int,
benchmark_name: str | None = None,
) -> list[dict[str, Any]]:
"""Generate synthetic test cases.
When ``benchmark_name`` is given, the description is enriched with that
benchmark's good/bad guidance and the generated cases are checked for any
extra fields the benchmark requires (e.g. ``canary``, ``context``). When
it is ``None`` a generic, non-safety-biased context is used instead.
"""
enriched_description = _enrich_description(agent_description, benchmark_name)
response = litellm.completion(
model=self._model,
api_base=self._api_base,
api_key=self._api_key,
max_tokens=4096,
messages=[
{"role": "system", "content": _SYSTEM_PROMPT},
{
"role": "user",
"content": _GENERATION_PROMPT.format(
agent_description=enriched_description,
n_good=n_good,
n_bad=n_bad,
),
},
],
)
content = response.choices[0].message.content
if not isinstance(content, str):
raise ValueError(f"Unexpected response content type: {type(content)}")
raw = json.loads(content.strip())
if not isinstance(raw, list):
raise ValueError(f"Expected a JSON array from the model, got {type(raw).__name__}")
cases: list[dict[str, Any]] = raw
expected = n_good + n_bad
if len(cases) != expected:
warnings.warn(
f"Requested {expected} cases ({n_good} good, {n_bad} bad) but "
f"model returned {len(cases)}",
UserWarning,
stacklevel=2,
)
for i, case in enumerate(cases):
missing = _REQUIRED_FIELDS - case.keys()
if missing:
raise ValueError(f"Case {i} missing fields: {missing}")
if case["expected_behavior"] not in ("good", "bad"):
raise ValueError(
f"Case {i} has invalid expected_behavior: {case['expected_behavior']}"
)
extra_fields = _BENCHMARK_EXTRA_FIELDS.get(benchmark_name or "", set())
if extra_fields:
for i, case in enumerate(cases):
missing = extra_fields - case.keys()
if missing:
raise ValueError(f"Case {i} missing benchmark-required fields: {missing}")
return cases
def save(
self,
cases: list[dict[str, Any]],
output_path: Path,
append: bool = False,
) -> list[dict[str, Any]]:
"""Write cases to disk. Returns the final list written (merged if append=True)."""
output_path.parent.mkdir(parents=True, exist_ok=True)
if append and output_path.exists():
existing: list[dict[str, Any]] = json.loads(output_path.read_text())
cases = existing + cases
output_path.write_text(json.dumps(cases, indent=2))
return cases