From cea391a041ac6f83b1d45a928c6b1aae354c80e0 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Tue, 9 Jan 2024 17:17:44 -0800 Subject: [PATCH 1/9] add generate and chat responses --- ollama/__init__.py | 9 ++++++++- ollama/_client.py | 46 +++++++++++++++++++++++++++++++++++++++------- ollama/_types.py | 35 ++++++++++++++++++++++++++++++++--- 3 files changed, 79 insertions(+), 11 deletions(-) diff --git a/ollama/__init__.py b/ollama/__init__.py index 2d0e94f..d48df1b 100644 --- a/ollama/__init__.py +++ b/ollama/__init__.py @@ -1,9 +1,16 @@ from ollama._client import Client, AsyncClient -from ollama._types import Message, Options +from ollama._types import ( + GenerateResponse, + ChatResponse, + Message, + Options, +) __all__ = [ 'Client', 'AsyncClient', + 'GenerateResponse', + 'ChatResponse', 'Message', 'Options', 'generate', diff --git a/ollama/_client.py b/ollama/_client.py index 19ec847..cc444f4 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -7,7 +7,7 @@ from pathlib import Path from hashlib import sha256 from base64 import b64encode -from typing import Any, AnyStr, Union, Optional, List, Mapping +from typing import Any, AnyStr, Union, Optional, Sequence, Mapping import sys @@ -56,13 +56,17 @@ class Client(BaseClient): prompt: str = '', system: str = '', template: str = '', - context: Optional[List[int]] = None, + context: Optional[Sequence[int]] = None, stream: bool = False, raw: bool = False, format: str = '', - images: Optional[List[AnyStr]] = None, + images: Optional[Sequence[AnyStr]] = None, options: Optional[Options] = None, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: + """ + Returns `GenerateResponse` if `stream` is `False`, otherwise returns a `GenerateResponse` generator. + """ + if not model: raise Exception('must provide a model') @@ -87,11 +91,15 @@ class Client(BaseClient): def chat( self, model: str = '', - messages: Optional[List[Message]] = None, + messages: Optional[Sequence[Message]] = None, stream: bool = False, format: str = '', options: Optional[Options] = None, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: + """ + Returns `ChatResponse` if `stream` is `False`, otherwise returns a `ChatResponse` generator. + """ + if not model: raise Exception('must provide a model') @@ -124,6 +132,9 @@ class Client(BaseClient): insecure: bool = False, stream: bool = False, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: + """ + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ return self._request_stream( 'POST', '/api/pull', @@ -141,6 +152,9 @@ class Client(BaseClient): insecure: bool = False, stream: bool = False, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: + """ + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ return self._request_stream( 'POST', '/api/push', @@ -159,6 +173,9 @@ class Client(BaseClient): modelfile: Optional[str] = None, stream: bool = False, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: + """ + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ if (realpath := _as_path(path)) and realpath.exists(): modelfile = self._parse_modelfile(realpath.read_text(), base=realpath.parent) elif modelfile: @@ -267,13 +284,16 @@ class AsyncClient(BaseClient): prompt: str = '', system: str = '', template: str = '', - context: Optional[List[int]] = None, + context: Optional[Sequence[int]] = None, stream: bool = False, raw: bool = False, format: str = '', - images: Optional[List[AnyStr]] = None, + images: Optional[Sequence[AnyStr]] = None, options: Optional[Options] = None, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: + """ + Returns `GenerateResponse` if `stream` is `False`, otherwise returns an asynchronous `GenerateResponse` generator. + """ if not model: raise Exception('must provide a model') @@ -298,11 +318,14 @@ class AsyncClient(BaseClient): async def chat( self, model: str = '', - messages: Optional[List[Message]] = None, + messages: Optional[Sequence[Message]] = None, stream: bool = False, format: str = '', options: Optional[Options] = None, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: + """ + Returns `ChatResponse` if `stream` is `False`, otherwise returns an asynchronous `ChatResponse` generator. + """ if not model: raise Exception('must provide a model') @@ -335,6 +358,9 @@ class AsyncClient(BaseClient): insecure: bool = False, stream: bool = False, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: + """ + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ return await self._request_stream( 'POST', '/api/pull', @@ -352,6 +378,9 @@ class AsyncClient(BaseClient): insecure: bool = False, stream: bool = False, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: + """ + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ return await self._request_stream( 'POST', '/api/push', @@ -370,6 +399,9 @@ class AsyncClient(BaseClient): modelfile: Optional[str] = None, stream: bool = False, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: + """ + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ if (realpath := _as_path(path)) and realpath.exists(): modelfile = await self._parse_modelfile(realpath.read_text(), base=realpath.parent) elif modelfile: diff --git a/ollama/_types.py b/ollama/_types.py index d263269..059dd1e 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -1,4 +1,4 @@ -from typing import Any, TypedDict, List +from typing import Any, TypedDict, Sequence import sys @@ -8,10 +8,39 @@ else: from typing import NotRequired +class BaseGenerateResponse(TypedDict): + model: str + created_at: str + done: bool + + total_duration: int + load_duration: int + prompt_eval_count: int + prompt_eval_duration: int + eval_count: int + eval_duration: int + + +class GenerateResponse(BaseGenerateResponse): + response: str + context: Sequence[int] + + class Message(TypedDict): role: str content: str - images: NotRequired[List[Any]] + images: NotRequired[Sequence[Any]] + + +class ChatResponse(BaseGenerateResponse): + message: Message + + +class ProgressResponse(TypedDict): + status: str + completed: int + total: int + digest: str class Options(TypedDict, total=False): @@ -50,4 +79,4 @@ class Options(TypedDict, total=False): mirostat_tau: float mirostat_eta: float penalize_newline: bool - stop: List[str] + stop: Sequence[str] From 2804a03d82af8cc86c29c076af2240d48d3943c8 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 09:49:09 -0800 Subject: [PATCH 2/9] httpx: kwargs --- ollama/_client.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/ollama/_client.py b/ollama/_client.py index cc444f4..09b4a47 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -20,14 +20,25 @@ from ollama._types import Message, Options class BaseClient: - def __init__(self, client, base_url: Optional[str] = None) -> None: - base_url = base_url or os.getenv('OLLAMA_HOST', 'http://127.0.0.1:11434') - self._client = client(base_url=base_url, follow_redirects=True, timeout=None) + def __init__( + self, + client, + base_url: Optional[str] = None, + follow_redirects: bool = True, + timeout: Any = None, + **kwargs, + ) -> None: + self._client = client( + base_url=base_url or os.getenv('OLLAMA_HOST', 'http://127.0.0.1:11434'), + follow_redirects=follow_redirects, + timeout=timeout, + **kwargs, + ) class Client(BaseClient): - def __init__(self, base_url: Optional[str] = None) -> None: - super().__init__(httpx.Client, base_url) + def __init__(self, base_url: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.Client, base_url, **kwargs) def _request(self, method: str, url: str, **kwargs) -> httpx.Response: response = self._client.request(method, url, **kwargs) @@ -247,8 +258,8 @@ class Client(BaseClient): class AsyncClient(BaseClient): - def __init__(self, base_url: Optional[str] = None) -> None: - super().__init__(httpx.AsyncClient, base_url) + def __init__(self, base_url: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.AsyncClient, base_url, **kwargs) async def _request(self, method: str, url: str, **kwargs) -> httpx.Response: response = await self._client.request(method, url, **kwargs) From 7601947a35d2ba84dc0698efe349aef3f317050e Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 11:31:58 -0800 Subject: [PATCH 3/9] request/response errors --- ollama/__init__.py | 6 +++ ollama/_client.py | 98 ++++++++++++++++++++++++++++++++++++---------- ollama/_types.py | 25 ++++++++++++ 3 files changed, 109 insertions(+), 20 deletions(-) diff --git a/ollama/__init__.py b/ollama/__init__.py index d48df1b..c4ebe2b 100644 --- a/ollama/__init__.py +++ b/ollama/__init__.py @@ -2,8 +2,11 @@ from ollama._client import Client, AsyncClient from ollama._types import ( GenerateResponse, ChatResponse, + ProgressResponse, Message, Options, + RequestError, + ResponseError, ) __all__ = [ @@ -11,8 +14,11 @@ __all__ = [ 'AsyncClient', 'GenerateResponse', 'ChatResponse', + 'ProgressResponse', 'Message', 'Options', + 'RequestError', + 'ResponseError', 'generate', 'chat', 'pull', diff --git a/ollama/_client.py b/ollama/_client.py index 09b4a47..d3b2a4a 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -16,7 +16,7 @@ if sys.version_info < (3, 9): else: from collections.abc import Iterator, AsyncIterator -from ollama._types import Message, Options +from ollama._types import Message, Options, RequestError, ResponseError class BaseClient: @@ -42,15 +42,26 @@ class Client(BaseClient): def _request(self, method: str, url: str, **kwargs) -> httpx.Response: response = self._client.request(method, url, **kwargs) - response.raise_for_status() + + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + raise ResponseError(e.response.text, e.response.status_code) from None + return response def _stream(self, method: str, url: str, **kwargs) -> Iterator[Mapping[str, Any]]: with self._client.stream(method, url, **kwargs) as r: + try: + r.raise_for_status() + except httpx.HTTPStatusError as e: + e.response.read() + raise ResponseError(e.response.text, e.response.status_code) from None + for line in r.iter_lines(): partial = json.loads(line) if e := partial.get('error'): - raise Exception(e) + raise ResponseError(e) yield partial def _request_stream( @@ -75,11 +86,17 @@ class Client(BaseClient): options: Optional[Options] = None, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: """ + Create a response using the requested model. + + Raises `RequestError` if a model is not provided. + + Raises `ResponseError` if the request could not be fulfilled. + Returns `GenerateResponse` if `stream` is `False`, otherwise returns a `GenerateResponse` generator. """ if not model: - raise Exception('must provide a model') + raise RequestError('must provide a model') return self._request_stream( 'POST', @@ -108,19 +125,25 @@ class Client(BaseClient): options: Optional[Options] = None, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: """ + Create a chat response using the requested model. + + Raises `RequestError` if a model is not provided. + + Raises `ResponseError` if the request could not be fulfilled. + Returns `ChatResponse` if `stream` is `False`, otherwise returns a `ChatResponse` generator. """ if not model: - raise Exception('must provide a model') + raise RequestError('must provide a model') for message in messages or []: if not isinstance(message, dict): raise TypeError('messages must be a list of strings') if not (role := message.get('role')) or role not in ['system', 'user', 'assistant']: - raise Exception('messages must contain a role and it must be one of "system", "user", or "assistant"') + raise RequestError('messages must contain a role and it must be one of "system", "user", or "assistant"') if not message.get('content'): - raise Exception('messages must contain content') + raise RequestError('messages must contain content') if images := message.get('images'): message['images'] = [_encode_image(image) for image in images] @@ -144,6 +167,8 @@ class Client(BaseClient): stream: bool = False, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: """ + Raises `ResponseError` if the request could not be fulfilled. + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. """ return self._request_stream( @@ -164,6 +189,8 @@ class Client(BaseClient): stream: bool = False, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: """ + Raises `ResponseError` if the request could not be fulfilled. + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. """ return self._request_stream( @@ -185,6 +212,8 @@ class Client(BaseClient): stream: bool = False, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: """ + Raises `ResponseError` if the request could not be fulfilled. + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. """ if (realpath := _as_path(path)) and realpath.exists(): @@ -192,7 +221,7 @@ class Client(BaseClient): elif modelfile: modelfile = self._parse_modelfile(modelfile) else: - raise Exception('must provide either path or modelfile') + raise RequestError('must provide either path or modelfile') return self._request_stream( 'POST', @@ -233,8 +262,8 @@ class Client(BaseClient): try: self._request('HEAD', f'/api/blobs/{digest}') - except httpx.HTTPStatusError as e: - if e.response.status_code != 404: + except ResponseError as e: + if e.status_code != 404: raise with open(path, 'rb') as r: @@ -263,16 +292,27 @@ class AsyncClient(BaseClient): async def _request(self, method: str, url: str, **kwargs) -> httpx.Response: response = await self._client.request(method, url, **kwargs) - response.raise_for_status() + + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + raise ResponseError(e.response.text, e.response.status_code) from None + return response async def _stream(self, method: str, url: str, **kwargs) -> AsyncIterator[Mapping[str, Any]]: async def inner(): async with self._client.stream(method, url, **kwargs) as r: + try: + r.raise_for_status() + except httpx.HTTPStatusError as e: + e.response.read() + raise ResponseError(e.response.text, e.response.status_code) from None + async for line in r.aiter_lines(): partial = json.loads(line) if e := partial.get('error'): - raise Exception(e) + raise ResponseError(e) yield partial return inner() @@ -303,10 +343,16 @@ class AsyncClient(BaseClient): options: Optional[Options] = None, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: """ + Create a response using the requested model. + + Raises `RequestError` if a model is not provided. + + Raises `ResponseError` if the request could not be fulfilled. + Returns `GenerateResponse` if `stream` is `False`, otherwise returns an asynchronous `GenerateResponse` generator. """ if not model: - raise Exception('must provide a model') + raise RequestError('must provide a model') return await self._request_stream( 'POST', @@ -335,18 +381,24 @@ class AsyncClient(BaseClient): options: Optional[Options] = None, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: """ + Create a chat response using the requested model. + + Raises `RequestError` if a model is not provided. + + Raises `ResponseError` if the request could not be fulfilled. + Returns `ChatResponse` if `stream` is `False`, otherwise returns an asynchronous `ChatResponse` generator. """ if not model: - raise Exception('must provide a model') + raise RequestError('must provide a model') for message in messages or []: if not isinstance(message, dict): raise TypeError('messages must be a list of strings') if not (role := message.get('role')) or role not in ['system', 'user', 'assistant']: - raise Exception('messages must contain a role and it must be one of "system", "user", or "assistant"') + raise RequestError('messages must contain a role and it must be one of "system", "user", or "assistant"') if not message.get('content'): - raise Exception('messages must contain content') + raise RequestError('messages must contain content') if images := message.get('images'): message['images'] = [_encode_image(image) for image in images] @@ -370,6 +422,8 @@ class AsyncClient(BaseClient): stream: bool = False, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: """ + Raises `ResponseError` if the request could not be fulfilled. + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. """ return await self._request_stream( @@ -390,6 +444,8 @@ class AsyncClient(BaseClient): stream: bool = False, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: """ + Raises `ResponseError` if the request could not be fulfilled. + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. """ return await self._request_stream( @@ -411,6 +467,8 @@ class AsyncClient(BaseClient): stream: bool = False, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: """ + Raises `ResponseError` if the request could not be fulfilled. + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. """ if (realpath := _as_path(path)) and realpath.exists(): @@ -418,7 +476,7 @@ class AsyncClient(BaseClient): elif modelfile: modelfile = await self._parse_modelfile(modelfile) else: - raise Exception('must provide either path or modelfile') + raise RequestError('must provide either path or modelfile') return await self._request_stream( 'POST', @@ -459,8 +517,8 @@ class AsyncClient(BaseClient): try: await self._request('HEAD', f'/api/blobs/{digest}') - except httpx.HTTPStatusError as e: - if e.response.status_code != 404: + except ResponseError as e: + if e.status_code != 404: raise async def upload_bytes(): @@ -498,7 +556,7 @@ def _encode_image(image) -> str: elif b := _as_bytesio(image): b64 = b64encode(b.read()) else: - raise Exception('images must be a list of bytes, path-like objects, or file-like objects') + raise RequestError('images must be a list of bytes, path-like objects, or file-like objects') return b64.decode('utf-8') diff --git a/ollama/_types.py b/ollama/_types.py index 059dd1e..226626e 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -80,3 +80,28 @@ class Options(TypedDict, total=False): mirostat_eta: float penalize_newline: bool stop: Sequence[str] + + +class RequestError(Exception): + """ + Common class for request errors. + """ + + def __init__(self, content: str): + super().__init__(content) + self.content = content + "Reason for the error." + + +class ResponseError(Exception): + """ + Common class for response errors. + """ + + def __init__(self, content: str, status_code: int = -1): + super().__init__(content) + self.content = content + "Reason for the error." + + self.status_code = status_code + "HTTP status code of the response." From f7e6980b1ba9c900ff15e795aecdc0f666476fe4 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 16:28:48 -0800 Subject: [PATCH 4/9] literal --- ollama/_client.py | 10 +++++----- ollama/_types.py | 5 +++-- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/ollama/_client.py b/ollama/_client.py index d3b2a4a..1ac5372 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -7,7 +7,7 @@ from pathlib import Path from hashlib import sha256 from base64 import b64encode -from typing import Any, AnyStr, Union, Optional, Sequence, Mapping +from typing import Any, AnyStr, Union, Optional, Sequence, Mapping, Literal import sys @@ -81,7 +81,7 @@ class Client(BaseClient): context: Optional[Sequence[int]] = None, stream: bool = False, raw: bool = False, - format: str = '', + format: Literal['', 'json'] = '', images: Optional[Sequence[AnyStr]] = None, options: Optional[Options] = None, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: @@ -121,7 +121,7 @@ class Client(BaseClient): model: str = '', messages: Optional[Sequence[Message]] = None, stream: bool = False, - format: str = '', + format: Literal['', 'json'] = '', options: Optional[Options] = None, ) -> Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]: """ @@ -338,7 +338,7 @@ class AsyncClient(BaseClient): context: Optional[Sequence[int]] = None, stream: bool = False, raw: bool = False, - format: str = '', + format: Literal['', 'json'] = '', images: Optional[Sequence[AnyStr]] = None, options: Optional[Options] = None, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: @@ -377,7 +377,7 @@ class AsyncClient(BaseClient): model: str = '', messages: Optional[Sequence[Message]] = None, stream: bool = False, - format: str = '', + format: Literal['', 'json'] = '', options: Optional[Options] = None, ) -> Union[Mapping[str, Any], AsyncIterator[Mapping[str, Any]]]: """ diff --git a/ollama/_types.py b/ollama/_types.py index 226626e..8258825 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -1,4 +1,4 @@ -from typing import Any, TypedDict, Sequence +from typing import Any, TypedDict, Sequence, Literal import sys @@ -27,7 +27,8 @@ class GenerateResponse(BaseGenerateResponse): class Message(TypedDict): - role: str + role: Literal['user', 'assistant', 'system'] + content: str images: NotRequired[Sequence[Any]] From 7c2ec01d2f9c0938fe8d61035ea34b302a6c371a Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 17:02:36 -0800 Subject: [PATCH 5/9] docstrings --- ollama/_client.py | 10 ++++++++++ ollama/_types.py | 45 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/ollama/_client.py b/ollama/_client.py index 1ac5372..6c2c9b3 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -28,6 +28,16 @@ class BaseClient: timeout: Any = None, **kwargs, ) -> None: + """ + Creates a httpx client. Default parameters are the same as those defined in httpx + except for the following: + + - `base_url`: http://127.0.0.1:11434 + - `follow_redirects`: True + - `timeout`: None + + `kwargs` are passed to the httpx client. + """ self._client = client( base_url=base_url or os.getenv('OLLAMA_HOST', 'http://127.0.0.1:11434'), follow_redirects=follow_redirects, diff --git a/ollama/_types.py b/ollama/_types.py index 8258825..b8fa06c 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -10,31 +10,76 @@ else: class BaseGenerateResponse(TypedDict): model: str + "Model used to generate response." + created_at: str + "Time when the request was created." + done: bool + "True if response is complete, otherwise False. Useful for streaming to detect the final response." total_duration: int + "Total duration in nanoseconds." + load_duration: int + "Load duration in nanoseconds." + prompt_eval_count: int + "Number of tokens evaluated in the prompt." + prompt_eval_duration: int + "Duration of evaluating the prompt in nanoseconds." + eval_count: int + "Number of tokens evaluated in inference." + eval_duration: int + "Duration of evaluating inference in nanoseconds." class GenerateResponse(BaseGenerateResponse): + """ + Response returned by generate requests. + """ + response: str + "Response content. When streaming, this contains a fragment of the response." + context: Sequence[int] + "Tokenized history up to the point of the response." class Message(TypedDict): + """ + Chat message. + """ + role: Literal['user', 'assistant', 'system'] + "Assumed role of the message. Response messages always has role 'assistant'." content: str + "Content of the message. Response messages contains message fragments when streaming." + images: NotRequired[Sequence[Any]] + """ + Optional list of image data for multimodal models. + + Valid input types are: + + - `str` or path-like object: path to image file + - `bytes` or bytes-like object: raw image data + + Valid image formats depend on the model. See the model card for more information. + """ class ChatResponse(BaseGenerateResponse): + """ + Response returned by chat requests. + """ + message: Message + "Response message." class ProgressResponse(TypedDict): From 64127ee821f11b8812e2b9a2b8ddd31bd5518caa Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 17:02:49 -0800 Subject: [PATCH 6/9] parse json error --- README.md | 15 +++++++++++++++ ollama/_types.py | 8 ++++++++ 2 files changed, 23 insertions(+) diff --git a/README.md b/README.md index ec7801c..770cc6a 100644 --- a/README.md +++ b/README.md @@ -68,3 +68,18 @@ async def chat(): asyncio.run(chat()) ``` + +## Handling Errors + +Errors are raised if requests return an error status or if an error is detected while streaming. + +```python +model = 'does-not-yet-exist' + +try: + ollama.chat(model) +except ollama.ResponseError as e: + print('Error:', e.content) + if e.status_code == 404: + ollama.pull(model) +``` diff --git a/ollama/_types.py b/ollama/_types.py index b8fa06c..fca614a 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -1,3 +1,4 @@ +import json from typing import Any, TypedDict, Sequence, Literal import sys @@ -145,6 +146,13 @@ class ResponseError(Exception): """ def __init__(self, content: str, status_code: int = -1): + try: + # try to parse content as JSON and extract 'error' + # fallback to raw content if JSON parsing fails + content = json.loads(content).get('error', content) + except json.JSONDecodeError: + ... + super().__init__(content) self.content = content "Reason for the error." From 008a6a6b00d6df0536d1595b85e7b6dae1028273 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 17:20:23 -0800 Subject: [PATCH 7/9] add embeddings --- ollama/__init__.py | 2 ++ ollama/_client.py | 24 ++++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/ollama/__init__.py b/ollama/__init__.py index c4ebe2b..048bb14 100644 --- a/ollama/__init__.py +++ b/ollama/__init__.py @@ -21,6 +21,7 @@ __all__ = [ 'ResponseError', 'generate', 'chat', + 'embeddings', 'pull', 'push', 'create', @@ -34,6 +35,7 @@ _client = Client() generate = _client.generate chat = _client.chat +embeddings = _client.embeddings pull = _client.pull push = _client.push create = _client.create diff --git a/ollama/_client.py b/ollama/_client.py index 6c2c9b3..e4aceac 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -170,6 +170,17 @@ class Client(BaseClient): stream=stream, ) + def embeddings(self, model: str = '', prompt: str = '', options: Optional[Options] = None) -> Sequence[float]: + return self._request( + 'POST', + '/api/embeddings', + json={ + 'model': model, + 'prompt': prompt, + 'options': options or {}, + }, + ).json() + def pull( self, model: str, @@ -425,6 +436,19 @@ class AsyncClient(BaseClient): stream=stream, ) + async def embeddings(self, model: str = '', prompt: str = '', options: Optional[Options] = None) -> Sequence[float]: + response = await self._request( + 'POST', + '/api/embeddings', + json={ + 'model': model, + 'prompt': prompt, + 'options': options or {}, + }, + ) + + return response.json() + async def pull( self, model: str, From d2ea72ae0753735116bd498967c7681581c62cbb Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 10 Jan 2024 17:20:35 -0800 Subject: [PATCH 8/9] return full list objects --- ollama/_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ollama/_client.py b/ollama/_client.py index e4aceac..35f6f2c 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -297,7 +297,7 @@ class Client(BaseClient): return {'status': 'success' if response.status_code == 200 else 'error'} def list(self) -> Mapping[str, Any]: - return self._request('GET', '/api/tags').json().get('models', []) + return self._request('GET', '/api/tags').json() def copy(self, source: str, target: str) -> Mapping[str, Any]: response = self._request('POST', '/api/copy', json={'source': source, 'destination': target}) @@ -573,7 +573,7 @@ class AsyncClient(BaseClient): async def list(self) -> Mapping[str, Any]: response = await self._request('GET', '/api/tags') - return response.json().get('models', []) + return response.json() async def copy(self, source: str, target: str) -> Mapping[str, Any]: response = await self._request('POST', '/api/copy', json={'source': source, 'destination': target}) From 1f68bae4830fc040e7e01f97bc3bd1f84df2c241 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Thu, 11 Jan 2024 09:35:10 -0800 Subject: [PATCH 9/9] wording --- ollama/_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ollama/_client.py b/ollama/_client.py index 35f6f2c..14608ba 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -149,7 +149,7 @@ class Client(BaseClient): for message in messages or []: if not isinstance(message, dict): - raise TypeError('messages must be a list of strings') + raise TypeError('messages must be a list of Message or dict-like objects') if not (role := message.get('role')) or role not in ['system', 'user', 'assistant']: raise RequestError('messages must contain a role and it must be one of "system", "user", or "assistant"') if not message.get('content'):