diff --git a/ollama/_types.py b/ollama/_types.py index ad041013..8e11871a 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -266,6 +266,9 @@ class TokenLogprob(SubscriptableBaseModel): logprob: float 'Log probability for the token.' + bytes: Optional[Sequence[int]] = None + 'Raw UTF-8 bytes of the token. Useful when a token is not a standalone valid string.' + class Logprob(TokenLogprob): top_logprobs: Optional[Sequence[TokenLogprob]] = None diff --git a/tests/test_type_serialization.py b/tests/test_type_serialization.py index e1af3512..3c0f981d 100644 --- a/tests/test_type_serialization.py +++ b/tests/test_type_serialization.py @@ -4,7 +4,16 @@ import pytest -from ollama._types import ChatRequest, CreateRequest, GenerateRequest, Image +from ollama._types import ( + ChatRequest, + ChatResponse, + CreateRequest, + GenerateRequest, + GenerateResponse, + Image, + Logprob, + TokenLogprob, +) def test_image_serialization_bytes(): @@ -122,3 +131,59 @@ def test_think_boolean_serialization(): assert ChatRequest(model='test-model', think=True).model_dump(exclude_none=True)['think'] is True assert ChatRequest(model='test-model', think=False).model_dump(exclude_none=True)['think'] is False assert 'think' not in ChatRequest(model='test-model', think=None).model_dump(exclude_none=True) + + +def test_generate_response_keeps_logprob_bytes(): + # Servers send raw token bytes next to token/logprob; they used to be dropped silently. + response = GenerateResponse.model_validate( + { + 'model': 'test-model', + 'response': 'Hello', + 'done': True, + 'logprobs': [{'token': 'Hello', 'logprob': -0.25, 'bytes': [72, 101, 108, 108, 111]}], + } + ) + + logprob = response.logprobs[0] + assert logprob.bytes == [72, 101, 108, 108, 111] + assert bytes(logprob.bytes).decode() == 'Hello' + # Kept as a list of ints so the model stays JSON-serializable. + assert response.model_dump(exclude_none=True)['logprobs'][0]['bytes'] == [72, 101, 108, 108, 111] + + +def test_chat_response_keeps_logprob_bytes(): + response = ChatResponse.model_validate( + { + 'model': 'test-model', + 'message': {'role': 'assistant', 'content': 'Hello'}, + 'done': True, + 'logprobs': [{'token': 'Hello', 'logprob': -0.25, 'bytes': [72, 101, 108, 108, 111]}], + } + ) + + assert response.logprobs[0].bytes == [72, 101, 108, 108, 111] + + +def test_logprob_bytes_is_optional(): + # Servers predating the field must keep parsing, with the attribute present but unset. + for logprob in ( + TokenLogprob.model_validate({'token': 'hi', 'logprob': -1.0}), + Logprob.model_validate({'token': 'hi', 'logprob': -1.0}), + ): + assert logprob.bytes is None + assert 'bytes' not in logprob.model_dump(exclude_none=True) + + +def test_top_logprobs_keep_bytes(): + # top_logprobs reuses TokenLogprob, so it has to carry bytes too. + logprob = Logprob.model_validate( + { + 'token': ' there', + 'logprob': -0.003, + 'bytes': [32, 116, 104, 101, 114, 101], + 'top_logprobs': [{'token': ' there', 'logprob': -0.003, 'bytes': [32, 116, 104, 101, 114, 101]}], + } + ) + + assert logprob.bytes == [32, 116, 104, 101, 114, 101] + assert logprob.top_logprobs[0].bytes == [32, 116, 104, 101, 114, 101]