aiecs 1.0.1__py3-none-any.whl → 1.7.6__py3-none-any.whl

aiecs/llm/clients/base_client.py

@@ -0,0 +1,324 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, List, AsyncGenerator
+from dataclasses import dataclass
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Lazy import to avoid circular dependency
+
+
+def _get_config_loader():
+    """Lazy import of config loader to avoid circular dependency"""
+    from aiecs.llm.config import get_llm_config_loader
+
+    return get_llm_config_loader()
+
+
+@dataclass
+class CacheControl:
+    """
+    Cache control marker for message content.
+
+    Used to indicate that a message or message block should be cached
+    by providers that support prompt caching (e.g., Anthropic, Google).
+    """
+
+    type: str = "ephemeral"  # Cache type - "ephemeral" for session-scoped caching
+
+
+@dataclass
+class LLMMessage:
+    """
+    Represents a message in an LLM conversation.
+
+    Attributes:
+        role: Message role - "system", "user", "assistant", or "tool"
+        content: Text content of the message (None when using tool calls)
+        tool_calls: Tool call information for assistant messages
+        tool_call_id: Tool call ID for tool response messages
+        cache_control: Cache control marker for prompt caching support
+    """
+
+    role: str  # "system", "user", "assistant", "tool"
+    content: Optional[str] = None  # None when using tool calls
+    tool_calls: Optional[List[Dict[str, Any]]] = None  # For assistant messages with tool calls
+    tool_call_id: Optional[str] = None  # For tool messages
+    cache_control: Optional[CacheControl] = None  # Cache control for prompt caching
+
+
+@dataclass
+class LLMResponse:
+    """
+    Response from an LLM provider.
+
+    Attributes:
+        content: Generated text content
+        provider: Name of the LLM provider (e.g., "openai", "google", "vertex")
+        model: Model name used for generation
+        tokens_used: Total tokens used (prompt + completion)
+        prompt_tokens: Number of tokens in the prompt
+        completion_tokens: Number of tokens in the completion
+        cost_estimate: Estimated cost in USD
+        response_time: Response time in seconds
+        metadata: Additional provider-specific metadata
+        cache_creation_tokens: Tokens used to create a new cache entry
+        cache_read_tokens: Tokens read from cache (indicates cache hit)
+        cache_hit: Whether the request hit a cached prompt prefix
+    """
+
+    content: str
+    provider: str
+    model: str
+    tokens_used: Optional[int] = None
+    prompt_tokens: Optional[int] = None
+    completion_tokens: Optional[int] = None
+    cost_estimate: Optional[float] = None
+    response_time: Optional[float] = None
+    metadata: Optional[Dict[str, Any]] = None
+    # Cache metadata for prompt caching observability
+    cache_creation_tokens: Optional[int] = None
+    cache_read_tokens: Optional[int] = None
+    cache_hit: Optional[bool] = None
+
+    def __post_init__(self):
+        """Ensure consistency of token data"""
+        # If there are detailed token information but no total, calculate the
+        # total
+        if self.prompt_tokens is not None and self.completion_tokens is not None and self.tokens_used is None:
+            self.tokens_used = self.prompt_tokens + self.completion_tokens
+
+        # If only total is available but no detailed information, try to
+        # estimate (cannot accurately allocate in this case)
+        elif self.tokens_used is not None and self.prompt_tokens is None and self.completion_tokens is None:
+            # In this case we cannot accurately allocate, keep as is
+            pass
+
+
+class LLMClientError(Exception):
+    """Base exception for LLM client errors"""
+
+
+class ProviderNotAvailableError(LLMClientError):
+    """Raised when a provider is not available or misconfigured"""
+
+
+class RateLimitError(LLMClientError):
+    """Raised when rate limit is exceeded"""
+
+
+class SafetyBlockError(LLMClientError):
+    """Raised when content is blocked by safety filters"""
+    
+    def __init__(
+        self,
+        message: str,
+        block_reason: Optional[str] = None,
+        block_type: Optional[str] = None,  # "prompt" or "response"
+        safety_ratings: Optional[List[Dict[str, Any]]] = None,
+    ):
+        """
+        Initialize SafetyBlockError with detailed information.
+        
+        Args:
+            message: Error message
+            block_reason: Reason for blocking (e.g., SAFETY, RECITATION, JAILBREAK)
+            block_type: Type of block - "prompt" if input was blocked, "response" if output was blocked
+            safety_ratings: List of safety ratings with category, severity, etc.
+        """
+        super().__init__(message)
+        self.block_reason = block_reason
+        self.block_type = block_type
+        self.safety_ratings = safety_ratings or []
+    
+    def __str__(self) -> str:
+        """Return detailed error message"""
+        msg = super().__str__()
+        if self.block_reason:
+            msg += f" (Block reason: {self.block_reason})"
+        if self.block_type:
+            msg += f" (Block type: {self.block_type})"
+        if self.safety_ratings:
+            categories = [r.get("category", "UNKNOWN") for r in self.safety_ratings if r.get("blocked")]
+            if categories:
+                msg += f" (Categories: {', '.join(categories)})"
+        return msg
+
+
+class BaseLLMClient(ABC):
+    """Abstract base class for all LLM provider clients"""
+
+    def __init__(self, provider_name: str):
+        self.provider_name = provider_name
+        self.logger = logging.getLogger(f"{__name__}.{provider_name}")
+
+    @abstractmethod
+    async def generate_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> LLMResponse:
+        """Generate text using the provider's API"""
+
+    @abstractmethod
+    async def stream_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> AsyncGenerator[str, None]:
+        """Stream text generation using the provider's API"""
+
+    @abstractmethod
+    async def close(self):
+        """Clean up resources"""
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    def _count_tokens_estimate(self, text: str) -> int:
+        """Rough token count estimation (4 chars ≈ 1 token for English)"""
+        return len(text) // 4
+
+    def _apply_cache_control(
+        self,
+        messages: List[LLMMessage],
+        enable_caching: bool = True,
+    ) -> List[LLMMessage]:
+        """
+        Apply cache control markers to cacheable messages.
+
+        Marks system messages and the first message in the conversation
+        with cache_control for providers that support prompt caching.
+
+        Args:
+            messages: List of LLM messages
+            enable_caching: Whether to enable caching (default: True)
+
+        Returns:
+            List of messages with cache_control applied where appropriate
+        """
+        if not enable_caching:
+            return messages
+
+        result = []
+        for msg in messages:
+            if msg.role == "system" and msg.cache_control is None:
+                # Mark system messages as cacheable
+                result.append(
+                    LLMMessage(
+                        role=msg.role,
+                        content=msg.content,
+                        tool_calls=msg.tool_calls,
+                        tool_call_id=msg.tool_call_id,
+                        cache_control=CacheControl(type="ephemeral"),
+                    )
+                )
+            else:
+                result.append(msg)
+        return result
+
+    def _extract_cache_metadata(
+        self,
+        usage: Any,
+    ) -> Dict[str, Any]:
+        """
+        Extract cache metadata from provider response usage data.
+
+        Override in subclasses for provider-specific extraction.
+
+        Args:
+            usage: Usage data from provider response
+
+        Returns:
+            Dictionary with cache_creation_tokens, cache_read_tokens, cache_hit
+        """
+        return {
+            "cache_creation_tokens": None,
+            "cache_read_tokens": None,
+            "cache_hit": None,
+        }
+
+    def _estimate_cost(
+        self,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        token_costs: Dict,
+    ) -> float:
+        """
+        Estimate the cost of the API call.
+
+        DEPRECATED: Use _estimate_cost_from_config instead for config-based cost estimation.
+        This method is kept for backward compatibility.
+        """
+        if model in token_costs:
+            costs = token_costs[model]
+            return (input_tokens * costs["input"] + output_tokens * costs["output"]) / 1000
+        return 0.0
+
+    def _estimate_cost_from_config(self, model_name: str, input_tokens: int, output_tokens: int) -> float:
+        """
+        Estimate the cost using configuration-based pricing.
+
+        Args:
+            model_name: Name of the model
+            input_tokens: Number of input tokens
+            output_tokens: Number of output tokens
+
+        Returns:
+            Estimated cost in USD
+        """
+        try:
+            loader = _get_config_loader()
+            model_config = loader.get_model_config(self.provider_name, model_name)
+
+            if model_config and model_config.costs:
+                input_cost = (input_tokens * model_config.costs.input) / 1000
+                output_cost = (output_tokens * model_config.costs.output) / 1000
+                return input_cost + output_cost
+            else:
+                self.logger.warning(f"No cost configuration found for model {model_name} " f"in provider {self.provider_name}")
+                return 0.0
+        except Exception as e:
+            self.logger.warning(f"Failed to estimate cost from config: {e}")
+            return 0.0
+
+    def _get_model_config(self, model_name: str):
+        """
+        Get model configuration from the config loader.
+
+        Args:
+            model_name: Name of the model
+
+        Returns:
+            ModelConfig if found, None otherwise
+        """
+        try:
+            loader = _get_config_loader()
+            return loader.get_model_config(self.provider_name, model_name)
+        except Exception as e:
+            self.logger.warning(f"Failed to get model config: {e}")
+            return None
+
+    def _get_default_model(self) -> Optional[str]:
+        """
+        Get the default model for this provider from configuration.
+
+        Returns:
+            Default model name if configured, None otherwise
+        """
+        try:
+            loader = _get_config_loader()
+            return loader.get_default_model(self.provider_name)
+        except Exception as e:
+            self.logger.warning(f"Failed to get default model: {e}")
+            return None