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

1""" 

2Progress tracking for document ingestion. 

3 

4Handles progress updates and tracking across ingestion jobs. 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10from typing import Any 

11 

12from lexigram.ai.workers.document_ingestion.types import ( 

13 IngestionProgress, 

14 IngestionStatus, 

15) 

16 

17 

18class ProgressTracker: 

19 """Tracks progress of document ingestion jobs.""" 

20 

21 def __init__(self) -> None: 

22 """Initialize progress tracker.""" 

23 self._progress: dict[str, IngestionProgress] = {} 

24 self._progress_lock = asyncio.Lock() 

25 

26 async def initialize_progress(self, job_id: str, document_id: str) -> None: 

27 """Initialize progress tracking for a job.""" 

28 async with self._progress_lock: 

29 self._progress[job_id] = IngestionProgress( 

30 document_id=document_id, 

31 status=IngestionStatus.PENDING, 

32 ) 

33 

34 async def get_progress(self, job_id: str) -> IngestionProgress | None: 

35 """Get ingestion progress for job.""" 

36 async with self._progress_lock: 

37 return self._progress.get(job_id) 

38 

39 async def update_progress( 

40 self, 

41 document_id: str, 

42 status: IngestionStatus | None = None, 

43 total_chunks: int | None = None, 

44 chunks_processed: int | None = None, 

45 error: str | None = None, 

46 ) -> None: 

47 """Update ingestion progress.""" 

48 async with self._progress_lock: 

49 # Find progress by document_id 

50 for progress in self._progress.values(): 

51 if progress.document_id == document_id: 

52 if total_chunks is not None: 

53 progress.total_chunks = total_chunks 

54 progress.update( 

55 status=status, 

56 chunks_processed=chunks_processed, 

57 error=error, 

58 ) 

59 break 

60 

61 async def get_all_progress(self) -> dict[str, IngestionProgress]: 

62 """Get all progress entries.""" 

63 async with self._progress_lock: 

64 return self._progress.copy() 

65 

66 def get_stats(self) -> dict[str, Any]: 

67 """Get progress tracking statistics.""" 

68 return { 

69 "active_jobs": len(self._progress), 

70 "completed_jobs": sum( 

71 1 

72 for p in self._progress.values() 

73 if p.status == IngestionStatus.COMPLETED 

74 ), 

75 "failed_jobs": sum( 

76 1 for p in self._progress.values() if p.status == IngestionStatus.FAILED 

77 ), 

78 }