Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-workers/src/lexigram/ai/workers/document_ingestion/types.py: 76%

1""" 

2Document ingestion types and data models. 

3 

4Contains all the core data structures used in document ingestion. 

5""" 

6 

7from __future__ import annotations 

8 

9from dataclasses import dataclass, field 

10from datetime import UTC, datetime 

11from enum import StrEnum 

12from typing import TYPE_CHECKING, Any 

13 

14if TYPE_CHECKING: 

15 from pathlib import Path 

16 

17 from lexigram.ai.workers.document_ingestion.parser import DocumentParser 

18 

19ChunkingConfigDict = dict[str, Any] 

20 

21 

22class IngestionStatus(StrEnum): 

23 """Document ingestion status.""" 

24 

25 PENDING = "pending" 

26 PARSING = "parsing" 

27 CHUNKING = "chunking" 

28 EMBEDDING = "embedding" 

29 STORING = "storing" 

30 COMPLETED = "completed" 

31 FAILED = "failed" 

32 

33 

34@dataclass 

35class Document: 

36 """Simple document container for ingestion.""" 

37 

38 content: str 

39 metadata: dict[str, Any] = field(default_factory=dict) 

40 

41 

42@dataclass 

43class IngestionProgress: 

44 """Track progress of document ingestion.""" 

45 

46 document_id: str 

47 status: IngestionStatus 

48 total_pages: int = 0 

49 pages_processed: int = 0 

50 total_chunks: int = 0 

51 chunks_processed: int = 0 

52 started_at: datetime = field(default_factory=lambda: datetime.now(UTC)) 

53 updated_at: datetime = field(default_factory=lambda: datetime.now(UTC)) 

54 error: str | None = None 

55 

56 @property 

57 def progress_percent(self) -> float: 

58 """Calculate overall progress percentage.""" 

59 if self.total_chunks == 0: 

60 return 0.0 

61 return (self.chunks_processed / self.total_chunks) * 100 

62 

63 def update( 

64 self, 

65 status: IngestionStatus | None = None, 

66 pages_processed: int | None = None, 

67 chunks_processed: int | None = None, 

68 error: str | None = None, 

69 ) -> None: 

70 """Update progress tracking.""" 

71 if status is not None: 

72 self.status = status 

73 if pages_processed is not None: 

74 self.pages_processed = pages_processed 

75 if chunks_processed is not None: 

76 self.chunks_processed = chunks_processed 

77 if error is not None: 

78 self.error = error 

79 self.status = IngestionStatus.FAILED 

80 self.updated_at = datetime.now(UTC) 

81 

82 def to_dict(self) -> dict[str, Any]: 

83 """Convert to dictionary for serialization.""" 

84 return { 

85 "document_id": self.document_id, 

86 "status": self.status.value, 

87 "total_pages": self.total_pages, 

88 "pages_processed": self.pages_processed, 

89 "total_chunks": self.total_chunks, 

90 "chunks_processed": self.chunks_processed, 

91 "progress_percent": self.progress_percent, 

92 "started_at": self.started_at.isoformat(), 

93 "updated_at": self.updated_at.isoformat(), 

94 "error": self.error, 

95 } 

96 

97 

98@dataclass 

99class IngestionResult: 

100 """Result of document ingestion.""" 

101 

102 document_id: str 

103 success: bool 

104 chunks_created: int = 0 

105 duration_seconds: float = 0.0 

106 error: str | None = None 

107 metadata: dict[str, Any] = field(default_factory=dict) 

108 

109 @classmethod 

110 def success_result( 

111 cls, 

112 document_id: str, 

113 chunks_created: int, 

114 duration: float, 

115 metadata: dict[str, Any] | None = None, 

116 ) -> IngestionResult: 

117 """Create successful ingestion result.""" 

118 return cls( 

119 document_id=document_id, 

120 success=True, 

121 chunks_created=chunks_created, 

122 duration_seconds=duration, 

123 metadata=metadata or {}, 

124 ) 

125 

126 @classmethod 

127 def failure_result( 

128 cls, 

129 document_id: str, 

130 error: str, 

131 duration: float, 

132 metadata: dict[str, Any] | None = None, 

133 ) -> IngestionResult: 

134 """Create failed ingestion result.""" 

135 return cls( 

136 document_id=document_id, 

137 success=False, 

138 error=error, 

139 duration_seconds=duration, 

140 metadata=metadata or {}, 

141 ) 

142 

143 def to_dict(self) -> dict[str, Any]: 

144 """Convert to dictionary for serialization.""" 

145 return { 

146 "document_id": self.document_id, 

147 "success": self.success, 

148 "chunks_created": self.chunks_created, 

149 "duration_seconds": self.duration_seconds, 

150 "error": self.error, 

151 "metadata": self.metadata, 

152 } 

153 

154 

155@dataclass 

156class DocumentIngestionJob: 

157 """JobProtocol configuration for document ingestion.""" 

158 

159 document_id: str 

160 file_path: Path 

161 collection_name: str 

162 parser: DocumentParser | None = None 

163 chunking_config: ChunkingConfigDict | None = None 

164 metadata: dict[str, Any] = field(default_factory=dict) 

165 batch_size: int = 50 # Number of chunks to process in batch 

166 

167 def to_job_kwargs(self) -> dict[str, Any]: 

168 """Convert to JobProtocol kwargs.""" 

169 return { 

170 "document_id": self.document_id, 

171 "file_path": str(self.file_path), 

172 "collection_name": self.collection_name, 

173 "metadata": self.metadata, 

174 "batch_size": self.batch_size, 

175 } 

176 

177 

178# Forward references for type hints 

179if False: # TYPE_CHECKING block for forward references 

180 pass