From 17530710c76d5cfd6f8617b58c194bfa362cd7df Mon Sep 17 00:00:00 2001 From: Tin Date: Wed, 28 Oct 2026 10:50:42 -0700 Subject: [PATCH] fix(github_copilot): preserve cache usage across response paths --- litellm/cost_calculator.py | 8 +- .../litellm_core_utils/streaming_handler.py | 3 + .../github_copilot/chat/transformation.py | 33 ++++-- litellm/llms/github_copilot/common_utils.py | 51 +++++++- .../test_streaming_handler.py | 71 +++++++++++ .../test_github_copilot_transformation.py | 111 +++++++++++++++++- tests/unit/test_cost_calculator.py | 47 ++++++++ 7 files changed, 308 insertions(+), 16 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 86b22069a39..7ce92772b91 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -519,7 +519,11 @@ def cost_per_token( # Anthropic reports prompt_tokens as input_tokens (excluding cache tokens). # Adjust so the helper's "prompt_tokens includes cache tokens" invariant holds. _normalized_prompt_tokens = float(prompt_tokens) - if _is_anthropic_style: + model_is_str: Final = isinstance(model, str) + _copilot_inclusive_prompt: Final = custom_llm_provider == "github_copilot" or ( + custom_llm_provider is None and model_is_str and model.startswith("github_copilot/") + ) + if _is_anthropic_style and not _copilot_inclusive_prompt: _normalized_prompt_tokens += _cache_read_tokens + _cache_creation_tokens response_cost: Final = _cost_per_token_custom_pricing_helper( @@ -546,8 +550,6 @@ def cost_per_token( # non-string objects. Only run the string-based dedup/prefix-join when it is # actually a string — e.g. a MagicMock's `.startswith()` is always truthy and # its slices return new mocks, which would spin the dedup loop forever. - model_is_str: Final = isinstance(model, str) - # Router/proxy deployments may repeat the provider segment (e.g. model_name # "openai/openai/gpt-5.5"). Strip duplicated `{provider}/` chains before joining. if caller_supplied_provider and model_is_str: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b64e146ca08..d0da5eb9675 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -36,6 +36,7 @@ from litellm.litellm_core_utils.model_response_utils import ( ) from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.llms.github_copilot.common_utils import normalize_copilot_chunk from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams @@ -1606,6 +1607,8 @@ class CustomStreamWrapper: return _ProviderChunkParsed(response_obj) def chunk_creator(self, chunk: Any): + if isinstance(chunk, (ChatCompletionChunk, ModelResponseStream)): + chunk = normalize_copilot_chunk(chunk) # rebind-ok: normalize a copy before both usage consumers if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields: self.logging_obj.record_native_bedrock_guardrail(chunk.provider_specific_fields) if self.custom_llm_provider != "cached_response" and isinstance( diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 83bc8df23a6..88ba6a4019f 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -8,7 +8,7 @@ from pydantic import TypeAdapter from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.llms.openai.common_utils import decoded_json_object +from litellm.llms.openai.common_utils import as_str_mapping, decoded_json_object from litellm.llms.openai.openai import OpenAIConfig from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.utils import ModelResponse @@ -18,6 +18,7 @@ from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, GetAPIKeyError, get_copilot_default_headers, + normalize_copilot_usage, ) from ..per_user_auth import require_github_copilot_user_session @@ -234,7 +235,7 @@ class GithubCopilotConfig(OpenAIConfig): return normalized @classmethod - def _synthesize_choices_for_anthropic_native(cls, response_json: dict) -> dict: + def _synthesize_choices_for_anthropic_native(cls, response_json: dict) -> dict[str, object]: """ Synthesize a `choices` array from an Anthropic-native Copilot response. @@ -282,21 +283,26 @@ class GithubCopilotConfig(OpenAIConfig): if thinking_blocks: message["thinking_blocks"] = thinking_blocks - synthesized: Final = { + usage: Final = response_json.get("usage") + usage_fields: Final = {"usage": cls._normalize_anthropic_usage(usage)} if isinstance(usage, dict) else {} + return { **response_json, "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + **usage_fields, } - usage: Final = response_json.get("usage") - if isinstance(usage, dict): - synthesized["usage"] = cls._normalize_anthropic_usage(usage) - return synthesized - def transform_parsed_response_dict(self, parsed_response: dict) -> dict: + def transform_parsed_response_dict(self, parsed_response: dict) -> dict[str, object]: """ Repair the OpenAI-SDK-parsed response on the handler path that bypasses transform_response. See: https://github.com/BerriAI/litellm/issues/30927 """ - return self._synthesize_choices_for_anthropic_native(parsed_response) + synthesized: Final = self._synthesize_choices_for_anthropic_native(parsed_response) + usage: Final = as_str_mapping(synthesized.get("usage")) + if usage is not None: + normalized: Final = normalize_copilot_usage(usage, synthesized.get("copilot_usage")) + if normalized is not usage: + return {**synthesized, "usage": normalized} + return synthesized def transform_response( self, @@ -330,8 +336,13 @@ class GithubCopilotConfig(OpenAIConfig): ) response_json: Final = decoded_json_object(decoded_response) - if not response_json.get("choices"): - synthesized: Final = self._synthesize_choices_for_anthropic_native(dict(response_json)) + raw_usage: Final = as_str_mapping(response_json.get("usage")) + if ( + not response_json.get("choices") + or "copilot_usage" in response_json + or (raw_usage is not None and "copilot_usage" in raw_usage) + ): + synthesized: Final = self.transform_parsed_response_dict(dict(response_json)) raw_response = httpx.Response( status_code=raw_response.status_code, headers=raw_response.headers, diff --git a/litellm/llms/github_copilot/common_utils.py b/litellm/llms/github_copilot/common_utils.py index 515b7bf4fb5..ca3774189c8 100644 --- a/litellm/llms/github_copilot/common_utils.py +++ b/litellm/llms/github_copilot/common_utils.py @@ -2,13 +2,16 @@ Constants for Copilot integration """ -from collections.abc import MutableMapping +from collections.abc import Mapping, MutableMapping, Sequence from typing import Final from uuid import uuid4 import httpx +from pydantic import BaseModel, TypeAdapter +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.utils import PromptTokensDetailsWrapper, Usage # Constants COPILOT_VERSION: Final = "0.26.7" @@ -17,6 +20,52 @@ USER_AGENT: Final = f"GitHubCopilotChat/{COPILOT_VERSION}" API_VERSION: Final = "2025-04-01" DEFAULT_GITHUB_COPILOT_API_BASE: Final = "https://api.githubcopilot.com" +_TOKEN_DETAILS: Final = TypeAdapter(tuple[object, ...]) + + +def _copilot_cache_count(entries: Sequence[object], token_type: str) -> int | None: + mappings: Final = tuple(as_str_mapping(entry) for entry in entries) + counts: Final = tuple( + entry.get("token_count") for entry in mappings if entry is not None and entry.get("token_type") == token_type + ) + valid_counts: Final = tuple( + count for count in counts if isinstance(count, int) and not isinstance(count, bool) and count >= 0 + ) + return sum(valid_counts) if valid_counts else None + + +def normalize_copilot_usage(usage: Mapping[str, object], copilot_usage: object = None) -> Mapping[str, object]: + outer: Final = as_str_mapping(copilot_usage) + metadata: Final = outer if outer is not None else as_str_mapping(usage.get("copilot_usage")) + token_details: Final = metadata.get("token_details") if metadata is not None else None + if not isinstance(token_details, list): + return usage + entries: Final = _TOKEN_DETAILS.validate_python(token_details) + details: Final = PromptTokensDetailsWrapper.model_validate( + as_str_mapping(usage.get("prompt_tokens_details")) or {} + ).model_dump(exclude_unset=True) + candidates: Final = ( + ("cache_read_input_tokens", "cached_tokens", _copilot_cache_count(entries, "cache_read")), + ("cache_creation_input_tokens", "cache_write_tokens", _copilot_cache_count(entries, "cache_write")), + ) + counts: Final = { + field: details.get(detail_field) if details.get(detail_field) is not None else count + for field, detail_field, count in candidates + if usage.get(field) is None and (details.get(detail_field) is not None or count is not None) + } + return {**usage, **counts} + + +def normalize_copilot_chunk(chunk: BaseModel) -> BaseModel: + usage: Final[object] = getattr(chunk, "usage", None) + if not isinstance(usage, BaseModel): + return chunk + copilot_usage: Final[object] = (chunk.model_extra or {}).get("copilot_usage") + if copilot_usage is None and "copilot_usage" not in (usage.model_extra or {}): + return chunk + normalized: Final = normalize_copilot_usage(usage.model_dump(exclude_unset=True), copilot_usage) + return chunk.model_copy(update={"usage": Usage.model_validate(normalized)}) + class GithubCopilotError(BaseLLMException): def __init__( diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6d700bae514..7a04f615eb6 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7943,3 +7943,74 @@ def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details( assert usage.prompt_tokens_details.cached_tokens == 80 assert usage.completion_tokens_details is not None assert usage.completion_tokens_details.reasoning_tokens == 5 + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("sdk", [False, True]) +@pytest.mark.parametrize("usage_only", [False, True]) +@pytest.mark.parametrize("include_usage", [False, True]) +@pytest.mark.parametrize("top_level", [False, True]) +@pytest.mark.parametrize("write_alias", [None, "cache_creation_tokens", "cache_creation_input_tokens"]) +@pytest.mark.asyncio +async def test_copilot_cache_usage_survives_stream_and_assembly( + async_stream: bool, sdk: bool, usage_only: bool, include_usage: bool, + top_level: bool, write_alias: str | None, +) -> None: + from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator + + metadata: Final = {"token_details": [ + {"token_type": "cache_read", "token_count": 70}, + {"token_type": "cache_write", "token_count": 20}, + ]} + reported: Final = { + "usage": { + "prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110, + **({} if top_level else {"copilot_usage": metadata}), + **({"prompt_tokens_details": {write_alias: 0}} if write_alias else {}), + }, + **({"copilot_usage": metadata} if top_level else {}), + } + bodies: Final = [ + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi"}}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + **({} if usage_only else reported)}, + *([{"choices": [], **reported}] if usage_only else []), + ] + parser: Final = OpenAIChatCompletionResponseIterator(streaming_response=None, sync_stream=True) + payloads: Final = [ + {"id": "chatcmpl-cache", "created": 1, "model": "claude-sonnet-5", + "object": "chat.completion.chunk", **body} + for body in bodies + ] + chunks: Final = [ + ChatCompletionChunk(**payload) if sdk else parser.chunk_parser(payload) + for payload in payloads + ] + before: Final = [chunk.model_dump() for chunk in chunks] + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="github_copilot/claude-sonnet-5", custom_llm_provider="openai", + stream_options={"include_usage": True} if include_usage else None, + logging_obj=Logging( + model="github_copilot/claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], stream=True, + call_type="acompletion", start_time=datetime.datetime(2025, 1, 1), + litellm_call_id="copilot-cache", function_id="copilot-cache", + ), + ) + received: Final = [chunk async for chunk in wrapper] if async_stream else list(wrapper) + final_usage: Final = ( + received[-1].usage if include_usage else received[-1]._hidden_params["usage"] + ) + assembled: Final = litellm.stream_chunk_builder(chunks=received) + assert assembled is not None + assert assembled.choices[0].message.content == "hi" + for usage in (final_usage, assembled.usage): + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (100, 10, 110) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 70 + assert usage.prompt_tokens_details.cache_write_tokens == (0 if write_alias else 20) + assert getattr(usage, "cache_read_input_tokens", None) == 70 + assert getattr(usage, "cache_creation_input_tokens", None) == (0 if write_alias else 20) + assert [chunk.model_dump() for chunk in chunks] == before diff --git a/tests/unit/llms/github_copilot/test_github_copilot_transformation.py b/tests/unit/llms/github_copilot/test_github_copilot_transformation.py index 2dbfae07526..66ea36695b9 100644 --- a/tests/unit/llms/github_copilot/test_github_copilot_transformation.py +++ b/tests/unit/llms/github_copilot/test_github_copilot_transformation.py @@ -1,7 +1,7 @@ import asyncio import json from datetime import datetime, timedelta -from typing import AsyncGenerator +from typing import AsyncGenerator, Final from unittest.mock import AsyncMock, MagicMock, mock_open, patch import pytest @@ -966,3 +966,112 @@ def test_transform_response_carries_upstream_usage(): assert result.usage.prompt_tokens == 11 assert result.usage.completion_tokens == 7 assert result.usage.total_tokens == 18 + + +def _copilot_normalized_usage(payload: dict[str, object], raw: bool) -> Usage: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_model_response_object, + ) + + config: Final = GithubCopilotConfig() + result: Final = ( + config.transform_response( + model="claude-sonnet-5", raw_response=httpx.Response(200, json=payload), + model_response=ModelResponse(), request_data={}, messages=[], + optional_params={}, litellm_params={}, encoding=None, + logging_obj=Logging( + model="claude-sonnet-5", messages=[], stream=False, + call_type="completion", start_time=datetime(2025, 1, 1), + litellm_call_id="copilot-cache", function_id="copilot-cache", + ), + ) + if raw else convert_to_model_response_object( + response_object=config.transform_parsed_response_dict(payload), + model_response_object=ModelResponse(), + ) + ) + assert isinstance(result, ModelResponse) + assert isinstance(result.usage, Usage) + return result.usage + + +@pytest.mark.parametrize("raw", [False, True]) +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("top_level", [False, True]) +@pytest.mark.parametrize( + "details,standard,expected", + [ + ([{"token_type": "cache_read", "token_count": 50}, + {"token_type": "cache_read", "token_count": 20}, + {"token_type": "cache_write", "token_count": 20}], {}, (70, 20)), + ([{"token_type": "cache_read", "token_count": 70}, + {"token_type": "cache_write", "token_count": 20}], + {"cache_read_input_tokens": 0, "cache_creation_input_tokens": 3}, (0, 3)), + ([{"token_type": "cache_read", "token_count": 70}, + {"token_type": "cache_write", "token_count": 20}], + {"prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 3}}, (0, 3)), + ([{"token_type": "cache_read", "token_count": 70}, + {"token_type": "cache_write", "token_count": 20}], + {"prompt_tokens_details": {"cache_creation_tokens": 0}}, (70, 0)), + ([{"token_type": "cache_read", "token_count": 70}, + {"token_type": "cache_write", "token_count": 20}], + {"prompt_tokens_details": {"cache_creation_input_tokens": 3}}, (70, 3)), + ([{"token_type": "cache_read", "token_count": 0}, + {"token_type": "cache_write", "token_count": 0}], {}, (0, 0)), + ([None, "bad", {}, {"token_type": "other", "token_count": 10}, + {"token_type": "cache_read", "token_count": True}, + {"token_type": "cache_read", "token_count": -1}, + {"token_type": "cache_write", "token_count": "20"}], {}, (None, None)), + (None, {}, (None, None)), + ({}, {}, (None, None)), + ], +) +def test_copilot_cache_usage_survives_response_conversion( + raw: bool, native: bool, top_level: bool, details: object, + standard: dict[str, object], expected: tuple[int | None, int | None], +) -> None: + metadata: Final = {"token_details": details} + usage: Final = { + **({"input_tokens": 100, "output_tokens": 10} if native else { + "prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110, + }), + "copilot_usage": {"token_details": [{"token_type": "cache_read", "token_count": 99}]} + if top_level else metadata, + **standard, + } + payload: Final = { + "id": "chatcmpl-cache", "created": 1, "model": "claude-sonnet-5", + "object": "chat.completion", "usage": usage, + **({"content": [{"type": "text", "text": "hi"}], "stop_reason": "end_turn"} + if native else {"choices": [{"index": 0, "message": { + "role": "assistant", "content": "hi"}, "finish_reason": "stop"}]}), + **({"copilot_usage": metadata} if top_level else {}), + } + before: Final = json.dumps(payload, sort_keys=True) + result: Final = _copilot_normalized_usage(payload, raw) + + assert (result.prompt_tokens, result.completion_tokens, result.total_tokens) == (100, 10, 110) + assert getattr(result, "cache_read_input_tokens", None) == expected[0] + assert getattr(result, "cache_creation_input_tokens", None) == expected[1] + assert getattr(result.prompt_tokens_details, "cached_tokens", None) == expected[0] + assert getattr(result.prompt_tokens_details, "cache_write_tokens", None) == expected[1] + assert json.dumps(payload, sort_keys=True) == before + + +@pytest.mark.parametrize("metadata", [None, [], "malformed"]) +@pytest.mark.parametrize("usage", [None, {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}]) +def test_copilot_cache_normalization_preserves_absent_metadata(usage: object, metadata: object) -> None: + payload: Final = { + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}}], + "usage": usage, "copilot_usage": metadata, + } + assert GithubCopilotConfig().transform_parsed_response_dict(payload) is payload + + +def test_copilot_cache_metadata_does_not_create_usage() -> None: + payload: Final = { + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}}], + "copilot_usage": {"token_details": [{"token_type": "cache_read", "token_count": 70}]}, + } + assert GithubCopilotConfig().transform_parsed_response_dict(payload) is payload diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 84f6dab9eff..46b00c1fd66 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -3700,6 +3700,53 @@ def test_custom_pricing_anthropic_style_cache_tokens_not_double_counted(): assert cost == pytest.approx(expected) +@pytest.mark.parametrize( + "model,provider", + [("claude-sonnet-5", "github_copilot"), ("github_copilot/claude-sonnet-5", None)], +) +@pytest.mark.parametrize("cached,created", [(70, 0), (0, 20), (70, 20)]) +@pytest.mark.parametrize("native", [False, True]) +def test_copilot_custom_pricing_keeps_cache_tokens_inside_prompt_total( + model: str, provider: str | None, cached: int, created: int, native: bool, +) -> None: + from litellm.llms.github_copilot.chat.transformation import GithubCopilotConfig + from litellm.types.utils import CostPerToken + + payload: Final = { + "id": "copilot-priced", "created": 1, "model": model, + "choices": [] if native else [{"index": 0, "message": { + "role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "content": [{"type": "text", "text": "hi"}], "stop_reason": "end_turn", + "usage": {"input_tokens": 100, "output_tokens": 10} if native else { + "prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110, + }, + "copilot_usage": {"token_details": [ + {"token_type": "cache_read", "token_count": cached}, + {"token_type": "cache_write", "token_count": created}, + ]}, + } + response: Final = convert_to_model_response_object( + response_object=GithubCopilotConfig().transform_parsed_response_dict(payload), + model_response_object=ModelResponse(), + ) + assert isinstance(response, ModelResponse) + assert isinstance(response.usage, Usage) + rates: Final = CostPerToken( + input_cost_per_token=1.0, output_cost_per_token=2.0, + cache_read_input_token_cost=0.1, cache_creation_input_token_cost=1.25, + ) + expected_input: Final = (100 - cached - created) * 1.0 + cached * 0.1 + created * 1.25 + assert cost_per_token( + model=model, custom_llm_provider=provider, prompt_tokens=response.usage.prompt_tokens, + completion_tokens=response.usage.completion_tokens, usage_object=response.usage, + custom_cost_per_token=rates, + ) == pytest.approx((expected_input, 20.0)) + assert completion_cost( + completion_response=response, model=model, custom_llm_provider=provider, + custom_cost_per_token=rates, + ) == pytest.approx(expected_input + 20.0) + + def test_custom_pricing_without_cache_keys_preserves_legacy_behavior(): """ Backward compatibility: when custom_cost_per_token omits both cache rates,