langsmith_evals/dataset.py

"""Create and manage the harmonic analysis LangSmith dataset.

The dataset is a stable artifact. Running this script is idempotent — it creates
the dataset on first run; on subsequent runs it syncs outputs and appends any
new examples.

Usage:
    LANGCHAIN_API_KEY=<key> python langsmith_evals/dataset.py
"""

import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

from dotenv import load_dotenv
from langsmith import Client

load_dotenv()

DATASET_NAME = "harmonic-analysis-chord-sequences"

_EXAMPLES = [
    {
        "inputs": {"user_input": "Analyse the chord progression C G Am F"},
        "outputs": {"tool_calls": [
            {"chord_text": "C G Am F", "limit": 10},
        ]},
    },
    {
        "inputs": {"user_input": "Analyse the chord progression Am F C G"},
        "outputs": {"tool_calls": [
            {"chord_text": "Am F C G", "limit": 10},
        ]},
    },
    {
        "inputs": {"user_input": "Analyse the chord progression D A Bm G"},
        "outputs": {"tool_calls": [
            {"chord_text": "D A Bm G", "limit": 10},
            {"chord_text": "D A Bm G", "limit": 20},
        ]},
    },
    {
        "inputs": {"user_input": "Analyse the chord progression Daug C9 Emaj7#11"},
        "outputs": {"tool_calls": [
            {"chord_text": "Daug C9 Emaj7#11", "limit": 10},
            {"chord_text": "Daug C9 Emaj7#11", "limit": 20},
        ]},
    },
]


def create(client: Client | None = None) -> None:
    """Create the dataset, or sync it with ``_EXAMPLES`` if it already exists.

    Idempotent: existing examples have their outputs updated; new examples are
    appended; nothing is deleted.

    :param client: LangSmith client. Creates one from environment if not provided.
    """
    client = client or Client()
    existing = list(client.list_datasets(dataset_name=DATASET_NAME))

    if existing:
        dataset = existing[0]
        for e in client.list_examples(dataset_id=dataset.id):
            client.delete_example(e.id)
        client.create_examples(dataset_id=dataset.id, examples=_EXAMPLES)
        print(f"Dataset '{DATASET_NAME}': reset to {len(_EXAMPLES)} example(s)")
        return

    dataset = client.create_dataset(
        dataset_name=DATASET_NAME,
        description="Chord sequence inputs with varying limits for evaluating the harmonic analysis agent.",
    )
    client.create_examples(dataset_id=dataset.id, examples=_EXAMPLES)
    print(f"Created dataset '{DATASET_NAME}' with {len(_EXAMPLES)} examples")
    print(f"URL: {dataset.url}")


if __name__ == "__main__":
    create()