diff --git a/ollama/_client.py b/ollama/_client.py index 18cb0fb4..8dfce824 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -209,7 +209,7 @@ def generate( template: str = '', context: Optional[Sequence[int]] = None, stream: Literal[False] = False, - think: Optional[bool] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, raw: bool = False, @@ -233,7 +233,7 @@ def generate( template: str = '', context: Optional[Sequence[int]] = None, stream: Literal[True] = True, - think: Optional[bool] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, raw: bool = False, @@ -256,7 +256,7 @@ def generate( template: Optional[str] = None, context: Optional[Sequence[int]] = None, stream: bool = False, - think: Optional[bool] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, raw: Optional[bool] = None, diff --git a/tests/test_client.py b/tests/test_client.py index 34657513..7b7ab38e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,4 +1,5 @@ import base64 +import inspect import json import os import re @@ -1486,3 +1487,13 @@ async def test_async_client_context_manager(): assert not client._client.is_closed assert client._client.is_closed + + +def test_generate_think_annotation_matches_chat(): + # The `think` parameter accepts bool or the 'low'/'medium'/'high' string levels. + # Client.generate must keep the same annotation as Client.chat and + # AsyncClient.generate so passing a string level does not raise a false type + # error (regression guard for the sync generate overloads/implementation). + expected = inspect.signature(Client.chat).parameters['think'].annotation + assert inspect.signature(Client.generate).parameters['think'].annotation == expected + assert inspect.signature(AsyncClient.generate).parameters['think'].annotation == expected