« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""
2Type definitions for batch embedding operations.
3"""
5from __future__ import annotations
7from dataclasses import dataclass, field
8from datetime import UTC, datetime
9from enum import StrEnum
10from typing import Any
13class EmbeddingStatus(StrEnum):
14 """Embedding job status."""
16 PENDING = "pending"
17 PROCESSING = "processing"
18 CACHING = "caching"
19 STORING = "storing"
20 COMPLETED = "completed"
21 FAILED = "failed"
24@dataclass
25class BatchEmbeddingProgress:
26 """Track progress of batch embedding job."""
28 job_id: str
29 status: EmbeddingStatus
30 total_texts: int = 0
31 texts_processed: int = 0
32 cache_hits: int = 0
33 cache_misses: int = 0
34 started_at: datetime = field(default_factory=lambda: datetime.now(UTC))
35 updated_at: datetime = field(default_factory=lambda: datetime.now(UTC))
36 error: str | None = None
38 @property
39 def progress_percent(self) -> float:
40 """Calculate progress percentage."""
41 if self.total_texts == 0:
42 return 0.0
43 return (self.texts_processed / self.total_texts) * 100
45 @property
46 def cache_hit_rate(self) -> float:
47 """Calculate cache hit rate."""
48 total = self.cache_hits + self.cache_misses
49 if total == 0:
50 return 0.0
51 return (self.cache_hits / total) * 100
53 def update(
54 self,
55 status: EmbeddingStatus | None = None,
56 texts_processed: int | None = None,
57 cache_hits: int | None = None,
58 cache_misses: int | None = None,
59 error: str | None = None,
60 ) -> None:
61 """Update progress tracking."""
62 if status is not None:
63 self.status = status
64 if texts_processed is not None:
65 self.texts_processed = texts_processed
66 if cache_hits is not None:
67 self.cache_hits += cache_hits
68 if cache_misses is not None:
69 self.cache_misses += cache_misses
70 if error is not None:
71 self.error = error
72 self.status = EmbeddingStatus.FAILED
73 self.updated_at = datetime.now(UTC)
75 def to_dict(self) -> dict[str, Any]:
76 """Convert to dictionary for serialization."""
77 return {
78 "job_id": self.job_id,
79 "status": self.status.value,
80 "total_texts": self.total_texts,
81 "texts_processed": self.texts_processed,
82 "cache_hits": self.cache_hits,
83 "cache_misses": self.cache_misses,
84 "cache_hit_rate": self.cache_hit_rate,
85 "progress_percent": self.progress_percent,
86 "started_at": self.started_at.isoformat(),
87 "updated_at": self.updated_at.isoformat(),
88 "error": self.error,
89 }
92@dataclass
93class BatchEmbeddingResult:
94 """Result of batch embedding job."""
96 job_id: str
97 success: bool
98 embeddings_generated: int = 0
99 cache_hits: int = 0
100 duration_seconds: float = 0.0
101 error: str | None = None
102 metadata: dict[str, Any] = field(default_factory=dict)
104 @classmethod
105 def success_result(
106 cls,
107 job_id: str,
108 embeddings_generated: int,
109 cache_hits: int,
110 duration: float,
111 metadata: dict[str, Any] | None = None,
112 ) -> BatchEmbeddingResult:
113 """Create successful embedding result."""
114 return cls(
115 job_id=job_id,
116 success=True,
117 embeddings_generated=embeddings_generated,
118 cache_hits=cache_hits,
119 duration_seconds=duration,
120 metadata=metadata or {},
121 )
123 @classmethod
124 def failure_result(
125 cls,
126 job_id: str,
127 error: str,
128 duration: float,
129 metadata: dict[str, Any] | None = None,
130 ) -> BatchEmbeddingResult:
131 """Create failed embedding result."""
132 return cls(
133 job_id=job_id,
134 success=False,
135 error=error,
136 duration_seconds=duration,
137 metadata=metadata or {},
138 )
140 def to_dict(self) -> dict[str, Any]:
141 """Convert to dictionary for serialization."""
142 return {
143 "job_id": self.job_id,
144 "success": self.success,
145 "embeddings_generated": self.embeddings_generated,
146 "cache_hits": self.cache_hits,
147 "duration_seconds": self.duration_seconds,
148 "error": self.error,
149 "metadata": self.metadata,
150 }
153@dataclass
154class BatchEmbeddingJob:
155 """JobProtocol configuration for batch embedding."""
157 chunks: list[Chunk] # type: ignore[name-defined] # Forward reference
158 collection_name: str
159 model_name: str | None = None
160 batch_size: int = 100 # Process 100 texts per API call
161 use_cache: bool = True
163 def to_job_kwargs(self) -> dict[str, Any]:
164 """Convert to JobProtocol kwargs."""
165 model_name = self.model_name
166 if model_name is None:
167 model_name = "text-embedding-ada-002"
169 return {
170 "chunks": [{"text": c.text, "metadata": c.metadata} for c in self.chunks],
171 "collection_name": self.collection_name,
172 "model_name": model_name,
173 "batch_size": self.batch_size,
174 "use_cache": self.use_cache,
175 }