follow ups

Author: ankrgyl

py/src/braintrust/dataset_pipeline.py

@@ -4,7 +4,6 @@
 
 from typing_extensions import NotRequired, TypedDict
 
-from .generated_types import ObjectReference
 from .logger import Metadata
 from .trace import Trace
 
@@ -19,7 +18,6 @@
     "DatasetPipelineTransform",
     "DatasetPipelineTransformArgs",
     "DatasetPipelineTransformResult",
-    "get_registered_dataset_pipelines",
 ]
 
 
@@ -49,13 +47,13 @@ class DatasetPipelineRow(TypedDict, total=False):
     expected: Any | None
     tags: Sequence[str] | None
     metadata: Metadata | None
-    origin: ObjectReference
 
 
 Row = TypeVar("Row", bound=DatasetPipelineRow, covariant=True)
 
 
 class DatasetPipelineTransformArgs(TypedDict, total=False):
+    id: str
     input: Any | None
     output: Any | None
     metadata: Metadata | None
@@ -69,6 +67,7 @@ class DatasetPipelineTransformArgs(TypedDict, total=False):
 class DatasetPipelineTransform(Protocol[Row]):
     def __call__(
         self,
+        id: str | None = None,
         input: Any | None = None,
         output: Any | None = None,
         metadata: Metadata | None = None,
@@ -88,10 +87,6 @@ class DatasetPipelineDefinition(Generic[Row]):
 _DATASET_PIPELINES: list[DatasetPipelineDefinition[Any]] = []
 
 
-def get_registered_dataset_pipelines() -> list[DatasetPipelineDefinition[Any]]:
-    return list(_DATASET_PIPELINES)
-
-
 def DatasetPipeline(
     name: str | None = None,
     *,