From f91f396f96883cd1110c7268cb9c4a48e9eed4fe Mon Sep 17 00:00:00 2001 From: Willy Douhard Date: Wed, 24 Jul 2024 11:48:11 +0200 Subject: [PATCH] feat: make datasetId optional for experiments --- literalai/api/__init__.py | 10 +++++----- literalai/api/dataset_helpers.py | 4 ++-- literalai/api/gql.py | 4 ++-- literalai/dataset_experiment.py | 18 +++++++----------- 4 files changed, 16 insertions(+), 20 deletions(-) diff --git a/literalai/api/__init__.py b/literalai/api/__init__.py index 2696b86..e0f61dd 100644 --- a/literalai/api/__init__.py +++ b/literalai/api/__init__.py @@ -1104,8 +1104,8 @@ def delete_dataset(self, id: str): def create_experiment( self, - dataset_id: str, name: str, + dataset_id: Optional[str] = None, prompt_id: Optional[str] = None, params: Optional[Dict] = None, ) -> "DatasetExperiment": @@ -1113,8 +1113,8 @@ def create_experiment( Creates a new experiment associated with a specific dataset. Args: - dataset_id (str): The unique identifier of the dataset. name (str): The name of the experiment. + dataset_id (Optional[str]): The unique identifier of the dataset. prompt_id (Optional[str]): The identifier of the prompt associated with the experiment. params (Optional[Dict]): Additional parameters for the experiment. @@ -1122,7 +1122,7 @@ def create_experiment( DatasetExperiment: The newly created experiment object. """ return self.gql_helper( - *create_experiment_helper(self, dataset_id, name, prompt_id, params) + *create_experiment_helper(self, name, dataset_id, prompt_id, params) ) def create_experiment_item( @@ -2315,15 +2315,15 @@ async def delete_dataset(self, id: str): async def create_experiment( self, - dataset_id: str, name: str, + dataset_id: Optional[str] = None, prompt_id: Optional[str] = None, params: Optional[Dict] = None, ) -> "DatasetExperiment": sync_api = LiteralAPI(self.api_key, self.url) return await self.gql_helper( - *create_experiment_helper(sync_api, dataset_id, name, prompt_id, params) + *create_experiment_helper(sync_api, name, dataset_id, prompt_id, params) ) create_experiment.__doc__ = LiteralAPI.create_experiment.__doc__ diff --git a/literalai/api/dataset_helpers.py b/literalai/api/dataset_helpers.py index 65f6694..7e03323 100644 --- a/literalai/api/dataset_helpers.py +++ b/literalai/api/dataset_helpers.py @@ -95,8 +95,8 @@ def process_response(response): def create_experiment_helper( api: "LiteralAPI", - dataset_id: str, name: str, + dataset_id: Optional[str] = None, prompt_id: Optional[str] = None, params: Optional[Dict] = None, ): @@ -119,7 +119,7 @@ def process_response(response): def create_experiment_item_helper( dataset_experiment_id: str, - dataset_item_id: str, + dataset_item_id: Optional[str] = None, input: Optional[Dict] = None, output: Optional[Dict] = None, ): diff --git a/literalai/api/gql.py b/literalai/api/gql.py index 5284ed8..e2adcc1 100644 --- a/literalai/api/gql.py +++ b/literalai/api/gql.py @@ -833,7 +833,7 @@ CREATE_EXPERIMENT = """ mutation CreateDatasetExperiment( $name: String! - $datasetId: String! + $datasetId: String $promptId: String $params: Json ) { @@ -854,7 +854,7 @@ CREATE_EXPERIMENT_ITEM = """ mutation CreateDatasetExperimentItem( $datasetExperimentId: String! - $datasetItemId: String! + $datasetItemId: String $input: Json $output: Json ) { diff --git a/literalai/dataset_experiment.py b/literalai/dataset_experiment.py index 6c85751..7254e00 100644 --- a/literalai/dataset_experiment.py +++ b/literalai/dataset_experiment.py @@ -1,6 +1,5 @@ from dataclasses import dataclass, field -from typing import Dict, List, Optional, TypedDict -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, List, Optional, TypedDict from literalai.my_types import ScoreDict, Utils @@ -11,7 +10,7 @@ class DatasetExperimentItemDict(TypedDict, total=False): id: str datasetExperimentId: str - datasetItemId: str + datasetItemId: Optional[str] scores: List[ScoreDict] input: Optional[Dict] output: Optional[Dict] @@ -21,7 +20,7 @@ class DatasetExperimentItemDict(TypedDict, total=False): class DatasetExperimentItem(Utils): id: str dataset_experiment_id: str - dataset_item_id: str + dataset_item_id: Optional[str] scores: List[ScoreDict] input: Optional[Dict] output: Optional[Dict] @@ -41,7 +40,7 @@ def from_dict(cls, item: DatasetExperimentItemDict) -> "DatasetExperimentItem": return cls( id=item.get("id", ""), dataset_experiment_id=item.get("datasetExperimentId", ""), - dataset_item_id=item.get("datasetItemId", ""), + dataset_item_id=item.get("datasetItemId"), scores=item.get("scores", []), input=item.get("input"), output=item.get("output"), @@ -64,7 +63,7 @@ class DatasetExperiment(Utils): id: str created_at: str name: str - dataset_id: str + dataset_id: Optional[str] params: Optional[Dict] prompt_id: Optional[str] = None items: List[DatasetExperimentItem] = field(default_factory=lambda: []) @@ -73,7 +72,7 @@ def log(self, item_dict: DatasetExperimentItemDict) -> DatasetExperimentItem: dataset_experiment_item = DatasetExperimentItem.from_dict( { "datasetExperimentId": self.id, - "datasetItemId": item_dict.get("datasetItemId", ""), + "datasetItemId": item_dict.get("datasetItemId"), "input": item_dict.get("input", {}), "output": item_dict.get("output", {}), "scores": item_dict.get("scores", []), @@ -110,8 +109,5 @@ def from_dict( dataset_id=dataset_experiment.get("datasetId", ""), params=dataset_experiment.get("params"), prompt_id=dataset_experiment.get("promptId"), - items=[ - DatasetExperimentItem.from_dict(item) - for item in items - ], + items=[DatasetExperimentItem.from_dict(item) for item in items], )