Merge pull request #58 from ollama/mxyng/python-user-agent

python user agent

examples/create/main.py

@@ -0,0 +1,20 @@
+import sys
+
+from ollama import create
+
+
+args = sys.argv[1:]
+if len(args) == 2:
+  # create from local file
+  path = args[1]
+else:
+  print('usage: python main.py <name> <filepath>')
+  sys.exit(1)
+
+# TODO: update to real Modelfile values
+modelfile = f"""
+FROM {path}
+"""
+
+for response in create(model=args[0], modelfile=modelfile, stream=True):
+  print(response['status'])

ollama/_client.py

@@ -2,11 +2,13 @@ import os
 import io
 import json
 import httpx
+import binascii
+import platform
 import urllib.parse
 from os import PathLike
 from pathlib import Path
 from hashlib import sha256
-from base64 import b64encode
+from base64 import b64encode, b64decode
 
 from typing import Any, AnyStr, Union, Optional, Sequence, Mapping, Literal
 
@@ -17,6 +19,13 @@ if sys.version_info < (3, 9):
 else:
   from collections.abc import Iterator, AsyncIterator
 
+from importlib import metadata
+
+try:
+  __version__ = metadata.version('ollama')
+except metadata.PackageNotFoundError:
+  __version__ = '0.0.0'
+
 from ollama._types import Message, Options, RequestError, ResponseError
 
 
@@ -36,10 +45,17 @@ class BaseClient:
     - `timeout`: None
     `kwargs` are passed to the httpx client.
     """
+
+    headers = kwargs.pop('headers', {})
+    headers['Content-Type'] = 'application/json'
+    headers['Accept'] = 'application/json'
+    headers['User-Agent'] = f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}'
+
     self._client = client(
       base_url=_parse_host(host or os.getenv('OLLAMA_HOST')),
       follow_redirects=follow_redirects,
       timeout=timeout,
+      headers=headers,
       **kwargs,
     )
 
@@ -92,6 +108,7 @@ class Client(BaseClient):
     format: Literal['', 'json'] = '',
     images: Optional[Sequence[AnyStr]] = None,
     options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
   ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]:
     """
     Create a response using the requested model.
@@ -120,6 +137,7 @@ class Client(BaseClient):
         'images': [_encode_image(image) for image in images or []],
         'format': format,
         'options': options or {},
+        'keep_alive': keep_alive,
       },
       stream=stream,
     )
@@ -131,6 +149,7 @@ class Client(BaseClient):
     stream: bool = False,
     format: Literal['', 'json'] = '',
     options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
   ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]:
     """
     Create a chat response using the requested model.
@@ -164,11 +183,18 @@ class Client(BaseClient):
         'stream': stream,
         'format': format,
         'options': options or {},
+        'keep_alive': keep_alive,
       },
       stream=stream,
     )
 
-  def embeddings(self, model: str = '', prompt: str = '', options: Optional[Options] = None) -> Sequence[float]:
+  def embeddings(
+    self,
+    model: str = '',
+    prompt: str = '',
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Sequence[float]:
     return self._request(
       'POST',
       '/api/embeddings',
@@ -176,6 +202,7 @@ class Client(BaseClient):
         'model': model,
         'prompt': prompt,
         'options': options or {},
+        'keep_alive': keep_alive,
       },
     ).json()
 
@@ -363,6 +390,7 @@ class AsyncClient(BaseClient):
     format: Literal['', 'json'] = '',
     images: Optional[Sequence[AnyStr]] = None,
     options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
   ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]:
     """
     Create a response using the requested model.
@@ -390,6 +418,7 @@ class AsyncClient(BaseClient):
         'images': [_encode_image(image) for image in images or []],
         'format': format,
         'options': options or {},
+        'keep_alive': keep_alive,
       },
       stream=stream,
     )
@@ -401,6 +430,7 @@ class AsyncClient(BaseClient):
     stream: bool = False,
     format: Literal['', 'json'] = '',
     options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
   ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]:
     """
     Create a chat response using the requested model.
@@ -433,11 +463,18 @@ class AsyncClient(BaseClient):
         'stream': stream,
         'format': format,
         'options': options or {},
+        'keep_alive': keep_alive,
       },
       stream=stream,
     )
 
-  async def embeddings(self, model: str = '', prompt: str = '', options: Optional[Options] = None) -> Sequence[float]:
+  async def embeddings(
+    self,
+    model: str = '',
+    prompt: str = '',
+    options: Optional[Options] = None,
+    keep_alive: Optional[Union[float, str]] = None,
+  ) -> Sequence[float]:
     response = await self._request(
       'POST',
       '/api/embeddings',
@@ -445,6 +482,7 @@ class AsyncClient(BaseClient):
         'model': model,
         'prompt': prompt,
         'options': options or {},
+        'keep_alive': keep_alive,
       },
     )
 
@@ -589,19 +627,43 @@ class AsyncClient(BaseClient):
 
 
 def _encode_image(image) -> str:
-  if p := _as_path(image):
-    b64 = b64encode(p.read_bytes())
-  elif b := _as_bytesio(image):
-    b64 = b64encode(b.read())
-  else:
-    raise RequestError('images must be a list of bytes, path-like objects, or file-like objects')
+  """
+  >>> _encode_image(b'ollama')
+  'b2xsYW1h'
+  >>> _encode_image(io.BytesIO(b'ollama'))
+  'b2xsYW1h'
+  >>> _encode_image('LICENSE')
+  'TUlUIExpY2Vuc2UKCkNvcHlyaWdodCAoYykgT2xsYW1hCgpQZXJtaXNzaW9uIGlzIGhlcmVieSBncmFudGVkLCBmcmVlIG9mIGNoYXJnZSwgdG8gYW55IHBlcnNvbiBvYnRhaW5pbmcgYSBjb3B5Cm9mIHRoaXMgc29mdHdhcmUgYW5kIGFzc29jaWF0ZWQgZG9jdW1lbnRhdGlvbiBmaWxlcyAodGhlICJTb2Z0d2FyZSIpLCB0byBkZWFsCmluIHRoZSBTb2Z0d2FyZSB3aXRob3V0IHJlc3RyaWN0aW9uLCBpbmNsdWRpbmcgd2l0aG91dCBsaW1pdGF0aW9uIHRoZSByaWdodHMKdG8gdXNlLCBjb3B5LCBtb2RpZnksIG1lcmdlLCBwdWJsaXNoLCBkaXN0cmlidXRlLCBzdWJsaWNlbnNlLCBhbmQvb3Igc2VsbApjb3BpZXMgb2YgdGhlIFNvZnR3YXJlLCBhbmQgdG8gcGVybWl0IHBlcnNvbnMgdG8gd2hvbSB0aGUgU29mdHdhcmUgaXMKZnVybmlzaGVkIHRvIGRvIHNvLCBzdWJqZWN0IHRvIHRoZSBmb2xsb3dpbmcgY29uZGl0aW9uczoKClRoZSBhYm92ZSBjb3B5cmlnaHQgbm90aWNlIGFuZCB0aGlzIHBlcm1pc3Npb24gbm90aWNlIHNoYWxsIGJlIGluY2x1ZGVkIGluIGFsbApjb3BpZXMgb3Igc3Vic3RhbnRpYWwgcG9ydGlvbnMgb2YgdGhlIFNvZnR3YXJlLgoKVEhFIFNPRlRXQVJFIElTIFBST1ZJREVEICJBUyBJUyIsIFdJVEhPVVQgV0FSUkFOVFkgT0YgQU5ZIEtJTkQsIEVYUFJFU1MgT1IKSU1QTElFRCwgSU5DTFVESU5HIEJVVCBOT1QgTElNSVRFRCBUTyBUSEUgV0FSUkFOVElFUyBPRiBNRVJDSEFOVEFCSUxJVFksCkZJVE5FU1MgRk9SIEEgUEFSVElDVUxBUiBQVVJQT1NFIEFORCBOT05JTkZSSU5HRU1FTlQuIElOIE5PIEVWRU5UIFNIQUxMIFRIRQpBVVRIT1JTIE9SIENPUFlSSUdIVCBIT0xERVJTIEJFIExJQUJMRSBGT1IgQU5ZIENMQUlNLCBEQU1BR0VTIE9SIE9USEVSCkxJQUJJTElUWSwgV0hFVEhFUiBJTiBBTiBBQ1RJT04gT0YgQ09OVFJBQ1QsIFRPUlQgT1IgT1RIRVJXSVNFLCBBUklTSU5HIEZST00sCk9VVCBPRiBPUiBJTiBDT05ORUNUSU9OIFdJVEggVEhFIFNPRlRXQVJFIE9SIFRIRSBVU0UgT1IgT1RIRVIgREVBTElOR1MgSU4gVEhFClNPRlRXQVJFLgo='
+  >>> _encode_image(Path('LICENSE'))
+  'TUlUIExpY2Vuc2UKCkNvcHlyaWdodCAoYykgT2xsYW1hCgpQZXJtaXNzaW9uIGlzIGhlcmVieSBncmFudGVkLCBmcmVlIG9mIGNoYXJnZSwgdG8gYW55IHBlcnNvbiBvYnRhaW5pbmcgYSBjb3B5Cm9mIHRoaXMgc29mdHdhcmUgYW5kIGFzc29jaWF0ZWQgZG9jdW1lbnRhdGlvbiBmaWxlcyAodGhlICJTb2Z0d2FyZSIpLCB0byBkZWFsCmluIHRoZSBTb2Z0d2FyZSB3aXRob3V0IHJlc3RyaWN0aW9uLCBpbmNsdWRpbmcgd2l0aG91dCBsaW1pdGF0aW9uIHRoZSByaWdodHMKdG8gdXNlLCBjb3B5LCBtb2RpZnksIG1lcmdlLCBwdWJsaXNoLCBkaXN0cmlidXRlLCBzdWJsaWNlbnNlLCBhbmQvb3Igc2VsbApjb3BpZXMgb2YgdGhlIFNvZnR3YXJlLCBhbmQgdG8gcGVybWl0IHBlcnNvbnMgdG8gd2hvbSB0aGUgU29mdHdhcmUgaXMKZnVybmlzaGVkIHRvIGRvIHNvLCBzdWJqZWN0IHRvIHRoZSBmb2xsb3dpbmcgY29uZGl0aW9uczoKClRoZSBhYm92ZSBjb3B5cmlnaHQgbm90aWNlIGFuZCB0aGlzIHBlcm1pc3Npb24gbm90aWNlIHNoYWxsIGJlIGluY2x1ZGVkIGluIGFsbApjb3BpZXMgb3Igc3Vic3RhbnRpYWwgcG9ydGlvbnMgb2YgdGhlIFNvZnR3YXJlLgoKVEhFIFNPRlRXQVJFIElTIFBST1ZJREVEICJBUyBJUyIsIFdJVEhPVVQgV0FSUkFOVFkgT0YgQU5ZIEtJTkQsIEVYUFJFU1MgT1IKSU1QTElFRCwgSU5DTFVESU5HIEJVVCBOT1QgTElNSVRFRCBUTyBUSEUgV0FSUkFOVElFUyBPRiBNRVJDSEFOVEFCSUxJVFksCkZJVE5FU1MgRk9SIEEgUEFSVElDVUxBUiBQVVJQT1NFIEFORCBOT05JTkZSSU5HRU1FTlQuIElOIE5PIEVWRU5UIFNIQUxMIFRIRQpBVVRIT1JTIE9SIENPUFlSSUdIVCBIT0xERVJTIEJFIExJQUJMRSBGT1IgQU5ZIENMQUlNLCBEQU1BR0VTIE9SIE9USEVSCkxJQUJJTElUWSwgV0hFVEhFUiBJTiBBTiBBQ1RJT04gT0YgQ09OVFJBQ1QsIFRPUlQgT1IgT1RIRVJXSVNFLCBBUklTSU5HIEZST00sCk9VVCBPRiBPUiBJTiBDT05ORUNUSU9OIFdJVEggVEhFIFNPRlRXQVJFIE9SIFRIRSBVU0UgT1IgT1RIRVIgREVBTElOR1MgSU4gVEhFClNPRlRXQVJFLgo='
+  >>> _encode_image('YWJj')
+  'YWJj'
+  >>> _encode_image(b'YWJj')
+  'YWJj'
+  """
 
-  return b64.decode('utf-8')
+  if p := _as_path(image):
+    return b64encode(p.read_bytes()).decode('utf-8')
+
+  try:
+    b64decode(image, validate=True)
+    return image if isinstance(image, str) else image.decode('utf-8')
+  except (binascii.Error, TypeError):
+    ...
+
+  if b := _as_bytesio(image):
+    return b64encode(b.read()).decode('utf-8')
+
+  raise RequestError('image must be bytes, path-like object, or file-like object')
 
 
 def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]:
   if isinstance(s, str) or isinstance(s, Path):
-    return Path(s)
+    try:
+      if (p := Path(s)).exists():
+        return p
+    except Exception:
+      ...
   return None
 
 

tests/test_client.py

@@ -29,6 +29,7 @@ def test_client_chat(httpserver: HTTPServer):
       'stream': False,
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json(
     {
@@ -75,6 +76,7 @@ def test_client_chat_stream(httpserver: HTTPServer):
       'stream': True,
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_handler(stream_handler)
 
@@ -103,6 +105,7 @@ def test_client_chat_images(httpserver: HTTPServer):
       'stream': False,
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json(
     {
@@ -139,6 +142,7 @@ def test_client_generate(httpserver: HTTPServer):
       'images': [],
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json(
     {
@@ -183,6 +187,7 @@ def test_client_generate_stream(httpserver: HTTPServer):
       'images': [],
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_handler(stream_handler)
 
@@ -210,6 +215,7 @@ def test_client_generate_images(httpserver: HTTPServer):
       'images': ['iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC'],
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json(
     {
@@ -513,6 +519,7 @@ async def test_async_client_chat(httpserver: HTTPServer):
       'stream': False,
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json({})
 
@@ -550,6 +557,7 @@ async def test_async_client_chat_stream(httpserver: HTTPServer):
       'stream': True,
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_handler(stream_handler)
 
@@ -579,6 +587,7 @@ async def test_async_client_chat_images(httpserver: HTTPServer):
       'stream': False,
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json({})
 
@@ -606,6 +615,7 @@ async def test_async_client_generate(httpserver: HTTPServer):
       'images': [],
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json({})
 
@@ -645,6 +655,7 @@ async def test_async_client_generate_stream(httpserver: HTTPServer):
       'images': [],
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_handler(stream_handler)
 
@@ -673,6 +684,7 @@ async def test_async_client_generate_images(httpserver: HTTPServer):
       'images': ['iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC'],
       'format': '',
       'options': {},
+      'keep_alive': None,
     },
   ).respond_with_json({})