diff --git a/ollama/_client.py b/ollama/_client.py index 8dfce82..779b988 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[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = 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[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = 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[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, raw: Optional[bool] = None, @@ -313,7 +313,7 @@ def chat( *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, @@ -329,7 +329,7 @@ def chat( *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, @@ -344,7 +344,7 @@ def chat( *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, @@ -842,7 +842,7 @@ async def generate( template: str = '', context: Optional[Sequence[int]] = None, stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, raw: bool = False, @@ -866,7 +866,7 @@ async def generate( template: str = '', context: Optional[Sequence[int]] = None, stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, raw: bool = False, @@ -889,7 +889,7 @@ async def generate( template: Optional[str] = None, context: Optional[Sequence[int]] = None, stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, raw: Optional[bool] = None, @@ -945,7 +945,7 @@ async def chat( *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, @@ -961,7 +961,7 @@ async def chat( *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, @@ -976,7 +976,7 @@ async def chat( *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, diff --git a/ollama/_types.py b/ollama/_types.py index 96529d6..b66cb67 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -207,7 +207,7 @@ class GenerateRequest(BaseGenerateRequest): images: Optional[Sequence[Image]] = None 'Image data for multimodal models.' - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None 'Enable thinking mode (for thinking models).' logprobs: Optional[bool] = None @@ -400,7 +400,7 @@ def serialize_model(self, nxt): tools: Optional[Sequence[Tool]] = None 'Tools to use for the chat.' - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None + think: Optional[Union[bool, Literal['low', 'medium', 'high', 'max']]] = None 'Enable thinking mode (for thinking models).' logprobs: Optional[bool] = None diff --git a/tests/test_client.py b/tests/test_client.py index 7b7ab38..498f152 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1490,10 +1490,28 @@ async def test_async_client_context_manager(): def test_generate_think_annotation_matches_chat(): - # The `think` parameter accepts bool or the 'low'/'medium'/'high' string levels. + # The `think` parameter accepts bool or the 'low'/'medium'/'high'/'max' 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 + + +def test_think_max_level_accepted(): + # The 'max' think level must be accepted by all generate/chat methods without + # raising a type error (mirrors the Go ThinkValue which accepts 'max'). + import typing + + for method in (Client.generate, Client.chat, AsyncClient.generate, AsyncClient.chat): + annotation = typing.get_type_hints(method).get('think') + args = typing.get_args(annotation) + # Flatten nested Unions/Literals to find the literal values + literals = set() + for arg in args: + if arg is bool or arg is bool: + continue + sub = typing.get_args(arg) + literals.update(sub) + assert 'max' in literals, f'{method.__qualname__} think annotation missing max: {literals}'