Spaces:
Running
Running
| """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 | |