push

Author: hassiebp

langfuse/_client/client.py

@@ -57,9 +57,13 @@
     LANGFUSE_TRACING_ENVIRONMENT,
 )
 from langfuse._client.experiments import (
+    EvaluatorFunction,
+    ExperimentData,
     ExperimentItem,
     ExperimentItemResult,
     ExperimentResult,
+    RunEvaluatorFunction,
+    TaskFunction,
     _run_evaluator,
     _run_task,
 )
@@ -2458,15 +2462,11 @@ def run_experiment(
         *,
         name: str,
         description: Optional[str] = None,
-        data: Union[
-            List[Union[ExperimentItem, dict, DatasetItem]], List[DatasetItemClient]
-        ],
-        task: Callable[
-            [Union[ExperimentItem, dict, DatasetItem, DatasetItemClient]], Any
-        ],
-        evaluators: Optional[List[Callable]] = None,
-        run_evaluators: Optional[List[Callable]] = None,
-        max_concurrency: Optional[int] = None,
+        data: ExperimentData,
+        task: TaskFunction,
+        evaluators: List[EvaluatorFunction] = [],
+        run_evaluators: List[RunEvaluatorFunction] = [],
+        max_concurrency: int = 50,
         metadata: Optional[Dict[str, Any]] = None,
     ) -> ExperimentResult:
         """Run an experiment on a dataset with automatic tracing and evaluation.
@@ -2524,27 +2524,20 @@ async def _run_experiment_async(
         *,
         name: str,
         description: Optional[str],
-        data: Union[
-            List[Union[ExperimentItem, dict, DatasetItem]], List[DatasetItemClient]
-        ],
-        task: Callable,
-        evaluators: List[Callable],
-        run_evaluators: List[Callable],
-        max_concurrency: Optional[int],
+        data: ExperimentData,
+        task: TaskFunction,
+        evaluators: List[EvaluatorFunction],
+        run_evaluators: List[RunEvaluatorFunction],
+        max_concurrency: int,
         metadata: Dict[str, Any],
     ) -> ExperimentResult:
         langfuse_logger.debug(f"Starting experiment '{name}' with {len(data)} items")
 
         # Set up concurrency control
-        max_workers = (
-            max_concurrency if max_concurrency is not None else min(len(data), 10)
-        )
-        semaphore = asyncio.Semaphore(max_workers)
+        semaphore = asyncio.Semaphore(max_concurrency)
 
         # Process all items
-        async def process_item(
-            item: Union[ExperimentItem, dict, DatasetItem, DatasetItemClient],
-        ) -> dict:
+        async def process_item(item: ExperimentItem) -> dict:
             async with semaphore:
                 return await self._process_experiment_item(
                     item, task, evaluators, name, description, metadata
@@ -2620,7 +2613,7 @@ async def process_item(
 
     async def _process_experiment_item(
         self,
-        item: Union[ExperimentItem, dict, DatasetItem, DatasetItemClient],
+        item: ExperimentItem,
         task: Callable,
         evaluators: List[Callable],
         experiment_name: str,

langfuse/_client/datasets.py

@@ -1,17 +1,23 @@
 import datetime as dt
 import logging
-from .span import LangfuseSpan
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional
 
 from opentelemetry.util._decorator import _agnosticcontextmanager
 
+from langfuse._client.experiments import (
+    EvaluatorFunction,
+    RunEvaluatorFunction,
+    TaskFunction,
+)
 from langfuse.model import (
     CreateDatasetRunItemRequest,
     Dataset,
     DatasetItem,
     DatasetStatus,
 )
 
+from .span import LangfuseSpan
+
 if TYPE_CHECKING:
     from langfuse._client.client import Langfuse
 
@@ -194,10 +200,10 @@ def run_experiment(
         *,
         name: str,
         description: Optional[str] = None,
-        task: Any,
-        evaluators: Optional[List[Any]] = None,
-        run_evaluators: Optional[List[Any]] = None,
-        max_concurrency: Optional[int] = None,
+        task: TaskFunction,
+        evaluators: List[EvaluatorFunction] = [],
+        run_evaluators: List[RunEvaluatorFunction] = [],
+        max_concurrency: int = 50,
         metadata: Optional[Dict[str, Any]] = None,
     ) -> Any:
         """Run an experiment on this dataset.

langfuse/_client/experiments.py

@@ -19,13 +19,11 @@
     Union,
 )
 
-from langfuse.model import DatasetItem
-
 if TYPE_CHECKING:
     from langfuse._client.datasets import DatasetItemClient
 
 
-class ExperimentItem(TypedDict, total=False):
+class LocalExperimentItem(TypedDict, total=False):
     """Structure for experiment data items.
 
     Args:
@@ -39,6 +37,10 @@ class ExperimentItem(TypedDict, total=False):
     metadata: Optional[Dict[str, Any]]
 
 
+ExperimentItem = Union[LocalExperimentItem, DatasetItemClient]
+ExperimentData = Union[List[LocalExperimentItem], List[DatasetItemClient]]
+
+
 class Evaluation(TypedDict, total=False):
     """Structure for evaluation results.
 
@@ -66,7 +68,7 @@ class ExperimentItemResult(TypedDict):
         dataset_run_id: Dataset run ID if this item was part of a Langfuse dataset
     """
 
-    item: Union[ExperimentItem, DatasetItem]
+    item: ExperimentItem
     output: Any
     evaluations: List[Evaluation]
     trace_id: Optional[str]
@@ -93,7 +95,10 @@ class TaskFunction(Protocol):
     """Protocol for experiment task functions."""
 
     def __call__(
-        self, item: Union[ExperimentItem, dict, DatasetItem, "DatasetItemClient"]
+        self,
+        *,
+        item: ExperimentItem,
+        **kwargs: Dict[str, Any],
     ) -> Union[Any, Awaitable[Any]]:
         """Execute the task on an experiment item.
 
@@ -116,6 +121,7 @@ def __call__(
         output: Any,
         expected_output: Any = None,
         metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Dict[str, Any],
     ) -> Union[
         Evaluation, List[Evaluation], Awaitable[Union[Evaluation, List[Evaluation]]]
     ]:
@@ -137,7 +143,10 @@ class RunEvaluatorFunction(Protocol):
     """Protocol for run-level evaluator functions."""
 
     def __call__(
-        self, *, item_results: List[ExperimentItemResult]
+        self,
+        *,
+        item_results: List[ExperimentItemResult],
+        **kwargs: Dict[str, Any],
     ) -> Union[
         Evaluation, List[Evaluation], Awaitable[Union[Evaluation, List[Evaluation]]]
     ]:
@@ -286,7 +295,7 @@ def _format_value(value: Any) -> str:
 
 
 async def _run_evaluator(
-    evaluator: EvaluatorFunction, **kwargs: Any
+    evaluator: Union[EvaluatorFunction, RunEvaluatorFunction], **kwargs: Any
 ) -> List[Evaluation]:
     """Run an evaluator function and normalize the result."""
     try:
@@ -299,8 +308,10 @@ async def _run_evaluator(
         # Normalize to list
         if isinstance(result, dict):
             return [result]
+
         elif isinstance(result, list):
             return result
+
         else:
             return []
 
@@ -310,12 +321,9 @@ async def _run_evaluator(
         return []
 
 
-async def _run_task(
-    task: TaskFunction,
-    item: Union[ExperimentItem, dict, DatasetItem, "DatasetItemClient"],
-) -> Any:
+async def _run_task(task: TaskFunction, item: ExperimentItem) -> Any:
     """Run a task function and handle sync/async."""
-    result = task(item)
+    result = task(item=item)
 
     # Handle async tasks
     if asyncio.iscoroutine(result):