Coverage for src/qdrant_loader/config/models.py: 98%

66 statements  

« prev     ^ index     » next       coverage.py v7.8.2, created at 2025-06-04 05:50 +0000

1"""Configuration models for multi-project support. 

2 

3This module defines the data structures used for multi-project configuration, 

4including project contexts, project configurations, and related models. 

5""" 

6 

7from dataclasses import dataclass 

8from typing import Any, Dict, List, Optional 

9from datetime import datetime 

10 

11from pydantic import BaseModel, Field 

12 

13from .sources import SourcesConfig 

14 

15 

16@dataclass 

17class ProjectContext: 

18 """Project context information passed through the pipeline.""" 

19 

20 project_id: str 

21 display_name: str 

22 description: Optional[str] 

23 collection_name: str 

24 config_overrides: Dict[str, Any] 

25 

26 def __post_init__(self): 

27 """Validate project context after initialization.""" 

28 if not self.project_id: 

29 raise ValueError("project_id cannot be empty") 

30 if not self.display_name: 

31 raise ValueError("display_name cannot be empty") 

32 if not self.collection_name: 

33 raise ValueError("collection_name cannot be empty") 

34 

35 

36class ProjectConfig(BaseModel): 

37 """Configuration for a single project.""" 

38 

39 project_id: str = Field(..., description="Unique project identifier") 

40 display_name: str = Field(..., description="Human-readable project name") 

41 description: Optional[str] = Field(None, description="Project description") 

42 sources: SourcesConfig = Field( 

43 default_factory=SourcesConfig, description="Project-specific sources" 

44 ) 

45 overrides: Dict[str, Any] = Field( 

46 default_factory=dict, description="Project-specific configuration overrides" 

47 ) 

48 

49 def get_effective_collection_name(self, global_collection_name: str) -> str: 

50 """Get the effective collection name for this project. 

51 

52 Args: 

53 global_collection_name: The global collection name from configuration 

54 

55 Returns: 

56 The collection name to use for this project (always the global collection name) 

57 """ 

58 # Always use the global collection name for all projects 

59 return global_collection_name 

60 

61 

62class ProjectsConfig(BaseModel): 

63 """Configuration for multiple projects.""" 

64 

65 projects: Dict[str, ProjectConfig] = Field( 

66 default_factory=dict, description="Project configurations" 

67 ) 

68 

69 def get_project(self, project_id: str) -> Optional[ProjectConfig]: 

70 """Get a project configuration by ID. 

71 

72 Args: 

73 project_id: The project identifier 

74 

75 Returns: 

76 The project configuration if found, None otherwise 

77 """ 

78 return self.projects.get(project_id) 

79 

80 def list_project_ids(self) -> List[str]: 

81 """Get a list of all project IDs. 

82 

83 Returns: 

84 List of project identifiers 

85 """ 

86 return list(self.projects.keys()) 

87 

88 def add_project(self, project_config: ProjectConfig) -> None: 

89 """Add a project configuration. 

90 

91 Args: 

92 project_config: The project configuration to add 

93 

94 Raises: 

95 ValueError: If project ID already exists 

96 """ 

97 if project_config.project_id in self.projects: 

98 raise ValueError(f"Project '{project_config.project_id}' already exists") 

99 

100 self.projects[project_config.project_id] = project_config 

101 

102 def to_dict(self) -> Dict[str, Any]: 

103 """Convert the projects configuration to a dictionary. 

104 

105 Returns: 

106 Dict[str, Any]: Projects configuration as a dictionary 

107 """ 

108 return { 

109 project_id: project_config.model_dump() 

110 for project_id, project_config in self.projects.items() 

111 } 

112 

113 

114@dataclass 

115class ParsedConfig: 

116 """Result of parsing a configuration file.""" 

117 

118 global_config: Any # Will be GlobalConfig, but avoiding circular import 

119 projects_config: ProjectsConfig 

120 

121 def get_all_projects(self) -> List[ProjectConfig]: 

122 """Get all project configurations. 

123 

124 Returns: 

125 List of all project configurations 

126 """ 

127 return list(self.projects_config.projects.values()) 

128 

129 

130class ProjectStats(BaseModel): 

131 """Statistics for a project.""" 

132 

133 project_id: str = Field(..., description="Project identifier") 

134 document_count: int = Field(default=0, description="Number of documents in project") 

135 source_count: int = Field(default=0, description="Number of sources in project") 

136 last_updated: Optional[datetime] = Field(None, description="Last update timestamp") 

137 storage_size: Optional[int] = Field(None, description="Storage size in bytes") 

138 

139 class Config: 

140 json_encoders = {datetime: lambda v: v.isoformat()} 

141 

142 

143class ProjectInfo(BaseModel): 

144 """Detailed information about a project.""" 

145 

146 id: str = Field(..., description="Project identifier") 

147 display_name: str = Field(..., description="Project display name") 

148 description: Optional[str] = Field(None, description="Project description") 

149 collection_name: str = Field(..., description="QDrant collection name") 

150 source_count: int = Field(default=0, description="Number of sources") 

151 document_count: int = Field(default=0, description="Number of documents") 

152 last_updated: Optional[datetime] = Field(None, description="Last update timestamp") 

153 

154 class Config: 

155 json_encoders = {datetime: lambda v: v.isoformat()} 

156 

157 

158class ProjectDetail(ProjectInfo): 

159 """Detailed project information including sources and statistics.""" 

160 

161 sources: List[Dict[str, Any]] = Field( 

162 default_factory=list, description="Source information" 

163 ) 

164 statistics: Dict[str, Any] = Field( 

165 default_factory=dict, description="Project statistics" 

166 )