diff --git a/docs/docs/usage-guide/additional_configurations.md b/docs/docs/usage-guide/additional_configurations.md index c93de699f9..415ac3d945 100644 --- a/docs/docs/usage-guide/additional_configurations.md +++ b/docs/docs/usage-guide/additional_configurations.md @@ -30,6 +30,14 @@ To see which model actually answered, how many tokens the run consumed, and how /review --config.output_run_details=true ``` +API-cost collection is a separate, default-off option controlled by `config.output_run_cost`. Enable both flags to collect it and add it inside the run-details section: + +``` +/review --config.output_run_details=true --config.output_run_cost=true +``` + +`config.output_run_details` remains the public-output gate: setting only `config.output_run_cost=true` collects run-level cost data but never adds it to a PR comment. + On providers that support GitHub-Flavored Markdown this appends a collapsible section to the generated comment; elsewhere `/review` and `/describe` append the same information as plain text: ``` @@ -38,10 +46,17 @@ On providers that support GitHub-Flavored Markdown this appends a collapsible se - Tokens: 12,340 in / 1,205 out / 13,545 total - Time cost: 8.2s - AI calls: 1 +- Estimated API cost: $0.08 USD + - anthropic/claude-opus-5: $0.07 USD + - anthropic/claude-sonnet-5: $0.01 USD ``` `Model` shows the model that produced the answer, marked `(fallback)` when the primary model failed and a fallback took over. The `Tokens` line appears only when the model provider reports usage. `AI calls` counts the successful LLM invocations made during the run. The flag is disabled by default. +`Estimated API cost` is derived synchronously from each completed LiteLLM response and its finalized usage. LiteLLM can account for cache reads, cache writes, reasoning tokens, and provider-specific usage categories when the response and its pricing data include them. Multi-model runs show a compact breakdown of the known costs. Exact `Decimal` values are retained for aggregation, while public currency output is rounded to two decimal places; a tiny positive value that would round to zero is shown as `<$0.01` instead of `$0.00`. If only some successful calls can be priced, the total is marked `partial` with the priced-call count; if none can be priced, the line reports `unavailable`. Missing pricing is never rendered as `$0`. + +The amount is an estimate based on LiteLLM's pricing data, not provider-invoice-authoritative billing. Reconcile it with the provider's billing records before using it for accounting or chargeback. Streaming responses are priced only after finalized usage is available; asynchronous callbacks and transient `response_cost` callback metadata are not treated as the sole source of truth. The public section contains only aggregate costs and configured model names, never prompts, response bodies, API keys, or provider request IDs. + Notes: - `/improve` appends the section only when it publishes a summary comment. If the provider lacks GFM support or `pr_code_suggestions.commitable_code_suggestions` is enabled, `/improve` posts inline comments instead, so no run details section appears. diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index 56ab0e7212..5230581e51 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -18,9 +18,9 @@ USER_MESSAGE_ONLY_MODELS) from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.litellm_helpers import ( - MockResponse, _get_azure_ad_token, _handle_streaming_response, - _process_litellm_extra_body) -from pr_agent.algo.run_details import record_ai_call + _get_azure_ad_token, _handle_streaming_response, + _process_litellm_extra_body, _response_field) +from pr_agent.algo.run_details import _as_decimal_cost, record_ai_call from pr_agent.algo.utils import ReasoningEffort, get_version from pr_agent.config_loader import get_settings from pr_agent.log import get_logger @@ -377,12 +377,73 @@ def prepare_logs(self, response, system, user, resp, finish_reason): return response_log @staticmethod - def _record_completion_metadata(response) -> None: - """Count the call and accumulate token usage when the provider reports it. + def _record_completion_metadata(response, model=None, display_model=None) -> None: + """Count a successful call and synchronously collect usage-based cost when possible.""" + usage = _response_field(response, "usage") + + cost_usd = None + if get_settings().get("config.output_run_cost", False): + # The guard covers the whole cost block, not just completion_cost: + # reading inline costs and probing usage call model_dump() on + # provider-specific objects, and a cost estimate must never fail a + # call that already succeeded and was billed. + try: + cost_usd = LiteLLMAIHandler._read_positive_response_cost(response, usage) + if cost_usd is None and model and LiteLLMAIHandler._has_priceable_usage(usage): + # Preserve LiteLLM's full usage object so completion_cost can price cache, + # reasoning, and provider-specific categories. Convert the small completed + # stream wrapper to a dictionary while retaining `response.usage`. + cost_response = response + if not isinstance(response, dict) and not hasattr(response, "model_dump"): + cost_response = response.dict() + cost_usd = litellm.completion_cost(completion_response=cost_response, model=model) + except Exception as e: + # Treat missing model pricing or insufficient usage as an unavailable call cost. + # Retain the successful call so the collector marks the aggregate safely. + get_logger().debug(f"Unable to estimate API cost for model {model}: {type(e).__name__}") + + recorded_model = display_model if display_model is not None else model + record_ai_call(usage, model=recorded_model, cost_usd=cost_usd) - Streaming models return a MockResponse without `usage`, so tokens stay unset. + @staticmethod + def _read_positive_response_cost(response, usage): + """Read a finalized inline cost, rejecting zero placeholders and invalid values.""" + candidates = [ + _response_field(usage, "response_cost"), + _response_field(usage, "cost"), + ] + + hidden_params = _response_field(response, "_hidden_params") + if hasattr(hidden_params, "model_dump"): + hidden_params = hidden_params.model_dump() + if isinstance(hidden_params, dict): + candidates.append(hidden_params.get("response_cost")) + + for candidate in candidates: + decimal_cost = _as_decimal_cost(candidate) + if decimal_cost is not None: + return decimal_cost + return None + + @staticmethod + def _has_priceable_usage(usage) -> bool: + """Return true when finalized usage reports a positive token count. + + Only token counters gate pricing: provider extras such as Groq's timing + floats (queue_time, prompt_time) are not billable quantities, and letting + them pass would send zero-token usage to completion_cost, which prices + it as 0.0 instead of raising. """ - record_ai_call(getattr(response, "usage", None)) + if usage is None: + return False + return any( + isinstance(count, int) and not isinstance(count, bool) and count > 0 + for count in ( + _response_field(usage, "prompt_tokens"), + _response_field(usage, "completion_tokens"), + _response_field(usage, "total_tokens"), + ) + ) def _configure_claude_extended_thinking(self, model: str, kwargs: dict) -> dict: """ @@ -913,19 +974,23 @@ def _as_int(value): body=None, ) from e - get_logger().debug(f"\nAI response:\n{resp}") + # Post-response bookkeeping happens outside the Bedrock IMDS lock above: it + # touches no os.environ credentials, and in IMDS mode the lock serializes + # every concurrent call, so holding it through logging and cost pricing + # would make each waiting coroutine pay for them serially. + get_logger().debug(f"\nAI response:\n{resp}") - # log the full response for debugging - response_log = self.prepare_logs(response_obj, system, user, resp, finish_reason) - get_logger().debug("Full_response", artifact=response_log) + # log the full response for debugging + response_log = self.prepare_logs(response_obj, system, user, resp, finish_reason) + get_logger().debug("Full_response", artifact=response_log) - # for CLI debugging - if get_settings().config.verbosity_level >= 2: - get_logger().info(f"\nAI response:\n{resp}") + # for CLI debugging + if get_settings().config.verbosity_level >= 2: + get_logger().info(f"\nAI response:\n{resp}") - self._record_completion_metadata(response_obj) + self._record_completion_metadata(response_obj, model=model, display_model=user_model) - return resp, finish_reason + return resp, finish_reason async def _get_completion(self, **kwargs): """ @@ -947,6 +1012,7 @@ async def _get_completion(self, **kwargs): # response normalization. Streaming avoids that conversion path. if model in self.streaming_required_models or force_streaming: kwargs["stream"] = True + kwargs["stream_options"] = {"include_usage": True} if force_streaming and model not in self.streaming_required_models: get_logger().info( f"Using streaming mode for model {model} " @@ -955,10 +1021,7 @@ async def _get_completion(self, **kwargs): else: get_logger().info(f"Using streaming mode for model {model}") response = await acompletion(**kwargs) - resp, finish_reason = await _handle_streaming_response(response) - # Create MockResponse for streaming since we don't have the full response object - mock_response = MockResponse(resp, finish_reason) - return resp, finish_reason, mock_response + return await _handle_streaming_response(response, model=model) else: response = await acompletion(**kwargs) if response is None or len(response["choices"]) == 0: diff --git a/pr_agent/algo/ai_handlers/litellm_helpers.py b/pr_agent/algo/ai_handlers/litellm_helpers.py index 483fb81883..dad733413b 100644 --- a/pr_agent/algo/ai_handlers/litellm_helpers.py +++ b/pr_agent/algo/ai_handlers/litellm_helpers.py @@ -43,7 +43,24 @@ } -async def _handle_streaming_response(response): +def _response_field(response, name): + if isinstance(response, dict): + return response.get(name) + return getattr(response, name, None) + + +def _stream_usage(chunk): + """Read finalized usage from a regular or metadata-only LiteLLM chunk.""" + usage = _response_field(chunk, "usage") + if usage is not None: + return usage + hidden_params = _response_field(chunk, "_hidden_params") + if isinstance(hidden_params, dict): + return hidden_params.get("usage") + return None + + +async def _handle_streaming_response(response, model=None): """ Handle streaming response from acompletion and collect the full response. @@ -51,13 +68,17 @@ async def _handle_streaming_response(response): response: The streaming response object from acompletion Returns: - tuple: (full_response_content, finish_reason) + tuple: (full_response_content, finish_reason, completed_response) """ full_response = "" finish_reason = None + finalized_usage = None try: async for chunk in response: + usage = _stream_usage(chunk) + if usage is not None: + finalized_usage = usage if chunk.choices and len(chunk.choices) > 0: choice = chunk.choices[0] delta = choice.delta @@ -76,13 +97,14 @@ async def _handle_streaming_response(response): elif not full_response and finish_reason: get_logger().debug(f"Streaming response resulted in empty content but completed with finish_reason: {finish_reason}") raise openai.APIError(f"Streaming response completed with finish_reason '{finish_reason}' but no content received") - return full_response, finish_reason + return full_response, finish_reason, MockResponse(full_response, finish_reason, finalized_usage, model) class MockResponse: - """Mock response object for streaming models to enable consistent logging.""" + """Represent a completed streaming response while retaining LiteLLM's finalized usage object.""" - def __init__(self, resp, finish_reason): + def __init__(self, resp, finish_reason, usage=None, model=None): + self.usage = usage self._data = { "choices": [ { @@ -91,9 +113,19 @@ def __init__(self, resp, finish_reason): } ] } + if model is not None: + self._data["model"] = model def dict(self): - return self._data + data = self._data.copy() + if self.usage is not None: + if hasattr(self.usage, "model_dump"): + data["usage"] = self.usage.model_dump() + elif isinstance(self.usage, dict): + data["usage"] = self.usage.copy() + else: + data["usage"] = vars(self.usage).copy() + return data def _get_azure_ad_token(): diff --git a/pr_agent/algo/ai_handlers/openai_ai_handler.py b/pr_agent/algo/ai_handlers/openai_ai_handler.py index 6850f01796..830b0e5dce 100644 --- a/pr_agent/algo/ai_handlers/openai_ai_handler.py +++ b/pr_agent/algo/ai_handlers/openai_ai_handler.py @@ -61,6 +61,10 @@ async def chat_completion(self, model: str, system: str, user: str, temperature: usage = chat_completion.usage get_logger().info("AI response", response=resp, messages=messages, finish_reason=finish_reason, model=model, usage=usage) + # Count the call and its tokens but no cost: this handler has no pricing + # source wired up, so with output_run_cost enabled its calls render as + # unpriced. Left as-is while the path stays cold — no setting selects + # this handler, it is reachable only by injecting it programmatically. record_ai_call(usage) return resp, finish_reason except openai.RateLimitError as e: diff --git a/pr_agent/algo/run_details.py b/pr_agent/algo/run_details.py index 78f5bd072e..aaafa062e6 100644 --- a/pr_agent/algo/run_details.py +++ b/pr_agent/algo/run_details.py @@ -10,6 +10,7 @@ import time from contextvars import ContextVar from dataclasses import dataclass, field +from decimal import Decimal, InvalidOperation from typing import Optional _run_details: ContextVar[Optional["RunDetails"]] = ContextVar( @@ -44,6 +45,13 @@ class RunDetails: total_tokens: int = 0 # Successful LLM invocations, counted even when their token usage is unavailable. num_ai_calls: int = 0 + # Accumulate costs only when cost output is enabled and LiteLLM can synchronously + # price a successful response with a positive amount. Use the known-call count to + # distinguish priced calls from missing pricing data. Retain per-model totals to + # keep fallback and multi-call runs auditable. + total_cost_usd: Decimal = field(default_factory=lambda: Decimal("0")) + known_cost_call_count: int = 0 + model_costs_usd: dict[str, Decimal] = field(default_factory=dict) # Monotonic reference taken when the collector is installed, i.e. at the top of the # tool's run(). Monotonic so that wall-clock adjustments cannot yield a negative duration. start_time: float = field(default_factory=time.monotonic) @@ -60,6 +68,15 @@ def has_token_usage(self) -> bool: or self.completion_tokens > 0 ) + @property + def cost_status(self) -> str: + """Return whether every, some, or none of the successful calls were priced.""" + if self.known_cost_call_count == 0: + return "unavailable" + if self.known_cost_call_count == self.num_ai_calls: + return "complete" + return "partial" + def init_run_details() -> RunDetails: """Install a fresh collector for the current run and return it.""" @@ -107,11 +124,36 @@ def add_token_usage(usage) -> None: details.total_tokens += total_tokens -def record_ai_call(usage=None) -> None: - """Count one AI call and accumulate token usage when available.""" +def _as_decimal_cost(cost_usd) -> Optional[Decimal]: + """Normalize a positive finite USD value without introducing float math. + + Zero is rejected on purpose: litellm.completion_cost returns 0.0 both for + zero-priced model entries (e.g. local/ollama models) and for usage without + billable tokens, so a zero here means "could not be priced", not "free" — + recording it would render a false "$0.00" with cost status complete. + """ + if cost_usd is None or isinstance(cost_usd, bool): + return None + try: + cost = cost_usd if isinstance(cost_usd, Decimal) else Decimal(str(cost_usd)) + except (InvalidOperation, TypeError, ValueError): + return None + if not cost.is_finite() or cost <= 0: + return None + return cost + + +def record_ai_call(usage=None, model: Optional[str] = None, cost_usd=None) -> None: + """Count one successful AI call and accumulate usage and known cost.""" details = get_run_details() if details is None: return details.num_ai_calls += 1 if usage is not None: add_token_usage(usage) + cost = _as_decimal_cost(cost_usd) + if cost is not None: + details.total_cost_usd += cost + details.known_cost_call_count += 1 + model_name = model or "unknown" + details.model_costs_usd[model_name] = details.model_costs_usd.get(model_name, Decimal("0")) + cost diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index ea69ce6de6..12bc956333 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -13,6 +13,7 @@ import time import traceback from datetime import datetime +from decimal import ROUND_HALF_UP, Decimal from enum import Enum from importlib.metadata import PackageNotFoundError, version from typing import Any, Iterable, List, Tuple, TypedDict @@ -1486,6 +1487,14 @@ def show_relevant_configurations(relevant_section: str) -> str: return markdown_text +def _format_usd(cost: Decimal) -> str: + """Format cost at two decimals without turning a positive amount into false zero.""" + rounded = cost.quantize(Decimal("0.01"), rounding=ROUND_HALF_UP) + if cost > 0 and rounded == 0: + return "<$0.01" + return f"${rounded:.2f}" + + def show_run_details(gfm_supported: bool) -> str: """Render the opt-in run-details section (model, tokens, time cost, AI calls). @@ -1508,6 +1517,21 @@ def show_run_details(gfm_supported: bool) -> str: lines.append(f"- Time cost: {details.duration_seconds:.1f}s") if details.num_ai_calls: lines.append(f"- AI calls: {details.num_ai_calls}") + if get_settings().get("config.output_run_cost", False) and details.num_ai_calls: + if details.cost_status == "unavailable": + # No causal claim here: pricing can be missing because usage was absent + # (streaming without a final usage chunk) or because the active handler + # never collects costs (openai/langchain handlers). + lines.append("- Estimated API cost: unavailable (no calls could be priced)") + else: + partial = "" + if details.cost_status == "partial": + partial = (f" (partial: {details.known_cost_call_count} of " + f"{details.num_ai_calls} successful calls priced)") + lines.append(f"- Estimated API cost: {_format_usd(details.total_cost_usd)} USD{partial}") + if len(details.model_costs_usd) > 1: + for model, cost in details.model_costs_usd.items(): + lines.append(f" - {model}: {_format_usd(cost)} USD") body = "\n".join(lines) if gfm_supported: diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml index 8f3647f0e5..31efec34f5 100644 --- a/pr_agent/settings/configuration.toml +++ b/pr_agent/settings/configuration.toml @@ -47,6 +47,7 @@ secret_provider="" # "" (disabled), "google_cloud_storage", or "aws_secrets_mana cli_mode=false output_relevant_configurations=false output_run_details=false # if true, append an agent run details section (model, tokens, time cost, AI calls) to generated PR comments +output_run_cost=false # if true, collect estimated LiteLLM API cost and include it inside the enabled run details section large_patch_policy = "clip" # "clip", "skip" duplicate_prompt_examples = false # Persistent inline comments (issue #2037): when true, the GitHub and GitLab providers diff --git a/tests/unittest/test_litellm_chat_completion_core.py b/tests/unittest/test_litellm_chat_completion_core.py index ae637ff3ed..f7d1c45eea 100644 --- a/tests/unittest/test_litellm_chat_completion_core.py +++ b/tests/unittest/test_litellm_chat_completion_core.py @@ -306,7 +306,11 @@ async def test_get_completion_uses_streaming_for_required_models(): patch("pr_agent.algo.ai_handlers.litellm_ai_handler._handle_streaming_response", new_callable=AsyncMock) as mock_stream: mock_call.return_value = "stream" - mock_stream.return_value = ("streamed text", "stop") + completed_response = MagicMock() + completed_response.dict.return_value = { + "choices": [{"message": {"content": "streamed text"}, "finish_reason": "stop"}] + } + mock_stream.return_value = ("streamed text", "stop", completed_response) resp, finish_reason, response_obj = await handler._get_completion( model="streaming-model", @@ -314,6 +318,8 @@ async def test_get_completion_uses_streaming_for_required_models(): ) assert mock_call.call_args.kwargs["stream"] is True + assert mock_call.call_args.kwargs["stream_options"] == {"include_usage": True} + mock_stream.assert_awaited_once_with("stream", model="streaming-model") assert resp == "streamed text" assert finish_reason == "stop" assert response_obj.dict()["choices"][0]["message"]["content"] == "streamed text" diff --git a/tests/unittest/test_litellm_custom_provider.py b/tests/unittest/test_litellm_custom_provider.py index b953de69d8..e8be216224 100644 --- a/tests/unittest/test_litellm_custom_provider.py +++ b/tests/unittest/test_litellm_custom_provider.py @@ -123,9 +123,9 @@ async def test_openai_compatible_endpoint_calls_force_streaming(monkeypatch): new_callable=AsyncMock, ) as mock_stream_handler, ): - mock_stream_handler.return_value = ("test", "stop") + mock_stream_handler.return_value = ("test", "stop", None) handler = LiteLLMAIHandler() - await handler._get_completion( + result = await handler._get_completion( model="claude-sonnet-4-5", messages=[], timeout=120, @@ -133,8 +133,11 @@ async def test_openai_compatible_endpoint_calls_force_streaming(monkeypatch): custom_llm_provider="openai", ) + assert result == ("test", "stop", None) call_kwargs = mock_completion.call_args[1] assert call_kwargs["stream"] is True + assert call_kwargs["stream_options"] == {"include_usage": True} + assert mock_stream_handler.call_args.kwargs["model"] == "claude-sonnet-4-5" @pytest.mark.asyncio @@ -152,9 +155,9 @@ async def test_openai_compatible_endpoint_normalizes_custom_provider_for_streami new_callable=AsyncMock, ) as mock_stream_handler, ): - mock_stream_handler.return_value = ("test", "stop") + mock_stream_handler.return_value = ("test", "stop", None) handler = LiteLLMAIHandler() - await handler._get_completion( + result = await handler._get_completion( model="claude-sonnet-4-5", messages=[], timeout=120, @@ -162,8 +165,10 @@ async def test_openai_compatible_endpoint_normalizes_custom_provider_for_streami custom_llm_provider=" OpenAI ", ) + assert result == ("test", "stop", None) call_kwargs = mock_completion.call_args[1] assert call_kwargs["stream"] is True + assert call_kwargs["stream_options"] == {"include_usage": True} @pytest.mark.asyncio diff --git a/tests/unittest/test_litellm_run_details.py b/tests/unittest/test_litellm_run_details.py index 71746b8e11..0d2355dce1 100644 --- a/tests/unittest/test_litellm_run_details.py +++ b/tests/unittest/test_litellm_run_details.py @@ -1,4 +1,7 @@ import asyncio +from decimal import Decimal +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch import pytest @@ -21,7 +24,17 @@ def __init__(self, usage): self.usage = usage def dict(self): - return {"choices": [{"message": {"content": "resp"}, "finish_reason": "stop"}]} + return { + "choices": [{"message": {"content": "resp"}, "finish_reason": "stop"}], + "usage": self.usage, + } + + +def _set_cost_collection(monkeypatch, enabled): + settings = SimpleNamespace( + get=lambda key, default=None: enabled if key == "config.output_run_cost" else default + ) + monkeypatch.setattr("pr_agent.algo.ai_handlers.litellm_ai_handler.get_settings", lambda: settings) def test_record_completion_metadata_accumulates_usage(): @@ -37,6 +50,116 @@ def test_record_completion_metadata_accumulates_usage(): assert details.total_tokens == 165 +def test_record_completion_metadata_collects_known_non_streaming_cost(monkeypatch): + usage = _Usage(100, 10, 110) + response = _Response(usage) + _set_cost_collection(monkeypatch, True) + init_run_details() + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + return_value=0.0842, + ) as completion_cost: + LiteLLMAIHandler._record_completion_metadata(response, model="model-a") + + details = get_run_details() + assert details.total_cost_usd == Decimal("0.0842") + assert details.known_cost_call_count == 1 + assert details.cost_status == "complete" + assert details.model_costs_usd == {"model-a": Decimal("0.0842")} + assert completion_cost.call_args.kwargs["completion_response"]["usage"] is usage + + +def test_record_completion_metadata_prices_routed_model_and_records_configured_model(monkeypatch): + usage = _Usage(100, 10, 110) + response = _Response(usage) + _set_cost_collection(monkeypatch, True) + init_run_details() + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + return_value=0.0842, + ) as completion_cost: + LiteLLMAIHandler._record_completion_metadata( + response, + model="azure/gpt-5", + display_model="gpt-5_thinking", + ) + + assert completion_cost.call_args.kwargs["model"] == "azure/gpt-5" + assert get_run_details().model_costs_usd == {"gpt-5_thinking": Decimal("0.0842")} + + +def test_record_completion_metadata_uses_positive_finalized_inline_cost(monkeypatch): + _set_cost_collection(monkeypatch, True) + init_run_details() + response = {"usage": {"cost": "0.0031"}} + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + ) as completion_cost: + LiteLLMAIHandler._record_completion_metadata(response, model="model-a") + + details = get_run_details() + assert details.total_cost_usd == Decimal("0.0031") + assert details.known_cost_call_count == 1 + completion_cost.assert_not_called() + + +def test_zero_inline_cost_without_priceable_usage_stays_unavailable(monkeypatch): + _set_cost_collection(monkeypatch, True) + init_run_details() + response = {"usage": {"response_cost": 0}} + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + ) as completion_cost: + LiteLLMAIHandler._record_completion_metadata(response, model="model-a") + + details = get_run_details() + assert details.total_cost_usd == Decimal("0") + assert details.known_cost_call_count == 0 + assert details.cost_status == "unavailable" + completion_cost.assert_not_called() + + +def test_zero_token_usage_with_provider_timing_floats_stays_unavailable(monkeypatch): + """Groq-style timing floats (queue_time, prompt_time) are not billable + quantities and must not send zero-token usage to completion_cost.""" + _set_cost_collection(monkeypatch, True) + init_run_details() + response = {"usage": {"prompt_tokens": 0, "completion_tokens": 0, + "queue_time": 0.019, "prompt_time": 0.004}} + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + return_value=0.0, + ) as completion_cost: + LiteLLMAIHandler._record_completion_metadata(response, model="model-a") + + details = get_run_details() + assert details.known_cost_call_count == 0 + assert details.cost_status == "unavailable" + completion_cost.assert_not_called() + + +def test_disabled_cost_collection_does_not_calculate_or_record_cost(monkeypatch): + _set_cost_collection(monkeypatch, False) + init_run_details() + response = _Response(_Usage(100, 10, 110)) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + ) as completion_cost: + LiteLLMAIHandler._record_completion_metadata(response, model="model-a") + + details = get_run_details() + assert details.num_ai_calls == 1 + assert details.known_cost_call_count == 0 + assert details.cost_status == "unavailable" + completion_cost.assert_not_called() + + def test_record_completion_metadata_counts_streaming_calls_without_tokens(): init_run_details() @@ -75,6 +198,92 @@ def _bare_handler(): return handler +def _streaming_handler(): + handler = LiteLLMAIHandler.__new__(LiteLLMAIHandler) + handler.streaming_required_models = ["streaming-model"] + handler.force_streaming_provider = "" + handler.force_streaming_api_base_substrings = [] + return handler + + +async def _async_chunks(*chunks): + for chunk in chunks: + yield chunk + + +@pytest.mark.asyncio +async def test_streamed_completion_preserves_finalized_usage_and_collects_known_cost(monkeypatch): + usage = { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "cache_read_input_tokens": 80, + "cache_creation_input_tokens": 10, + "completion_tokens_details": {"reasoning_tokens": 8}, + } + content_chunk = SimpleNamespace( + choices=[SimpleNamespace(delta=SimpleNamespace(content="streamed"), finish_reason="stop")], + usage=None, + ) + usage_chunk = SimpleNamespace(choices=[], usage=usage) + handler = _streaming_handler() + _set_cost_collection(monkeypatch, True) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as acompletion, patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + return_value=0.071, + ) as completion_cost: + acompletion.return_value = _async_chunks(content_chunk, usage_chunk) + init_run_details() + resp, finish_reason, response = await handler._get_completion( + model="streaming-model", + messages=[{"role": "user", "content": "hello"}], + ) + handler._record_completion_metadata(response, model="streaming-model") + + details = get_run_details() + assert (resp, finish_reason) == ("streamed", "stop") + assert response.usage is usage + assert details.total_tokens == 120 + assert details.total_cost_usd == Decimal("0.071") + assert details.cost_status == "complete" + assert acompletion.call_args.kwargs["stream_options"] == {"include_usage": True} + assert completion_cost.call_args.kwargs["completion_response"]["usage"] == usage + + +@pytest.mark.asyncio +async def test_streamed_completion_without_finalized_usage_marks_cost_unavailable(monkeypatch): + content_chunk = SimpleNamespace( + choices=[SimpleNamespace(delta=SimpleNamespace(content="streamed"), finish_reason="stop")], + usage=None, + ) + handler = _streaming_handler() + _set_cost_collection(monkeypatch, True) + + with patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.acompletion", + new_callable=AsyncMock, + ) as acompletion, patch( + "pr_agent.algo.ai_handlers.litellm_ai_handler.litellm.completion_cost", + ) as completion_cost: + acompletion.return_value = _async_chunks(content_chunk) + init_run_details() + _, _, response = await handler._get_completion(model="streaming-model", messages=[]) + handler._record_completion_metadata(response, model="streaming-model") + + details = get_run_details() + assert response.usage is None + assert details.num_ai_calls == 1 + assert details.has_token_usage is False + assert details.known_cost_call_count == 0 + assert details.cost_status == "unavailable" + assert details.total_cost_usd == Decimal("0") + completion_cost.assert_not_called() + + @pytest.mark.asyncio async def test_chat_completion_records_the_call_it_just_made(monkeypatch): """Guard the wiring, not just the recorder. @@ -100,6 +309,35 @@ async def fake_get_completion(**_kwargs): assert details.total_tokens == 110 +@pytest.mark.asyncio +async def test_chat_completion_preserves_configured_model_for_cost_breakdown(monkeypatch): + handler = _bare_handler() + handler.azure = True + response = _Response(_Usage(100, 10, 110)) + routed_models = [] + + async def fake_get_completion(**kwargs): + routed_models.append(kwargs["model"]) + return "resp", "stop", response + + monkeypatch.setattr(handler, "_get_completion", fake_get_completion) + + with patch.object( + handler, + "_record_completion_metadata", + wraps=handler._record_completion_metadata, + ) as record_completion_metadata: + init_run_details() + await handler.chat_completion(model="gpt-4.1", system="sys", user="usr") + + assert routed_models == ["azure/gpt-4.1"] + record_completion_metadata.assert_called_once_with( + response, + model="azure/gpt-4.1", + display_model="gpt-4.1", + ) + + @pytest.mark.asyncio async def test_chat_completion_does_not_record_when_the_call_fails(monkeypatch): """A failed model must not be counted, or fallback runs would inflate the totals.""" diff --git a/tests/unittest/test_run_details.py b/tests/unittest/test_run_details.py index 474800fe9d..eab5c79863 100644 --- a/tests/unittest/test_run_details.py +++ b/tests/unittest/test_run_details.py @@ -1,4 +1,5 @@ import asyncio +from decimal import Decimal import pytest @@ -26,6 +27,10 @@ def test_init_returns_fresh_instance_with_zeroed_counters(): assert details.completion_tokens == 0 assert details.total_tokens == 0 assert details.num_ai_calls == 0 + assert details.total_cost_usd == Decimal("0") + assert details.known_cost_call_count == 0 + assert details.model_costs_usd == {} + assert details.cost_status == "unavailable" assert details.has_token_usage is False assert details.duration_seconds >= 0 @@ -109,6 +114,61 @@ def test_record_ai_call_counts_calls_even_without_usage(): assert details.total_tokens == 12 +def test_record_ai_call_aggregates_decimal_costs_by_model(): + init_run_details() + record_model_used("model-b", is_fallback=True) + + record_ai_call(_Usage(10, 2, 12), model="model-a", cost_usd=Decimal("0.0710")) + record_ai_call(_Usage(5, 1, 6), model="model-b", cost_usd="0.0132") + record_ai_call(_Usage(1, 1, 2), model="model-a", cost_usd=0.00001) + + details = get_run_details() + assert details.total_cost_usd == Decimal("0.08421") + assert details.known_cost_call_count == 3 + assert details.cost_status == "complete" + assert details.fallback_used is True + assert details.model_costs_usd == { + "model-a": Decimal("0.07101"), + "model-b": Decimal("0.0132"), + } + + +def test_record_ai_call_treats_zero_cost_as_unpriced(): + """litellm.completion_cost returns 0.0 for unpriced models and empty usage; + recording it would render a false '$0.00' with cost status complete.""" + init_run_details() + + record_ai_call(_Usage(10, 2, 12), model="zero-priced", cost_usd=0.0) + record_ai_call(_Usage(10, 2, 12), model="zero-priced", cost_usd=Decimal("0")) + + details = get_run_details() + assert details.total_cost_usd == Decimal("0") + assert details.known_cost_call_count == 0 + assert details.cost_status == "unavailable" + assert details.model_costs_usd == {} + + +def test_record_ai_call_marks_partial_and_unavailable_cost_without_fabricating_zero(): + init_run_details() + + record_ai_call(_Usage(10, 2, 12), model="known", cost_usd=Decimal("0.0042")) + record_ai_call(None, model="unknown", cost_usd=None) + + details = get_run_details() + assert details.total_cost_usd == Decimal("0.0042") + assert details.known_cost_call_count == 1 + assert details.cost_status == "partial" + + init_run_details() + record_ai_call(None, model="unknown", cost_usd=None) + + details = get_run_details() + assert details.total_cost_usd == Decimal("0") + assert details.known_cost_call_count == 0 + assert details.cost_status == "unavailable" + assert details.model_costs_usd == {} + + @pytest.mark.asyncio async def test_concurrent_child_tasks_accumulate_into_parent_collector(): init_run_details() diff --git a/tests/unittest/test_run_details_wiring.py b/tests/unittest/test_run_details_wiring.py index d62535edfa..484d4b9e52 100644 --- a/tests/unittest/test_run_details_wiring.py +++ b/tests/unittest/test_run_details_wiring.py @@ -3,15 +3,18 @@ import pytest -from pr_agent.algo.run_details import init_run_details, record_ai_call, record_model_used +from pr_agent.algo.run_details import (init_run_details, record_ai_call, + record_model_used) from pr_agent.config_loader import get_settings from pr_agent.tools.pr_code_suggestions import PRCodeSuggestions from pr_agent.tools.pr_description import PRDescription from pr_agent.tools.pr_reviewer import PRReviewer -from tests.unittest._settings_helpers import restore_settings, snapshot_settings +from tests.unittest._settings_helpers import (restore_settings, + snapshot_settings) _TRACKED_KEYS_REVIEW = ( "config.output_run_details", + "config.output_run_cost", "config.publish_output", "config.is_auto_command", "data", @@ -73,6 +76,7 @@ async def _noop_async(*_args, **_kwargs): def test_flag_defaults_to_false(): assert get_settings().config.get("output_run_details", None) is False + assert get_settings().config.get("output_run_cost", None) is False @pytest.mark.asyncio @@ -105,6 +109,7 @@ async def test_pr_reviewer_appends_run_details_only_when_enabled(monkeypatch): get_settings().set("config.is_auto_command", False) get_settings().pr_reviewer.enable_help_text = False + get_settings().set("config.output_run_cost", True) get_settings().set("config.output_run_details", False) await reviewer.run() without_details = get_settings().data["artifact"] @@ -114,7 +119,9 @@ async def test_pr_reviewer_appends_run_details_only_when_enabled(monkeypatch): with_details = get_settings().data["artifact"] assert "⚙️ Agent run details" not in without_details + assert "Estimated API cost" not in without_details assert "⚙️ Agent run details" in with_details + assert "Estimated API cost: unavailable" in with_details finally: restore_settings(snapshot) diff --git a/tests/unittest/test_show_run_details.py b/tests/unittest/test_show_run_details.py index 4f7de379c5..dfa3c813a6 100644 --- a/tests/unittest/test_show_run_details.py +++ b/tests/unittest/test_show_run_details.py @@ -1,8 +1,13 @@ import re +from decimal import Decimal +import pytest + +from pr_agent.algo import run_details from pr_agent.algo.run_details import (get_run_details, init_run_details, record_ai_call, record_model_used) from pr_agent.algo.utils import show_run_details +from pr_agent.config_loader import get_settings class _Usage: @@ -12,6 +17,15 @@ def __init__(self, prompt_tokens, completion_tokens, total_tokens): self.total_tokens = total_tokens +@pytest.fixture(autouse=True) +def disable_cost_output_by_default(): + settings = get_settings() + previous = settings.config.get("output_run_cost", False) + settings.set("config.output_run_cost", False) + yield + settings.set("config.output_run_cost", previous) + + def test_renders_all_fields_in_a_details_block_when_gfm_supported(): init_run_details() record_model_used("openai/gpt-5.4", is_fallback=False) @@ -101,3 +115,72 @@ def test_returns_empty_string_when_collector_not_initialized(): assert show_run_details(gfm_supported=True) == "" finally: run_details._run_details.reset(token) + + +def test_disabled_cost_output_preserves_existing_run_details_byte_for_byte(monkeypatch): + monkeypatch.setattr(run_details.time, "monotonic", lambda: 108.2) + details = init_run_details() + details.start_time = 100.0 + record_model_used("model-a", is_fallback=False) + record_ai_call(_Usage(10, 2, 12), model="model-a", cost_usd=Decimal("0.0842")) + + output = show_run_details(gfm_supported=True) + + assert output == ( + "\n