diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 1636890707..5b9840d95b 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -772,12 +772,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): def chunk_parser(self, chunk: dict) -> ModelResponseStream: try: - return ModelResponseStream( - id=chunk["id"], - object="chat.completion.chunk", - created=chunk.get("created"), - model=chunk.get("model"), - choices=chunk.get("choices", []), - ) + kwargs = { + "id": chunk["id"], + "object": "chat.completion.chunk", + "created": chunk.get("created"), + "model": chunk.get("model"), + "choices": chunk.get("choices", []), + } + if "usage" in chunk and chunk["usage"] is not None: + kwargs["usage"] = chunk["usage"] + return ModelResponseStream(**kwargs) except Exception as e: raise e diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 62a06468e9..e003841ac0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -71,9 +71,7 @@ try: from mcp.shared.tool_name_validation import ( validate_tool_name, # pyright: ignore[reportAssignmentType] ) - from mcp.shared.tool_name_validation import ( - SEP_986_URL, - ) + from mcp.shared.tool_name_validation import SEP_986_URL except ImportError: from pydantic import BaseModel diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index c0695bf358..e6ab199168 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -9,7 +9,10 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) -from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIGPTConfig, + OpenAIChatCompletionStreamingHandler, +) class TestOpenAIGPTConfig: @@ -136,6 +139,75 @@ class TestGetOptionalParamsIntegration: assert regular_params.get("user") == "my-end-user" assert responses_params.get("user") == "my-end-user" + +class TestOpenAIChatCompletionStreamingHandler: + """Tests for OpenAIChatCompletionStreamingHandler.chunk_parser()""" + + def test_chunk_parser_preserves_usage(self): + """ + Test that chunk_parser preserves the usage field from streaming chunks. + + """ + handler = OpenAIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + usage_chunk = { + "id": "gen-123", + "created": 1234567890, + "model": "openai/gpt-4o-mini", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": ""}, + "finish_reason": None, + } + ], + "usage": { + "prompt_tokens": 13797, + "completion_tokens": 350, + "total_tokens": 14147, + }, + } + + result = handler.chunk_parser(usage_chunk) + + assert result.usage is not None + assert result.usage.prompt_tokens == 13797 + assert result.usage.completion_tokens == 350 + assert result.usage.total_tokens == 14147 + + def test_chunk_parser_without_usage(self): + """Test that chunk_parser works normally for chunks without usage.""" + handler = OpenAIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + chunk = { + "id": "gen-123", + "created": 1234567890, + "model": "openai/gpt-4o-mini", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hello"}, + "finish_reason": None, + } + ], + } + + result = handler.chunk_parser(chunk) + + assert result.id == "gen-123" + assert result.choices[0].delta.content == "Hello" + assert not hasattr(result, "usage") or result.usage is None + + +class TestPromptCacheKeyIntegration: + """Tests for prompt_cache_key support""" + def test_prompt_cache_key_in_optional_params(self): """Test that 'prompt_cache_key' flows through get_optional_params for OpenAI models.""" from litellm.utils import get_optional_params