dslighting 1.3.9__py3-none-any.whl

dsat/services/llm.py

@@ -0,0 +1,486 @@
+# dsat/services/llm.py
+
+"""
+Unified, asynchronous LLM service powered by LiteLLM.
+Provides a simple interface for standard calls, structured JSON output,
+and automatic cost tracking.
+"""
+import logging
+import asyncio
+import yaml
+import copy
+import time
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Type, Optional, Any, Dict, List
+
+import litellm
+from pydantic import BaseModel, ValidationError
+
+from dsat.config import LLMConfig # Use the main pydantic config
+from dsat.common.exceptions import LLMError
+
+logger = logging.getLogger(__name__)
+
+# Configure LiteLLM globally
+litellm.telemetry = False # Disable anonymous telemetry
+litellm.input_callbacks = []
+litellm.success_callbacks = []
+litellm.failure_callbacks = []
+
+# Load custom model pricing from YAML configuration file
+def _load_custom_model_pricing():
+    """Load custom model pricing configuration from config.yaml file."""
+    try:
+        # Get the path to the config.yaml file relative to this module
+        current_dir = Path(__file__).parent
+        framework_dir = current_dir.parent.parent  # Go up to ds_agent_framework
+        config_yaml_path = framework_dir / "config.yaml"
+        
+        if config_yaml_path.exists():
+            with open(config_yaml_path, 'r', encoding='utf-8') as f:
+                config = yaml.safe_load(f)
+                return config.get('custom_model_pricing', {})
+        else:
+            # Changed to debug to avoid confusing warnings for pip-installed packages
+            logger.debug(f"Config file not found at {config_yaml_path} (this is expected for pip-installed packages)")
+            return {}
+    except Exception as e:
+        logger.error(f"Failed to load cost configuration: {e}")
+        return {}
+
+# Load and apply custom model pricing
+CUSTOM_MODEL_PRICING = _load_custom_model_pricing()
+if CUSTOM_MODEL_PRICING:
+    litellm.model_cost.update(CUSTOM_MODEL_PRICING)
+    logger.info(f"Loaded custom model pricing for {len(CUSTOM_MODEL_PRICING)} models")
+
+class LLMService:
+    """
+    A robust wrapper around LiteLLM that handles requests, structured formatting,
+    and cost tracking. It's configured via the main DSATConfig's LLMConfig.
+    """
+    def __init__(self, config: LLMConfig):
+        self.config = config
+        self.total_cost = 0.0
+        self.total_prompt_tokens = 0
+        self.total_completion_tokens = 0
+        self.total_prompt_cost = 0.0
+        self.total_completion_cost = 0.0
+        self.call_history: List[Dict[str, Any]] = []
+
+        # Parse API key pool from string list format
+        import json
+        try:
+            self.api_keys = json.loads(config.api_key) if config.api_key and config.api_key.startswith('[') else [config.api_key]
+        except:
+            self.api_keys = [config.api_key]
+        self.current_key_index = 0
+        logger.info(f"Initialized LLM service with {len(self.api_keys)} API keys")
+
+    def _get_current_api_key(self) -> str:
+        """Get current API key from the pool."""
+        return self.api_keys[self.current_key_index]
+
+    def _switch_api_key(self) -> bool:
+        """Switch to next API key in pool. Returns True if switched, False if no more keys."""
+        if len(self.api_keys) <= 1:
+            return False
+        self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
+        logger.warning(f"Switching to API key {self.current_key_index + 1}/{len(self.api_keys)}")
+        return True
+
+    @staticmethod
+    def _safe_float(value: Any) -> Optional[float]:
+        try:
+            return float(value)
+        except (TypeError, ValueError):
+            return None
+
+    @staticmethod
+    def _safe_int(value: Any) -> Optional[int]:
+        try:
+            return int(value)
+        except (TypeError, ValueError):
+            return None
+
+    def _is_retryable_error(self, error: Exception) -> bool:
+        """
+        Determines whether an error is retryable based on litellm exception types.
+        Only network timeouts, rate limits, and temporary service issues should be retried.
+        Authentication errors, invalid requests, and other permanent failures should not be retried.
+        """
+        # Import litellm exceptions locally to avoid import issues
+        try:
+            import litellm.exceptions as litellm_exceptions
+        except ImportError:
+            # If litellm exceptions module is not available, be conservative and retry
+            return True
+            
+        # Non-retryable errors - fail immediately
+        non_retryable_errors = (
+            litellm_exceptions.AuthenticationError,  # API key issues
+            litellm_exceptions.InvalidRequestError,  # Request format/parameter issues  
+            litellm_exceptions.PermissionDeniedError,  # Insufficient permissions
+            litellm_exceptions.NotFoundError,  # Model/endpoint not found
+        )
+        
+        # Retryable errors - can be retried
+        retryable_errors = (
+            litellm_exceptions.RateLimitError,  # Rate limit exceeded
+            litellm_exceptions.ServiceUnavailableError,  # Temporary service issues
+            litellm_exceptions.Timeout,  # Network timeout
+            litellm_exceptions.APIConnectionError,  # Connection issues
+            litellm_exceptions.InternalServerError,  # Server-side temporary issues
+        )
+        
+        # Check for specific error types
+        if isinstance(error, non_retryable_errors):
+            return False
+        elif isinstance(error, retryable_errors):
+            return True
+        else:
+            # For unknown errors, be conservative and retry
+            # This handles generic network errors, etc.
+            return True
+
+    def _supports_response_format(self) -> bool:
+        """
+        Whether it's safe to pass `response_format` through LiteLLM for this model.
+
+        Some OpenAI reasoning models (e.g. `o4-mini-*`) reject `response_format` and
+        require JSON-only behavior to be enforced via prompt instead.
+        """
+        raw_model = (self.config.model or "").strip()
+        model = raw_model.split("/")[-1].strip()
+        model_lower = model.lower()
+        if model_lower.startswith("o4-mini-") or model_lower == "o4-mini":
+            return False
+        if model_lower == "kimi-k2-instruct-0905":
+            return False
+        return True
+
+    @staticmethod
+    def _is_insufficient_balance_error(error: Exception) -> bool:
+        """
+        Detect insufficient balance/quota errors that should trigger API key rotation.
+        """
+        message = str(error).lower()
+        return (
+            "insufficient" in message and "balance" in message
+        ) or "insufficient_quota" in message or "insufficient quota" in message
+
+    async def _make_llm_call_with_retries(
+        self, messages: list, response_format: Optional[dict] = None, max_retries: int = 3, base_delay: float = 1.0
+    ):
+        """
+        Internal method to make LLM calls with centralized retry logic and exponential backoff.
+        This method is the single point of contact with the LiteLLM library.
+
+        Args:
+            messages: List of message dictionaries for the LLM.
+            response_format: Optional response format specification.
+            max_retries: Maximum number of retry attempts.
+            base_delay: Base delay in seconds for exponential backoff.
+
+        Returns:
+            The raw LiteLLM response object upon success.
+
+        Raises:
+            LLMError: If all retry attempts fail due to API errors or empty responses.
+        """
+        if response_format and not self._supports_response_format():
+            logger.info(
+                "Dropping unsupported `response_format` for model %s; using prompt-enforced JSON instead.",
+                self.config.model,
+            )
+            response_format = None
+
+        logger.info(f"prompt: {messages[-1]['content']}")
+        last_exception = None
+        for attempt in range(max_retries):
+            call_id = uuid.uuid4().hex
+            call_started_at = datetime.utcnow()
+            perf_start = time.perf_counter()
+            try:
+                kwargs = {
+                    "model": self.config.model,
+                    "messages": messages,
+                    "temperature": self.config.temperature,
+                    "api_key": self._get_current_api_key(),
+                    "api_base": self.config.api_base
+                }
+                if self.config.provider:
+                    kwargs["custom_llm_provider"] = self.config.provider
+
+                if response_format:
+                    kwargs["response_format"] = response_format
+
+                response = await litellm.acompletion(**kwargs)
+
+                try:
+                    content = response.choices[0].message.content
+                    if content and content.strip():
+                        duration = time.perf_counter() - perf_start
+                        self._record_successful_call(
+                            call_id=call_id,
+                            call_started_at=call_started_at,
+                            duration=duration,
+                            messages=messages,
+                            response=response,
+                            content=content,
+                            response_format=response_format
+                        )
+                        return response  # Success!
+                    else:
+                        # Treat empty response as a failure to be retried
+                        logger.warning(f"LLM returned an empty response on attempt {attempt + 1}/{max_retries}.")
+                        last_exception = LLMError("LLM returned an empty response.")
+                except (IndexError, AttributeError) as content_error:
+                    logger.warning(f"Invalid response structure on attempt {attempt + 1}/{max_retries}: {content_error}")
+                    last_exception = LLMError(f"Invalid response structure: {content_error}")
+
+            except Exception as e:
+                # Special handling for RateLimitError - switch API key immediately
+                import litellm.exceptions as litellm_exceptions
+                if isinstance(e, litellm_exceptions.RateLimitError):
+                    logger.warning(f"RateLimitError on attempt {attempt + 1}/{max_retries}: {e}")
+                    if self._switch_api_key():
+                        logger.info(f"Switched to next API key, retrying immediately...")
+                        last_exception = e
+                        continue  # Skip delay, retry immediately with new key
+                    else:
+                        logger.warning(f"No more API keys to switch, will use retry delay")
+                        last_exception = e
+                elif isinstance(e, litellm_exceptions.APIError) and self._is_insufficient_balance_error(e):
+                    logger.warning(f"Insufficient balance on attempt {attempt + 1}/{max_retries}: {e}")
+                    if self._switch_api_key():
+                        logger.info("Switched to next API key after insufficient balance, retrying immediately...")
+                        last_exception = e
+                        continue
+                    else:
+                        logger.warning("No more API keys to switch after insufficient balance, will use retry delay")
+                        last_exception = e
+                # Check if this is a retryable error
+                elif self._is_retryable_error(e):
+                    logger.warning(f"Retryable LLM error on attempt {attempt + 1}/{max_retries}: {e}")
+                    last_exception = e
+                    logger.debug(f"Debug info - messages: {messages}, response_format: {response_format if response_format else 'None'}")
+                else:
+                    # Non-retryable error - fail immediately
+                    logger.error(f"Non-retryable LLM error: {e}")
+                    raise LLMError(f"LLM call failed with non-retryable error: {e}") from e
+
+            # If this was the last attempt, break the loop to raise the final error
+            if attempt == max_retries - 1:
+                break
+
+            # Exponential backoff with jitter
+            delay = base_delay * (3 ** attempt) + (asyncio.get_event_loop().time() % 1)
+            logger.info(f"Retrying LLM call in {delay:.2f} seconds ({attempt + 2}/{max_retries})...")
+            await asyncio.sleep(delay)
+
+        raise LLMError(f"LLM call failed after {max_retries} attempts. Last error: {last_exception}") from last_exception
+
+    def _record_successful_call(
+        self,
+        call_id: str,
+        call_started_at: datetime,
+        duration: float,
+        messages: list,
+        response: Any,
+        content: str,
+        response_format: Optional[dict],
+    ) -> None:
+        """
+        将一次成功的调用附加到历史中，并更新累计 token / 费用。
+        """
+        usage_payload = self._extract_usage(response)
+        try:
+            call_cost_raw = litellm.completion_cost(completion_response=response)
+            call_cost = float(call_cost_raw) if call_cost_raw is not None else 0.0
+        except Exception:
+            call_cost = 0.0
+
+        self.total_cost += call_cost
+
+        prompt_tokens = usage_payload.get("prompt_tokens") if usage_payload else None
+        completion_tokens = usage_payload.get("completion_tokens") if usage_payload else None
+
+        if prompt_tokens:
+            self.total_prompt_tokens += prompt_tokens
+        if completion_tokens:
+            self.total_completion_tokens += completion_tokens
+
+        prompt_cost_val = usage_payload.get("prompt_tokens_cost") if usage_payload else None
+        completion_cost_val = usage_payload.get("completion_tokens_cost") if usage_payload else None
+
+        if prompt_cost_val is not None:
+            self.total_prompt_cost += prompt_cost_val
+        if completion_cost_val is not None:
+            self.total_completion_cost += completion_cost_val
+
+        total_tokens = usage_payload.get("total_tokens") if usage_payload else None
+        cost_per_token = (call_cost / total_tokens) if total_tokens else None
+
+        history_entry = {
+            "call_id": call_id,
+            "model": self.config.model,
+            "provider": self.config.provider,
+            "timestamp_utc": call_started_at.isoformat() + "Z",
+            "duration_seconds": round(duration, 4),
+            "response_format": "json_object" if response_format else "text",
+            "messages": copy.deepcopy(messages),
+            "response": content,
+            "usage": usage_payload or None,
+            "cost": call_cost,
+            "cost_per_token": cost_per_token,
+        }
+        self.call_history.append(history_entry)
+        logger.info(f"LLM call complete. Model: {self.config.model}, Cost: ${call_cost:.6f}")
+
+    def _extract_usage(self, response: Any) -> Dict[str, Any]:
+        """
+        从 LiteLLM Response 中提取 token / 费用信息，确保可 JSON 序列化。
+        """
+        usage = getattr(response, "usage", None)
+        if not usage:
+            return {}
+
+        payload: Dict[str, Any] = {
+            "prompt_tokens": self._safe_int(getattr(usage, "prompt_tokens", None)),
+            "completion_tokens": self._safe_int(getattr(usage, "completion_tokens", None)),
+            "total_tokens": self._safe_int(getattr(usage, "total_tokens", None)),
+            "prompt_tokens_cost": self._safe_float(getattr(usage, "prompt_tokens_cost", None)),
+            "completion_tokens_cost": self._safe_float(getattr(usage, "completion_tokens_cost", None)),
+        }
+        total_tokens_cost = self._safe_float(getattr(usage, "total_tokens_cost", None))
+        if total_tokens_cost is None:
+            prompt_cost = payload.get("prompt_tokens_cost")
+            completion_cost = payload.get("completion_tokens_cost")
+            if prompt_cost is not None and completion_cost is not None:
+                total_tokens_cost = prompt_cost + completion_cost
+        payload["total_tokens_cost"] = total_tokens_cost
+        return payload
+
+    async def call(self, prompt: str, system_message: Optional[str] = None, max_retries: Optional[int] = None) -> str:
+        """
+        Makes a standard, asynchronous call to the LLM and returns the text response.
+        The retry logic is handled by the internal _make_llm_call_with_retries method.
+
+        Args:
+            prompt: The user's prompt.
+            system_message: An optional system message to guide the LLM's behavior.
+            max_retries: Maximum number of retry attempts (default: 3).
+
+        Returns:
+            The string content of the LLM's response.
+        """
+        retries = max_retries if max_retries is not None else self.config.max_retries
+        messages = []
+        if system_message:
+            messages.append({"role": "system", "content": system_message})
+        messages.append({"role": "user", "content": prompt})
+
+        logger.debug(f"Calling LLM ({self.config.model}) with prompt: {prompt[:100]}...")
+
+        response = await self._make_llm_call_with_retries(messages, max_retries=retries)
+        content = response.choices[0].message.content
+        logger.info(f"content: {content}")
+        return content
+
+    async def call_with_json(self, prompt: str, output_model: Type[BaseModel], max_retries: Optional[int] = None) -> BaseModel:
+        """
+        Calls the LLM and forces the output to be a JSON object conforming to the
+        provided Pydantic model. The retry logic is handled by the internal method.
+
+        Args:
+            prompt: The user's prompt.
+            output_model: The Pydantic model class for the desired output structure.
+            max_retries: Maximum number of retry attempts (default: 3).
+
+        Returns:
+            An instantiated Pydantic model with the LLM's response.
+        """
+        retries = max_retries if max_retries is not None else self.config.max_retries
+        system_message = (
+            "You are a helpful assistant that always responds with a JSON object "
+            "that strictly adheres to the provided JSON Schema. Do not add any "
+            "other text, explanations, or markdown formatting."
+        )
+
+        prompt_with_schema = (
+            f"{prompt}\n\n# RESPONSE JSON SCHEMA:\n"
+            f"```json\n{output_model.model_json_schema()}\n```"
+        )
+
+        messages = [
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": prompt_with_schema}
+        ]
+
+        logger.debug(f"Calling LLM ({self.config.model}) for structured JSON output...")
+
+        for attempt in range(max(1, retries)):
+            if self._supports_response_format():
+                response = await self._make_llm_call_with_retries(
+                    messages,
+                    response_format={"type": "json_object"},
+                    max_retries=retries,
+                )
+            else:
+                if attempt == 0:
+                    logger.info(
+                        "Model %s does not support `response_format`; falling back to prompt-enforced JSON.",
+                        self.config.model,
+                    )
+                response = await self._make_llm_call_with_retries(
+                    messages,
+                    response_format=None,
+                    max_retries=retries,
+                )
+
+            try:
+                response_content = response.choices[0].message.content
+                logger.info(f"content: {response_content}")
+            except (IndexError, AttributeError) as e:
+                raise LLMError(f"Invalid response structure from LLM: {e}") from e
+
+            try:
+                return output_model.model_validate_json(response_content)
+            except ValidationError as e:
+                logger.error(
+                    "Failed to validate LLM JSON response against Pydantic model (attempt %d/%d): %s",
+                    attempt + 1,
+                    retries,
+                    e,
+                )
+                logger.debug("Invalid JSON received: %s", response_content)
+                if attempt + 1 >= retries:
+                    raise LLMError(
+                        f"LLM returned invalid JSON that could not be parsed: {e}"
+                    ) from e
+
+    def get_total_cost(self) -> float:
+        """Returns the total accumulated cost for this LLM instance."""
+        return self.total_cost
+
+    def get_call_history(self) -> List[Dict[str, Any]]:
+        """Returns a deep copy of the call history for telemetry persistence."""
+        return copy.deepcopy(self.call_history)
+
+    def get_usage_summary(self) -> Dict[str, Any]:
+        """汇总本实例的 token/费用信息。"""
+        total_tokens = self.total_prompt_tokens + self.total_completion_tokens
+        summary = {
+            "prompt_tokens": self.total_prompt_tokens,
+            "completion_tokens": self.total_completion_tokens,
+            "total_tokens": total_tokens,
+            "prompt_tokens_cost": round(self.total_prompt_cost, 12),
+            "completion_tokens_cost": round(self.total_completion_cost, 12),
+            "total_cost": round(self.total_cost, 12),
+            "call_count": len(self.call_history),
+        }
+        summary["cost_per_token"] = (self.total_cost / total_tokens) if total_tokens else None
+        return summary