---
name: agentbyte-dataset
description: Guide for creating and using Agentbyte datasets with SQLite or JSON backends. Use this when asked to load datasets, create dataset structures, work with schema-driven data loading, register custom engines, or integrate datasets with eval pipelines.
---

Use this skill when you need to:
- **Load** an existing dataset using `load_dataset()`
- **Create** a new dataset with config.json, schema.json (SQLite), and data.json
- **Load** the same dataset structure from local disk or S3
- **Iterate** any dataset engine-agnostically using `fetch_all()`
- **Register** a custom engine using `register_engine()`
- **Display** dataset schemas in pretty or raw format
- **Save** eval datasets using `save_eval_dataset()` (SQLite or JSON engine)
- **Load eval tasks** from any registered engine using `tasks_from_eval_dataset()`
- **Publish** a local dataset artifact to S3 using `publish_dataset()`

## Dataset Architecture

Agentbyte datasets are file-based and engine-agnostic. The dataset contract stays the same whether the files live on local disk or in S3.

Local layout:

### File Structure
```
data/samples/{category}/{name}/
├── config.json      # Engine, table name, in-memory flag, description
├── schema.json      # Column definitions (SQLite only — absent for JSON)
└── data.json        # Records only ({"records": [...]})
```

Equivalent S3 layout:
```
s3://<bucket>/<optional-prefix>/{category}/{name}/
├── config.json
├── schema.json      # SQLite only
└── data.json
```

### Supported Engines
- **sqlite** (default): Structured data with SQL queries, strict schema, constraints
- **json**: Semi-structured data, flexible JSON format, inferred schema
- **Custom**: Register additional engines with `register_engine()`

## Using Datasets

### Load a Dataset
```python
from pathlib import Path

from agentbyte.dataset import LocalDatasetConfig, S3DatasetConfig, load_dataset

# Local source
local_config = LocalDatasetConfig(local_root=Path("data/samples"))
ds = load_dataset("category/name", local_config)  # Returns SQLiteDataset or JSONDataset
print(ds.engine)        # "sqlite" or "json"
print(ds.table_name)    # Table/collection name

# S3 source
s3_config = S3DatasetConfig(
    s3_bucket="datapsycho-datasets",
    s3_region="eu-west-1",
    s3_prefix="datasets",
)
ds_s3 = load_dataset("category/name", s3_config)
```

### Engine-Agnostic Access: `fetch_all()`

`fetch_all()` is the preferred method for reading rows — it works with any registered engine and returns `list[dict[str, Any]]`.

```python
# Works with SQLite, JSON, or any custom engine
rows = ds.fetch_all()
for row in rows:
    print(row["id"], row["name"])
```

### Display Schema

Both datasets support schema inspection:
```python
# Pretty format (default) - shows column names, types, constraints
ds.schema.show()

# Raw format - list of column dicts
from agentbyte.dataset import SchemaFormat
ds.schema.show(format=SchemaFormat.RAW)
```

**SQLite Only:** SqliteSchema has `.to_create_table_sql()` and `.get_column_names()` methods.

## Creating a Dataset

### 1. Create Directory
```python
from pathlib import Path

dataset_dir = Path("data/samples/mydata/mydataset")
dataset_dir.mkdir(parents=True, exist_ok=True)
```

### 2. Write config.json
Define engine, table name, and metadata:
```python
import json

config = {
    "engine": "sqlite",           # "sqlite" or "json"
    "table_name": "my_table",     # SQL table name (SQLite only; JSON uses it as a label)
    "in_memory": True,            # Use in-memory storage
    "description": "My dataset"   # Optional description
}

with open(dataset_dir / "config.json", "w") as f:
    json.dump(config, f, indent=2)
```

### 3. Write schema.json (SQLite Only)
List of column definitions with types and constraints:
```python
schema = [
    {"name": "id", "type": "TEXT", "primary_key": True, "not_null": True},
    {"name": "name", "type": "TEXT", "not_null": True},
    {"name": "value", "type": "REAL"},
    {"name": "status", "type": "TEXT"}
]

with open(dataset_dir / "schema.json", "w") as f:
    json.dump(schema, f, indent=2)
```

**SQLite Column Constraints (all optional):**
- `primary_key` (bool): Is PRIMARY KEY
- `not_null` (bool): Is NOT NULL
- `unique` (bool): Is UNIQUE
- JSON datasets do not require schema.json — schema is inferred from data

### 4. Write data.json
Records only; schema is separate:
```python
data = {
    "records": [
        {"id": "001", "name": "Item 1", "value": 100.0, "status": "active"},
        {"id": "002", "name": "Item 2", "value": 200.0, "status": "inactive"},
        {"id": "003", "name": "Item 3", "value": 150.0, "status": "active"}
    ]
}

with open(dataset_dir / "data.json", "w") as f:
    json.dump(data, f, indent=2)
```

## Engine Registry

### Register a Custom Engine
Third-party engines can be registered once at module import time without modifying `load_dataset()`:

```python
import json
from pathlib import Path
from typing import Any

import pandas as pd

from agentbyte.dataset import load_dataset, register_engine
from agentbyte.dataset.base import BaseDataset, DatasetConfig
from agentbyte.dataset.config import LocalDatasetConfig


class PandasDataset(BaseDataset):
    def __init__(
        self,
        dataset_id: str,
        config: DatasetConfig,
        data_text: str,
        source_uri: str,
        schema_text: str | None = None,
    ) -> None:
        self.dataset_id = dataset_id
        self.config = config
        self.data_text = data_text
        self.source_uri = source_uri
        self.schema_text = schema_text
        self._frame: pd.DataFrame | None = None

    @property
    def engine(self) -> str:
        return self.config.engine

    @property
    def table_name(self) -> str:
        return self.config.table_name

    def _load_frame(self) -> pd.DataFrame:
        payload = json.loads(self.data_text)
        records = payload.get("records", [])
        if not isinstance(records, list):
            raise ValueError(
                f"Data file must contain a 'records' list for dataset: {self.source_uri}/data.json"
            )
        return pd.DataFrame(records)

    def fetch_all(self) -> list[dict[str, Any]]:
        if self._frame is None:
            self._frame = self._load_frame()
        return self._frame.to_dict(orient="records")

register_engine("pandas", PandasDataset)

# Now usable — config.json must have "engine": "pandas"
ds = load_dataset("contracts/asmd", LocalDatasetConfig(local_root=Path("data")))
rows = ds.fetch_all()
```

A custom engine class must:
1. Subclass `BaseDataset`
2. Implement all abstract members on `BaseDataset`: `engine`, `table_name`, `fetch_all()`
3. Accept the standard constructor signature (including `schema_text: str | None = None`)

## Publisher Layer

Use publishing when a dataset is created locally and then needs to be shared
through the same canonical dataset layout in S3.

```python
from pathlib import Path

from agentbyte.dataset import publish_dataset
from agentbyte.dataset.config import LocalDatasetConfig
from agentbyte.dataset.publish_config import S3DatasetPublishConfig

result = publish_dataset(
    "contracts10/asmd",
    local_config=LocalDatasetConfig(local_root=Path("data/samples")),
    publish_config=S3DatasetPublishConfig(
        s3_bucket="my-bucket",
        s3_prefix="datasets",
        overwrite_existing=True,
    ),
)

print(result.destination_uri)
# s3://my-bucket/datasets/contracts10/asmd
print(result.uploaded_files)
# ["config.json", "data.json", "schema.json"]
```

Publisher rules:
- `publish_dataset()` validates `dataset_id` in `"category/name"` format
- the local dataset directory must exist under `LocalDatasetConfig.local_root`
- `config.json` and `data.json` are required
- `schema.json` is uploaded only when present
- publishing reuses the same artifact layout as local disk and S3 loading

## Eval Dataset Integration

### Save an Eval Dataset
`save_eval_dataset()` writes the six-column eval schema to any registered engine.

```python
from pathlib import Path
from agentbyte.eval import EvalDatasetRecord, save_eval_dataset

records = [
    EvalDatasetRecord(
        id="task_1",
        task="What is the status of contract CW001?",
        expected_response="Published.",
        eval_criteria=["Correct status", "No hallucinations"],
        context={
            "reference_sql": "SELECT status FROM contracts WHERE id='CW001'",
            "reference_table": "contracts",
        },
    ),
]

# SQLite engine (default) — writes config.json + schema.json + data.json
save_eval_dataset(records, root=Path("/data"), dataset_id="asmd/sql_eval")

# JSON engine — writes config.json + data.json only (no schema.json)
save_eval_dataset(records, root=Path("/data"), dataset_id="asmd/sql_eval", engine="json")
```

### Publish Eval Datasets After Saving

```python
from agentbyte.dataset import publish_dataset
from agentbyte.dataset.config import LocalDatasetConfig
from agentbyte.dataset.publish_config import S3DatasetPublishConfig

publish_result = publish_dataset(
    "asmd/sql_eval",
    local_config=LocalDatasetConfig(local_root=Path("/data")),
    publish_config=S3DatasetPublishConfig(
        s3_bucket="my-bucket",
        s3_prefix="eval-datasets",
    ),
)
```

### Load Eval Tasks from a Dataset
`tasks_from_eval_dataset()` uses `fetch_all()` internally — works with any registered engine.

```python
from agentbyte.dataset import load_dataset
from agentbyte.dataset.config import LocalDatasetConfig
from agentbyte.eval import tasks_from_eval_dataset

dataset = load_dataset("asmd/sql_eval", LocalDatasetConfig(local_root=Path("/data")))
tasks = tasks_from_eval_dataset(dataset)

print(tasks[0].name)                         # id column → EvalTask.name
print(tasks[0].input)                        # task column → EvalTask.input
print(tasks[0].expected_output)              # expected_response column
print(tasks[0].metadata["_criteria"])        # eval_criteria → metadata["_criteria"]
print(tasks[0].metadata["reference_sql"])    # context keys unpacked flat
```

### Eval Dataset Column Schema

| Column | Type | Maps to |
|--------|------|---------|
| `id` | TEXT (PK) | `EvalTask.name` |
| `task` | TEXT | `EvalTask.input` |
| `expected_response` | TEXT | `EvalTask.expected_output` |
| `eval_criteria` | TEXT | `EvalTask.metadata["_criteria"]` |
| `scenario` | TEXT | `EvalTask.metadata["scenario"]` |
| `context` | TEXT (JSON) | Unpacked flat into `EvalTask.metadata` |

`tasks_from_eval_dataset()` raises `ValueError` (not `KeyError`) when `id` or `task` columns are absent.

`EvalRunner.evaluate()` will automatically read `task.metadata["_criteria"]` (when present) as judge criteria unless explicit `criteria=...` is passed at call time.

## Examples

### Example 1: Load Existing Dataset
```python
from pathlib import Path

from agentbyte.dataset import LocalDatasetConfig, load_dataset

ds = load_dataset("contracts/asmd", LocalDatasetConfig(local_root=Path("data/samples")))
rows = ds.fetch_all()
for row in rows:
    print(row)
```

### Example 2: Engine-Agnostic Iteration with `fetch_all()`
```python
from pathlib import Path

from agentbyte.dataset import LocalDatasetConfig, load_dataset

# Works with SQLite, JSON, or any custom engine
ds = load_dataset("contracts/asmd", LocalDatasetConfig(local_root=Path("data/samples")))
rows = ds.fetch_all()

for row in rows:
    print(row["id"], row["author"])
```

### Example 3: Load & Iterate JSON Dataset
```python
from pathlib import Path

from agentbyte.dataset import LocalDatasetConfig, load_dataset

ds = load_dataset("documents/wiki", LocalDatasetConfig(local_root=Path("data/samples")))

# Engine-agnostic — same interface as SQLite
rows = ds.fetch_all()
for row in rows:
    print(row["title"])
```

### Example 4: Register a Custom Engine and Use It
```python
from agentbyte.dataset import register_engine, load_dataset, LocalDatasetConfig
from agentbyte.dataset.base import BaseDataset, DatasetConfig

class StubDataset(BaseDataset):
    def __init__(self, dataset_id, config, data_text, source_uri, schema_text=None):
        self._engine = config.engine
        self._table = config.table_name

    @property
    def engine(self): return self._engine

    @property
    def table_name(self): return self._table

    def fetch_all(self): return [{"id": "1", "task": "hello"}]

register_engine("stub", StubDataset)
# config.json must have "engine": "stub"
ds = load_dataset("custom/test", LocalDatasetConfig(local_root=Path("data")))
```

### Example 5: End-to-End Eval Dataset Workflow
```python
from pathlib import Path
from agentbyte.eval import (
    EvalDatasetRecord, save_eval_dataset, tasks_from_eval_dataset
)
from agentbyte.dataset import load_dataset
from agentbyte.dataset.config import LocalDatasetConfig

records = [
    EvalDatasetRecord(
        id="q1", task="List active contracts.", expected_response="A, B.",
        eval_criteria=["Complete list"],
        context={"reference_sql": "SELECT ...", "reference_table": "t"},
    )
]
save_eval_dataset(records, root=Path("/data"), dataset_id="asmd/eval")

dataset = load_dataset("asmd/eval", LocalDatasetConfig(local_root=Path("/data")))
tasks = tasks_from_eval_dataset(dataset)
# tasks[0].metadata["_criteria"] == ["Complete list"]
# tasks[0].metadata["reference_sql"] == "SELECT ..."
```

### Example 6: Load From S3
```python
from agentbyte.dataset import S3DatasetConfig, load_dataset

s3_config = S3DatasetConfig(
    s3_bucket="datapsycho-datasets",
    s3_region="eu-west-1",
    s3_prefix="datasets",
)

ds = load_dataset("contracts10/asmd", s3_config)
rows = ds.fetch_all()
print(rows[0])
```

## Guardrails

- **Do** use `fetch_all()` for engine-agnostic access; reserve `connect()` for SQL or engine-specific features
- **Do** use canonical dataset IDs in `load_dataset()` (e.g., `"contracts/asmd"`, not full file paths)
- **Do** pass a source config object to `load_dataset()` (`LocalDatasetConfig` or `S3DatasetConfig`)
- **Do** keep `dataset_id` logical (`"category/name"`) and configure location through the source config
- **Do** keep schema.json separate from data.json for clarity and reusability (SQLite only)
- **Do** use `dict(row)` to convert sqlite3.Row objects to dicts for JSON serialization
- **Do** call `ds.connect()` only once; it caches the connection and schema on first call
- **Do** register third-party engines once at module import time, before any `load_dataset()` call
- **Don't** hard-code values in datasets (use config.json for metadata)
- **Don't** store bucket names or S3 prefixes inside `config.json`
- **Don't** forget `.json` extensions for all three files
- **Don't** put schema definitions inside data.json (separate files for clarity)
- **Don't** assume JSON datasets support SQL — use `fetch_all()` or dict iteration instead
- **Don't** use `DataStore` — it has been removed; use `BaseDataset` + `fetch_all()` instead
- **Don't** require `connect()` on custom engine APIs unless you need engine-specific features; `fetch_all()` is the stable cross-engine contract

## API Reference

### Factory Functions
```python
load_dataset(
    dataset_id: str,
    config: DatasetSourceConfig,
) -> BaseDataset
```
Loads dataset using the provided source config. Returns `SQLiteDataset` or `JSONDataset` (or any registered engine) based on `config.engine`.

```python
register_engine(name: str, cls: type[BaseDataset]) -> None
```
Registers a custom engine class under the given name. Must be called before `load_dataset()`.

### Eval Integration
```python
save_eval_dataset(
    records: list[EvalDatasetRecord],
    root: Path,
    dataset_id: str,
    engine: str = "sqlite",     # "sqlite" writes schema.json; "json" skips it
    table_name: str = "eval",
) -> None
```

```python
tasks_from_eval_dataset(dataset: BaseDataset) -> list[EvalTask]
```

### Source Config Types
- **DatasetSourceConfig**: Abstract config contract
- **LocalDatasetConfig**: Local source config with `local_root`
- **S3DatasetConfig**: S3 source config with `s3_bucket`, `s3_region`, `s3_prefix`

### Classes (Generic)
- **BaseDataset**: Abstract base; properties `engine`, `table_name`; abstract method `fetch_all()`. `connect()` is on each concrete engine class, not on the base.
- **Column**: name (str), type (str)
- **Schema**: columns (list[Column])
- **DatasetConfig**: engine (str), table_name (str), in_memory (bool), description (str)
- **SchemaFormat**: RAW or PRETTY enum for display formats

### Classes (SQLite-Specific)
- **SqliteColumn**: Extends Column; adds primary_key, not_null, unique; method `to_sql_def()`
- **SqliteSchema**: Extends Schema; methods `to_create_table_sql()`, `get_column_names()`, `show(format)`
- **SQLiteDataset**: Implements BaseDataset; properties: engine, table_name, schema, in_memory

### Classes (JSON-Specific)
- **JSONDataset**: Implements BaseDataset; properties: engine, table_name, schema (inferred)

## Troubleshooting

| Issue | Fix |
|-------|-----|
| `FileNotFoundError: Schema file not found` | Ensure `schema.json` exists in dataset directory (SQLite only) |
| `FileNotFoundError: Data file not found` | Ensure `data.json` exists in dataset directory |
| `FileNotFoundError: Dataset object not found: s3://...` | Ensure bucket, prefix, dataset_id, and object names are correct |
| `ValidationError: s3_bucket field required` | Set `AGENTBYTE_DATASET_S3_BUCKET` or pass `S3DatasetConfig(s3_bucket=...)` |
| `KeyError: 'records'` | data.json must have `{"records": [...]}` structure |
| `sqlite3.Row has no .dict() method` | Use `dict(row)` instead of `row.dict()` |
| `NotImplementedError: Engine 'X' is not registered` | Call `register_engine("X", MyDataset)` before `load_dataset()` |
| `ValueError: missing required eval columns` | Dataset was not saved with `save_eval_dataset()` — ensure `id` and `task` columns are present |
| Notebook paths fail in `load_dataset()` | Pass an explicit `LocalDatasetConfig(local_root=Path("data/samples"))` |
