diff --git a/src/social_hook/llm/litellm_client.py b/src/social_hook/llm/litellm_client.py index be5bcf1..46b9dcb 100644 --- a/src/social_hook/llm/litellm_client.py +++ b/src/social_hook/llm/litellm_client.py @@ -116,10 +116,10 @@ def complete( # (token-field naming, reasoning params) instead of erroring. "drop_params": True, } - # The pipeline always passes exactly one tool — force it. (A provider - # that rejects forced tool_choice is caught by the JSON-from-text - # fallback in _extract_content.) "required" covers the rare multi-tool - # case. + # The pipeline always passes exactly one tool — force it. drop_params + # lets litellm drop a forced tool_choice the model can't honor; a model + # that then answers in prose is recovered by the JSON-from-text fallback + # in _extract_content. "required" covers the rare multi-tool case. if len(tools) == 1: kwargs["tool_choice"] = {"type": "function", "function": {"name": tools[0]["name"]}} elif tools: @@ -183,7 +183,8 @@ def _extract_content(self, response: Any, tools: list[dict[str, Any]]) -> list[A if blocks: return blocks - text = (getattr(message, "content", None) or "").strip() + content = getattr(message, "content", None) + text = content.strip() if isinstance(content, str) else "" if len(tools) == 1 and text: try: parsed = extract_json_object(text) @@ -197,10 +198,14 @@ def _extract_content(self, response: Any, tools: list[dict[str, Any]]) -> list[A def _extract_usage(self, response: Any) -> NormalizedUsage: usage_obj = getattr(response, "usage", None) - input_tokens = int(getattr(usage_obj, "prompt_tokens", 0) or 0) + prompt_tokens = int(getattr(usage_obj, "prompt_tokens", 0) or 0) output_tokens = int(getattr(usage_obj, "completion_tokens", 0) or 0) ptd = getattr(usage_obj, "prompt_tokens_details", None) cache_read = int((getattr(ptd, "cached_tokens", 0) if ptd is not None else 0) or 0) + # OpenAI-style prompt_tokens INCLUDES cached tokens; normalize to the + # Anthropic-style disjoint accounting (non-cached input + separate + # cache_read) so cost and usage logging are consistent across providers. + input_tokens = max(prompt_tokens - cache_read, 0) cost_cents, cost_source = self._resolve_cost( response, input_tokens, output_tokens, cache_read diff --git a/tests/test_litellm_client.py b/tests/test_litellm_client.py index 5d1af67..9fc2903 100644 --- a/tests/test_litellm_client.py +++ b/tests/test_litellm_client.py @@ -148,6 +148,28 @@ def test_cost_falls_back_to_catalog(self, mock_completion, mock_cost): assert abs(resp.usage.cost_cents - 93.0) < 0.01 assert resp.usage.cost_source == "registry" + @patch("litellm.completion_cost") + @patch("litellm.completion") + def test_cached_tokens_not_double_billed(self, mock_completion, mock_cost): + # OpenAI-style prompt_tokens INCLUDES cached tokens; the catalog + # fallback must bill the cached portion once (at 0.1x), not twice. + mock_completion.return_value = _mock_response( + tool_calls=[_tool_call("route_action", "{}")], + prompt_tokens=100_000, + cached_tokens=90_000, + completion_tokens=0, + response_cost=None, + ) + mock_cost.side_effect = Exception("model not in litellm cost map") + resp = self._client().complete(messages=[{"role": "user", "content": "hi"}], tools=[TOOL]) + # non-cached input = 10k; GLM-5.2 $0.93/M in, cache read at 0.1x: + # (10000*0.93 + 90000*0.93*0.1)/1e6*100 = 0.93 + 0.837 = 1.767c + assert abs(resp.usage.cost_cents - 1.767) < 0.01 + assert resp.usage.cost_source == "registry" + # input_tokens is normalized to non-cached (Anthropic-style disjoint) + assert resp.usage.input_tokens == 10_000 + assert resp.usage.cache_read_input_tokens == 90_000 + @patch("litellm.completion") def test_provider_error_wrapped(self, mock_completion): from social_hook.errors import MalformedResponseError