From cea391a041ac6f83b1d45a928c6b1aae354c80e0 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Tue, 9 Jan 2024 17:17:44 -0800 Subject: [PATCH] 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]