Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/pipeline/types.py: 75%
105 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Core types for RAG pipeline execution.
3This module provides the foundational types used throughout the pipeline,
4including context, errors, and configuration structures.
5"""
7from __future__ import annotations
9from dataclasses import dataclass, field
10from datetime import UTC, datetime
11from enum import StrEnum
12from typing import Any
13import uuid
15from lexigram.ai.rag.chunking import Chunk
16from lexigram.ai.rag.query import TransformedQuery
17from lexigram.ai.rag.routing import QueryIntent, RoutingDecision
18from lexigram.ai.rag.synthesis import (
19 ContextChunk,
20 QualityMetrics,
21 SynthesisResult,
22)
25class ErrorStrategy(StrEnum):
26 """Strategy for handling errors in pipeline stages."""
28 RETRY = "retry" # Retry with exponential backoff
29 FALLBACK = "fallback" # Use fallback implementation
30 SKIP = "skip" # Skip stage and continue
31 FAIL_FAST = "fail_fast" # Raise exception immediately
32 GRACEFUL = "graceful" # Return partial results with warnings
35class StageStatus(StrEnum):
36 """Status of a pipeline stage."""
38 PENDING = "pending"
39 RUNNING = "running"
40 COMPLETED = "completed"
41 FAILED = "failed"
42 SKIPPED = "skipped"
45@dataclass
46class PipelineError:
47 """Represents an error that occurred during pipeline execution."""
49 stage: str
50 error: Exception
51 timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
52 recoverable: bool = True
53 retry_count: int = 0
56@dataclass
57class StageMetrics:
58 """Metrics for a single pipeline stage."""
60 stage_name: str
61 status: StageStatus
62 start_time: datetime | None = None
63 end_time: datetime | None = None
64 duration_ms: float = 0.0
65 error: PipelineError | None = None
67 @property
68 def duration_seconds(self) -> float:
69 """Get duration in seconds."""
70 return self.duration_ms / 1000.0
73@dataclass
74class PipelineContext:
75 """Context object that flows through pipeline stages.
77 This context carries all state through the pipeline, accumulating
78 results from each stage and tracking metrics/errors.
79 """
81 # Required input
82 query: str
84 # Optional input documents
85 documents: list[str] = field(default_factory=list)
86 document_paths: list[str] = field(default_factory=list)
88 # Stage results - populated by pipeline stages
89 transformed_queries: list[TransformedQuery] = field(default_factory=list)
90 chunks: list[Chunk] = field(default_factory=list)
91 retrieved_chunks: list[ContextChunk] = field(default_factory=list)
92 optimized_chunks: list[ContextChunk] = field(default_factory=list)
93 synthesis_result: SynthesisResult | None = None
94 quality_metrics: QualityMetrics | None = None
96 # Routing and intent
97 routing_decision: RoutingDecision | None = None
98 intent: QueryIntent | None = None
100 # Request tracking
101 request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
102 created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
104 # Custom metadata
105 metadata: dict[str, Any] = field(default_factory=dict)
107 # Performance metrics
108 stage_metrics: list[StageMetrics] = field(default_factory=list)
109 token_usage: dict[str, int] = field(default_factory=dict)
110 cache_hits: dict[str, bool] = field(default_factory=dict)
112 # Error tracking
113 errors: list[PipelineError] = field(default_factory=list)
114 warnings: list[str] = field(default_factory=list)
116 @property
117 def has_errors(self) -> bool:
118 """Check if any errors occurred."""
119 return len(self.errors) > 0
121 @property
122 def has_warnings(self) -> bool:
123 """Check if any warnings were issued."""
124 return len(self.warnings) > 0
126 @property
127 def total_duration_ms(self) -> float:
128 """Calculate total pipeline duration in milliseconds."""
129 return sum(m.duration_ms for m in self.stage_metrics)
131 @property
132 def total_duration_seconds(self) -> float:
133 """Calculate total pipeline duration in seconds."""
134 return self.total_duration_ms / 1000.0
136 @property
137 def failed_stages(self) -> list[str]:
138 """Get list of failed stage names."""
139 return [
140 m.stage_name for m in self.stage_metrics if m.status == StageStatus.FAILED
141 ]
143 @property
144 def completed_stages(self) -> list[str]:
145 """Get list of completed stage names."""
146 return [
147 m.stage_name
148 for m in self.stage_metrics
149 if m.status == StageStatus.COMPLETED
150 ]
152 def add_error(
153 self,
154 stage: str,
155 error: Exception,
156 recoverable: bool = True,
157 retry_count: int = 0,
158 ) -> None:
159 """Add an error to the context.
161 Args:
162 stage: Name of the stage where error occurred
163 error: The exception that was raised
164 recoverable: Whether the error is recoverable
165 retry_count: Number of times this error has been retried
166 """
167 pipeline_error = PipelineError(
168 stage=stage,
169 error=error,
170 recoverable=recoverable,
171 retry_count=retry_count,
172 )
173 self.errors.append(pipeline_error)
175 def add_warning(self, message: str) -> None:
176 """Add a warning message to the context.
178 Args:
179 message: Warning message
180 """
181 self.warnings.append(message)
183 def start_stage(self, stage_name: str) -> None:
184 """Mark a stage as started.
186 Args:
187 stage_name: Name of the stage
188 """
189 metrics = StageMetrics(
190 stage_name=stage_name,
191 status=StageStatus.RUNNING,
192 start_time=datetime.now(UTC),
193 )
194 self.stage_metrics.append(metrics)
196 def complete_stage(
197 self,
198 stage_name: str,
199 status: StageStatus = StageStatus.COMPLETED,
200 error: PipelineError | None = None,
201 ) -> None:
202 """Mark a stage as completed.
204 Args:
205 stage_name: Name of the stage
206 status: Final status of the stage
207 error: Optional error if stage failed
208 """
209 for metrics in reversed(self.stage_metrics):
210 if (
211 metrics.stage_name == stage_name
212 and metrics.status == StageStatus.RUNNING
213 ):
214 metrics.status = status
215 metrics.end_time = datetime.now(UTC)
216 if metrics.start_time:
217 duration = (metrics.end_time - metrics.start_time).total_seconds()
218 metrics.duration_ms = duration * 1000.0
219 metrics.error = error
220 break
222 def record_token_usage(self, operation: str, tokens: int) -> None:
223 """Record token usage for an operation.
225 Args:
226 operation: Name of the operation (e.g., "embedding", "llm_call")
227 tokens: Number of tokens used
228 """
229 if operation in self.token_usage:
230 self.token_usage[operation] += tokens
231 else:
232 self.token_usage[operation] = tokens
234 def record_cache_hit(self, cache_key: str, hit: bool) -> None:
235 """Record a cache hit or miss.
237 Args:
238 cache_key: Key for the cache lookup
239 hit: Whether the cache was hit
240 """
241 self.cache_hits[cache_key] = hit
243 def to_dict(self) -> dict[str, Any]:
244 """Convert context to dictionary for serialization.
246 Returns:
247 Dictionary representation of the context
248 """
249 return {
250 "request_id": self.request_id,
251 "query": self.query,
252 "created_at": self.created_at.isoformat(),
253 "total_duration_ms": self.total_duration_ms,
254 "completed_stages": self.completed_stages,
255 "failed_stages": self.failed_stages,
256 "has_errors": self.has_errors,
257 "has_warnings": self.has_warnings,
258 "error_count": len(self.errors),
259 "warning_count": len(self.warnings),
260 "token_usage": self.token_usage,
261 "cache_hits": dict(self.cache_hits),
262 "metadata": self.metadata,
263 }