- Cleaned up unnecessary comments in `__all__` definitions for better readability. - Adjusted indentation and formatting across modules for improved clarity (e.g., long lines, logical grouping). - Simplified conditional expressions and inline comments for context scoring and ranking. - Replaced some hard-coded values with type-safe annotations (e.g., `ClassVar`). - Removed unused imports and ensured consistent usage across test files. - Updated `test_score_not_cached_on_context` to clarify caching behavior. - Improved truncation strategy logic and marker handling.
355 lines
9.5 KiB
Python
355 lines
9.5 KiB
Python
"""
|
|
Context Management Engine Exceptions.
|
|
|
|
Provides a hierarchy of exceptions for context assembly,
|
|
token budget management, and related operations.
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
|
|
class ContextError(Exception):
|
|
"""
|
|
Base exception for all context management errors.
|
|
|
|
All context-related exceptions should inherit from this class
|
|
to allow for catch-all handling when needed.
|
|
"""
|
|
|
|
def __init__(self, message: str, details: dict[str, Any] | None = None) -> None:
|
|
"""
|
|
Initialize context error.
|
|
|
|
Args:
|
|
message: Human-readable error message
|
|
details: Optional dict with additional error context
|
|
"""
|
|
self.message = message
|
|
self.details = details or {}
|
|
super().__init__(message)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""Convert exception to dictionary for logging/serialization."""
|
|
return {
|
|
"error_type": self.__class__.__name__,
|
|
"message": self.message,
|
|
"details": self.details,
|
|
}
|
|
|
|
|
|
class BudgetExceededError(ContextError):
|
|
"""
|
|
Raised when token budget is exceeded.
|
|
|
|
This occurs when the assembled context would exceed the
|
|
allocated token budget for a specific context type or total.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Token budget exceeded",
|
|
allocated: int = 0,
|
|
requested: int = 0,
|
|
context_type: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize budget exceeded error.
|
|
|
|
Args:
|
|
message: Error message
|
|
allocated: Tokens allocated for this context type
|
|
requested: Tokens requested
|
|
context_type: Type of context that exceeded budget
|
|
"""
|
|
details: dict[str, Any] = {
|
|
"allocated": allocated,
|
|
"requested": requested,
|
|
"overage": requested - allocated,
|
|
}
|
|
if context_type:
|
|
details["context_type"] = context_type
|
|
|
|
super().__init__(message, details)
|
|
self.allocated = allocated
|
|
self.requested = requested
|
|
self.context_type = context_type
|
|
|
|
|
|
class TokenCountError(ContextError):
|
|
"""
|
|
Raised when token counting fails.
|
|
|
|
This typically occurs when the LLM Gateway token counting
|
|
service is unavailable or returns an error.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Failed to count tokens",
|
|
model: str | None = None,
|
|
text_length: int | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize token count error.
|
|
|
|
Args:
|
|
message: Error message
|
|
model: Model for which counting was attempted
|
|
text_length: Length of text that failed to count
|
|
"""
|
|
details: dict[str, Any] = {}
|
|
if model:
|
|
details["model"] = model
|
|
if text_length is not None:
|
|
details["text_length"] = text_length
|
|
|
|
super().__init__(message, details)
|
|
self.model = model
|
|
self.text_length = text_length
|
|
|
|
|
|
class CompressionError(ContextError):
|
|
"""
|
|
Raised when context compression fails.
|
|
|
|
This can occur when summarization or truncation cannot
|
|
reduce content to fit within the budget.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Failed to compress context",
|
|
original_tokens: int | None = None,
|
|
target_tokens: int | None = None,
|
|
achieved_tokens: int | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize compression error.
|
|
|
|
Args:
|
|
message: Error message
|
|
original_tokens: Tokens before compression
|
|
target_tokens: Target token count
|
|
achieved_tokens: Tokens achieved after compression attempt
|
|
"""
|
|
details: dict[str, Any] = {}
|
|
if original_tokens is not None:
|
|
details["original_tokens"] = original_tokens
|
|
if target_tokens is not None:
|
|
details["target_tokens"] = target_tokens
|
|
if achieved_tokens is not None:
|
|
details["achieved_tokens"] = achieved_tokens
|
|
|
|
super().__init__(message, details)
|
|
self.original_tokens = original_tokens
|
|
self.target_tokens = target_tokens
|
|
self.achieved_tokens = achieved_tokens
|
|
|
|
|
|
class AssemblyTimeoutError(ContextError):
|
|
"""
|
|
Raised when context assembly exceeds time limit.
|
|
|
|
Context assembly must complete within a configurable
|
|
time limit to maintain responsiveness.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Context assembly timed out",
|
|
timeout_ms: int = 0,
|
|
elapsed_ms: float = 0.0,
|
|
stage: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize assembly timeout error.
|
|
|
|
Args:
|
|
message: Error message
|
|
timeout_ms: Configured timeout in milliseconds
|
|
elapsed_ms: Actual elapsed time in milliseconds
|
|
stage: Pipeline stage where timeout occurred
|
|
"""
|
|
details: dict[str, Any] = {
|
|
"timeout_ms": timeout_ms,
|
|
"elapsed_ms": round(elapsed_ms, 2),
|
|
}
|
|
if stage:
|
|
details["stage"] = stage
|
|
|
|
super().__init__(message, details)
|
|
self.timeout_ms = timeout_ms
|
|
self.elapsed_ms = elapsed_ms
|
|
self.stage = stage
|
|
|
|
|
|
class ScoringError(ContextError):
|
|
"""
|
|
Raised when context scoring fails.
|
|
|
|
This occurs when relevance, recency, or priority scoring
|
|
encounters an error.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Failed to score context",
|
|
scorer_type: str | None = None,
|
|
context_id: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize scoring error.
|
|
|
|
Args:
|
|
message: Error message
|
|
scorer_type: Type of scorer that failed
|
|
context_id: ID of context being scored
|
|
"""
|
|
details: dict[str, Any] = {}
|
|
if scorer_type:
|
|
details["scorer_type"] = scorer_type
|
|
if context_id:
|
|
details["context_id"] = context_id
|
|
|
|
super().__init__(message, details)
|
|
self.scorer_type = scorer_type
|
|
self.context_id = context_id
|
|
|
|
|
|
class FormattingError(ContextError):
|
|
"""
|
|
Raised when context formatting fails.
|
|
|
|
This occurs when converting assembled context to
|
|
model-specific format fails.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Failed to format context",
|
|
model: str | None = None,
|
|
adapter: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize formatting error.
|
|
|
|
Args:
|
|
message: Error message
|
|
model: Target model
|
|
adapter: Adapter that failed
|
|
"""
|
|
details: dict[str, Any] = {}
|
|
if model:
|
|
details["model"] = model
|
|
if adapter:
|
|
details["adapter"] = adapter
|
|
|
|
super().__init__(message, details)
|
|
self.model = model
|
|
self.adapter = adapter
|
|
|
|
|
|
class CacheError(ContextError):
|
|
"""
|
|
Raised when cache operations fail.
|
|
|
|
This is typically non-fatal and should be handled
|
|
gracefully by falling back to recomputation.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Cache operation failed",
|
|
operation: str | None = None,
|
|
cache_key: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize cache error.
|
|
|
|
Args:
|
|
message: Error message
|
|
operation: Cache operation that failed (get, set, delete)
|
|
cache_key: Key involved in the failed operation
|
|
"""
|
|
details: dict[str, Any] = {}
|
|
if operation:
|
|
details["operation"] = operation
|
|
if cache_key:
|
|
details["cache_key"] = cache_key
|
|
|
|
super().__init__(message, details)
|
|
self.operation = operation
|
|
self.cache_key = cache_key
|
|
|
|
|
|
class ContextNotFoundError(ContextError):
|
|
"""
|
|
Raised when expected context is not found.
|
|
|
|
This occurs when required context sources return
|
|
no results or are unavailable.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Required context not found",
|
|
source: str | None = None,
|
|
query: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize context not found error.
|
|
|
|
Args:
|
|
message: Error message
|
|
source: Source that returned no results
|
|
query: Query used to search
|
|
"""
|
|
details: dict[str, Any] = {}
|
|
if source:
|
|
details["source"] = source
|
|
if query:
|
|
details["query"] = query
|
|
|
|
super().__init__(message, details)
|
|
self.source = source
|
|
self.query = query
|
|
|
|
|
|
class InvalidContextError(ContextError):
|
|
"""
|
|
Raised when context data is invalid.
|
|
|
|
This occurs when context content or metadata
|
|
fails validation.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str = "Invalid context data",
|
|
field: str | None = None,
|
|
value: Any | None = None,
|
|
reason: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize invalid context error.
|
|
|
|
Args:
|
|
message: Error message
|
|
field: Field that is invalid
|
|
value: Invalid value (may be redacted for security)
|
|
reason: Reason for invalidity
|
|
"""
|
|
details: dict[str, Any] = {}
|
|
if field:
|
|
details["field"] = field
|
|
if value is not None:
|
|
# Avoid logging potentially sensitive values
|
|
details["value_type"] = type(value).__name__
|
|
if reason:
|
|
details["reason"] = reason
|
|
|
|
super().__init__(message, details)
|
|
self.field = field
|
|
self.value = value
|
|
self.reason = reason
|