Coverage for agentlib_flexquant/data_structures/market.py: 100%

1import pydantic 

2from typing import Union, Literal 

3from agentlib_flexquant.data_structures.globals import FlexibilityDirections 

4 

5 

6class RandomOptions(pydantic.BaseModel): 

7 type: Literal["random"] 

8 random_seed: int = pydantic.Field( 

9 name="random_seed", default=None, 

10 description="Random seed for reproducing experiments" 

11 ) 

12 

13 pos_neg_rate: float = pydantic.Field( 

14 name="pos_neg_rate", default=0, 

15 description="Determines the likelihood positive and the negative flexibility." 

16 "A higher rate means that more positive offers will be accepted.", 

17 le=1, ge=0 

18 ) 

19 

20 offer_acceptance_rate: float = pydantic.Field( 

21 name="offer_acceptance_rate", default=0.5, 

22 description="Determines the likelihood of an accepted offer", 

23 le=1, ge=0 

24 ) 

25 

26 

27class SingleOptions(pydantic.BaseModel): 

28 type: Literal["single"] 

29 start_time: float = pydantic.Field(description="After this time, the first available flex offer" 

30 " is accepted") 

31 direction: FlexibilityDirections = pydantic.Field(default="positive", description="Direction of the flexibility") 

32 

33 

34class CustomOptions(pydantic.BaseModel): 

35 type: Literal["custom"] 

36 

37 model_config = pydantic.ConfigDict(extra="allow") 

38 

39 

40class MarketSpecifications(pydantic.BaseModel): 

41 type: str = pydantic.Field( 

42 default=None, 

43 description="Name of market type" 

44 ) 

45 

46 cooldown: int = pydantic.Field( 

47 name="cooldown", default=6, 

48 description="cooldown time (no timesteps) after a provision" 

49 ) 

50 

51 minimum_average_flex: float = pydantic.Field( 

52 name="minimum_average_flex", default=0, 

53 unit="W", 

54 description="minimum average of an accepted offer" 

55 ) 

56 

57 options: Union[RandomOptions, SingleOptions, CustomOptions] = pydantic.Field( 

58 ..., 

59 description="Market options, changes depending on 'type'", 

60 discriminator='type' 

61 ) 

62 

63 # Root validator to automatically populate the options.type from the top-level type 

64 @pydantic.model_validator(mode='before') 

65 @classmethod 

66 def set_options_type(cls, values): 

67 market_type = values.get('type') 

68 options = values.get('options', {}) 

69 

70 # Ensure the options dict contains the correct 'type' field 

71 if isinstance(options, dict) and 'type' not in options: 

72 options['type'] = market_type 

73 values['options'] = options 

74 

75 return values 

76 

77 

78class RandomMarket(MarketSpecifications): 

79 type: str = "random"