openhands-sdk 1.7.3__py3-none-any.whl

openhands/sdk/llm/utils/retry_mixin.py

@@ -0,0 +1,128 @@
+from collections.abc import Callable, Iterable
+from typing import Any, cast
+
+from tenacity import (
+    RetryCallState,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from openhands.sdk.llm.exceptions import LLMNoResponseError
+from openhands.sdk.logger import get_logger
+
+
+logger = get_logger(__name__)
+
+# Helpful alias for listener signature: (attempt_number, max_retries) -> None
+RetryListener = Callable[[int, int, BaseException | None], None]
+
+
+class RetryMixin:
+    """Mixin class for retry logic."""
+
+    def retry_decorator(
+        self,
+        num_retries: int = 5,
+        retry_exceptions: tuple[type[BaseException], ...] = (LLMNoResponseError,),
+        retry_min_wait: int = 8,
+        retry_max_wait: int = 64,
+        retry_multiplier: float = 2.0,
+        retry_listener: RetryListener | None = None,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """
+        Create a LLM retry decorator with customizable parameters.
+        This is used for 429 errors, and a few other exceptions in LLM classes.
+        """
+
+        def before_sleep(retry_state: RetryCallState) -> None:
+            # Log first (also validates outcome as part of logging)
+            self.log_retry_attempt(retry_state)
+
+            if retry_listener is not None:
+                exc = (
+                    retry_state.outcome.exception()
+                    if retry_state.outcome is not None
+                    else None
+                )
+                retry_listener(retry_state.attempt_number, num_retries, exc)
+
+            # If there is no outcome or no exception, nothing to tweak.
+            if retry_state.outcome is None:
+                return
+            exc = retry_state.outcome.exception()
+            if exc is None:
+                return
+
+            # Only adjust temperature for LLMNoResponseError
+            if isinstance(exc, LLMNoResponseError):
+                kwargs = getattr(retry_state, "kwargs", None)
+                if isinstance(kwargs, dict):
+                    current_temp = kwargs.get("temperature", 0)
+                    if current_temp == 0:
+                        kwargs["temperature"] = 1.0
+                        logger.warning(
+                            "LLMNoResponseError with temperature=0, "
+                            "setting temperature to 1.0 for next attempt."
+                        )
+                    else:
+                        logger.warning(
+                            f"LLMNoResponseError with temperature={current_temp}, "
+                            "keeping original temperature"
+                        )
+
+        retry_decorator: Callable[[Callable[..., Any]], Callable[..., Any]] = retry(
+            before_sleep=before_sleep,
+            stop=stop_after_attempt(num_retries),
+            reraise=True,
+            retry=retry_if_exception_type(retry_exceptions),
+            wait=wait_exponential(
+                multiplier=retry_multiplier,
+                min=retry_min_wait,
+                max=retry_max_wait,
+            ),
+        )
+        return retry_decorator
+
+    def log_retry_attempt(self, retry_state: RetryCallState) -> None:
+        """Log retry attempts."""
+
+        if retry_state.outcome is None:
+            logger.error(
+                "retry_state.outcome is None. "
+                "This should not happen, please check the retry logic."
+            )
+            return
+
+        exc = retry_state.outcome.exception()
+        if exc is None:
+            logger.error("retry_state.outcome.exception() returned None.")
+            return
+
+        # Try to get max attempts from the stop condition if present
+        max_attempts: int | None = None
+        retry_obj = getattr(retry_state, "retry_object", None)
+        stop_condition = getattr(retry_obj, "stop", None)
+        if stop_condition is not None:
+            # stop_any has .stops, single stop does not
+            stops: Iterable[Any]
+            if hasattr(stop_condition, "stops"):
+                stops = stop_condition.stops  # type: ignore[attr-defined]
+            else:
+                stops = [stop_condition]
+            for stop_func in stops:
+                if hasattr(stop_func, "max_attempts"):
+                    max_attempts = getattr(stop_func, "max_attempts")
+                    break
+
+        # Attach dynamic fields for downstream consumers (keep existing behavior)
+        setattr(cast(Any, exc), "retry_attempt", retry_state.attempt_number)
+        if max_attempts is not None:
+            setattr(cast(Any, exc), "max_retries", max_attempts)
+
+        logger.error(
+            "%s. Attempt #%d | You can customize retry values in the configuration.",
+            exc,
+            retry_state.attempt_number,
+        )

openhands/sdk/llm/utils/telemetry.py

@@ -0,0 +1,362 @@
+import json
+import os
+import time
+import uuid
+import warnings
+from collections.abc import Callable
+from typing import Any, ClassVar
+
+from litellm.cost_calculator import completion_cost as litellm_completion_cost
+from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
+from litellm.types.utils import CostPerToken, ModelResponse, Usage
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+
+from openhands.sdk.llm.utils.metrics import Metrics
+from openhands.sdk.logger import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class Telemetry(BaseModel):
+    """
+    Handles latency, token/cost accounting, and optional logging.
+    All runtime state (like start times) lives in private attrs.
+    """
+
+    # --- Config fields ---
+    model_name: str = Field(default="unknown", description="Name of the LLM model")
+    log_enabled: bool = Field(default=False, description="Whether to log completions")
+    log_dir: str | None = Field(
+        default=None, description="Directory to write logs if enabled"
+    )
+    input_cost_per_token: float | None = Field(
+        default=None, ge=0, description="Custom Input cost per token (USD)"
+    )
+    output_cost_per_token: float | None = Field(
+        default=None, ge=0, description="Custom Output cost per token (USD)"
+    )
+
+    metrics: Metrics = Field(..., description="Metrics collector instance")
+
+    # --- Runtime fields (not serialized) ---
+    _req_start: float = PrivateAttr(default=0.0)
+    _req_ctx: dict[str, Any] = PrivateAttr(default_factory=dict)
+    _last_latency: float = PrivateAttr(default=0.0)
+    _log_completions_callback: Callable[[str, str], None] | None = PrivateAttr(
+        default=None
+    )
+    _stats_update_callback: Callable[[], None] | None = PrivateAttr(default=None)
+
+    model_config: ClassVar[ConfigDict] = ConfigDict(
+        extra="forbid", arbitrary_types_allowed=True
+    )
+
+    # ---------- Lifecycle ----------
+    def set_log_completions_callback(
+        self, callback: Callable[[str, str], None] | None
+    ) -> None:
+        """Set a callback function for logging instead of writing to file.
+
+        Args:
+            callback: A function that takes (filename, log_data) and handles the log.
+                     Used for streaming logs in remote execution contexts.
+        """
+        self._log_completions_callback = callback
+
+    def set_stats_update_callback(self, callback: Callable[[], None] | None) -> None:
+        """Set a callback function to be notified when stats are updated.
+
+        Args:
+            callback: A function called whenever metrics are updated.
+                     Used for streaming stats updates in remote execution contexts.
+        """
+        self._stats_update_callback = callback
+
+    def on_request(self, log_ctx: dict | None) -> None:
+        self._req_start = time.time()
+        self._req_ctx = log_ctx or {}
+
+    def on_response(
+        self,
+        resp: ModelResponse | ResponsesAPIResponse,
+        raw_resp: ModelResponse | None = None,
+    ) -> Metrics:
+        """
+        Side-effects:
+          - records latency, tokens, cost into Metrics
+          - optionally writes a JSON log file
+        """
+        # 1) latency
+        self._last_latency = time.time() - (self._req_start or time.time())
+        response_id = resp.id
+        self.metrics.add_response_latency(self._last_latency, response_id)
+
+        # 2) cost
+        cost = self._compute_cost(resp)
+        # Intentionally skip logging zero-cost (0.0) responses; only record
+        # positive cost
+        if cost:
+            self.metrics.add_cost(cost)
+
+        # 3) tokens - use typed usage field when available
+        usage = getattr(resp, "usage", None)
+
+        if usage and self._has_meaningful_usage(usage):
+            self._record_usage(
+                usage, response_id, self._req_ctx.get("context_window", 0)
+            )
+
+        # 4) optional logging
+        if self.log_enabled:
+            self.log_llm_call(resp, cost, raw_resp=raw_resp)
+
+        # 5) notify about stats update
+        if self._stats_update_callback is not None:
+            try:
+                self._stats_update_callback()
+            except Exception:
+                logger.exception("Stats update callback failed", exc_info=True)
+
+        return self.metrics.deep_copy()
+
+    def on_error(self, _err: BaseException) -> None:
+        # Stub for error tracking / counters
+        return
+
+    # ---------- Helpers ----------
+    def _has_meaningful_usage(self, usage: Usage | ResponseAPIUsage | None) -> bool:
+        """Check if usage has meaningful (non-zero) token counts.
+
+        Supports both Chat Completions Usage and Responses API Usage shapes.
+        """
+        if usage is None:
+            return False
+        try:
+            prompt_tokens = getattr(usage, "prompt_tokens", None)
+            if prompt_tokens is None:
+                prompt_tokens = getattr(usage, "input_tokens", 0)
+            completion_tokens = getattr(usage, "completion_tokens", None)
+            if completion_tokens is None:
+                completion_tokens = getattr(usage, "output_tokens", 0)
+
+            pt = int(prompt_tokens or 0)
+            ct = int(completion_tokens or 0)
+            return pt > 0 or ct > 0
+        except Exception:
+            return False
+
+    def _record_usage(
+        self, usage: Usage | ResponseAPIUsage, response_id: str, context_window: int
+    ) -> None:
+        """
+        Record token usage, supporting both Chat Completions Usage and
+        Responses API Usage.
+
+        Chat shape:
+          - prompt_tokens, completion_tokens
+          - prompt_tokens_details.cached_tokens
+          - completion_tokens_details.reasoning_tokens
+          - _cache_creation_input_tokens for cache_write
+        Responses shape:
+          - input_tokens, output_tokens
+          - input_tokens_details.cached_tokens
+          - output_tokens_details.reasoning_tokens
+        """
+        prompt_tokens = int(
+            getattr(usage, "prompt_tokens", None)
+            or getattr(usage, "input_tokens", 0)
+            or 0
+        )
+        completion_tokens = int(
+            getattr(usage, "completion_tokens", None)
+            or getattr(usage, "output_tokens", 0)
+            or 0
+        )
+
+        cache_read = 0
+        p_details = getattr(usage, "prompt_tokens_details", None) or getattr(
+            usage, "input_tokens_details", None
+        )
+        if p_details is not None:
+            cache_read = int(getattr(p_details, "cached_tokens", 0) or 0)
+
+        # Kimi-K2-thinking populate usage.cached_tokens field
+        if not cache_read and hasattr(usage, "cached_tokens"):
+            cache_read = int(getattr(usage, "cached_tokens", 0) or 0)
+
+        reasoning_tokens = 0
+        c_details = getattr(usage, "completion_tokens_details", None) or getattr(
+            usage, "output_tokens_details", None
+        )
+        if c_details is not None:
+            reasoning_tokens = int(getattr(c_details, "reasoning_tokens", 0) or 0)
+
+        # Chat-specific: litellm may set a hidden cache write field
+        cache_write = int(getattr(usage, "_cache_creation_input_tokens", 0) or 0)
+
+        self.metrics.add_token_usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            cache_read_tokens=cache_read,
+            cache_write_tokens=cache_write,
+            reasoning_tokens=reasoning_tokens,
+            context_window=context_window,
+            response_id=response_id,
+        )
+
+    def _compute_cost(self, resp: ModelResponse | ResponsesAPIResponse) -> float | None:
+        """Try provider header → litellm direct. Return None on failure."""
+        extra_kwargs = {}
+        if (
+            self.input_cost_per_token is not None
+            and self.output_cost_per_token is not None
+        ):
+            cost_per_token = CostPerToken(
+                input_cost_per_token=self.input_cost_per_token,
+                output_cost_per_token=self.output_cost_per_token,
+            )
+            logger.debug(f"Using custom cost per token: {cost_per_token}")
+            extra_kwargs["custom_cost_per_token"] = cost_per_token
+
+        try:
+            hidden = getattr(resp, "_hidden_params", {}) or {}
+            cost = hidden.get("additional_headers", {}).get(
+                "llm_provider-x-litellm-response-cost"
+            )
+            if cost is not None:
+                return float(cost)
+        except Exception as e:
+            logger.debug(f"Failed to get cost from LiteLLM headers: {e}")
+
+        # move on to litellm cost calculator
+        # Handle model name properly - if it doesn't contain "/", use as-is
+        model_parts = self.model_name.split("/")
+        if len(model_parts) > 1:
+            extra_kwargs["model"] = "/".join(model_parts[1:])
+        else:
+            extra_kwargs["model"] = self.model_name
+        try:
+            return float(
+                litellm_completion_cost(completion_response=resp, **extra_kwargs)
+            )
+        except Exception as e:
+            warnings.warn(f"Cost calculation failed: {e}")
+            return None
+
+    def log_llm_call(
+        self,
+        resp: ModelResponse | ResponsesAPIResponse,
+        cost: float | None,
+        raw_resp: ModelResponse | ResponsesAPIResponse | None = None,
+    ) -> None:
+        # Skip if neither file logging nor callback is configured
+        if not self.log_dir and not self._log_completions_callback:
+            return
+        try:
+            # Prepare filename and log data
+            filename = (
+                f"{self.model_name.replace('/', '__')}-"
+                f"{time.time():.3f}-"
+                f"{uuid.uuid4().hex[:4]}.json"
+            )
+
+            data = self._req_ctx.copy()
+            data["response"] = (
+                resp  # ModelResponse | ResponsesAPIResponse;
+                # serialized via _safe_json
+            )
+            data["cost"] = float(cost or 0.0)
+            data["timestamp"] = time.time()
+            data["latency_sec"] = self._last_latency
+
+            # Usage summary (prompt, completion, reasoning tokens) for quick inspection
+            try:
+                usage = getattr(resp, "usage", None)
+                if usage:
+                    prompt_tokens = int(
+                        getattr(usage, "prompt_tokens", None)
+                        or getattr(usage, "input_tokens", 0)
+                        or 0
+                    )
+                    completion_tokens = int(
+                        getattr(usage, "completion_tokens", None)
+                        or getattr(usage, "output_tokens", 0)
+                        or 0
+                    )
+                    details = getattr(
+                        usage, "completion_tokens_details", None
+                    ) or getattr(usage, "output_tokens_details", None)
+                    reasoning_tokens = (
+                        int(getattr(details, "reasoning_tokens", 0) or 0)
+                        if details
+                        else 0
+                    )
+                    p_details = getattr(
+                        usage, "prompt_tokens_details", None
+                    ) or getattr(usage, "input_tokens_details", None)
+                    cache_read_tokens = (
+                        int(getattr(p_details, "cached_tokens", 0) or 0)
+                        if p_details
+                        else 0
+                    )
+
+                    data["usage_summary"] = {
+                        "prompt_tokens": prompt_tokens,
+                        "completion_tokens": completion_tokens,
+                        "reasoning_tokens": reasoning_tokens,
+                        "cache_read_tokens": cache_read_tokens,
+                    }
+            except Exception:
+                # Best-effort only; don't fail logging
+                pass
+
+            # Raw response *before* nonfncall -> call conversion
+            if raw_resp:
+                data["raw_response"] = (
+                    raw_resp  # ModelResponse | ResponsesAPIResponse;
+                    # serialized via _safe_json
+                )
+            # Pop duplicated tools to avoid logging twice
+            if (
+                "tools" in data
+                and isinstance(data.get("kwargs"), dict)
+                and "tools" in data["kwargs"]
+            ):
+                data["kwargs"].pop("tools")
+
+            log_data = json.dumps(data, default=_safe_json, ensure_ascii=False)
+
+            # Use callback if set (for remote execution), otherwise write to file
+            if self._log_completions_callback:
+                self._log_completions_callback(filename, log_data)
+            elif self.log_dir:
+                # Create log directory if it doesn't exist
+                os.makedirs(self.log_dir, exist_ok=True)
+                if not os.access(self.log_dir, os.W_OK):
+                    raise PermissionError(f"log_dir is not writable: {self.log_dir}")
+
+                fname = os.path.join(self.log_dir, filename)
+                with open(fname, "w", encoding="utf-8") as f:
+                    f.write(log_data)
+        except Exception as e:
+            warnings.warn(f"Telemetry logging failed: {e}")
+
+
+def _safe_json(obj: Any) -> Any:
+    # Centralized serializer for telemetry logs.
+    # Prefer robust serialization for Pydantic models first to avoid cycles.
+    # Typed LiteLLM responses
+    if isinstance(obj, ModelResponse) or isinstance(obj, ResponsesAPIResponse):
+        return obj.model_dump(mode="json", exclude_none=True)
+
+    # Any Pydantic BaseModel (e.g., ToolDefinition, ChatCompletionToolParam, etc.)
+    if isinstance(obj, BaseModel):
+        # Use Pydantic's serializer which respects field exclusions (e.g., executors)
+        return obj.model_dump(mode="json", exclude_none=True)
+
+    # Fallbacks for other non-serializable objects used elsewhere in the log payload
+    try:
+        return obj.__dict__
+    except Exception:
+        return str(obj)

openhands/sdk/llm/utils/unverified_models.py

@@ -0,0 +1,156 @@
+import importlib
+
+import litellm
+from pydantic import SecretStr
+
+from openhands.sdk.llm.utils.verified_models import VERIFIED_MODELS
+from openhands.sdk.logger import get_logger
+
+
+def _get_boto3():
+    """Get boto3 module if available, otherwise return None."""
+    try:
+        return importlib.import_module("boto3")
+    except ModuleNotFoundError:
+        return None
+
+
+logger = get_logger(__name__)
+
+
+def _list_bedrock_foundation_models(
+    aws_region_name: str, aws_access_key_id: str, aws_secret_access_key: str
+) -> list[str]:
+    boto3 = _get_boto3()
+    if boto3 is None:
+        logger.warning(
+            "boto3 is not installed. To use Bedrock models,"
+            "install with: openhands-sdk[boto3]"
+        )
+        return []
+
+    try:
+        # The AWS bedrock model id is not queried, if no AWS parameters are configured.
+        client = boto3.client(
+            service_name="bedrock",
+            region_name=aws_region_name,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+        foundation_models_list = client.list_foundation_models(
+            byOutputModality="TEXT", byInferenceType="ON_DEMAND"
+        )
+        model_summaries = foundation_models_list["modelSummaries"]
+        return ["bedrock/" + model["modelId"] for model in model_summaries]
+    except Exception as err:
+        logger.warning(
+            "%s. Please config AWS_REGION_NAME AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY"
+            " if you want use bedrock model.",
+            err,
+        )
+        return []
+
+
+def get_supported_llm_models(
+    aws_region_name: str | None = None,
+    aws_access_key_id: SecretStr | None = None,
+    aws_secret_access_key: SecretStr | None = None,
+) -> list[str]:
+    """Get all models supported by LiteLLM.
+
+    This function combines models from litellm and Bedrock, removing any
+    error-prone Bedrock models.
+
+    Returns:
+        list[str]: A sorted list of unique model names.
+    """
+    litellm_model_list = litellm.model_list + list(litellm.model_cost.keys())
+    litellm_model_list_without_bedrock = list(
+        filter(lambda m: not m.startswith("bedrock"), litellm_model_list)
+    )
+    bedrock_model_list = []
+    if aws_region_name and aws_access_key_id and aws_secret_access_key:
+        bedrock_model_list = _list_bedrock_foundation_models(
+            aws_region_name,
+            aws_access_key_id.get_secret_value(),
+            aws_secret_access_key.get_secret_value(),
+        )
+    model_list = litellm_model_list_without_bedrock + bedrock_model_list
+    return model_list
+
+
+def _split_is_actually_version(split: list[str]) -> bool:
+    return (
+        len(split) > 1
+        and bool(split[1])
+        and bool(split[1][0])
+        and split[1][0].isdigit()
+    )
+
+
+def _extract_model_and_provider(model: str) -> tuple[str, str, str]:
+    """
+    Extract provider and model information from a model identifier.
+    """
+    separator = "/"
+    split = model.split(separator)
+
+    if len(split) == 1:
+        # no "/" separator found, try with "."
+        separator = "."
+        split = model.split(separator)
+        if _split_is_actually_version(split):
+            split = [separator.join(split)]  # undo the split
+
+    if len(split) == 1:
+        matched_provider = ""
+        for provider, models in VERIFIED_MODELS.items():
+            if split[0] in models:
+                matched_provider = provider
+                break
+
+        if matched_provider:
+            return matched_provider, split[0], "/"
+
+        return matched_provider, model, ""
+
+    provider = split[0]
+    model_id = separator.join(split[1:])
+    return provider, model_id, separator
+
+
+def get_unverified_models(
+    aws_region_name: str | None = None,
+    aws_access_key_id: SecretStr | None = None,
+    aws_secret_access_key: SecretStr | None = None,
+) -> dict[str, list[str]]:
+    """
+    Organize a mapping of unverified model identifiers by provider.
+    """
+    result_dict: dict[str, list[str]] = {}
+
+    models = get_supported_llm_models(
+        aws_region_name, aws_access_key_id, aws_secret_access_key
+    )
+    for model in models:
+        provider, model_id, separator = _extract_model_and_provider(model)
+
+        # Ignore "anthropic" providers with a separator of "."
+        # These are outdated and incompatible providers.
+        if provider == "anthropic" and separator == ".":
+            continue
+
+        # Dedup verified models
+        if provider in VERIFIED_MODELS and model_id in VERIFIED_MODELS[provider]:
+            continue
+
+        key = provider or "other"
+        if key not in result_dict:
+            result_dict[key] = []
+
+        result_dict[key].append(model_id)
+
+    return result_dict
+
+
+UNVERIFIED_MODELS_EXCLUDING_BEDROCK = get_unverified_models()