Merge pull request #775 from parea-ai/PAI-1063-extend-experiment-list-endpoint

feat: extend experiments list endpoint

Author: joschkabraun

parea/client.py

@@ -25,6 +25,7 @@
     CreateTestCases,
     ExperimentSchema,
     ExperimentStatsSchema,
+    ExperimentWithPinnedStatsSchema,
     FeedbackRequest,
     FinishExperimentRequestSchema,
     ListExperimentUUIDsFilters,
@@ -350,13 +351,13 @@ async def aget_trace_log(self, trace_id: str) -> TraceLog:
         response = await self._client.request_async("GET", GET_TRACE_LOG_ENDPOINT.format(trace_id=trace_id))
         return structure_trace_log_from_api(response.json())
 
-    def list_experiment_uuids(self, filter_conditions: Optional[ListExperimentUUIDsFilters] = ListExperimentUUIDsFilters()) -> List[str]:
+    def list_experiments(self, filter_conditions: Optional[ListExperimentUUIDsFilters] = ListExperimentUUIDsFilters()) -> List[ExperimentWithPinnedStatsSchema]:
         response = self._client.request("POST", LIST_EXPERIMENTS_ENDPOINT, data=asdict(filter_conditions))
-        return response.json()
+        return structure(response.json(), List[ExperimentWithPinnedStatsSchema])
 
-    async def alist_experiment_uuids(self, filter_conditions: Optional[ListExperimentUUIDsFilters] = ListExperimentUUIDsFilters()) -> List[str]:
+    async def alist_experiments(self, filter_conditions: Optional[ListExperimentUUIDsFilters] = ListExperimentUUIDsFilters()) -> List[ExperimentWithPinnedStatsSchema]:
         response = await self._client.request_async("POST", LIST_EXPERIMENTS_ENDPOINT, data=asdict(filter_conditions))
-        return response.json()
+        return structure(response.json(), List[ExperimentWithPinnedStatsSchema])
 
     def get_experiment_trace_logs(self, experiment_uuid: str, filters: TraceLogFilters = TraceLogFilters()) -> List[TraceLog]:
         response = self._client.request("POST", GET_EXPERIMENT_LOGS_ENDPOINT.format(experiment_uuid=experiment_uuid), data=asdict(filters))

parea/cookbook/list_experiments.py

@@ -9,7 +9,8 @@
 
 p = Parea(api_key=os.getenv("PAREA_API_KEY"))
 
-exp_uuids = p.list_experiment_uuids(ListExperimentUUIDsFilters(experiment_name_filter="Greeting"))
-print(f"Num. experiments: {len(exp_uuids)}")
-trace_logs = p.get_experiment_trace_logs(exp_uuids[0])
+experiments = p.list_experiments(ListExperimentUUIDsFilters(experiment_name_filter="Greeting"))
+print(f"Num. experiments: {len(experiments)}")
+trace_logs = p.get_experiment_trace_logs(experiments[0].uuid)
 print(f"Num. trace logs: {len(trace_logs)}")
+print(f"Trace log: {trace_logs[0]}")

parea/schemas/models.py

@@ -309,6 +309,50 @@ class ListExperimentUUIDsFilters:
     run_name_filter: Optional[str] = None
 
 
+class ExperimentStatus(str, Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class StatisticOperation(str, Enum):
+    MEAN = "mean"
+    MEDIAN = "median"
+    VARIANCE = "variance"
+    STANDARD_DEVIATION = "standard_deviation"
+    MIN = "min"
+    MAX = "max"
+    MSE = "mse"
+    MAE = "mae"
+    CORRELATION = "correlation"
+    SPEARMAN_CORRELATION = "spearman_correlation"
+    ACCURACY = "accuracy"
+    CUSTOM = "custom"
+
+
+@define
+class ExperimentPinnedStatistic:
+    var1: str
+    operation: StatisticOperation
+    value: float
+    var2: Optional[str] = None
+
+
+@define
+class ExperimentWithPinnedStatsSchema:
+    name: str
+    uuid: str
+    created_at: str
+    run_name: str
+    project_uuid: str
+    status: ExperimentStatus
+    is_public: bool = False
+    metadata: Optional[Dict[str, str]] = None
+    pinned_stats: list[ExperimentPinnedStatistic] = []
+    num_samples: Optional[int] = None
+
+
 class FilterOperator(str, Enum):
     EQUALS = "equals"
     NOT_EQUALS = "not_equals"

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "parea-ai"
 packages = [{ include = "parea" }]
-version = "0.2.134"
+version = "0.2.135"
 description = "Parea python sdk"
 readme = "README.md"
 authors = ["joel-parea-ai <[email protected]>"]