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

1""" 

2Progress tracking for batch embedding operations. 

3""" 

4 

5from __future__ import annotations 

6 

7import asyncio 

8from typing import Any 

9 

10from lexigram.ai.workers.batch_embedding.types import ( 

11 BatchEmbeddingProgress, 

12 EmbeddingStatus, 

13) 

14 

15 

16class ProgressTracker: 

17 """Thread-safe progress tracking for embedding jobs.""" 

18 

19 def __init__(self) -> None: 

20 """Initialize progress tracker.""" 

21 self._progress: dict[str, BatchEmbeddingProgress] = {} 

22 self._lock = asyncio.Lock() 

23 

24 async def initialize_job( 

25 self, 

26 job_id: str, 

27 total_texts: int, 

28 status: EmbeddingStatus = EmbeddingStatus.PENDING, 

29 ) -> None: 

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

31 async with self._lock: 

32 self._progress[job_id] = BatchEmbeddingProgress( 

33 job_id=job_id, 

34 status=status, 

35 total_texts=total_texts, 

36 ) 

37 

38 async def update_progress( 

39 self, 

40 job_id: str, 

41 status: EmbeddingStatus | None = None, 

42 texts_processed: int | None = None, 

43 cache_hits: int | None = None, 

44 cache_misses: int | None = None, 

45 error: str | None = None, 

46 ) -> None: 

47 """Update progress for a job.""" 

48 async with self._lock: 

49 if job_id in self._progress: 

50 self._progress[job_id].update( 

51 status=status, 

52 texts_processed=texts_processed, 

53 cache_hits=cache_hits, 

54 cache_misses=cache_misses, 

55 error=error, 

56 ) 

57 

58 async def get_progress(self, job_id: str) -> BatchEmbeddingProgress | None: 

59 """Get progress for a job.""" 

60 async with self._lock: 

61 return self._progress.get(job_id) 

62 

63 async def remove_job(self, job_id: str) -> None: 

64 """Remove progress tracking for a job.""" 

65 async with self._lock: 

66 self._progress.pop(job_id, None) 

67 

68 def get_active_jobs(self) -> list[str]: 

69 """Get list of active job IDs.""" 

70 return list(self._progress.keys()) 

71 

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

73 """Get progress tracking statistics.""" 

74 return { 

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

76 "jobs_by_status": { 

77 status.value: sum( 

78 1 for p in self._progress.values() if p.status == status 

79 ) 

80 for status in EmbeddingStatus 

81 }, 

82 }