mini-swe-agent 1.16.0__py3-none-any.whl

minisweagent/models/__init__.py

@@ -0,0 +1,114 @@
+"""This file provides convenience functions for selecting models.
+You can ignore this file completely if you explicitly set your model in your run script.
+"""
+
+import copy
+import importlib
+import os
+import threading
+
+from minisweagent import Model
+
+
+class GlobalModelStats:
+    """Global model statistics tracker with optional limits."""
+
+    def __init__(self):
+        self._cost = 0.0
+        self._n_calls = 0
+        self._lock = threading.Lock()
+        self.cost_limit = float(os.getenv("MSWEA_GLOBAL_COST_LIMIT", "0"))
+        self.call_limit = int(os.getenv("MSWEA_GLOBAL_CALL_LIMIT", "0"))
+        if (self.cost_limit > 0 or self.call_limit > 0) and not os.getenv("MSWEA_SILENT_STARTUP"):
+            print(f"Global cost/call limit: ${self.cost_limit:.4f} / {self.call_limit}")
+
+    def add(self, cost: float) -> None:
+        """Add a model call with its cost, checking limits."""
+        with self._lock:
+            self._cost += cost
+            self._n_calls += 1
+        if 0 < self.cost_limit < self._cost or 0 < self.call_limit < self._n_calls + 1:
+            raise RuntimeError(f"Global cost/call limit exceeded: ${self._cost:.4f} / {self._n_calls + 1}")
+
+    @property
+    def cost(self) -> float:
+        return self._cost
+
+    @property
+    def n_calls(self) -> int:
+        return self._n_calls
+
+
+GLOBAL_MODEL_STATS = GlobalModelStats()
+
+
+def get_model(input_model_name: str | None = None, config: dict | None = None) -> Model:
+    """Get an initialized model object from any kind of user input or settings."""
+    resolved_model_name = get_model_name(input_model_name, config)
+    if config is None:
+        config = {}
+    config = copy.deepcopy(config)
+    config["model_name"] = resolved_model_name
+
+    model_class = get_model_class(resolved_model_name, config.pop("model_class", ""))
+
+    if (from_env := os.getenv("MSWEA_MODEL_API_KEY")) and not str(type(model_class)).endswith("DeterministicModel"):
+        config.setdefault("model_kwargs", {})["api_key"] = from_env
+
+    if (
+        any(s in resolved_model_name.lower() for s in ["anthropic", "sonnet", "opus", "claude"])
+        and "set_cache_control" not in config
+    ):
+        # Select cache control for Anthropic models by default
+        config["set_cache_control"] = "default_end"
+
+    return model_class(**config)
+
+
+def get_model_name(input_model_name: str | None = None, config: dict | None = None) -> str:
+    """Get a model name from any kind of user input or settings."""
+    if config is None:
+        config = {}
+    if input_model_name:
+        return input_model_name
+    if from_config := config.get("model_name"):
+        return from_config
+    if from_env := os.getenv("MSWEA_MODEL_NAME"):
+        return from_env
+    raise ValueError("No default model set. Please run `mini-extra config setup` to set one.")
+
+
+_MODEL_CLASS_MAPPING = {
+    "anthropic": "minisweagent.models.anthropic.AnthropicModel",
+    "litellm": "minisweagent.models.litellm_model.LitellmModel",
+    "litellm_response": "minisweagent.models.litellm_response_api_model.LitellmResponseAPIModel",
+    "openrouter": "minisweagent.models.openrouter_model.OpenRouterModel",
+    "portkey": "minisweagent.models.portkey_model.PortkeyModel",
+    "portkey_response": "minisweagent.models.portkey_response_api_model.PortkeyResponseAPIModel",
+    "requesty": "minisweagent.models.requesty_model.RequestyModel",
+    "deterministic": "minisweagent.models.test_models.DeterministicModel",
+}
+
+
+def get_model_class(model_name: str, model_class: str = "") -> type:
+    """Select the best model class.
+
+    If a model_class is provided (as shortcut name, or as full import path,
+    e.g., "anthropic" or "minisweagent.models.anthropic.AnthropicModel"),
+    it takes precedence over the `model_name`.
+    Otherwise, the model_name is used to select the best model class.
+    """
+    if model_class:
+        full_path = _MODEL_CLASS_MAPPING.get(model_class, model_class)
+        try:
+            module_name, class_name = full_path.rsplit(".", 1)
+            module = importlib.import_module(module_name)
+            return getattr(module, class_name)
+        except (ValueError, ImportError, AttributeError):
+            msg = f"Unknown model class: {model_class} (resolved to {full_path}, available: {_MODEL_CLASS_MAPPING})"
+            raise ValueError(msg)
+
+    # Default to LitellmModel
+    from minisweagent.models.litellm_model import LitellmModel
+
+    return LitellmModel

minisweagent/models/anthropic.py

@@ -0,0 +1,35 @@
+import os
+import warnings
+from typing import Literal
+
+from minisweagent.models.litellm_model import LitellmModel, LitellmModelConfig
+from minisweagent.models.utils.cache_control import set_cache_control
+from minisweagent.models.utils.key_per_thread import get_key_per_thread
+
+
+class AnthropicModelConfig(LitellmModelConfig):
+    set_cache_control: Literal["default_end"] | None = "default_end"
+    """Set explicit cache control markers, for example for Anthropic models"""
+
+
+class AnthropicModel(LitellmModel):
+    """This class is now only a thin wrapper around the LitellmModel class.
+    It is largely kept for backwards compatibility.
+    It will not be selected by `get_model` and `get_model_class` unless explicitly specified.
+    """
+
+    def __init__(self, *, config_class: type = AnthropicModelConfig, **kwargs):
+        super().__init__(config_class=config_class, **kwargs)
+
+    def query(self, messages: list[dict], **kwargs) -> dict:
+        api_key = None
+        # Legacy only
+        if rotating_keys := os.getenv("ANTHROPIC_API_KEYS"):
+            warnings.warn(
+                "ANTHROPIC_API_KEYS is deprecated and will be removed in the future. "
+                "Simply use the ANTHROPIC_API_KEY environment variable instead. "
+                "Key rotation is no longer required."
+            )
+            api_key = get_key_per_thread(rotating_keys.split("::"))
+        messages = set_cache_control(messages, mode="default_end")
+        return super().query(messages, api_key=api_key, **kwargs)

minisweagent/models/extra/roulette.py

@@ -0,0 +1,61 @@
+import random
+from dataclasses import asdict, dataclass
+
+from minisweagent import Model
+from minisweagent.models import get_model
+
+
+@dataclass
+class RouletteModelConfig:
+    model_kwargs: list[dict]
+    """The models to choose from"""
+    model_name: str = "roulette"
+
+
+class RouletteModel:
+    def __init__(self, *, config_class: type = RouletteModelConfig, **kwargs):
+        """This "meta"-model randomly selects one of the models at every call"""
+        self.config = config_class(**kwargs)
+        self.models = [get_model(config=config) for config in self.config.model_kwargs]
+
+    @property
+    def cost(self) -> float:
+        return sum(model.cost for model in self.models)
+
+    @property
+    def n_calls(self) -> int:
+        return sum(model.n_calls for model in self.models)
+
+    def get_template_vars(self) -> dict:
+        return asdict(self.config) | {"n_model_calls": self.n_calls, "model_cost": self.cost}
+
+    def select_model(self) -> Model:
+        return random.choice(self.models)
+
+    def query(self, *args, **kwargs) -> dict:
+        model = self.select_model()
+        response = model.query(*args, **kwargs)
+        response["model_name"] = model.config.model_name
+        return response
+
+
+@dataclass
+class InterleavingModelConfig:
+    model_kwargs: list[dict]
+    sequence: list[int] | None = None
+    """If set to 0, 0, 1, we will return the first model 2 times, then the second model 1 time,
+    then the first model again, etc."""
+    model_name: str = "interleaving"
+
+
+class InterleavingModel(RouletteModel):
+    def __init__(self, *, config_class: type = InterleavingModelConfig, **kwargs):
+        """This "meta"-model alternates between the models in the sequence for every call"""
+        super().__init__(config_class=config_class, **kwargs)
+
+    def select_model(self) -> Model:
+        if self.config.sequence is None:
+            i_model = self.n_calls % len(self.models)
+        else:
+            i_model = self.config.sequence[self.n_calls % len(self.config.sequence)]
+        return self.models[i_model]

minisweagent/models/litellm_model.py

@@ -0,0 +1,100 @@
+import json
+import logging
+import os
+from collections.abc import Callable
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+import litellm
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_not_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from minisweagent.models import GLOBAL_MODEL_STATS
+from minisweagent.models.utils.cache_control import set_cache_control
+
+logger = logging.getLogger("litellm_model")
+
+
+@dataclass
+class LitellmModelConfig:
+    model_name: str
+    model_kwargs: dict[str, Any] = field(default_factory=dict)
+    litellm_model_registry: Path | str | None = os.getenv("LITELLM_MODEL_REGISTRY_PATH")
+    set_cache_control: Literal["default_end"] | None = None
+    """Set explicit cache control markers, for example for Anthropic models"""
+    cost_tracking: Literal["default", "ignore_errors"] = os.getenv("MSWEA_COST_TRACKING", "default")
+    """Cost tracking mode for this model. Can be "default" or "ignore_errors" (ignore errors/missing cost info)"""
+
+
+class LitellmModel:
+    def __init__(self, *, config_class: Callable = LitellmModelConfig, **kwargs):
+        self.config = config_class(**kwargs)
+        self.cost = 0.0
+        self.n_calls = 0
+        if self.config.litellm_model_registry and Path(self.config.litellm_model_registry).is_file():
+            litellm.utils.register_model(json.loads(Path(self.config.litellm_model_registry).read_text()))
+
+    @retry(
+        stop=stop_after_attempt(int(os.getenv("MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT", "10"))),
+        wait=wait_exponential(multiplier=1, min=4, max=60),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        retry=retry_if_not_exception_type(
+            (
+                litellm.exceptions.UnsupportedParamsError,
+                litellm.exceptions.NotFoundError,
+                litellm.exceptions.PermissionDeniedError,
+                litellm.exceptions.ContextWindowExceededError,
+                litellm.exceptions.APIError,
+                litellm.exceptions.AuthenticationError,
+                KeyboardInterrupt,
+            )
+        ),
+    )
+    def _query(self, messages: list[dict[str, str]], **kwargs):
+        try:
+            return litellm.completion(
+                model=self.config.model_name, messages=messages, **(self.config.model_kwargs | kwargs)
+            )
+        except litellm.exceptions.AuthenticationError as e:
+            e.message += " You can permanently set your API key with `mini-extra config set KEY VALUE`."
+            raise e
+
+    def query(self, messages: list[dict[str, str]], **kwargs) -> dict:
+        if self.config.set_cache_control:
+            messages = set_cache_control(messages, mode=self.config.set_cache_control)
+        response = self._query(messages, **kwargs)
+        try:
+            cost = litellm.cost_calculator.completion_cost(response)
+            if cost <= 0.0:
+                raise ValueError(f"Cost must be > 0.0, got {cost}")
+        except Exception as e:
+            cost = 0.0
+            if self.config.cost_tracking != "ignore_errors":
+                msg = (
+                    f"Error calculating cost for model {self.config.model_name}: {e}, perhaps it's not registered? "
+                    "You can ignore this issue from your config file with cost_tracking: 'ignore_errors' or "
+                    "globally with export MSWEA_COST_TRACKING='ignore_errors'. "
+                    "Alternatively check the 'Cost tracking' section in the documentation at "
+                    "https://klieret.short.gy/mini-local-models. "
+                    " Still stuck? Please open a github issue at https://github.com/SWE-agent/mini-swe-agent/issues/new/choose!"
+                )
+                logger.critical(msg)
+                raise RuntimeError(msg) from e
+        self.n_calls += 1
+        self.cost += cost
+        GLOBAL_MODEL_STATS.add(cost)
+        return {
+            "content": response.choices[0].message.content or "",  # type: ignore
+            "extra": {
+                "response": response.model_dump(),
+            },
+        }
+
+    def get_template_vars(self) -> dict[str, Any]:
+        return asdict(self.config) | {"n_model_calls": self.n_calls, "model_cost": self.cost}

minisweagent/models/litellm_response_api_model.py

@@ -0,0 +1,80 @@
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import litellm
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_not_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from minisweagent.models.litellm_model import LitellmModel, LitellmModelConfig
+from minisweagent.models.utils.openai_utils import coerce_responses_text
+
+logger = logging.getLogger("litellm_response_api_model")
+
+
+@dataclass
+class LitellmResponseAPIModelConfig(LitellmModelConfig):
+    pass
+
+
+class LitellmResponseAPIModel(LitellmModel):
+    def __init__(self, *, config_class: Callable = LitellmResponseAPIModelConfig, **kwargs):
+        super().__init__(config_class=config_class, **kwargs)
+        self._previous_response_id: str | None = None
+
+    @retry(
+        stop=stop_after_attempt(10),
+        wait=wait_exponential(multiplier=1, min=4, max=60),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        retry=retry_if_not_exception_type(
+            (
+                litellm.exceptions.UnsupportedParamsError,
+                litellm.exceptions.NotFoundError,
+                litellm.exceptions.PermissionDeniedError,
+                litellm.exceptions.ContextWindowExceededError,
+                litellm.exceptions.APIError,
+                litellm.exceptions.AuthenticationError,
+                KeyboardInterrupt,
+            )
+        ),
+    )
+    def _query(self, messages: list[dict[str, str]], **kwargs):
+        try:
+            resp = litellm.responses(
+                model=self.config.model_name,
+                input=messages if self._previous_response_id is None else messages[-1:],
+                previous_response_id=self._previous_response_id,
+                **(self.config.model_kwargs | kwargs),
+            )
+            self._previous_response_id = getattr(resp, "id", None)
+            return resp
+        except litellm.exceptions.AuthenticationError as e:
+            e.message += " You can permanently set your API key with `mini-extra config set KEY VALUE`."
+            raise e
+
+    def query(self, messages: list[dict[str, str]], **kwargs) -> dict:
+        response = self._query(messages, **kwargs)
+        print(response)
+        text = coerce_responses_text(response)
+        try:
+            cost = litellm.cost_calculator.completion_cost(response)
+        except Exception as e:
+            logger.critical(
+                f"Error calculating cost for model {self.config.model_name}: {e}. "
+                "Please check the 'Updating the model registry' section in the documentation. "
+                "http://bit.ly/4p31bi4 Still stuck? Please open a github issue for help!"
+            )
+            raise
+        self.n_calls += 1
+        self.cost += cost
+        from minisweagent.models import GLOBAL_MODEL_STATS
+
+        GLOBAL_MODEL_STATS.add(cost)
+        return {
+            "content": text,
+        }

minisweagent/models/openrouter_model.py

@@ -0,0 +1,125 @@
+import json
+import logging
+import os
+from dataclasses import asdict, dataclass, field
+from typing import Any, Literal
+
+import requests
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_not_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from minisweagent.models import GLOBAL_MODEL_STATS
+from minisweagent.models.utils.cache_control import set_cache_control
+
+logger = logging.getLogger("openrouter_model")
+
+
+@dataclass
+class OpenRouterModelConfig:
+    model_name: str
+    model_kwargs: dict[str, Any] = field(default_factory=dict)
+    set_cache_control: Literal["default_end"] | None = None
+    """Set explicit cache control markers, for example for Anthropic models"""
+    cost_tracking: Literal["default", "ignore_errors"] = os.getenv("MSWEA_COST_TRACKING", "default")
+    """Cost tracking mode for this model. Can be "default" or "ignore_errors" (ignore errors/missing cost info)"""
+
+
+class OpenRouterAPIError(Exception):
+    """Custom exception for OpenRouter API errors."""
+
+    pass
+
+
+class OpenRouterAuthenticationError(Exception):
+    """Custom exception for OpenRouter authentication errors."""
+
+    pass
+
+
+class OpenRouterRateLimitError(Exception):
+    """Custom exception for OpenRouter rate limit errors."""
+
+    pass
+
+
+class OpenRouterModel:
+    def __init__(self, **kwargs):
+        self.config = OpenRouterModelConfig(**kwargs)
+        self.cost = 0.0
+        self.n_calls = 0
+        self._api_url = "https://openrouter.ai/api/v1/chat/completions"
+        self._api_key = os.getenv("OPENROUTER_API_KEY", "")
+
+    @retry(
+        stop=stop_after_attempt(int(os.getenv("MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT", "10"))),
+        wait=wait_exponential(multiplier=1, min=4, max=60),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        retry=retry_if_not_exception_type(
+            (
+                OpenRouterAuthenticationError,
+                KeyboardInterrupt,
+            )
+        ),
+    )
+    def _query(self, messages: list[dict[str, str]], **kwargs):
+        headers = {
+            "Authorization": f"Bearer {self._api_key}",
+            "Content-Type": "application/json",
+        }
+
+        payload = {
+            "model": self.config.model_name,
+            "messages": messages,
+            "usage": {"include": True},
+            **(self.config.model_kwargs | kwargs),
+        }
+
+        try:
+            response = requests.post(self._api_url, headers=headers, data=json.dumps(payload), timeout=60)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.HTTPError as e:
+            if response.status_code == 401:
+                error_msg = "Authentication failed. You can permanently set your API key with `mini-extra config set OPENROUTER_API_KEY YOUR_KEY`."
+                raise OpenRouterAuthenticationError(error_msg) from e
+            elif response.status_code == 429:
+                raise OpenRouterRateLimitError("Rate limit exceeded") from e
+            else:
+                raise OpenRouterAPIError(f"HTTP {response.status_code}: {response.text}") from e
+        except requests.exceptions.RequestException as e:
+            raise OpenRouterAPIError(f"Request failed: {e}") from e
+
+    def query(self, messages: list[dict[str, str]], **kwargs) -> dict:
+        if self.config.set_cache_control:
+            messages = set_cache_control(messages, mode=self.config.set_cache_control)
+        response = self._query(messages, **kwargs)
+
+        usage = response.get("usage", {})
+        cost = usage.get("cost", 0.0)
+        if cost <= 0.0 and self.config.cost_tracking != "ignore_errors":
+            raise RuntimeError(
+                f"No valid cost information available from OpenRouter API for model {self.config.model_name}: "
+                f"Usage {usage}, cost {cost}. Cost must be > 0.0. Set cost_tracking: 'ignore_errors' in your config file or "
+                "export MSWEA_COST_TRACKING='ignore_errors' to ignore cost tracking errors "
+                "(for example for free/local models), more information at https://klieret.short.gy/mini-local-models "
+                "for more details. Still stuck? Please open a github issue at https://github.com/SWE-agent/mini-swe-agent/issues/new/choose!"
+            )
+
+        self.n_calls += 1
+        self.cost += cost
+        GLOBAL_MODEL_STATS.add(cost)
+
+        return {
+            "content": response["choices"][0]["message"]["content"] or "",
+            "extra": {
+                "response": response,  # already is json
+            },
+        }
+
+    def get_template_vars(self) -> dict[str, Any]:
+        return asdict(self.config) | {"n_model_calls": self.n_calls, "model_cost": self.cost}

minisweagent/models/portkey_model.py

@@ -0,0 +1,154 @@
+import json
+import logging
+import os
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+import litellm
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_not_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from minisweagent.models import GLOBAL_MODEL_STATS
+from minisweagent.models.utils.cache_control import set_cache_control
+
+logger = logging.getLogger("portkey_model")
+
+try:
+    from portkey_ai import Portkey
+except ImportError:
+    raise ImportError(
+        "The portkey-ai package is required to use PortkeyModel. Please install it with: pip install portkey-ai"
+    )
+
+
+@dataclass
+class PortkeyModelConfig:
+    model_name: str
+    model_kwargs: dict[str, Any] = field(default_factory=dict)
+    litellm_model_registry: Path | str | None = os.getenv("LITELLM_MODEL_REGISTRY_PATH")
+    """We currently use litellm to calculate costs. Here you can register additional models to litellm's model registry.
+    Note that this might change if we get better support for Portkey and change how we calculate costs.
+    """
+    litellm_model_name_override: str = ""
+    """We currently use litellm to calculate costs. Here you can override the model name to use for litellm in case it
+    doesn't match the Portkey model name.
+    Note that this might change if we get better support for Portkey and change how we calculate costs.
+    """
+    set_cache_control: Literal["default_end"] | None = None
+    """Set explicit cache control markers, for example for Anthropic models"""
+    cost_tracking: Literal["default", "ignore_errors"] = os.getenv("MSWEA_COST_TRACKING", "default")
+    """Cost tracking mode for this model. Can be "default" or "ignore_errors" (ignore errors/missing cost info)"""
+
+
+class PortkeyModel:
+    def __init__(self, *, config_class: type = PortkeyModelConfig, **kwargs):
+        self.config = config_class(**kwargs)
+        self.cost = 0.0
+        self.n_calls = 0
+        if self.config.litellm_model_registry and Path(self.config.litellm_model_registry).is_file():
+            litellm.utils.register_model(json.loads(Path(self.config.litellm_model_registry).read_text()))
+
+        # Get API key from environment or raise error
+        self._api_key = os.getenv("PORTKEY_API_KEY")
+        if not self._api_key:
+            raise ValueError(
+                "Portkey API key is required. Set it via the "
+                "PORTKEY_API_KEY environment variable. You can permanently set it with "
+                "`mini-extra config set PORTKEY_API_KEY YOUR_KEY`."
+            )
+
+        # Get virtual key from environment
+        virtual_key = os.getenv("PORTKEY_VIRTUAL_KEY")
+
+        # Initialize Portkey client
+        client_kwargs = {"api_key": self._api_key}
+        if virtual_key:
+            client_kwargs["virtual_key"] = virtual_key
+
+        self.client = Portkey(**client_kwargs)
+
+    @retry(
+        stop=stop_after_attempt(int(os.getenv("MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT", "10"))),
+        wait=wait_exponential(multiplier=1, min=4, max=60),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        retry=retry_if_not_exception_type((KeyboardInterrupt, TypeError, ValueError)),
+    )
+    def _query(self, messages: list[dict[str, str]], **kwargs):
+        # return self.client.with_options(metadata={"request_id": request_id}).chat.completions.create(
+        return self.client.chat.completions.create(
+            model=self.config.model_name,
+            messages=messages,
+            **(self.config.model_kwargs | kwargs),
+        )
+
+    def query(self, messages: list[dict[str, str]], **kwargs) -> dict:
+        if self.config.set_cache_control:
+            messages = set_cache_control(messages, mode=self.config.set_cache_control)
+        response = self._query(messages, **kwargs)
+        cost = self._calculate_cost(response)
+        self.n_calls += 1
+        self.cost += cost
+        GLOBAL_MODEL_STATS.add(cost)
+        return {
+            "content": response.choices[0].message.content or "",
+            "extra": {
+                "response": response.model_dump(),
+                "cost": cost,
+            },
+        }
+
+    def get_template_vars(self) -> dict[str, Any]:
+        return asdict(self.config) | {"n_model_calls": self.n_calls, "model_cost": self.cost}
+
+    def _calculate_cost(self, response) -> float:
+        response_for_cost_calc = response.model_copy()
+        if self.config.litellm_model_name_override:
+            if response_for_cost_calc.model:
+                response_for_cost_calc.model = self.config.litellm_model_name_override
+        prompt_tokens = response_for_cost_calc.usage.prompt_tokens
+        if prompt_tokens is None:
+            logger.warning(
+                f"Prompt tokens are None for model {self.config.model_name}. Setting to 0. Full response: {response_for_cost_calc.model_dump()}"
+            )
+            prompt_tokens = 0
+        total_tokens = response_for_cost_calc.usage.total_tokens
+        completion_tokens = response_for_cost_calc.usage.completion_tokens
+        if completion_tokens is None:
+            logger.warning(
+                f"Completion tokens are None for model {self.config.model_name}. Setting to 0. Full response: {response_for_cost_calc.model_dump()}"
+            )
+            completion_tokens = 0
+        if total_tokens - prompt_tokens - completion_tokens != 0:
+            # This is most likely related to how portkey treats cached tokens: It doesn't count them towards the prompt tokens (?)
+            logger.warning(
+                f"WARNING: Total tokens - prompt tokens - completion tokens != 0: {response_for_cost_calc.model_dump()}."
+                " This is probably a portkey bug or incompatibility with litellm cost tracking. "
+                "Setting prompt tokens based on total tokens and completion tokens. You might want to double check your costs. "
+                f"Full response: {response_for_cost_calc.model_dump()}"
+            )
+            response_for_cost_calc.usage.prompt_tokens = total_tokens - completion_tokens
+        try:
+            cost = litellm.cost_calculator.completion_cost(
+                response_for_cost_calc, model=self.config.litellm_model_name_override or None
+            )
+            assert cost >= 0.0, f"Cost is negative: {cost}"
+        except Exception as e:
+            cost = 0.0
+            if self.config.cost_tracking != "ignore_errors":
+                msg = (
+                    f"Error calculating cost for model {self.config.model_name} based on {response_for_cost_calc.model_dump()}: {e}. "
+                    "You can ignore this issue from your config file with cost_tracking: 'ignore_errors' or "
+                    "globally with export MSWEA_COST_TRACKING='ignore_errors' to ignore this error. "
+                    "Alternatively check the 'Cost tracking' section in the documentation at "
+                    "https://klieret.short.gy/mini-local-models. "
+                    "Still stuck? Please open a github issue at https://github.com/SWE-agent/mini-swe-agent/issues/new/choose!"
+                )
+                logger.critical(msg)
+                raise RuntimeError(msg) from e
+        return cost