chemgraph-loop / src /chemgraph /eval /datasets.py
rockyaaos's picture
ChemGraph Loop: guarded real-agent API (EMT/TBLite single-point energy)
c509967 verified
Raw
History Blame Contribute Delete
4.79 kB
"""Ground-truth dataset loading and validation for ChemGraph evaluation."""
import json
from pathlib import Path
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field
from chemgraph.utils.logging_config import setup_logger
logger = setup_logger(__name__)
# Path to the bundled default ground-truth dataset.
_DEFAULT_DATASET = Path(__file__).parent / "data" / "ground_truth.json"
def default_dataset_path() -> str:
"""Return the absolute path to the bundled default ground-truth dataset.
The dataset ships with the ``chemgraph`` package under
``chemgraph/eval/data/ground_truth.json`` and contains 14
evaluation queries covering single-tool, multi-step, and
reaction-energy calculations.
Returns
-------
str
Absolute path to the default ``ground_truth.json``.
"""
return str(_DEFAULT_DATASET.resolve())
class GroundTruthItem(BaseModel):
"""A single evaluation query with its expected tool-call sequence"""
id: str = Field(description="Unique identifier for the query.")
query: str = Field(description="The natural-language query to send to the agent.")
expected_tool_calls: list = Field(
description="Ordered list of expected tool-call dicts."
)
expected_result: Any = Field(
default="",
description="Optional expected final result (string or list of step dicts).",
)
expected_structured_output: Optional[Dict[str, Any]] = Field(
default=None,
description=(
"Expected structured output in ResponseFormatter format. "
"When present, the deterministic structured-output judge "
"can compare field-by-field against the agent's output."
),
)
category: str = Field(
default="",
description="Optional category / experiment tag.",
)
def load_dataset(path: str) -> List[GroundTruthItem]:
"""Load a ground-truth dataset from a JSON file.
Automatically detects the two formats used in ChemGraph:
1. **List format** -- a JSON array of ``{id, query, answer}`` objects
(used by the bundled ``data/ground_truth.json``).
2. **Dict format** -- a JSON object keyed by query/name, each
containing ``manual_workflow`` with ``tool_calls`` and ``result``
(used by legacy ``run_manual/`` baselines).
Parameters
----------
path : str
Path to the JSON file.
Returns
-------
list[GroundTruthItem]
Validated list of ground-truth items.
Raises
------
ValueError
If the file cannot be parsed into either known format.
FileNotFoundError
If the file does not exist.
"""
p = Path(path)
if not p.exists():
raise FileNotFoundError(f"Dataset file not found: {path}")
with open(p, "r", encoding="utf-8") as f:
raw = json.load(f)
items: List[GroundTruthItem] = []
if isinstance(raw, list):
# List format: [{id, query, category?, answer: {tool_calls, result, structured_output?}}, ...]
for idx, entry in enumerate(raw):
answer = entry.get("answer", {})
items.append(
GroundTruthItem(
id=str(entry.get("id", idx)),
query=entry["query"],
expected_tool_calls=answer.get("tool_calls", []),
expected_result=answer.get("result", ""),
expected_structured_output=answer.get("structured_output"),
category=entry.get("category", ""),
)
)
elif isinstance(raw, dict):
# Dict format: {name: {manual_workflow: {tool_calls, result}, ...}, ...}
for idx, (name, data) in enumerate(raw.items()):
workflow = data.get("manual_workflow", data.get("llm_workflow", {}))
tool_calls = workflow.get("tool_calls", [])
result = workflow.get("result", "")
# For dict format, the key is typically the molecule/reaction
# name which also serves as the query. If a "query" field
# exists at the top level, prefer it.
query = data.get("query", name)
items.append(
GroundTruthItem(
id=str(idx),
query=query,
expected_tool_calls=tool_calls,
expected_result=result if result else "",
expected_structured_output=workflow.get("structured_output"),
category=name,
)
)
else:
raise ValueError(
f"Unrecognised dataset format in {path}. Expected a JSON list or dict."
)
logger.info(f"Loaded {len(items)} ground-truth items from {path}")
return items