include tool calls

* Add type overloads for chat() method in _client.py

* Overloading

* Fix Overload Overlap

* Fix async chat

* Lint

* Reverse

---------

Co-authored-by: Simon Ottenhaus <simon.ottenhaus@kenbun.de>

README.md

@@ -2,24 +2,6 @@
 
 The Ollama Python library provides the easiest way to integrate Python 3.8+ projects with [Ollama](https://github.com/ollama/ollama).
 
-## Prerequisites
-
-You need to have a local ollama server running to be able to continue. To do this:
-
-- Download: https://ollama.com/
-- Run an LLM: https://ollama.com/library
-    - Example: `ollama run llama2`
-    - Example: `ollama run llama2:70b`
-
-Then:
-
-```sh
-curl https://ollama.ai/install.sh | sh
-ollama serve
-```
-
-Next you can go ahead with `ollama-python`.
-
 ## Install
 
 ```sh
@@ -125,6 +107,12 @@ ollama.push('user/llama3')
 ollama.embeddings(model='llama3', prompt='The sky is blue because of rayleigh scattering')
 ```
 
+### Ps
+
+```python
+ollama.ps()
+```
+
 ## Custom client
 
 A custom client can be created with the following fields:

examples/ps/main.py

@@ -0,0 +1,31 @@
+from ollama import ps, pull, chat
+
+response = pull('mistral', stream=True)
+progress_states = set()
+for progress in response:
+  if progress.get('status') in progress_states:
+    continue
+  progress_states.add(progress.get('status'))
+  print(progress.get('status'))
+
+print('\n')
+
+response = chat('mistral', messages=[{'role': 'user', 'content': 'Hello!'}])
+print(response['message']['content'])
+
+print('\n')
+
+response = ps()
+
+name = response['models'][0]['name']
+size = response['models'][0]['size']
+size_vram = response['models'][0]['size_vram']
+
+if size == size_vram:
+  print(f'{name}: 100% GPU')
+elif not size_vram:
+  print(f'{name}: 100% CPU')
+else:
+  size_cpu = size - size_vram
+  cpu_percent = round(size_cpu / size * 100)
+  print(f'{name}: {cpu_percent}% CPU/{100 - cpu_percent}% GPU')

ollama/__init__.py

@@ -29,6 +29,7 @@ __all__ = [
   'list',
   'copy',
   'show',
+  'ps',
 ]
 
 _client = Client()
@@ -43,3 +44,4 @@ delete = _client.delete
 list = _client.list
 copy = _client.copy
 show = _client.show
+ps = _client.ps

ollama/_client.py

@@ -7,10 +7,11 @@ import platform
 import urllib.parse
 from os import PathLike
 from pathlib import Path
+from copy import deepcopy
 from hashlib import sha256
 from base64 import b64encode, b64decode
 
-from typing import Any, AnyStr, Union, Optional, Sequence, Mapping, Literal
+from typing import Any, AnyStr, Union, Optional, Sequence, Mapping, Literal, overload
 
 import sys
 
@@ -96,6 +97,40 @@ class Client(BaseClient):
   ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]:
     return self._stream(*args, **kwargs) if stream else self._request(*args, **kwargs).json()
 
+  @overload
+  def generate(
+    self,
+    model: str = '',
+    prompt: str = '',
+    system: str = '',
+    template: str = '',
+    context: Optional[Sequence[int]] = None,
+    stream: Literal[False] = False,
+    raw: bool = False,
+    format: Literal['', 'json'] = '',
+    images: Optional[Sequence[AnyStr]] = None,
+    tools: Optional[Sequence[Any]] = None,
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  def generate(
+    self,
+    model: str = '',
+    prompt: str = '',
+    system: str = '',
+    template: str = '',
+    context: Optional[Sequence[int]] = None,
+    stream: Literal[True] = True,
+    raw: bool = False,
+    format: Literal['', 'json'] = '',
+    images: Optional[Sequence[AnyStr]] = None,
+    tools: Optional[Sequence[Any]] = None,
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Iterator[Mapping[str, Any]]: ...
+
   def generate(
     self,
     model: str = '',
@@ -107,6 +142,7 @@ class Client(BaseClient):
     raw: bool = False,
     format: Literal['', 'json'] = '',
     images: Optional[Sequence[AnyStr]] = None,
+    tools: Optional[Sequence[Any]] = None,
     options: Optional[Options] = None,
     keep_alive: Optional[Union[float, str]] = None,
   ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]:
@@ -135,6 +171,7 @@ class Client(BaseClient):
         'stream': stream,
         'raw': raw,
         'images': [_encode_image(image) for image in images or []],
+        'tools': tools or [],
         'format': format,
         'options': options or {},
         'keep_alive': keep_alive,
@@ -142,6 +179,28 @@ class Client(BaseClient):
       stream=stream,
     )
 
+  @overload
+  def chat(
+    self,
+    model: str = '',
+    messages: Optional[Sequence[Message]] = None,
+    stream: Literal[False] = False,
+    format: Literal['', 'json'] = '',
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  def chat(
+    self,
+    model: str = '',
+    messages: Optional[Sequence[Message]] = None,
+    stream: Literal[True] = True,
+    format: Literal['', 'json'] = '',
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Iterator[Mapping[str, Any]]: ...
+
   def chat(
     self,
     model: str = '',
@@ -164,6 +223,8 @@ class Client(BaseClient):
     if not model:
       raise RequestError('must provide a model')
 
+    messages = deepcopy(messages)
+
     for message in messages or []:
       if not isinstance(message, dict):
         raise TypeError('messages must be a list of Message or dict-like objects')
@@ -206,6 +267,22 @@ class Client(BaseClient):
       },
     ).json()
 
+  @overload
+  def pull(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[False] = False,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  def pull(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[True] = True,
+  ) -> Iterator[Mapping[str, Any]]: ...
+
   def pull(
     self,
     model: str,
@@ -228,6 +305,22 @@ class Client(BaseClient):
       stream=stream,
     )
 
+  @overload
+  def push(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[False] = False,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  def push(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[True] = True,
+  ) -> Iterator[Mapping[str, Any]]: ...
+
   def push(
     self,
     model: str,
@@ -250,6 +343,26 @@ class Client(BaseClient):
       stream=stream,
     )
 
+  @overload
+  def create(
+    self,
+    model: str,
+    path: Optional[Union[str, PathLike]] = None,
+    modelfile: Optional[str] = None,
+    quantize: Optional[str] = None,
+    stream: Literal[False] = False,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  def create(
+    self,
+    model: str,
+    path: Optional[Union[str, PathLike]] = None,
+    modelfile: Optional[str] = None,
+    quantize: Optional[str] = None,
+    stream: Literal[True] = True,
+  ) -> Iterator[Mapping[str, Any]]: ...
+
   def create(
     self,
     model: str,
@@ -336,6 +449,9 @@ class Client(BaseClient):
   def show(self, model: str) -> Mapping[str, Any]:
     return self._request('POST', '/api/show', json={'name': model}).json()
 
+  def ps(self) -> Mapping[str, Any]:
+    return self._request('GET', '/api/ps').json()
+
 
 class AsyncClient(BaseClient):
   def __init__(self, host: Optional[str] = None, **kwargs) -> None:
@@ -380,6 +496,38 @@ class AsyncClient(BaseClient):
     response = await self._request(*args, **kwargs)
     return response.json()
 
+  @overload
+  async def generate(
+    self,
+    model: str = '',
+    prompt: str = '',
+    system: str = '',
+    template: str = '',
+    context: Optional[Sequence[int]] = None,
+    stream: Literal[False] = False,
+    raw: bool = False,
+    format: Literal['', 'json'] = '',
+    images: Optional[Sequence[AnyStr]] = None,
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  async def generate(
+    self,
+    model: str = '',
+    prompt: str = '',
+    system: str = '',
+    template: str = '',
+    context: Optional[Sequence[int]] = None,
+    stream: Literal[True] = True,
+    raw: bool = False,
+    format: Literal['', 'json'] = '',
+    images: Optional[Sequence[AnyStr]] = None,
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> AsyncIterator[Mapping[str, Any]]: ...
+
   async def generate(
     self,
     model: str = '',
@@ -425,6 +573,28 @@ class AsyncClient(BaseClient):
       stream=stream,
     )
 
+  @overload
+  async def chat(
+    self,
+    model: str = '',
+    messages: Optional[Sequence[Message]] = None,
+    stream: Literal[False] = False,
+    format: Literal['', 'json'] = '',
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  async def chat(
+    self,
+    model: str = '',
+    messages: Optional[Sequence[Message]] = None,
+    stream: Literal[True] = True,
+    format: Literal['', 'json'] = '',
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> AsyncIterator[Mapping[str, Any]]: ...
+
   async def chat(
     self,
     model: str = '',
@@ -446,6 +616,8 @@ class AsyncClient(BaseClient):
     if not model:
       raise RequestError('must provide a model')
 
+    messages = deepcopy(messages)
+
     for message in messages or []:
       if not isinstance(message, dict):
         raise TypeError('messages must be a list of strings')
@@ -490,6 +662,22 @@ class AsyncClient(BaseClient):
 
     return response.json()
 
+  @overload
+  async def pull(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[False] = False,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  async def pull(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[True] = True,
+  ) -> AsyncIterator[Mapping[str, Any]]: ...
+
   async def pull(
     self,
     model: str,
@@ -512,6 +700,22 @@ class AsyncClient(BaseClient):
       stream=stream,
     )
 
+  @overload
+  async def push(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[False] = False,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  async def push(
+    self,
+    model: str,
+    insecure: bool = False,
+    stream: Literal[True] = True,
+  ) -> AsyncIterator[Mapping[str, Any]]: ...
+
   async def push(
     self,
     model: str,
@@ -534,6 +738,26 @@ class AsyncClient(BaseClient):
       stream=stream,
     )
 
+  @overload
+  async def create(
+    self,
+    model: str,
+    path: Optional[Union[str, PathLike]] = None,
+    modelfile: Optional[str] = None,
+    quantize: Optional[str] = None,
+    stream: Literal[False] = False,
+  ) -> Mapping[str, Any]: ...
+
+  @overload
+  async def create(
+    self,
+    model: str,
+    path: Optional[Union[str, PathLike]] = None,
+    modelfile: Optional[str] = None,
+    quantize: Optional[str] = None,
+    stream: Literal[True] = True,
+  ) -> AsyncIterator[Mapping[str, Any]]: ...
+
   async def create(
     self,
     model: str,
@@ -629,6 +853,10 @@ class AsyncClient(BaseClient):
     response = await self._request('POST', '/api/show', json={'name': model})
     return response.json()
 
+  async def ps(self) -> Mapping[str, Any]:
+    response = await self._request('GET', '/api/ps')
+    return response.json()
+
 
 def _encode_image(image) -> str:
   """

ollama/_types.py

@@ -19,6 +19,9 @@ class BaseGenerateResponse(TypedDict):
   done: bool
   'True if response is complete, otherwise False. Useful for streaming to detect the final response.'
 
+  done_reason: str
+  'Reason for completion. Only present when done is True.'
+
   total_duration: int
   'Total duration in nanoseconds.'
 
@@ -49,6 +52,8 @@ class GenerateResponse(BaseGenerateResponse):
   context: Sequence[int]
   'Tokenized history up to the point of the response.'
 
+  tool_calls: Sequence[Any]
+  'List of tool calls made by the model.'
 
 class Message(TypedDict):
   """

poetry.lock

@@ -386,13 +386,13 @@ testing = ["pytest", "pytest-benchmark"]
 
 [[package]]
 name = "pytest"
-version = "8.2.0"
+version = "8.2.2"
 description = "pytest: simple powerful testing with Python"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "pytest-8.2.0-py3-none-any.whl", hash = "sha256:1733f0620f6cda4095bbf0d9ff8022486e91892245bb9e7d5542c018f612f233"},
-    {file = "pytest-8.2.0.tar.gz", hash = "sha256:d507d4482197eac0ba2bae2e9babf0672eb333017bcedaa5fb1a3d42c1174b3f"},
+    {file = "pytest-8.2.2-py3-none-any.whl", hash = "sha256:c434598117762e2bd304e526244f67bf66bbd7b5d6cf22138be51ff661980343"},
+    {file = "pytest-8.2.2.tar.gz", hash = "sha256:de4bb8104e201939ccdc688b27a89a7be2079b22e2bd2b07f806b6ba71117977"},
 ]
 
 [package.dependencies]
@@ -408,13 +408,13 @@ dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments
 
 [[package]]
 name = "pytest-asyncio"
-version = "0.23.6"
+version = "0.23.7"
 description = "Pytest support for asyncio"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "pytest-asyncio-0.23.6.tar.gz", hash = "sha256:ffe523a89c1c222598c76856e76852b787504ddb72dd5d9b6617ffa8aa2cde5f"},
-    {file = "pytest_asyncio-0.23.6-py3-none-any.whl", hash = "sha256:68516fdd1018ac57b846c9846b954f0393b26f094764a28c955eabb0536a4e8a"},
+    {file = "pytest_asyncio-0.23.7-py3-none-any.whl", hash = "sha256:009b48127fbe44518a547bddd25611551b0e43ccdbf1e67d12479f569832c20b"},
+    {file = "pytest_asyncio-0.23.7.tar.gz", hash = "sha256:5f5c72948f4c49e7db4f29f2521d4031f1c27f86e57b046126654083d4770268"},
 ]
 
 [package.dependencies]
@@ -458,28 +458,28 @@ Werkzeug = ">=2.0.0"
 
 [[package]]
 name = "ruff"
-version = "0.4.3"
+version = "0.4.7"
 description = "An extremely fast Python linter and code formatter, written in Rust."
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "ruff-0.4.3-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:b70800c290f14ae6fcbb41bbe201cf62dfca024d124a1f373e76371a007454ce"},
-    {file = "ruff-0.4.3-py3-none-macosx_11_0_arm64.whl", hash = "sha256:08a0d6a22918ab2552ace96adeaca308833873a4d7d1d587bb1d37bae8728eb3"},
-    {file = "ruff-0.4.3-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:eba1f14df3c758dd7de5b55fbae7e1c8af238597961e5fb628f3de446c3c40c5"},
-    {file = "ruff-0.4.3-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:819fb06d535cc76dfddbfe8d3068ff602ddeb40e3eacbc90e0d1272bb8d97113"},
-    {file = "ruff-0.4.3-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0bfc9e955e6dc6359eb6f82ea150c4f4e82b660e5b58d9a20a0e42ec3bb6342b"},
-    {file = "ruff-0.4.3-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:510a67d232d2ebe983fddea324dbf9d69b71c4d2dfeb8a862f4a127536dd4cfb"},
-    {file = "ruff-0.4.3-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:dc9ff11cd9a092ee7680a56d21f302bdda14327772cd870d806610a3503d001f"},
-    {file = "ruff-0.4.3-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:29efff25bf9ee685c2c8390563a5b5c006a3fee5230d28ea39f4f75f9d0b6f2f"},
-    {file = "ruff-0.4.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:18b00e0bcccf0fc8d7186ed21e311dffd19761cb632241a6e4fe4477cc80ef6e"},
-    {file = "ruff-0.4.3-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:262f5635e2c74d80b7507fbc2fac28fe0d4fef26373bbc62039526f7722bca1b"},
-    {file = "ruff-0.4.3-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:7363691198719c26459e08cc17c6a3dac6f592e9ea3d2fa772f4e561b5fe82a3"},
-    {file = "ruff-0.4.3-py3-none-musllinux_1_2_i686.whl", hash = "sha256:eeb039f8428fcb6725bb63cbae92ad67b0559e68b5d80f840f11914afd8ddf7f"},
-    {file = "ruff-0.4.3-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:927b11c1e4d0727ce1a729eace61cee88a334623ec424c0b1c8fe3e5f9d3c865"},
-    {file = "ruff-0.4.3-py3-none-win32.whl", hash = "sha256:25cacda2155778beb0d064e0ec5a3944dcca9c12715f7c4634fd9d93ac33fd30"},
-    {file = "ruff-0.4.3-py3-none-win_amd64.whl", hash = "sha256:7a1c3a450bc6539ef00da6c819fb1b76b6b065dec585f91456e7c0d6a0bbc725"},
-    {file = "ruff-0.4.3-py3-none-win_arm64.whl", hash = "sha256:71ca5f8ccf1121b95a59649482470c5601c60a416bf189d553955b0338e34614"},
-    {file = "ruff-0.4.3.tar.gz", hash = "sha256:ff0a3ef2e3c4b6d133fbedcf9586abfbe38d076041f2dc18ffb2c7e0485d5a07"},
+    {file = "ruff-0.4.7-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:e089371c67892a73b6bb1525608e89a2aca1b77b5440acf7a71dda5dac958f9e"},
+    {file = "ruff-0.4.7-py3-none-macosx_11_0_arm64.whl", hash = "sha256:10f973d521d910e5f9c72ab27e409e839089f955be8a4c8826601a6323a89753"},
+    {file = "ruff-0.4.7-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:59c3d110970001dfa494bcd95478e62286c751126dfb15c3c46e7915fc49694f"},
+    {file = "ruff-0.4.7-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:fa9773c6c00f4958f73b317bc0fd125295110c3776089f6ef318f4b775f0abe4"},
+    {file = "ruff-0.4.7-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:07fc80bbb61e42b3b23b10fda6a2a0f5a067f810180a3760c5ef1b456c21b9db"},
+    {file = "ruff-0.4.7-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:fa4dafe3fe66d90e2e2b63fa1591dd6e3f090ca2128daa0be33db894e6c18648"},
+    {file = "ruff-0.4.7-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a7c0083febdec17571455903b184a10026603a1de078428ba155e7ce9358c5f6"},
+    {file = "ruff-0.4.7-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ad1b20e66a44057c326168437d680a2166c177c939346b19c0d6b08a62a37589"},
+    {file = "ruff-0.4.7-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cbf5d818553add7511c38b05532d94a407f499d1a76ebb0cad0374e32bc67202"},
+    {file = "ruff-0.4.7-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:50e9651578b629baec3d1513b2534de0ac7ed7753e1382272b8d609997e27e83"},
+    {file = "ruff-0.4.7-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:8874a9df7766cb956b218a0a239e0a5d23d9e843e4da1e113ae1d27ee420877a"},
+    {file = "ruff-0.4.7-py3-none-musllinux_1_2_i686.whl", hash = "sha256:b9de9a6e49f7d529decd09381c0860c3f82fa0b0ea00ea78409b785d2308a567"},
+    {file = "ruff-0.4.7-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:13a1768b0691619822ae6d446132dbdfd568b700ecd3652b20d4e8bc1e498f78"},
+    {file = "ruff-0.4.7-py3-none-win32.whl", hash = "sha256:769e5a51df61e07e887b81e6f039e7ed3573316ab7dd9f635c5afaa310e4030e"},
+    {file = "ruff-0.4.7-py3-none-win_amd64.whl", hash = "sha256:9e3ab684ad403a9ed1226894c32c3ab9c2e0718440f6f50c7c5829932bc9e054"},
+    {file = "ruff-0.4.7-py3-none-win_arm64.whl", hash = "sha256:10f2204b9a613988e3484194c2c9e96a22079206b22b787605c255f130db5ed7"},
+    {file = "ruff-0.4.7.tar.gz", hash = "sha256:2331d2b051dc77a289a653fcc6a42cce357087c5975738157cd966590b18b5e1"},
 ]
 
 [[package]]