---
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, or work with schema-driven data loading.
---

Use this skill when you need to:
- **Load** an existing dataset using `load_dataset()`
- **Create** a new dataset with config.json, schema.json, and data.json
- **Load** the same dataset structure from local disk or S3
- **Query** SQLite datasets with SQL or iterate JSON datasets
- **Display** dataset schemas in pretty or raw format
- **Work with** multiple data engines (SQLite for structured, JSON for semi-structured)

## 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 (name, type, SQL constraints)
└── data.json        # Records only ({"records": [...]})
```

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

### Supported Engines
- **sqlite** (default): Structured data with SQL queries, strict schema, constraints
- **json**: Semi-structured data, flexible JSON format, inferred schema
- **Future**: duckdb, lance, qdrant

## 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)
    "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
- All other fields ignored for JSON datasets (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)
```

## 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)
```

### Connect & Get Data

**SQLite Dataset:**
```python
conn = ds.connect()  # Returns sqlite3.Connection
rows = conn.execute("SELECT * FROM my_table").fetchall()
for row in rows:
    print(dict(row))  # Convert sqlite3.Row to dict
```

**JSON Dataset:**
```python
data = ds.connect()  # Returns dict with {"records": [...]}
for record in data.get("records", []):
    print(record)
```

### 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.

## Examples

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

from agentbyte.dataset import LocalDatasetConfig, load_dataset

ds = load_dataset("contracts/asmd", LocalDatasetConfig(local_root=Path("data/samples")))
conn = ds.connect()

# Display schema
ds.schema.show()

# Query
active = conn.execute(
    "SELECT id, name FROM contracts_asmd WHERE status='active'"
).fetchall()

for row in active:
    print(dict(row))
```

### Example 2: 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")))
data = ds.connect()  # dict with {"records": [...]}

for i, record in enumerate(data["records"], 1):
    print(f"{i}. {record['title']}: {record['url']}")
```

### Example 3: Create Dataset Programmatically
See "Creating a Dataset" section above for step-by-step code.

### Example 4: 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)
conn = ds.connect()
rows = conn.execute("SELECT * FROM contracts_asmd").fetchall()
```

## Guardrails

- **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 S3 location through `S3DatasetConfig`.
- **Do** keep schema.json separate from data.json for clarity and reusability.
- **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.
- **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 dict iteration instead.

## API Reference

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

### 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)
- **Column**: name (str), type (str)
- **Schema**: columns (list[Column])
- **DataRecord**: id (str), data (Any), metadata (dict)
- **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 |
| `FileNotFoundError: Data file not found` | Ensure `data.json` exists in dataset directory |
| `FileNotFoundError: Dataset object not found: s3://...` | Ensure the 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()` |
| Notebook paths fail in `load_dataset()` | Pass an explicit `LocalDatasetConfig(local_root=Path("data/samples"))` |
