chore(llmobs): add base experiments classes (#13930)

Adds skeleton code for Experiments, experiment tasks, and experiment
evaluator classes/decorators. Implementation of experiment run() has
been left out for a follow-up PR.

Basic structure of Experiments:
- Call LLMObs.experiment_task as a decorator to wrap a task function
(must have `input` as an arg)
- Call LLMObs.experiment_evaluator as a decorator to wrap an evaluator
function (must have `input/output/expected_output` as args)
- Create a Dataset
- Create an Experiment(name: str, task, dataset, evaluators,
description, config)
- Call experiment.run(...)

Some concerns:
- Should experiment task/evaluator decorators support async/generator
methods? Currently (and based on #13314) it only supports sync methods.
- The ExperimentTask wrapper class requires `input` as an arg name,
which shadows Python builtins.

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: Yun-Kim

ddtrace/llmobs/_experiment.py

@@ -1,4 +1,5 @@
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -13,7 +14,7 @@
 
 
 class DatasetRecord(TypedDict):
-    input: NonNoneJSONType
+    input_data: NonNoneJSONType
     expected_output: JSONType
     metadata: Dict[str, Any]
     record_id: NotRequired[Optional[str]]
@@ -28,3 +29,35 @@ def __init__(self, name: str, dataset_id: str, data: List[DatasetRecord]) -> Non
         self.name = name
         self._id = dataset_id
         self._data = data
+
+
+class Experiment:
+    def __init__(
+        self,
+        name: str,
+        task: Callable[[Dict[str, NonNoneJSONType]], JSONType],
+        dataset: Dataset,
+        evaluators: List[Callable[[NonNoneJSONType, JSONType, JSONType], JSONType]],
+        description: str = "",
+        config: Optional[Dict[str, Any]] = None,
+        _llmobs: Optional[Any] = None,  # LLMObs service (cannot import here due to circular dependency)
+    ) -> None:
+        self.name = name
+        self._task = task
+        self._dataset = dataset
+        self._evaluators = evaluators
+        self._description = description
+        self._config: Dict[str, Any] = config or {}
+        self._llmobs = _llmobs
+        self._id: Optional[str] = None
+
+    def run(self, jobs: int = 1, raise_errors: bool = False, sample_size: Optional[int] = None) -> None:
+        task_results = self._run_task(jobs, raise_errors, sample_size)
+        self._run_evaluators(task_results, raise_errors=raise_errors)
+        return
+
+    def _run_task(self, jobs: int, raise_errors: bool = False, sample_size: Optional[int] = None) -> List[Any]:
+        return []
+
+    def _run_evaluators(self, task_results, raise_errors: bool = False) -> None:
+        pass

ddtrace/llmobs/_llmobs.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 from dataclasses import field
+import inspect
 import json
 import os
 import time
@@ -74,6 +75,9 @@
 from ddtrace.llmobs._context import LLMObsContextProvider
 from ddtrace.llmobs._evaluators.runner import EvaluatorRunner
 from ddtrace.llmobs._experiment import Dataset
+from ddtrace.llmobs._experiment import Experiment
+from ddtrace.llmobs._experiment import JSONType
+from ddtrace.llmobs._experiment import NonNoneJSONType
 from ddtrace.llmobs._utils import AnnotationContext
 from ddtrace.llmobs._utils import LinkTracker
 from ddtrace.llmobs._utils import ToolCallTracker
@@ -573,6 +577,42 @@ def create_dataset(cls, name: str, description: str) -> Dataset:
     def _delete_dataset(cls, dataset_id: str) -> None:
         return cls._instance._dne_client.dataset_delete(dataset_id)
 
+    @classmethod
+    def experiment(
+        cls,
+        name: str,
+        task: Callable[[Dict[str, NonNoneJSONType]], JSONType],
+        dataset: Dataset,
+        evaluators: List[Callable[[NonNoneJSONType, JSONType, JSONType], JSONType]],
+        description: str = "",
+    ) -> Experiment:
+        """Initializes an Experiment to run a task on a Dataset and evaluators.
+
+        :param name: The name of the experiment.
+        :param task: The task function to run. Must accept a parameter ``input_data`` and optionally ``config``.
+        :param dataset: The dataset to run the experiment on, created with LLMObs.pull/create_dataset().
+        :param evaluators: A list of evaluator functions to evaluate the task output.
+                           Must accept parameters ``input_data``, ``output_data``, and ``expected_output``.
+        :param description: A description of the experiment.
+        """
+        if not callable(task):
+            raise TypeError("task must be a callable function.")
+        sig = inspect.signature(task)
+        params = sig.parameters
+        if "input_data" not in params:
+            raise TypeError("Task function must have an 'input_data' parameter.")
+        if not isinstance(dataset, Dataset):
+            raise TypeError("Dataset must be an LLMObs Dataset object.")
+        if not evaluators or not all(callable(evaluator) for evaluator in evaluators):
+            raise TypeError("Evaluators must be a list of callable functions.")
+        for evaluator in evaluators:
+            sig = inspect.signature(evaluator)
+            params = sig.parameters
+            required_params = ("input_data", "output_data", "expected_output")
+            if not all(param in params for param in required_params):
+                raise TypeError("Evaluator function must have parameters {}.".format(required_params))
+        return Experiment(name, task, dataset, evaluators, description=description, _llmobs=cls)
+
     @classmethod
     def register_processor(cls, processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None) -> None:
         """Register a processor to be called on each LLMObs span.

ddtrace/llmobs/_writer.py

@@ -356,7 +356,7 @@ def dataset_pull(self, name: str) -> Dataset:
             class_records.append(
                 {
                     "record_id": record["id"],
-                    "input": attrs["input"],
+                    "input_data": attrs["input"],
                     "expected_output": attrs["expected_output"],
                     "metadata": attrs.get("metadata", {}),
                 }

tests/llmobs/test_experiments.py

@@ -11,10 +11,19 @@
 """
 
 import os
+import re
 
 import pytest
 
 
+def dummy_task(input_data):
+    return input_data
+
+
+def dummy_evaluator(input_data, output_data, expected_output):
+    return output_data == expected_output
+
+
 @pytest.fixture
 def test_dataset(llmobs):
     ds = llmobs.create_dataset(name="test-dataset", description="A test dataset")
@@ -44,3 +53,59 @@ def test_dataset_pull_non_existent(llmobs):
 def test_dataset_pull(llmobs, test_dataset):
     dataset = llmobs.pull_dataset(name=test_dataset.name)
     assert dataset._id is not None
+
+
+def test_experiment_invalid_task_type_raises(llmobs, test_dataset):
+    with pytest.raises(TypeError, match="task must be a callable function."):
+        llmobs.experiment("test_experiment", 123, test_dataset, [dummy_evaluator])
+
+
+def test_experiment_invalid_task_signature_raises(llmobs, test_dataset):
+    with pytest.raises(TypeError, match="Task function must have an 'input_data' parameter."):
+
+        def my_task(not_input):
+            pass
+
+        llmobs.experiment("test_experiment", my_task, test_dataset, [dummy_evaluator])
+
+
+def test_experiment_invalid_dataset_raises(llmobs):
+    with pytest.raises(TypeError, match="Dataset must be an LLMObs Dataset object."):
+        llmobs.experiment("test_experiment", dummy_task, 123, [dummy_evaluator])
+
+
+def test_experiment_invalid_evaluators_type_raises(llmobs, test_dataset):
+    with pytest.raises(TypeError, match="Evaluators must be a list of callable functions"):
+        llmobs.experiment("test_experiment", dummy_task, test_dataset, [])
+    with pytest.raises(TypeError, match="Evaluators must be a list of callable functions"):
+        llmobs.experiment("test_experiment", dummy_task, test_dataset, [123])
+
+
+def test_experiment_invalid_evaluator_signature_raises(llmobs, test_dataset):
+    expected_err = "Evaluator function must have parameters ('input_data', 'output_data', 'expected_output')."
+    with pytest.raises(TypeError, match=re.escape(expected_err)):
+
+        def my_evaluator_missing_expected_output(input_data, output_data):
+            pass
+
+        llmobs.experiment("test_experiment", dummy_task, test_dataset, [my_evaluator_missing_expected_output])
+    with pytest.raises(TypeError, match=re.escape(expected_err)):
+
+        def my_evaluator_missing_input(output_data, expected_output):
+            pass
+
+        llmobs.experiment("test_experiment", dummy_task, test_dataset, [my_evaluator_missing_input])
+    with pytest.raises(TypeError, match=re.escape(expected_err)):
+
+        def my_evaluator_missing_output(input_data, expected_output):
+            pass
+
+        llmobs.experiment("test_experiment", dummy_task, test_dataset, [my_evaluator_missing_output])
+
+
+def test_experiment_create(llmobs, test_dataset):
+    exp = llmobs.experiment("test_experiment", dummy_task, test_dataset, [dummy_evaluator], description="lorem ipsum")
+    assert exp.name == "test_experiment"
+    assert exp._task == dummy_task
+    assert exp._dataset == test_dataset
+    assert exp._evaluators == [dummy_evaluator]