Coverage for bbconf/config_parser/models.py: 88%

1import logging 

2from pathlib import Path 

3from typing import Literal 

4 

5from pydantic import BaseModel, ConfigDict, computed_field, field_validator 

6from yacman import load_yaml 

7 

8from bbconf.config_parser.const import ( # DEFAULT_VEC2VEC_MODEL, 

9 DEFAULT_DB_DIALECT, 

10 DEFAULT_DB_DRIVER, 

11 DEFAULT_DB_NAME, 

12 DEFAULT_DB_PORT, 

13 DEFAULT_QDRANT_BIVEC_COLLECTION_NAME, 

14 DEFAULT_QDRANT_FILE_COLLECTION_NAME, 

15 DEFAULT_QDRANT_HYBRID_COLLECTION_NAME, 

16 DEFAULT_QDRANT_PORT, 

17 DEFAULT_REGION2_VEC_MODEL, 

18 DEFAULT_S3_BUCKET, 

19 DEFAULT_SERVER_HOST, 

20 DEFAULT_SERVER_PORT, 

21 DEFAULT_SPARSE_MODEL, 

22 DEFAULT_TEXT2VEC_MODEL, 

23) 

24 

25_LOGGER = logging.getLogger(__name__) 

26 

27 

28class ConfigDB(BaseModel): 

29 host: str 

30 port: int = DEFAULT_DB_PORT 

31 user: str 

32 password: str 

33 database: str = DEFAULT_DB_NAME 

34 dialect: str = DEFAULT_DB_DIALECT 

35 driver: str | None = DEFAULT_DB_DRIVER 

36 run_migrations: bool = False 

37 

38 model_config = ConfigDict(extra="forbid") 

39 

40 @computed_field 

41 @property 

42 def url(self) -> str: 

43 """ 

44 The URL of the database. 

45 

46 Returns: 

47 The URL of the database. 

48 """ 

49 return f"{self.dialect}+{self.driver}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}" 

50 

51 

52class ConfigQdrant(BaseModel): 

53 host: str 

54 port: int = DEFAULT_QDRANT_PORT 

55 api_key: str | None = None 

56 file_collection: str = DEFAULT_QDRANT_FILE_COLLECTION_NAME 

57 text_collection: str | None = DEFAULT_QDRANT_BIVEC_COLLECTION_NAME 

58 hybrid_collection: str | None = DEFAULT_QDRANT_HYBRID_COLLECTION_NAME 

59 

60 

61class ConfigServer(BaseModel): 

62 host: str = DEFAULT_SERVER_HOST 

63 port: int = DEFAULT_SERVER_PORT 

64 

65 

66class ConfigPath(BaseModel): 

67 region2vec: str = DEFAULT_REGION2_VEC_MODEL 

68 # vec2vec: str = DEFAULT_VEC2VEC_MODEL 

69 text2vec: str = DEFAULT_TEXT2VEC_MODEL 

70 sparse_model: str = DEFAULT_SPARSE_MODEL 

71 umap_model: str | None = None # Path or link to pre-trained UMAP model 

72 

73 

74class AccessMethodsStruct(BaseModel): 

75 type: str 

76 description: str | None = None 

77 prefix: str 

78 

79 

80class AccessMethods(BaseModel): 

81 http: AccessMethodsStruct | None = None 

82 s3: AccessMethodsStruct | None = None 

83 local: AccessMethodsStruct | None = None 

84 

85 

86class ConfigS3(BaseModel): 

87 endpoint_url: str | None = None 

88 aws_access_key_id: str | None = None 

89 aws_secret_access_key: str | None = None 

90 bucket: str | None = DEFAULT_S3_BUCKET 

91 

92 @field_validator("aws_access_key_id", "aws_secret_access_key") 

93 def validate_aws_credentials(cls, value): 

94 # Do this if AWS credentials are not provided 

95 if value in [ 

96 "AWS_SECRET_ACCESS_KEY", 

97 "AWS_ACCESS_KEY_ID", 

98 "", 

99 "$AWS_ACCESS_KEY_ID", 

100 "$AWS_SECRET_ACCESS_KEY", 

101 ]: 

102 return None 

103 return value 

104 

105 @computed_field 

106 @property 

107 def modify_access(self) -> bool: 

108 """ 

109 If the AWS credentials are provided, set the modify access to True. (create = True) 

110 

111 Returns: 

112 True if AWS credentials are provided, False otherwise. 

113 """ 

114 if self.aws_access_key_id and self.aws_secret_access_key: 

115 return True 

116 _LOGGER.warning( 

117 "AWS credentials are not provided. The S3 bucket will be read-only." 

118 ) 

119 return False 

120 

121 

122class ConfigAnalysis(BaseModel): 

123 """Analysis backend configuration. 

124 

125 Controls which statistics engine is used for BED file analysis. 

126 """ 

127 

128 backend: Literal["r", "gtars"] = "r" 

129 

130 model_config = ConfigDict(extra="forbid") 

131 

132 

133class ConfigFile(BaseModel): 

134 database: ConfigDB 

135 qdrant: ConfigQdrant | None = None 

136 server: ConfigServer 

137 path: ConfigPath 

138 access_methods: AccessMethods | None = None 

139 s3: ConfigS3 | None = None 

140 analysis: ConfigAnalysis | None = ConfigAnalysis() 

141 

142 model_config = ConfigDict(extra="allow") 

143 

144 @classmethod 

145 def from_yaml(cls, path: Path): 

146 """ 

147 Load the database configuration from a YAML file. 

148 

149 Args: 

150 path: The path to the YAML file. 

151 

152 Returns: 

153 The database configuration. 

154 """ 

155 return cls.model_validate(load_yaml(path.as_posix()))