scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

fix(github_copilot): preserve cache usage across response paths

Browse source
This commit is contained in:
Tin 2026-10-28 10:50:42 -07:00
parent 9da926b90f
commit 17530710c7
7 changed files with 308 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/cost_calculator.py
View file

@ -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:

3
litellm/litellm_core_utils/streaming_handler.py
View file

@ -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(

33
litellm/llms/github_copilot/chat/transformation.py
View file

@ -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,

51
litellm/llms/github_copilot/common_utils.py
View file

@ -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__(

71
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -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

111
tests/unit/llms/github_copilot/test_github_copilot_transformation.py
View file

@ -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

47
tests/unit/test_cost_calculator.py
View file

@ -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,