diff --git a/ollama/__init__.py b/ollama/__init__.py index 92bba280..e72e81f1 100644 --- a/ollama/__init__.py +++ b/ollama/__init__.py @@ -15,6 +15,8 @@ ShowResponse, StatusResponse, Tool, + ToolChoice, + ToolChoiceFunction, WebFetchResponse, WebSearchResponse, ) @@ -37,6 +39,8 @@ 'ShowResponse', 'StatusResponse', 'Tool', + 'ToolChoice', + 'ToolChoiceFunction', 'WebFetchResponse', 'WebSearchResponse', ] diff --git a/ollama/_client.py b/ollama/_client.py index 18cb0fb4..d529617b 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -67,6 +67,8 @@ ShowResponse, StatusResponse, Tool, + ToolChoice, + ToolChoiceFunction, WebFetchRequest, WebFetchResponse, WebSearchRequest, @@ -312,6 +314,7 @@ def chat( messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + tool_choice: Optional[Union[ToolChoice, Mapping[str, Any]]] = None, stream: Literal[False] = False, think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, @@ -328,6 +331,7 @@ def chat( messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + tool_choice: Optional[Union[ToolChoice, Mapping[str, Any]]] = None, stream: Literal[True] = True, think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, @@ -343,6 +347,7 @@ def chat( messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + tool_choice: Optional[Union[ToolChoice, Mapping[str, Any]]] = None, stream: bool = False, think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, @@ -392,6 +397,7 @@ def add_two_numbers(a: int, b: int) -> int: model=model, messages=list(_copy_messages(messages)), tools=list(_copy_tools(tools)), + tool_choice=tool_choice, stream=stream, think=think, logprobs=logprobs, @@ -944,6 +950,7 @@ async def chat( messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + tool_choice: Optional[Union[ToolChoice, Mapping[str, Any]]] = None, stream: Literal[False] = False, think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, @@ -960,6 +967,7 @@ async def chat( messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + tool_choice: Optional[Union[ToolChoice, Mapping[str, Any]]] = None, stream: Literal[True] = True, think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, @@ -975,6 +983,7 @@ async def chat( messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, *, tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + tool_choice: Optional[Union[ToolChoice, Mapping[str, Any]]] = None, stream: bool = False, think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, logprobs: Optional[bool] = None, @@ -1025,6 +1034,7 @@ def add_two_numbers(a: int, b: int) -> int: model=model, messages=list(_copy_messages(messages)), tools=list(_copy_tools(tools)), + tool_choice=tool_choice, stream=stream, think=think, logprobs=logprobs, diff --git a/ollama/_types.py b/ollama/_types.py index 96529d63..690be548 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -384,6 +384,28 @@ class Property(SubscriptableBaseModel): function: Optional[Function] = None + +class ToolChoiceFunction(SubscriptableBaseModel): + """ + Forces the model to call a specific tool by name. + + Example:: + + tool_choice=ToolChoiceFunction(function=ToolChoiceFunction.Function(name='get_weather')) + """ + + class Function(SubscriptableBaseModel): + name: str + 'Name of the function to call.' + + type: Literal['function'] = 'function' + function: Function + 'The function to force the model to call.' + + +ToolChoice = Union[Literal['none', 'auto', 'required'], ToolChoiceFunction] + + class ChatRequest(BaseGenerateRequest): @model_serializer(mode='wrap') def serialize_model(self, nxt): @@ -400,6 +422,9 @@ def serialize_model(self, nxt): tools: Optional[Sequence[Tool]] = None 'Tools to use for the chat.' + tool_choice: Optional[ToolChoice] = None + """Controls which tool the model calls: 'none', 'auto', 'required', or ToolChoiceFunction.""" + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None 'Enable thinking mode (for thinking models).'