diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5542249d2c7..682b2a2f538 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -40,6 +40,7 @@ from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + CHOICE_LOGPROBS_OR_JSON, CacheCreationTokenDetails, CompletionTokensDetailsWrapper, Delta, @@ -1201,7 +1202,10 @@ class CustomStreamWrapper: choice_json.pop( "finish_reason", None ) # for mistral etc. which return a value in their last chunk (not-openai compatible). - choices.append(StreamingChoices(**choice_json)) + choice_json.pop("logprobs", None) + choices.append( + StreamingChoices(**choice_json, logprobs=getattr(choice, "logprobs", None)) + ) except Exception: choices.append(StreamingChoices()) setattr(model_response, "choices", choices) @@ -1605,7 +1609,7 @@ class CustomStreamWrapper: model_response.system_fingerprint = response_obj["original_chunk"].system_fingerprint self.system_fingerprint = response_obj["original_chunk"].system_fingerprint if response_obj["logprobs"] is not None: - model_response.choices[0].logprobs = response_obj["logprobs"] + model_response.choices[0].logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(response_obj["logprobs"]) if response_obj["usage"] is not None: if isinstance(response_obj["usage"], dict): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6a758961c7b..780c9ef6824 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -7,6 +7,7 @@ from enum import Enum from types import MappingProxyType from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -38,6 +39,7 @@ from openai.types.moderation_create_response import ( ) from pydantic import ( BaseModel, + BeforeValidator, ConfigDict, Field, FieldSerializationInfo, @@ -45,6 +47,7 @@ from pydantic import ( JsonValue, PrivateAttr, SkipValidation, + TypeAdapter, ValidationError, field_serializer, field_validator, @@ -1199,7 +1202,7 @@ class ChatCompletionTokenLogprob(OpenAIObject): unlikely. """ - top_logprobs: list[TopLogprob] + top_logprobs: list[TopLogprob] = Field(default_factory=list) """List of the most likely tokens and their log probability, at this token position. @@ -1236,9 +1239,13 @@ class ChatCompletionTokenLogprob(OpenAIObject): class ChoiceLogprobs(OpenAIObject): + model_config = ConfigDict(defer_build=False) + content: list[ChatCompletionTokenLogprob] | None = None """A list of message content tokens with log probability information.""" + refusal: Sequence[ChatCompletionTokenLogprob] | None = None + def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) @@ -1252,6 +1259,21 @@ class ChoiceLogprobs(OpenAIObject): return getattr(self, key) +def _sdk_logprobs_as_dict(logprobs: object) -> object: + if isinstance(logprobs, ChoiceLogprobs) or not isinstance(logprobs, BaseModel): + return logprobs + return logprobs.model_dump(exclude_unset=True) + + +CHOICE_LOGPROBS_OR_JSON: Final[TypeAdapter[ChoiceLogprobs | JsonValue]] = TypeAdapter( + Annotated[ + ChoiceLogprobs | JsonValue, + Field(union_mode="left_to_right"), + BeforeValidator(_sdk_logprobs_as_dict), + ] +) + + class FunctionCall(OpenAIObject): arguments: str name: str | None = None @@ -1765,7 +1787,7 @@ class Choices(SafeAttributeModel, OpenAIObject): params["message"] = Message(**dump) if logprobs is not None: if isinstance(logprobs, dict): - params["logprobs"] = ChoiceLogprobs(**logprobs) + params["logprobs"] = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs) else: params["logprobs"] = logprobs else: @@ -2131,8 +2153,8 @@ class StreamingChoices(OpenAIObject): if enhancements is not None: self.enhancements = enhancements - if logprobs is not None and isinstance(logprobs, dict): - self.logprobs = ChoiceLogprobs(**logprobs) + if isinstance(logprobs, (dict, BaseModel)): + self.logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs) else: self.logprobs = logprobs diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index faecb8d8ef6..de8d6a7a8c5 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1,6 +1,7 @@ import asyncio import importlib import json +from collections.abc import Mapping from datetime import datetime, timedelta from typing import Final @@ -1111,6 +1112,44 @@ def test_convert_to_model_response_object_with_logprobs(): assert result.system_fingerprint is None assert result._hidden_params is not None + +@pytest.mark.parametrize( + "upstream_logprobs", + [ + pytest.param({"content": [{"token": "Hi", "logprob": None}]}, id="null-content-logprob"), + pytest.param( + { + "content": [{"token": "Hi", "logprob": -0.1, "top_logprobs": []}], + "refusal": [{"token": "I", "logprob": None}], + }, + id="null-refusal-logprob", + ), + ], +) +def test_convert_to_model_response_object_passes_nonconforming_logprobs_through( + upstream_logprobs: Mapping[str, object], +) -> None: + result: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1702685778, + "model": "test-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi"}, + "logprobs": upstream_logprobs, + "finish_reason": "stop", + } + ], + }, + model_response_object=ModelResponse(), + ) + + assert json.loads(result.model_dump_json())["choices"][0]["logprobs"] == upstream_logprobs + + def test_convert_to_model_response_object_error(): """Test error handling for None response object.""" with pytest.raises(Exception, match="Error in response object format"): diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 41e8c50fa85..b9c8feb223d 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2,12 +2,15 @@ import asyncio import datetime import itertools import json +import os import queue +import subprocess import sys import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from pathlib import Path from types import CodeType, FrameType, MappingProxyType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -33,6 +36,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( + ChoiceLogprobs, CompletionTokensDetailsWrapper, Delta, ModelResponse, @@ -1624,6 +1628,356 @@ async def test_stream_with_null_choices_chunk_finishes_normally( assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"] +def _token_logprobs(token: str) -> dict[str, object]: + return { + "content": [{"token": token, "logprob": -0.5, "bytes": list(token.encode()), "top_logprobs": []}], + "refusal": None, + } + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize( + ("finish_chunk_extra", "stream_options", "expected_usage_totals"), + [ + pytest.param({}, None, [], id="finish-chunk-with-logprobs"), + pytest.param( + {"usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}}, + {"include_usage": True}, + [6], + id="finish-chunk-with-logprobs-and-usage", + ), + ], +) +async def test_stream_finish_chunk_with_logprobs_reaches_caller_as_litellm_logprobs( + async_stream: bool, + transport: str, + finish_chunk_extra: dict[str, object], + stream_options: dict[str, bool] | None, + expected_usage_totals: list[int], + logging_obj: Logging, +) -> None: + upstream_tokens: Final = ("Hel", "lo", "<|return|>") + upstream_chunks: Final = ( + { + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hel"}, + "logprobs": _token_logprobs("Hel"), + "finish_reason": None, + } + ] + }, + { + "choices": [ + {"index": 0, "delta": {"content": "lo"}, "logprobs": _token_logprobs("lo"), "finish_reason": None} + ] + }, + { + "choices": [{"index": 0, "delta": {}, "logprobs": _token_logprobs("<|return|>"), "finish_reason": "stop"}], + **finish_chunk_extra, + }, + ) + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-logprobs", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + **chunk, + } + ) + + "\n\n" + for chunk in upstream_chunks + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + if transport == "sdk" + else OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", + stream_options=stream_options, + ) + + try: + chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + wire_choices: Final = tuple( + itertools.chain.from_iterable( + json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))["choices"] for chunk in chunks + ) + ) + finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")] + streamed_logprobs: Final = tuple( + choice.logprobs + for choice in itertools.chain.from_iterable(chunk.choices for chunk in chunks) + if choice.logprobs is not None + ) + streamed_tokens: Final = tuple( + entry.token for entry in itertools.chain.from_iterable(logprobs.content or () for logprobs in streamed_logprobs) + ) + + assert [choice["finish_reason"] for choice in finish_choices] == ["stop"] + assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == [upstream_tokens[-1]] + assert all(isinstance(logprobs, ChoiceLogprobs) for logprobs in streamed_logprobs) + assert streamed_tokens == upstream_tokens + assert [chunk.usage.total_tokens for chunk in chunks if getattr(chunk, "usage", None)] == expected_usage_totals + + +async def _sdk_stream_chunks( + upstream_chunks: Sequence[Mapping[str, object]], + async_stream: bool, + logging_obj: Logging, + stream_options: dict[str, bool] | None = None, +) -> list[ModelResponseStream]: + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-logprobs", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + **chunk, + } + ) + + "\n\n" + for chunk in upstream_chunks + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options=stream_options, + ) + try: + return [chunk async for chunk in response] if async_stream else list(response) + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize( + ("upstream_logprobs", "finish_chunk_extra", "stream_options", "expected_usage_totals"), + [ + pytest.param( + {"content": [{"token": "<|return|>", "bytes": None, "logprob": None, "top_logprobs": None}], "refusal": None}, + {}, + None, + [], + id="null-logprob", + ), + pytest.param( + {"content": [{"token": "<|return|>", "logprob": None}]}, + {}, + None, + [], + id="null-logprob-keeps-only-the-keys-upstream-sent", + ), + pytest.param(-0.25, {}, None, [], id="scalar-logprobs"), + pytest.param( + {"content": [{"token": "<|return|>", "bytes": None, "logprob": None, "top_logprobs": None}], "refusal": None}, + {"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}}, + {"include_usage": True}, + [5], + id="null-logprob-and-usage", + ), + ], +) +async def test_stream_with_nonstandard_logprobs_passes_them_through_and_finishes( + async_stream: bool, + upstream_logprobs: object, + finish_chunk_extra: dict[str, object], + stream_options: dict[str, bool] | None, + expected_usage_totals: list[int], + logging_obj: Logging, +) -> None: + chunks: Final = await _sdk_stream_chunks( + ( + { + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "hi"}, + "logprobs": upstream_logprobs, + "finish_reason": None, + } + ] + }, + { + "choices": [{"index": 0, "delta": {}, "logprobs": upstream_logprobs, "finish_reason": "stop"}], + **finish_chunk_extra, + }, + ), + async_stream, + logging_obj, + stream_options, + ) + + wire_choices: Final = tuple( + itertools.chain.from_iterable( + json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))["choices"] for chunk in chunks + ) + ) + logprobs_choices: Final = [choice for choice in wire_choices if "logprobs" in choice] + + assert [choice.get("finish_reason") for choice in logprobs_choices] == [None, "stop"] + assert [choice["logprobs"] for choice in logprobs_choices] == [upstream_logprobs, upstream_logprobs] + assert [chunk.usage.total_tokens for chunk in chunks if getattr(chunk, "usage", None)] == expected_usage_totals + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("logprobs_field", ["content", "refusal"]) +async def test_stream_logprobs_without_top_logprobs_reach_caller_as_litellm_logprobs( + async_stream: bool, logprobs_field: Literal["content", "refusal"], logging_obj: Logging +) -> None: + chunks: Final = await _sdk_stream_chunks( + tuple( + { + "choices": [ + { + "index": 0, + "delta": delta, + "logprobs": {logprobs_field: [{"token": token, "logprob": -0.5}]}, + "finish_reason": finish_reason, + } + ] + } + for token, delta, finish_reason in ( + ("hi", {"role": "assistant", "content": "hi"}, None), + ("<|return|>", {}, "stop"), + ) + ), + async_stream, + logging_obj, + ) + + streamed_logprobs: Final = tuple( + choice.logprobs + for choice in itertools.chain.from_iterable(chunk.choices for chunk in chunks) + if choice.logprobs is not None + ) + streamed_entries: Final = tuple( + itertools.chain.from_iterable( + (logprobs.content if logprobs_field == "content" else logprobs.refusal) or () + for logprobs in streamed_logprobs + ) + ) + + assert [(entry.token, entry.logprob, entry.top_logprobs) for entry in streamed_entries] == [ + ("hi", -0.5, []), + ("<|return|>", -0.5, []), + ] + + +FINISH_ONLY_LOGPROBS_STREAM_SCRIPT: Final = """ +import json +import sys + +sys.path.insert(0, sys.argv[1]) + +import httpx +from openai import OpenAI + +import litellm + +upstream_choices = ( + {"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}, + { + "index": 0, + "delta": {}, + "logprobs": {"content": [{"token": "<|return|>", "logprob": -0.5, "top_logprobs": []}]}, + "finish_reason": "stop", + }, +) +body = "".join( + "data: " + + json.dumps({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [choice]}) + + "\\n\\n" + for choice in upstream_choices +) + "data: [DONE]\\n\\n" +client = OpenAI( + api_key="unused", + base_url="https://provider.example/v1", + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"}) + ) + ), +) +stream = litellm.completion( + model="openai/m", messages=[{"role": "user", "content": "hi"}], stream=True, logprobs=True, client=client +) +print(json.dumps([json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in stream])) +""" + + +def test_stream_with_logprobs_only_on_finish_chunk_serializes_in_a_fresh_process() -> None: + result: Final = subprocess.run( + [sys.executable, "-I", "-c", FINISH_ONLY_LOGPROBS_STREAM_SCRIPT, str(Path(litellm.__file__).parent.parent)], + capture_output=True, + text=True, + timeout=120, + env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + ) + + assert result.returncode == 0, result.stdout + result.stderr + wire_choices: Final = tuple( + itertools.chain.from_iterable(chunk["choices"] for chunk in json.loads(result.stdout.strip().splitlines()[-1])) + ) + finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")] + assert [choice["finish_reason"] for choice in finish_choices] == ["stop"] + assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == ["<|return|>"] + + def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper: """A CustomStreamWrapper over the real OpenAI-compatible line iterator, fed an HTTP 200 SSE body that carries an in-body error payload the way diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1aec7d5b761..960616e2631 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29161,7 +29161,7 @@ export interface components { /** Token */ token: string; /** Top Logprobs */ - top_logprobs: components["schemas"]["TopLogprob"][]; + top_logprobs?: components["schemas"]["TopLogprob"][]; } & { [key: string]: unknown; }; @@ -29292,6 +29292,8 @@ export interface components { ChoiceLogprobs: { /** Content */ content?: components["schemas"]["ChatCompletionTokenLogprob"][] | null; + /** Refusal */ + refusal?: components["schemas"]["ChatCompletionTokenLogprob"][] | null; } & { [key: string]: unknown; };