From 6b3c5b23bfdab90bb48c7f86811aad0e887668ef Mon Sep 17 00:00:00 2001 From: r7mekmy4g67w6l Date: Thu, 6 Aug 2026 17:07:25 +0200 Subject: [PATCH 1/6] fix: preserve empty string message content in chat requests Signed-off-by: r7mekmy4g67w6l --- ollama/_client.py | 2861 +++++++++++++++++++++++---------------------- 1 file changed, 1431 insertions(+), 1430 deletions(-) diff --git a/ollama/_client.py b/ollama/_client.py index 8dfce824..156d4e59 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -1,1430 +1,1431 @@ -import contextlib -import ipaddress -import json -import os -import platform -import sys -import urllib.parse -from hashlib import sha256 -from os import PathLike -from pathlib import Path -from typing import ( - Any, - Callable, - Dict, - List, - Literal, - Mapping, - Optional, - Sequence, - Type, - TypeVar, - Union, - overload, -) - -import anyio -from pydantic.json_schema import JsonSchemaValue - -from ollama._utils import convert_function_to_tool - -if sys.version_info < (3, 9): - from typing import AsyncIterator, Iterator -else: - from collections.abc import AsyncIterator, Iterator - -from importlib import metadata - -try: - __version__ = metadata.version('ollama') -except metadata.PackageNotFoundError: - __version__ = '0.0.0' - -import httpx - -from ollama._types import ( - ChatRequest, - ChatResponse, - CopyRequest, - CreateRequest, - DeleteRequest, - EmbeddingsRequest, - EmbeddingsResponse, - EmbedRequest, - EmbedResponse, - GenerateRequest, - GenerateResponse, - Image, - ListResponse, - Message, - Options, - ProcessResponse, - ProgressResponse, - PullRequest, - PushRequest, - ResponseError, - ShowRequest, - ShowResponse, - StatusResponse, - Tool, - WebFetchRequest, - WebFetchResponse, - WebSearchRequest, - WebSearchResponse, -) - -T = TypeVar('T') - - -class BaseClient(contextlib.AbstractContextManager, contextlib.AbstractAsyncContextManager): - def __init__( - self, - client, - host: Optional[str] = None, - *, - follow_redirects: bool = True, - timeout: Any = None, - headers: Optional[Mapping[str, str]] = None, - **kwargs, - ) -> None: - """ - Creates a httpx client. Default parameters are the same as those defined in httpx - except for the following: - - `follow_redirects`: True - - `timeout`: None - `kwargs` are passed to the httpx client. - """ - - headers = { - k.lower(): v - for k, v in { - **(headers or {}), - 'Content-Type': 'application/json', - 'Accept': 'application/json', - 'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}', - }.items() - if v is not None - } - api_key = os.getenv('OLLAMA_API_KEY', None) - if not headers.get('authorization') and api_key: - headers['authorization'] = f'Bearer {api_key}' - - self._client = client( - base_url=_parse_host(host or os.getenv('OLLAMA_HOST')), - follow_redirects=follow_redirects, - timeout=timeout, - headers=headers, - **kwargs, - ) - - def __exit__(self, exc_type, exc_val, exc_tb): - self.close() - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.close() - - -CONNECTION_ERROR_MESSAGE = 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - - -class Client(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.Client, host, **kwargs) - - def close(self): - self._client.close() - - def _request_raw(self, *args, **kwargs): - try: - r = self._client.request(*args, **kwargs) - r.raise_for_status() - return r - except httpx.HTTPStatusError as e: - raise ResponseError(e.response.text, e.response.status_code) from None - except httpx.ConnectError: - raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None - - @overload - def _request( - self, - cls: Type[T], - *args, - stream: Literal[False] = False, - **kwargs, - ) -> T: ... - - @overload - def _request( - self, - cls: Type[T], - *args, - stream: Literal[True] = True, - **kwargs, - ) -> Iterator[T]: ... - - @overload - def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, Iterator[T]]: ... - - def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, Iterator[T]]: - if stream: - - def inner(): - with self._client.stream(*args, **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(): - part = json.loads(line) - if err := part.get('error'): - raise ResponseError(err) - yield cls(**part) - - return inner() - - return cls(**self._request_raw(*args, **kwargs).json()) - - @overload - def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> GenerateResponse: ... - - @overload - def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> Iterator[GenerateResponse]: ... - - def generate( - self, - model: str = '', - prompt: Optional[str] = None, - suffix: Optional[str] = None, - *, - system: Optional[str] = None, - template: Optional[str] = None, - context: Optional[Sequence[int]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: Optional[bool] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> Union[GenerateResponse, Iterator[GenerateResponse]]: - """ - 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. - """ - - return self._request( - GenerateResponse, - 'POST', - '/api/generate', - json=GenerateRequest( - model=model, - prompt=prompt, - suffix=suffix, - system=system, - template=template, - context=context, - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - raw=raw, - format=format, - images=list(_copy_images(images)) if images else None, - options=options, - keep_alive=keep_alive, - width=width, - height=height, - steps=steps, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> ChatResponse: ... - - @overload - def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> Iterator[ChatResponse]: ... - - def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> Union[ChatResponse, Iterator[ChatResponse]]: - """ - Create a chat response using the requested model. - - Args: - tools: - A JSON schema as a dict, an Ollama Tool or a Python Function. - Python functions need to follow Google style docstrings to be converted to an Ollama Tool. - For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings - stream: Whether to stream the response. - format: The format of the response. - - Example: - def add_two_numbers(a: int, b: int) -> int: - ''' - Add two numbers together. - - Args: - a: First number to add - b: Second number to add - - Returns: - int: The sum of a and b - ''' - return a + b - - client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) - - 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. - """ - return self._request( - ChatResponse, - 'POST', - '/api/chat', - json=ChatRequest( - model=model, - messages=list(_copy_messages(messages)), - tools=list(_copy_tools(tools)), - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - format=format, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - stream=stream, - ) - - def embed( - self, - model: str = '', - input: Union[str, Sequence[str]] = '', - truncate: Optional[bool] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - dimensions: Optional[int] = None, - ) -> EmbedResponse: - return self._request( - EmbedResponse, - 'POST', - '/api/embed', - json=EmbedRequest( - model=model, - input=input, - truncate=truncate, - options=options, - keep_alive=keep_alive, - dimensions=dimensions, - ).model_dump(exclude_none=True), - ) - - def embeddings( - self, - model: str = '', - prompt: Optional[str] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> EmbeddingsResponse: - """ - Deprecated in favor of `embed`. - """ - return self._request( - EmbeddingsResponse, - 'POST', - '/api/embeddings', - json=EmbeddingsRequest( - model=model, - prompt=prompt, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - ) - - @overload - def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> Iterator[ProgressResponse]: ... - - def pull( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return self._request( - ProgressResponse, - 'POST', - '/api/pull', - json=PullRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> Iterator[ProgressResponse]: ... - - def push( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return self._request( - ProgressResponse, - 'POST', - '/api/push', - json=PushRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[True] = True, - ) -> Iterator[ProgressResponse]: ... - - def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: bool = False, - ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return self._request( - ProgressResponse, - 'POST', - '/api/create', - json=CreateRequest( - model=model, - stream=stream, - quantize=quantize, - from_=from_, - files=files, - adapters=adapters, - license=license, - template=template, - system=system, - parameters=parameters, - messages=messages, - ).model_dump(exclude_none=True), - stream=stream, - ) - - def create_blob(self, path: Union[str, Path]) -> str: - sha256sum = sha256() - with open(path, 'rb') as r: - while True: - chunk = r.read(32 * 1024) - if not chunk: - break - sha256sum.update(chunk) - - digest = f'sha256:{sha256sum.hexdigest()}' - - with open(path, 'rb') as r: - self._request_raw('POST', f'/api/blobs/{digest}', content=r) - - return digest - - def list(self) -> ListResponse: - return self._request( - ListResponse, - 'GET', - '/api/tags', - ) - - def delete(self, model: str) -> StatusResponse: - r = self._request_raw( - 'DELETE', - '/api/delete', - json=DeleteRequest( - model=model, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - def copy(self, source: str, destination: str) -> StatusResponse: - r = self._request_raw( - 'POST', - '/api/copy', - json=CopyRequest( - source=source, - destination=destination, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - def show(self, model: str) -> ShowResponse: - return self._request( - ShowResponse, - 'POST', - '/api/show', - json=ShowRequest( - model=model, - ).model_dump(exclude_none=True), - ) - - def ps(self) -> ProcessResponse: - return self._request( - ProcessResponse, - 'GET', - '/api/ps', - ) - - def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: - """ - Performs a web search - - Args: - query: The query to search for - max_results: The maximum number of results to return (default: 3) - - Returns: - WebSearchResponse with the search results - Raises: - ValueError: If OLLAMA_API_KEY environment variable is not set - """ - if not self._client.headers.get('authorization', '').startswith('Bearer '): - raise ValueError('Authorization header with Bearer token is required for web search') - - return self._request( - WebSearchResponse, - 'POST', - 'https://ollama.com/api/web_search', - json=WebSearchRequest( - query=query, - max_results=max_results, - ).model_dump(exclude_none=True), - ) - - def web_fetch(self, url: str) -> WebFetchResponse: - """ - Fetches the content of a web page for the provided URL. - - Args: - url: The URL to fetch - - Returns: - WebFetchResponse with the fetched result - """ - if not self._client.headers.get('authorization', '').startswith('Bearer '): - raise ValueError('Authorization header with Bearer token is required for web fetch') - - return self._request( - WebFetchResponse, - 'POST', - 'https://ollama.com/api/web_fetch', - json=WebFetchRequest( - url=url, - ).model_dump(exclude_none=True), - ) - - -class AsyncClient(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.AsyncClient, host, **kwargs) - - async def close(self): - await self._client.aclose() - - async def _request_raw(self, *args, **kwargs): - try: - r = await self._client.request(*args, **kwargs) - r.raise_for_status() - return r - except httpx.HTTPStatusError as e: - raise ResponseError(e.response.text, e.response.status_code) from None - except httpx.ConnectError: - raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None - - @overload - async def _request( - self, - cls: Type[T], - *args, - stream: Literal[False] = False, - **kwargs, - ) -> T: ... - - @overload - async def _request( - self, - cls: Type[T], - *args, - stream: Literal[True] = True, - **kwargs, - ) -> AsyncIterator[T]: ... - - @overload - async def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, AsyncIterator[T]]: ... - - async def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, AsyncIterator[T]]: - if stream: - - async def inner(): - async with self._client.stream(*args, **kwargs) as r: - try: - r.raise_for_status() - except httpx.HTTPStatusError as e: - await e.response.aread() - raise ResponseError(e.response.text, e.response.status_code) from None - - async for line in r.aiter_lines(): - part = json.loads(line) - if err := part.get('error'): - raise ResponseError(err) - yield cls(**part) - - return inner() - - return cls(**(await self._request_raw(*args, **kwargs)).json()) - - async def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: - """ - Performs a web search - - Args: - query: The query to search for - max_results: The maximum number of results to return (default: 3) - - Returns: - WebSearchResponse with the search results - """ - return await self._request( - WebSearchResponse, - 'POST', - 'https://ollama.com/api/web_search', - json=WebSearchRequest( - query=query, - max_results=max_results, - ).model_dump(exclude_none=True), - ) - - async def web_fetch(self, url: str) -> WebFetchResponse: - """ - Fetches the content of a web page for the provided URL. - - Args: - url: The URL to fetch - - Returns: - WebFetchResponse with the fetched result - """ - return await self._request( - WebFetchResponse, - 'POST', - 'https://ollama.com/api/web_fetch', - json=WebFetchRequest( - url=url, - ).model_dump(exclude_none=True), - ) - - @overload - async def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> GenerateResponse: ... - - @overload - async def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> AsyncIterator[GenerateResponse]: ... - - async def generate( - self, - model: str = '', - prompt: Optional[str] = None, - suffix: Optional[str] = None, - *, - system: Optional[str] = None, - template: Optional[str] = None, - context: Optional[Sequence[int]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: Optional[bool] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> Union[GenerateResponse, AsyncIterator[GenerateResponse]]: - """ - 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. - """ - return await self._request( - GenerateResponse, - 'POST', - '/api/generate', - json=GenerateRequest( - model=model, - prompt=prompt, - suffix=suffix, - system=system, - template=template, - context=context, - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - raw=raw, - format=format, - images=list(_copy_images(images)) if images else None, - options=options, - keep_alive=keep_alive, - width=width, - height=height, - steps=steps, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - async def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> ChatResponse: ... - - @overload - async def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> AsyncIterator[ChatResponse]: ... - - async def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> Union[ChatResponse, AsyncIterator[ChatResponse]]: - """ - Create a chat response using the requested model. - - Args: - tools: - A JSON schema as a dict, an Ollama Tool or a Python Function. - Python functions need to follow Google style docstrings to be converted to an Ollama Tool. - For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings - stream: Whether to stream the response. - format: The format of the response. - - Example: - def add_two_numbers(a: int, b: int) -> int: - ''' - Add two numbers together. - - Args: - a: First number to add - b: Second number to add - - Returns: - int: The sum of a and b - ''' - return a + b - - await client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) - - 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. - """ - - return await self._request( - ChatResponse, - 'POST', - '/api/chat', - json=ChatRequest( - model=model, - messages=list(_copy_messages(messages)), - tools=list(_copy_tools(tools)), - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - format=format, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - stream=stream, - ) - - async def embed( - self, - model: str = '', - input: Union[str, Sequence[str]] = '', - truncate: Optional[bool] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - dimensions: Optional[int] = None, - ) -> EmbedResponse: - return await self._request( - EmbedResponse, - 'POST', - '/api/embed', - json=EmbedRequest( - model=model, - input=input, - truncate=truncate, - options=options, - keep_alive=keep_alive, - dimensions=dimensions, - ).model_dump(exclude_none=True), - ) - - async def embeddings( - self, - model: str = '', - prompt: Optional[str] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> EmbeddingsResponse: - """ - Deprecated in favor of `embed`. - """ - return await self._request( - EmbeddingsResponse, - 'POST', - '/api/embeddings', - json=EmbeddingsRequest( - model=model, - prompt=prompt, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - ) - - @overload - async def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - async def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> AsyncIterator[ProgressResponse]: ... - - async def pull( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return await self._request( - ProgressResponse, - 'POST', - '/api/pull', - json=PullRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - async def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - async def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> AsyncIterator[ProgressResponse]: ... - - async def push( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return await self._request( - ProgressResponse, - 'POST', - '/api/push', - json=PushRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - async def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - async def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[True] = True, - ) -> AsyncIterator[ProgressResponse]: ... - - async def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: bool = False, - ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - - return await self._request( - ProgressResponse, - 'POST', - '/api/create', - json=CreateRequest( - model=model, - stream=stream, - quantize=quantize, - from_=from_, - files=files, - adapters=adapters, - license=license, - template=template, - system=system, - parameters=parameters, - messages=messages, - ).model_dump(exclude_none=True), - stream=stream, - ) - - async def create_blob(self, path: Union[str, Path]) -> str: - sha256sum = sha256() - async with await anyio.open_file(path, 'rb') as r: - while True: - chunk = await r.read(32 * 1024) - if not chunk: - break - sha256sum.update(chunk) - - digest = f'sha256:{sha256sum.hexdigest()}' - - async def upload_bytes(): - async with await anyio.open_file(path, 'rb') as r: - while True: - chunk = await r.read(32 * 1024) - if not chunk: - break - yield chunk - - await self._request_raw('POST', f'/api/blobs/{digest}', content=upload_bytes()) - - return digest - - async def list(self) -> ListResponse: - return await self._request( - ListResponse, - 'GET', - '/api/tags', - ) - - async def delete(self, model: str) -> StatusResponse: - r = await self._request_raw( - 'DELETE', - '/api/delete', - json=DeleteRequest( - model=model, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - async def copy(self, source: str, destination: str) -> StatusResponse: - r = await self._request_raw( - 'POST', - '/api/copy', - json=CopyRequest( - source=source, - destination=destination, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - async def show(self, model: str) -> ShowResponse: - return await self._request( - ShowResponse, - 'POST', - '/api/show', - json=ShowRequest( - model=model, - ).model_dump(exclude_none=True), - ) - - async def ps(self) -> ProcessResponse: - return await self._request( - ProcessResponse, - 'GET', - '/api/ps', - ) - - -def _copy_images(images: Optional[Sequence[Union[Image, Any]]]) -> Iterator[Image]: - for image in images or []: - yield image if isinstance(image, Image) else Image(value=image) - - -def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]]]) -> Iterator[Message]: - for message in messages or []: - yield Message.model_validate( - {k: list(_copy_images(v)) if k == 'images' else v for k, v in dict(message).items() if v}, - ) - - -def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None) -> Iterator[Tool]: - for unprocessed_tool in tools or []: - yield convert_function_to_tool(unprocessed_tool) if callable(unprocessed_tool) else Tool.model_validate(unprocessed_tool) - - -def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]: - if isinstance(s, (str, Path)): - try: - if (p := Path(s)).exists(): - return p - except Exception: - ... - return None - - -def _parse_host(host: Optional[str]) -> str: - """ - >>> _parse_host(None) - 'http://127.0.0.1:11434' - >>> _parse_host('') - 'http://127.0.0.1:11434' - >>> _parse_host('1.2.3.4') - 'http://1.2.3.4:11434' - >>> _parse_host(':56789') - 'http://127.0.0.1:56789' - >>> _parse_host('1.2.3.4:56789') - 'http://1.2.3.4:56789' - >>> _parse_host('http://1.2.3.4') - 'http://1.2.3.4:80' - >>> _parse_host('https://1.2.3.4') - 'https://1.2.3.4:443' - >>> _parse_host('https://1.2.3.4:56789') - 'https://1.2.3.4:56789' - >>> _parse_host('example.com') - 'http://example.com:11434' - >>> _parse_host('example.com:56789') - 'http://example.com:56789' - >>> _parse_host('http://example.com') - 'http://example.com:80' - >>> _parse_host('https://example.com') - 'https://example.com:443' - >>> _parse_host('https://example.com:56789') - 'https://example.com:56789' - >>> _parse_host('example.com/') - 'http://example.com:11434' - >>> _parse_host('example.com:56789/') - 'http://example.com:56789' - >>> _parse_host('example.com/path') - 'http://example.com:11434/path' - >>> _parse_host('example.com:56789/path') - 'http://example.com:56789/path' - >>> _parse_host('https://example.com:56789/path') - 'https://example.com:56789/path' - >>> _parse_host('example.com:56789/path/') - 'http://example.com:56789/path' - >>> _parse_host('[0001:002:003:0004::1]') - 'http://[0001:002:003:0004::1]:11434' - >>> _parse_host('[0001:002:003:0004::1]:56789') - 'http://[0001:002:003:0004::1]:56789' - >>> _parse_host('http://[0001:002:003:0004::1]') - 'http://[0001:002:003:0004::1]:80' - >>> _parse_host('https://[0001:002:003:0004::1]') - 'https://[0001:002:003:0004::1]:443' - >>> _parse_host('https://[0001:002:003:0004::1]:56789') - 'https://[0001:002:003:0004::1]:56789' - >>> _parse_host('[0001:002:003:0004::1]/') - 'http://[0001:002:003:0004::1]:11434' - >>> _parse_host('[0001:002:003:0004::1]:56789/') - 'http://[0001:002:003:0004::1]:56789' - >>> _parse_host('[0001:002:003:0004::1]/path') - 'http://[0001:002:003:0004::1]:11434/path' - >>> _parse_host('[0001:002:003:0004::1]:56789/path') - 'http://[0001:002:003:0004::1]:56789/path' - >>> _parse_host('https://[0001:002:003:0004::1]:56789/path') - 'https://[0001:002:003:0004::1]:56789/path' - >>> _parse_host('[0001:002:003:0004::1]:56789/path/') - 'http://[0001:002:003:0004::1]:56789/path' - """ - - host, port = host or '', 11434 - scheme, _, hostport = host.partition('://') - if not hostport: - scheme, hostport = 'http', host - elif scheme == 'http': - port = 80 - elif scheme == 'https': - port = 443 - - split = urllib.parse.urlsplit(f'{scheme}://{hostport}') - host = split.hostname or '127.0.0.1' - port = split.port or port - - try: - if isinstance(ipaddress.ip_address(host), ipaddress.IPv6Address): - # Fix missing square brackets for IPv6 from urlsplit - host = f'[{host}]' - except ValueError: - ... - - if path := split.path.strip('/'): - return f'{scheme}://{host}:{port}/{path}' - - return f'{scheme}://{host}:{port}' +import contextlib +import ipaddress +import json +import os +import platform +import sys +import urllib.parse +from hashlib import sha256 +from os import PathLike +from pathlib import Path +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + Mapping, + Optional, + Sequence, + Type, + TypeVar, + Union, + overload, +) + +import anyio +from pydantic.json_schema import JsonSchemaValue + +from ollama._utils import convert_function_to_tool + +if sys.version_info < (3, 9): + from typing import AsyncIterator, Iterator +else: + from collections.abc import AsyncIterator, Iterator + +from importlib import metadata + +try: + __version__ = metadata.version('ollama') +except metadata.PackageNotFoundError: + __version__ = '0.0.0' + +import httpx + +from ollama._types import ( + ChatRequest, + ChatResponse, + CopyRequest, + CreateRequest, + DeleteRequest, + EmbeddingsRequest, + EmbeddingsResponse, + EmbedRequest, + EmbedResponse, + GenerateRequest, + GenerateResponse, + Image, + ListResponse, + Message, + Options, + ProcessResponse, + ProgressResponse, + PullRequest, + PushRequest, + ResponseError, + ShowRequest, + ShowResponse, + StatusResponse, + Tool, + WebFetchRequest, + WebFetchResponse, + WebSearchRequest, + WebSearchResponse, +) + +T = TypeVar('T') + + +class BaseClient(contextlib.AbstractContextManager, contextlib.AbstractAsyncContextManager): + def __init__( + self, + client, + host: Optional[str] = None, + *, + follow_redirects: bool = True, + timeout: Any = None, + headers: Optional[Mapping[str, str]] = None, + **kwargs, + ) -> None: + """ + Creates a httpx client. Default parameters are the same as those defined in httpx + except for the following: + - `follow_redirects`: True + - `timeout`: None + `kwargs` are passed to the httpx client. + """ + + headers = { + k.lower(): v + for k, v in { + **(headers or {}), + 'Content-Type': 'application/json', + 'Accept': 'application/json', + 'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}', + }.items() + if v is not None + } + api_key = os.getenv('OLLAMA_API_KEY', None) + if not headers.get('authorization') and api_key: + headers['authorization'] = f'Bearer {api_key}' + + self._client = client( + base_url=_parse_host(host or os.getenv('OLLAMA_HOST')), + follow_redirects=follow_redirects, + timeout=timeout, + headers=headers, + **kwargs, + ) + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.close() + + +CONNECTION_ERROR_MESSAGE = 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + + +class Client(BaseClient): + def __init__(self, host: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.Client, host, **kwargs) + + def close(self): + self._client.close() + + def _request_raw(self, *args, **kwargs): + try: + r = self._client.request(*args, **kwargs) + r.raise_for_status() + return r + except httpx.HTTPStatusError as e: + raise ResponseError(e.response.text, e.response.status_code) from None + except httpx.ConnectError: + raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None + + @overload + def _request( + self, + cls: Type[T], + *args, + stream: Literal[False] = False, + **kwargs, + ) -> T: ... + + @overload + def _request( + self, + cls: Type[T], + *args, + stream: Literal[True] = True, + **kwargs, + ) -> Iterator[T]: ... + + @overload + def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, Iterator[T]]: ... + + def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, Iterator[T]]: + if stream: + + def inner(): + with self._client.stream(*args, **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(): + part = json.loads(line) + if err := part.get('error'): + raise ResponseError(err) + yield cls(**part) + + return inner() + + return cls(**self._request_raw(*args, **kwargs).json()) + + @overload + def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> GenerateResponse: ... + + @overload + def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> Iterator[GenerateResponse]: ... + + def generate( + self, + model: str = '', + prompt: Optional[str] = None, + suffix: Optional[str] = None, + *, + system: Optional[str] = None, + template: Optional[str] = None, + context: Optional[Sequence[int]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: Optional[bool] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> Union[GenerateResponse, Iterator[GenerateResponse]]: + """ + 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. + """ + + return self._request( + GenerateResponse, + 'POST', + '/api/generate', + json=GenerateRequest( + model=model, + prompt=prompt, + suffix=suffix, + system=system, + template=template, + context=context, + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + raw=raw, + format=format, + images=list(_copy_images(images)) if images else None, + options=options, + keep_alive=keep_alive, + width=width, + height=height, + steps=steps, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> ChatResponse: ... + + @overload + def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> Iterator[ChatResponse]: ... + + def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> Union[ChatResponse, Iterator[ChatResponse]]: + """ + Create a chat response using the requested model. + + Args: + tools: + A JSON schema as a dict, an Ollama Tool or a Python Function. + Python functions need to follow Google style docstrings to be converted to an Ollama Tool. + For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings + stream: Whether to stream the response. + format: The format of the response. + + Example: + def add_two_numbers(a: int, b: int) -> int: + ''' + Add two numbers together. + + Args: + a: First number to add + b: Second number to add + + Returns: + int: The sum of a and b + ''' + return a + b + + client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) + + 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. + """ + return self._request( + ChatResponse, + 'POST', + '/api/chat', + json=ChatRequest( + model=model, + messages=list(_copy_messages(messages)), + tools=list(_copy_tools(tools)), + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + format=format, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + stream=stream, + ) + + def embed( + self, + model: str = '', + input: Union[str, Sequence[str]] = '', + truncate: Optional[bool] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + dimensions: Optional[int] = None, + ) -> EmbedResponse: + return self._request( + EmbedResponse, + 'POST', + '/api/embed', + json=EmbedRequest( + model=model, + input=input, + truncate=truncate, + options=options, + keep_alive=keep_alive, + dimensions=dimensions, + ).model_dump(exclude_none=True), + ) + + def embeddings( + self, + model: str = '', + prompt: Optional[str] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> EmbeddingsResponse: + """ + Deprecated in favor of `embed`. + """ + return self._request( + EmbeddingsResponse, + 'POST', + '/api/embeddings', + json=EmbeddingsRequest( + model=model, + prompt=prompt, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + ) + + @overload + def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> Iterator[ProgressResponse]: ... + + def pull( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return self._request( + ProgressResponse, + 'POST', + '/api/pull', + json=PullRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> Iterator[ProgressResponse]: ... + + def push( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return self._request( + ProgressResponse, + 'POST', + '/api/push', + json=PushRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[True] = True, + ) -> Iterator[ProgressResponse]: ... + + def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: bool = False, + ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return self._request( + ProgressResponse, + 'POST', + '/api/create', + json=CreateRequest( + model=model, + stream=stream, + quantize=quantize, + from_=from_, + files=files, + adapters=adapters, + license=license, + template=template, + system=system, + parameters=parameters, + messages=messages, + ).model_dump(exclude_none=True), + stream=stream, + ) + + def create_blob(self, path: Union[str, Path]) -> str: + sha256sum = sha256() + with open(path, 'rb') as r: + while True: + chunk = r.read(32 * 1024) + if not chunk: + break + sha256sum.update(chunk) + + digest = f'sha256:{sha256sum.hexdigest()}' + + with open(path, 'rb') as r: + self._request_raw('POST', f'/api/blobs/{digest}', content=r) + + return digest + + def list(self) -> ListResponse: + return self._request( + ListResponse, + 'GET', + '/api/tags', + ) + + def delete(self, model: str) -> StatusResponse: + r = self._request_raw( + 'DELETE', + '/api/delete', + json=DeleteRequest( + model=model, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + def copy(self, source: str, destination: str) -> StatusResponse: + r = self._request_raw( + 'POST', + '/api/copy', + json=CopyRequest( + source=source, + destination=destination, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + def show(self, model: str) -> ShowResponse: + return self._request( + ShowResponse, + 'POST', + '/api/show', + json=ShowRequest( + model=model, + ).model_dump(exclude_none=True), + ) + + def ps(self) -> ProcessResponse: + return self._request( + ProcessResponse, + 'GET', + '/api/ps', + ) + + def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: + """ + Performs a web search + + Args: + query: The query to search for + max_results: The maximum number of results to return (default: 3) + + Returns: + WebSearchResponse with the search results + Raises: + ValueError: If OLLAMA_API_KEY environment variable is not set + """ + if not self._client.headers.get('authorization', '').startswith('Bearer '): + raise ValueError('Authorization header with Bearer token is required for web search') + + return self._request( + WebSearchResponse, + 'POST', + 'https://ollama.com/api/web_search', + json=WebSearchRequest( + query=query, + max_results=max_results, + ).model_dump(exclude_none=True), + ) + + def web_fetch(self, url: str) -> WebFetchResponse: + """ + Fetches the content of a web page for the provided URL. + + Args: + url: The URL to fetch + + Returns: + WebFetchResponse with the fetched result + """ + if not self._client.headers.get('authorization', '').startswith('Bearer '): + raise ValueError('Authorization header with Bearer token is required for web fetch') + + return self._request( + WebFetchResponse, + 'POST', + 'https://ollama.com/api/web_fetch', + json=WebFetchRequest( + url=url, + ).model_dump(exclude_none=True), + ) + + +class AsyncClient(BaseClient): + def __init__(self, host: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.AsyncClient, host, **kwargs) + + async def close(self): + await self._client.aclose() + + async def _request_raw(self, *args, **kwargs): + try: + r = await self._client.request(*args, **kwargs) + r.raise_for_status() + return r + except httpx.HTTPStatusError as e: + raise ResponseError(e.response.text, e.response.status_code) from None + except httpx.ConnectError: + raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None + + @overload + async def _request( + self, + cls: Type[T], + *args, + stream: Literal[False] = False, + **kwargs, + ) -> T: ... + + @overload + async def _request( + self, + cls: Type[T], + *args, + stream: Literal[True] = True, + **kwargs, + ) -> AsyncIterator[T]: ... + + @overload + async def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, AsyncIterator[T]]: ... + + async def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, AsyncIterator[T]]: + if stream: + + async def inner(): + async with self._client.stream(*args, **kwargs) as r: + try: + r.raise_for_status() + except httpx.HTTPStatusError as e: + await e.response.aread() + raise ResponseError(e.response.text, e.response.status_code) from None + + async for line in r.aiter_lines(): + part = json.loads(line) + if err := part.get('error'): + raise ResponseError(err) + yield cls(**part) + + return inner() + + return cls(**(await self._request_raw(*args, **kwargs)).json()) + + async def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: + """ + Performs a web search + + Args: + query: The query to search for + max_results: The maximum number of results to return (default: 3) + + Returns: + WebSearchResponse with the search results + """ + return await self._request( + WebSearchResponse, + 'POST', + 'https://ollama.com/api/web_search', + json=WebSearchRequest( + query=query, + max_results=max_results, + ).model_dump(exclude_none=True), + ) + + async def web_fetch(self, url: str) -> WebFetchResponse: + """ + Fetches the content of a web page for the provided URL. + + Args: + url: The URL to fetch + + Returns: + WebFetchResponse with the fetched result + """ + return await self._request( + WebFetchResponse, + 'POST', + 'https://ollama.com/api/web_fetch', + json=WebFetchRequest( + url=url, + ).model_dump(exclude_none=True), + ) + + @overload + async def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> GenerateResponse: ... + + @overload + async def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> AsyncIterator[GenerateResponse]: ... + + async def generate( + self, + model: str = '', + prompt: Optional[str] = None, + suffix: Optional[str] = None, + *, + system: Optional[str] = None, + template: Optional[str] = None, + context: Optional[Sequence[int]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: Optional[bool] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> Union[GenerateResponse, AsyncIterator[GenerateResponse]]: + """ + 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. + """ + return await self._request( + GenerateResponse, + 'POST', + '/api/generate', + json=GenerateRequest( + model=model, + prompt=prompt, + suffix=suffix, + system=system, + template=template, + context=context, + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + raw=raw, + format=format, + images=list(_copy_images(images)) if images else None, + options=options, + keep_alive=keep_alive, + width=width, + height=height, + steps=steps, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + async def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> ChatResponse: ... + + @overload + async def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> AsyncIterator[ChatResponse]: ... + + async def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> Union[ChatResponse, AsyncIterator[ChatResponse]]: + """ + Create a chat response using the requested model. + + Args: + tools: + A JSON schema as a dict, an Ollama Tool or a Python Function. + Python functions need to follow Google style docstrings to be converted to an Ollama Tool. + For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings + stream: Whether to stream the response. + format: The format of the response. + + Example: + def add_two_numbers(a: int, b: int) -> int: + ''' + Add two numbers together. + + Args: + a: First number to add + b: Second number to add + + Returns: + int: The sum of a and b + ''' + return a + b + + await client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) + + 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. + """ + + return await self._request( + ChatResponse, + 'POST', + '/api/chat', + json=ChatRequest( + model=model, + messages=list(_copy_messages(messages)), + tools=list(_copy_tools(tools)), + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + format=format, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + stream=stream, + ) + + async def embed( + self, + model: str = '', + input: Union[str, Sequence[str]] = '', + truncate: Optional[bool] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + dimensions: Optional[int] = None, + ) -> EmbedResponse: + return await self._request( + EmbedResponse, + 'POST', + '/api/embed', + json=EmbedRequest( + model=model, + input=input, + truncate=truncate, + options=options, + keep_alive=keep_alive, + dimensions=dimensions, + ).model_dump(exclude_none=True), + ) + + async def embeddings( + self, + model: str = '', + prompt: Optional[str] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> EmbeddingsResponse: + """ + Deprecated in favor of `embed`. + """ + return await self._request( + EmbeddingsResponse, + 'POST', + '/api/embeddings', + json=EmbeddingsRequest( + model=model, + prompt=prompt, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + ) + + @overload + async def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + async def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> AsyncIterator[ProgressResponse]: ... + + async def pull( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return await self._request( + ProgressResponse, + 'POST', + '/api/pull', + json=PullRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + async def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + async def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> AsyncIterator[ProgressResponse]: ... + + async def push( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return await self._request( + ProgressResponse, + 'POST', + '/api/push', + json=PushRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + async def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + async def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[True] = True, + ) -> AsyncIterator[ProgressResponse]: ... + + async def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: bool = False, + ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + + return await self._request( + ProgressResponse, + 'POST', + '/api/create', + json=CreateRequest( + model=model, + stream=stream, + quantize=quantize, + from_=from_, + files=files, + adapters=adapters, + license=license, + template=template, + system=system, + parameters=parameters, + messages=messages, + ).model_dump(exclude_none=True), + stream=stream, + ) + + async def create_blob(self, path: Union[str, Path]) -> str: + sha256sum = sha256() + async with await anyio.open_file(path, 'rb') as r: + while True: + chunk = await r.read(32 * 1024) + if not chunk: + break + sha256sum.update(chunk) + + digest = f'sha256:{sha256sum.hexdigest()}' + + async def upload_bytes(): + async with await anyio.open_file(path, 'rb') as r: + while True: + chunk = await r.read(32 * 1024) + if not chunk: + break + yield chunk + + await self._request_raw('POST', f'/api/blobs/{digest}', content=upload_bytes()) + + return digest + + async def list(self) -> ListResponse: + return await self._request( + ListResponse, + 'GET', + '/api/tags', + ) + + async def delete(self, model: str) -> StatusResponse: + r = await self._request_raw( + 'DELETE', + '/api/delete', + json=DeleteRequest( + model=model, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + async def copy(self, source: str, destination: str) -> StatusResponse: + r = await self._request_raw( + 'POST', + '/api/copy', + json=CopyRequest( + source=source, + destination=destination, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + async def show(self, model: str) -> ShowResponse: + return await self._request( + ShowResponse, + 'POST', + '/api/show', + json=ShowRequest( + model=model, + ).model_dump(exclude_none=True), + ) + + async def ps(self) -> ProcessResponse: + return await self._request( + ProcessResponse, + 'GET', + '/api/ps', + ) + + +def _copy_images(images: Optional[Sequence[Union[Image, Any]]]) -> Iterator[Image]: + for image in images or []: + yield image if isinstance(image, Image) else Image(value=image) + + +def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]]]) -> Iterator[Message]: + for message in messages or []: + # Keep empty strings (e.g. tool results with content='') — only drop None. + yield Message.model_validate( + {k: list(_copy_images(v)) if k == 'images' else v for k, v in dict(message).items() if v is not None}, + ) + + +def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None) -> Iterator[Tool]: + for unprocessed_tool in tools or []: + yield convert_function_to_tool(unprocessed_tool) if callable(unprocessed_tool) else Tool.model_validate(unprocessed_tool) + + +def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]: + if isinstance(s, (str, Path)): + try: + if (p := Path(s)).exists(): + return p + except Exception: + ... + return None + + +def _parse_host(host: Optional[str]) -> str: + """ + >>> _parse_host(None) + 'http://127.0.0.1:11434' + >>> _parse_host('') + 'http://127.0.0.1:11434' + >>> _parse_host('1.2.3.4') + 'http://1.2.3.4:11434' + >>> _parse_host(':56789') + 'http://127.0.0.1:56789' + >>> _parse_host('1.2.3.4:56789') + 'http://1.2.3.4:56789' + >>> _parse_host('http://1.2.3.4') + 'http://1.2.3.4:80' + >>> _parse_host('https://1.2.3.4') + 'https://1.2.3.4:443' + >>> _parse_host('https://1.2.3.4:56789') + 'https://1.2.3.4:56789' + >>> _parse_host('example.com') + 'http://example.com:11434' + >>> _parse_host('example.com:56789') + 'http://example.com:56789' + >>> _parse_host('http://example.com') + 'http://example.com:80' + >>> _parse_host('https://example.com') + 'https://example.com:443' + >>> _parse_host('https://example.com:56789') + 'https://example.com:56789' + >>> _parse_host('example.com/') + 'http://example.com:11434' + >>> _parse_host('example.com:56789/') + 'http://example.com:56789' + >>> _parse_host('example.com/path') + 'http://example.com:11434/path' + >>> _parse_host('example.com:56789/path') + 'http://example.com:56789/path' + >>> _parse_host('https://example.com:56789/path') + 'https://example.com:56789/path' + >>> _parse_host('example.com:56789/path/') + 'http://example.com:56789/path' + >>> _parse_host('[0001:002:003:0004::1]') + 'http://[0001:002:003:0004::1]:11434' + >>> _parse_host('[0001:002:003:0004::1]:56789') + 'http://[0001:002:003:0004::1]:56789' + >>> _parse_host('http://[0001:002:003:0004::1]') + 'http://[0001:002:003:0004::1]:80' + >>> _parse_host('https://[0001:002:003:0004::1]') + 'https://[0001:002:003:0004::1]:443' + >>> _parse_host('https://[0001:002:003:0004::1]:56789') + 'https://[0001:002:003:0004::1]:56789' + >>> _parse_host('[0001:002:003:0004::1]/') + 'http://[0001:002:003:0004::1]:11434' + >>> _parse_host('[0001:002:003:0004::1]:56789/') + 'http://[0001:002:003:0004::1]:56789' + >>> _parse_host('[0001:002:003:0004::1]/path') + 'http://[0001:002:003:0004::1]:11434/path' + >>> _parse_host('[0001:002:003:0004::1]:56789/path') + 'http://[0001:002:003:0004::1]:56789/path' + >>> _parse_host('https://[0001:002:003:0004::1]:56789/path') + 'https://[0001:002:003:0004::1]:56789/path' + >>> _parse_host('[0001:002:003:0004::1]:56789/path/') + 'http://[0001:002:003:0004::1]:56789/path' + """ + + host, port = host or '', 11434 + scheme, _, hostport = host.partition('://') + if not hostport: + scheme, hostport = 'http', host + elif scheme == 'http': + port = 80 + elif scheme == 'https': + port = 443 + + split = urllib.parse.urlsplit(f'{scheme}://{hostport}') + host = split.hostname or '127.0.0.1' + port = split.port or port + + try: + if isinstance(ipaddress.ip_address(host), ipaddress.IPv6Address): + # Fix missing square brackets for IPv6 from urlsplit + host = f'[{host}]' + except ValueError: + ... + + if path := split.path.strip('/'): + return f'{scheme}://{host}:{port}/{path}' + + return f'{scheme}://{host}:{port}' From 8544799e7e537fe025bd5933558e8e5cb698fd96 Mon Sep 17 00:00:00 2001 From: r7mekmy4g67w6l Date: Thu, 6 Aug 2026 17:07:31 +0200 Subject: [PATCH 2/6] fix: preserve empty string message content in chat requests Signed-off-by: r7mekmy4g67w6l --- tests/test_client.py | 3014 +++++++++++++++++++++--------------------- 1 file changed, 1515 insertions(+), 1499 deletions(-) diff --git a/tests/test_client.py b/tests/test_client.py index 7b7ab38e..ddc3c233 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,1499 +1,1515 @@ -import base64 -import inspect -import json -import os -import re -import tempfile -from pathlib import Path -from typing import Any - -import pytest -from httpx import Response as httpxResponse -from pydantic import BaseModel -from pytest_httpserver import HTTPServer, URIPattern -from werkzeug.wrappers import Request, Response - -from ollama._client import CONNECTION_ERROR_MESSAGE, AsyncClient, Client, _copy_tools -from ollama._types import Image, Message - -PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC' -PNG_BYTES = base64.b64decode(PNG_BASE64) - -pytestmark = pytest.mark.anyio - - -@pytest.fixture -def anyio_backend(): - return 'asyncio' - - -class PrefixPattern(URIPattern): - def __init__(self, prefix: str): - self.prefix = prefix - - def match(self, uri): - return uri.startswith(self.prefix) - - -def test_client_chat(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -def test_client_chat_with_logprobs(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Hi'}], - 'tools': [], - 'stream': False, - 'logprobs': True, - 'top_logprobs': 3, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': 'Hello', - }, - 'logprobs': [ - { - 'token': 'Hello', - 'logprob': -0.1, - 'top_logprobs': [ - {'token': 'Hello', 'logprob': -0.1}, - {'token': 'Hi', 'logprob': -1.0}, - ], - } - ], - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Hi'}], logprobs=True, top_logprobs=3) - assert response['logprobs'][0]['token'] == 'Hello' - assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' - - -def test_client_chat_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['I ', "don't ", 'know.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': message, - }, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) - - it = iter(['I ', "don't ", 'know.']) - for part in response: - assert part['message']['role'] in 'assistant' - assert part['message']['content'] == next(it) - - -@pytest.mark.parametrize('message_format', ('dict', 'pydantic_model')) -@pytest.mark.parametrize('file_style', ('path', 'bytes')) -def test_client_chat_images(httpserver: HTTPServer, message_format: str, file_style: str, tmp_path): - from ollama._types import Image, Message - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [ - { - 'role': 'user', - 'content': 'Why is the sky blue?', - 'images': [PNG_BASE64], - }, - ], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = Client(httpserver.url_for('/')) - - if file_style == 'bytes': - image_content = PNG_BYTES - elif file_style == 'path': - image_path = tmp_path / 'transparent.png' - image_path.write_bytes(PNG_BYTES) - image_content = str(image_path) - - if message_format == 'pydantic_model': - messages = [Message(role='user', content='Why is the sky blue?', images=[Image(value=image_content)])] - elif message_format == 'dict': - messages = [{'role': 'user', 'content': 'Why is the sky blue?', 'images': [image_content]}] - else: - raise ValueError(f'Invalid message format: {message_format}') - - response = client.chat('dummy', messages=messages) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -def test_client_chat_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering"}', - }, - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' - - -def test_client_chat_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - }, - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -async def test_async_client_chat_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering"}', - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' - - -async def test_async_client_chat_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -def test_client_generate(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?') - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -def test_client_generate_with_logprobs(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why', - 'stream': False, - 'logprobs': True, - 'top_logprobs': 2, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Hello', - 'logprobs': [ - { - 'token': 'Hello', - 'logprob': -0.2, - 'top_logprobs': [ - {'token': 'Hello', 'logprob': -0.2}, - {'token': 'Hi', 'logprob': -1.5}, - ], - } - ], - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why', logprobs=True, top_logprobs=2) - assert response['logprobs'][0]['token'] == 'Hello' - assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' - - -def test_client_generate_with_image_type(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'What is in this image?', - 'stream': False, - 'images': [PNG_BASE64], - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'A blue sky.', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'What is in this image?', images=[Image(value=PNG_BASE64)]) - assert response['model'] == 'dummy' - assert response['response'] == 'A blue sky.' - - -def test_client_generate_with_invalid_image(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'What is in this image?', - 'stream': False, - 'images': ['invalid_base64'], - }, - ).respond_with_json({'error': 'Invalid image data'}, status=400) - - client = Client(httpserver.url_for('/')) - with pytest.raises(ValueError): - client.generate('dummy', 'What is in this image?', images=[Image(value='invalid_base64')]) - - -def test_client_generate_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['Because ', 'it ', 'is.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'response': message, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?', stream=True) - - it = iter(['Because ', 'it ', 'is.']) - for part in response: - assert part['model'] == 'dummy' - assert part['response'] == next(it) - - -def test_client_generate_images(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - 'images': [PNG_BASE64], - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as temp: - temp.write(PNG_BYTES) - temp.flush() - response = client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -def test_client_generate_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering"}', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?', format='json') - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' - - -def test_client_generate_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -async def test_async_client_generate_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering"}', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?', format='json') - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' - - -async def test_async_client_generate_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -def test_client_generate_image(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy-image', - 'prompt': 'a sunset over mountains', - 'stream': False, - 'width': 1024, - 'height': 768, - 'steps': 20, - }, - ).respond_with_json( - { - 'model': 'dummy-image', - 'image': PNG_BASE64, - 'done': True, - 'done_reason': 'stop', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy-image', 'a sunset over mountains', width=1024, height=768, steps=20) - assert response['model'] == 'dummy-image' - assert response['image'] == PNG_BASE64 - assert response['done'] is True - - -def test_client_generate_image_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - # Progress updates - for i in range(1, 4): - yield ( - json.dumps( - { - 'model': 'dummy-image', - 'completed': i, - 'total': 3, - 'done': False, - } - ) - + '\n' - ) - # Final response with image - yield ( - json.dumps( - { - 'model': 'dummy-image', - 'image': PNG_BASE64, - 'done': True, - 'done_reason': 'stop', - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy-image', - 'prompt': 'a sunset over mountains', - 'stream': True, - 'width': 512, - 'height': 512, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy-image', 'a sunset over mountains', stream=True, width=512, height=512) - - parts = list(response) - # Check progress updates - assert parts[0]['completed'] == 1 - assert parts[0]['total'] == 3 - assert parts[0]['done'] is False - # Check final response - assert parts[-1]['image'] == PNG_BASE64 - assert parts[-1]['done'] is True - - -async def test_async_client_generate_image(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy-image', - 'prompt': 'a robot painting', - 'stream': False, - 'width': 1024, - 'height': 1024, - }, - ).respond_with_json( - { - 'model': 'dummy-image', - 'image': PNG_BASE64, - 'done': True, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy-image', 'a robot painting', width=1024, height=1024) - assert response['model'] == 'dummy-image' - assert response['image'] == PNG_BASE64 - - -def test_client_pull(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - response = client.pull('dummy') - assert response['status'] == 'success' - - -def test_client_pull_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'pulling manifest'}) + '\n' - yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' - yield json.dumps({'status': 'writing manifest'}) + '\n' - yield json.dumps({'status': 'removing any unused layers'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.pull('dummy', stream=True) - - it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) - for part in response: - assert part['status'] == next(it) - - -def test_client_push(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - response = client.push('dummy') - assert response['status'] == 'success' - - -def test_client_push_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'retrieving manifest'}) + '\n' - yield json.dumps({'status': 'pushing manifest'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.push('dummy', stream=True) - - it = iter(['retrieving manifest', 'pushing manifest', 'success']) - for part in response: - assert part['status'] == next(it) - - -@pytest.fixture -def userhomedir(): - with tempfile.TemporaryDirectory() as temp: - home = os.getenv('HOME', '') - os.environ['HOME'] = temp - yield Path(temp) - os.environ['HOME'] = home - - -def test_client_create_with_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) - assert response['status'] == 'success' - - -def test_client_create_with_parameters_roundtrip(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'quantize': 'q4_k_m', - 'from': 'mymodel', - 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - 'license': 'this is my license', - 'system': '\nUse\nmultiline\nstrings.\n', - 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = client.create( - 'dummy', - quantize='q4_k_m', - from_='mymodel', - adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - license='this is my license', - system='\nUse\nmultiline\nstrings.\n', - parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - stream=False, - ) - assert response['status'] == 'success' - - -def test_client_create_from_library(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'from': 'llama2', - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - - response = client.create('dummy', from_='llama2') - assert response['status'] == 'success' - - -def test_client_create_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -def test_client_create_blob_exists(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -def test_client_delete(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) - client = Client(httpserver.url_for('/api/delete')) - response = client.delete('dummy') - assert response['status'] == 'success' - - -def test_client_copy(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) - client = Client(httpserver.url_for('/api/copy')) - response = client.copy('dum', 'dummer') - assert response['status'] == 'success' - - -async def test_async_client_chat(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -async def test_async_client_chat_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['I ', "don't ", 'know.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': message, - }, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) - - it = iter(['I ', "don't ", 'know.']) - async for part in response: - assert part['message']['role'] == 'assistant' - assert part['message']['content'] == next(it) - - -async def test_async_client_chat_images(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [ - { - 'role': 'user', - 'content': 'Why is the sky blue?', - 'images': [PNG_BASE64], - }, - ], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?', 'images': [PNG_BYTES]}]) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -async def test_async_client_generate(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?') - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -async def test_async_client_generate_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['Because ', 'it ', 'is.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'response': message, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?', stream=True) - - it = iter(['Because ', 'it ', 'is.']) - async for part in response: - assert part['model'] == 'dummy' - assert part['response'] == next(it) - - -async def test_async_client_generate_images(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - 'images': [PNG_BASE64], - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as temp: - temp.write(PNG_BYTES) - temp.flush() - response = await client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -async def test_async_client_pull(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.pull('dummy') - assert response['status'] == 'success' - - -async def test_async_client_pull_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'pulling manifest'}) + '\n' - yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' - yield json.dumps({'status': 'writing manifest'}) + '\n' - yield json.dumps({'status': 'removing any unused layers'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.pull('dummy', stream=True) - - it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) - async for part in response: - assert part['status'] == next(it) - - -async def test_async_client_push(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.push('dummy') - assert response['status'] == 'success' - - -async def test_async_client_push_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'retrieving manifest'}) + '\n' - yield json.dumps({'status': 'pushing manifest'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.push('dummy', stream=True) - - it = iter(['retrieving manifest', 'pushing manifest', 'success']) - async for part in response: - assert part['status'] == next(it) - - -async def test_async_client_create_with_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = await client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) - assert response['status'] == 'success' - - -async def test_async_client_create_with_parameters_roundtrip(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'quantize': 'q4_k_m', - 'from': 'mymodel', - 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - 'license': 'this is my license', - 'system': '\nUse\nmultiline\nstrings.\n', - 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = await client.create( - 'dummy', - quantize='q4_k_m', - from_='mymodel', - adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - license='this is my license', - system='\nUse\nmultiline\nstrings.\n', - parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - stream=False, - ) - assert response['status'] == 'success' - - -async def test_async_client_create_from_library(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'from': 'llama2', - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - - response = await client.create('dummy', from_='llama2') - assert response['status'] == 'success' - - -async def test_async_client_create_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = await client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -async def test_async_client_create_blob_exists(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = await client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -async def test_async_client_delete(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) - client = AsyncClient(httpserver.url_for('/api/delete')) - response = await client.delete('dummy') - assert response['status'] == 'success' - - -async def test_async_client_copy(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) - client = AsyncClient(httpserver.url_for('/api/copy')) - response = await client.copy('dum', 'dummer') - assert response['status'] == 'success' - - -def test_headers(): - client = Client() - assert client._client.headers['content-type'] == 'application/json' - assert client._client.headers['accept'] == 'application/json' - assert client._client.headers['user-agent'].startswith('ollama-python/') - - client = Client( - headers={ - 'X-Custom': 'value', - 'Content-Type': 'text/plain', - } - ) - assert client._client.headers['x-custom'] == 'value' - assert client._client.headers['content-type'] == 'application/json' - - -def test_copy_tools(): - def func1(x: int) -> str: - """Simple function 1. - Args: - x (integer): A number - """ - - def func2(y: str) -> int: - """Simple function 2. - Args: - y (string): A string - """ - - # Test with list of functions - tools = list(_copy_tools([func1, func2])) - assert len(tools) == 2 - assert tools[0].function.name == 'func1' - assert tools[1].function.name == 'func2' - - # Test with empty input - assert list(_copy_tools()) == [] - assert list(_copy_tools(None)) == [] - assert list(_copy_tools([])) == [] - - # Test with mix of functions and tool dicts - tool_dict = { - 'type': 'function', - 'function': { - 'name': 'test', - 'description': 'Test function', - 'parameters': { - 'type': 'object', - 'properties': {'x': {'type': 'string', 'description': 'A string', 'enum': ['a', 'b', 'c']}, 'y': {'type': ['integer', 'number'], 'description': 'An integer'}}, - 'required': ['x'], - }, - }, - } - - tools = list(_copy_tools([func1, tool_dict])) - assert len(tools) == 2 - assert tools[0].function.name == 'func1' - assert tools[1].function.name == 'test' - - -def test_tool_validation(): - arbitrary_tool = {'type': 'custom_type', 'function': {'name': 'test'}} - tools = list(_copy_tools([arbitrary_tool])) - assert len(tools) == 1 - assert tools[0].type == 'custom_type' - assert tools[0].function.name == 'test' - - -def test_client_connection_error(): - client = Client('http://localhost:1234') - - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.generate('model', 'prompt') - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.show('model') - - -async def test_async_client_connection_error(): - client = AsyncClient('http://localhost:1234') - with pytest.raises(ConnectionError) as exc_info: - await client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) - assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - with pytest.raises(ConnectionError) as exc_info: - await client.generate('model', 'prompt') - assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - with pytest.raises(ConnectionError) as exc_info: - await client.show('model') - assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - - -def test_arbitrary_roles_accepted_in_message(): - _ = Message(role='somerandomrole', content="I'm ok with you adding any role message now!") - - -def _mock_request(*args: Any, **kwargs: Any) -> Response: - return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") - - -def test_arbitrary_roles_accepted_in_message_request(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(Client, '_request', _mock_request) - - client = Client() - - client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) - - -async def _mock_request_async(*args: Any, **kwargs: Any) -> Response: - return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") - - -async def test_arbitrary_roles_accepted_in_message_request_async(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(AsyncClient, '_request', _mock_request_async) - - client = AsyncClient() - - await client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) - - -def test_client_web_search_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - - client = Client() - - with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'): - client.web_search('test query') - - -def test_client_web_fetch_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - - client = Client() - - with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'): - client.web_fetch('https://example.com') - - -def _mock_request_web_search(self, cls, method, url, json=None, **kwargs): - assert method == 'POST' - assert url == 'https://ollama.com/api/web_search' - assert json is not None and 'query' in json and 'max_results' in json - return httpxResponse(status_code=200, content='{"results": {}, "success": true}') - - -def _mock_request_web_fetch(self, cls, method, url, json=None, **kwargs): - assert method == 'POST' - assert url == 'https://ollama.com/api/web_fetch' - assert json is not None and 'url' in json - return httpxResponse(status_code=200, content='{"results": {}, "success": true}') - - -def test_client_web_search_with_env_api_key(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') - monkeypatch.setattr(Client, '_request', _mock_request_web_search) - - client = Client() - client.web_search('what is ollama?', max_results=2) - - -def test_client_web_fetch_with_env_api_key(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') - monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) - - client = Client() - client.web_fetch('https://example.com') - - -def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - monkeypatch.setattr(Client, '_request', _mock_request_web_search) - - client = Client(headers={'Authorization': 'Bearer custom-token'}) - client.web_search('what is ollama?', max_results=1) - - -def test_client_web_fetch_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) - - client = Client(headers={'Authorization': 'Bearer custom-token'}) - client.web_fetch('https://example.com') - - -def test_client_bearer_header_from_env(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') - - client = Client() - assert client._client.headers['authorization'] == 'Bearer env-token' - - -def test_client_explicit_bearer_header_overrides_env(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') - monkeypatch.setattr(Client, '_request', _mock_request_web_search) - - client = Client(headers={'Authorization': 'Bearer explicit-token'}) - assert client._client.headers['authorization'] == 'Bearer explicit-token' - client.web_search('override check') - - -def test_client_close(): - client = Client() - client.close() - assert client._client.is_closed - - -@pytest.mark.anyio -async def test_async_client_close(): - client = AsyncClient() - await client.close() - assert client._client.is_closed - - -def test_client_context_manager(): - with Client() as client: - assert isinstance(client, Client) - assert not client._client.is_closed - - assert client._client.is_closed - - -@pytest.mark.anyio -async def test_async_client_context_manager(): - async with AsyncClient() as client: - assert isinstance(client, AsyncClient) - assert not client._client.is_closed - - assert client._client.is_closed - - -def test_generate_think_annotation_matches_chat(): - # The `think` parameter accepts bool or the 'low'/'medium'/'high' string levels. - # Client.generate must keep the same annotation as Client.chat and - # AsyncClient.generate so passing a string level does not raise a false type - # error (regression guard for the sync generate overloads/implementation). - expected = inspect.signature(Client.chat).parameters['think'].annotation - assert inspect.signature(Client.generate).parameters['think'].annotation == expected - assert inspect.signature(AsyncClient.generate).parameters['think'].annotation == expected +import base64 +import inspect +import json +import os +import re +import tempfile +from pathlib import Path +from typing import Any + +import pytest +from httpx import Response as httpxResponse +from pydantic import BaseModel +from pytest_httpserver import HTTPServer, URIPattern +from werkzeug.wrappers import Request, Response + +from ollama._client import CONNECTION_ERROR_MESSAGE, AsyncClient, Client, _copy_tools +from ollama._types import Image, Message + +PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC' +PNG_BYTES = base64.b64decode(PNG_BASE64) + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend(): + return 'asyncio' + + +class PrefixPattern(URIPattern): + def __init__(self, prefix: str): + self.prefix = prefix + + def match(self, uri): + return uri.startswith(self.prefix) + + +def test_client_chat(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +def test_client_chat_with_logprobs(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Hi'}], + 'tools': [], + 'stream': False, + 'logprobs': True, + 'top_logprobs': 3, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': 'Hello', + }, + 'logprobs': [ + { + 'token': 'Hello', + 'logprob': -0.1, + 'top_logprobs': [ + {'token': 'Hello', 'logprob': -0.1}, + {'token': 'Hi', 'logprob': -1.0}, + ], + } + ], + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Hi'}], logprobs=True, top_logprobs=3) + assert response['logprobs'][0]['token'] == 'Hello' + assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' + + +def test_client_chat_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['I ', "don't ", 'know.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': message, + }, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) + + it = iter(['I ', "don't ", 'know.']) + for part in response: + assert part['message']['role'] in 'assistant' + assert part['message']['content'] == next(it) + + +@pytest.mark.parametrize('message_format', ('dict', 'pydantic_model')) +@pytest.mark.parametrize('file_style', ('path', 'bytes')) +def test_client_chat_images(httpserver: HTTPServer, message_format: str, file_style: str, tmp_path): + from ollama._types import Image, Message + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [ + { + 'role': 'user', + 'content': 'Why is the sky blue?', + 'images': [PNG_BASE64], + }, + ], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = Client(httpserver.url_for('/')) + + if file_style == 'bytes': + image_content = PNG_BYTES + elif file_style == 'path': + image_path = tmp_path / 'transparent.png' + image_path.write_bytes(PNG_BYTES) + image_content = str(image_path) + + if message_format == 'pydantic_model': + messages = [Message(role='user', content='Why is the sky blue?', images=[Image(value=image_content)])] + elif message_format == 'dict': + messages = [{'role': 'user', 'content': 'Why is the sky blue?', 'images': [image_content]}] + else: + raise ValueError(f'Invalid message format: {message_format}') + + response = client.chat('dummy', messages=messages) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +def test_client_chat_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering"}', + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' + + +def test_client_chat_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +async def test_async_client_chat_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering"}', + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' + + +async def test_async_client_chat_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +def test_client_generate(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?') + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +def test_client_generate_with_logprobs(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why', + 'stream': False, + 'logprobs': True, + 'top_logprobs': 2, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Hello', + 'logprobs': [ + { + 'token': 'Hello', + 'logprob': -0.2, + 'top_logprobs': [ + {'token': 'Hello', 'logprob': -0.2}, + {'token': 'Hi', 'logprob': -1.5}, + ], + } + ], + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why', logprobs=True, top_logprobs=2) + assert response['logprobs'][0]['token'] == 'Hello' + assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' + + +def test_client_generate_with_image_type(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'What is in this image?', + 'stream': False, + 'images': [PNG_BASE64], + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'A blue sky.', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'What is in this image?', images=[Image(value=PNG_BASE64)]) + assert response['model'] == 'dummy' + assert response['response'] == 'A blue sky.' + + +def test_client_generate_with_invalid_image(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'What is in this image?', + 'stream': False, + 'images': ['invalid_base64'], + }, + ).respond_with_json({'error': 'Invalid image data'}, status=400) + + client = Client(httpserver.url_for('/')) + with pytest.raises(ValueError): + client.generate('dummy', 'What is in this image?', images=[Image(value='invalid_base64')]) + + +def test_client_generate_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['Because ', 'it ', 'is.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'response': message, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?', stream=True) + + it = iter(['Because ', 'it ', 'is.']) + for part in response: + assert part['model'] == 'dummy' + assert part['response'] == next(it) + + +def test_client_generate_images(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + 'images': [PNG_BASE64], + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as temp: + temp.write(PNG_BYTES) + temp.flush() + response = client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +def test_client_generate_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering"}', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?', format='json') + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' + + +def test_client_generate_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +async def test_async_client_generate_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering"}', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?', format='json') + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' + + +async def test_async_client_generate_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +def test_client_generate_image(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy-image', + 'prompt': 'a sunset over mountains', + 'stream': False, + 'width': 1024, + 'height': 768, + 'steps': 20, + }, + ).respond_with_json( + { + 'model': 'dummy-image', + 'image': PNG_BASE64, + 'done': True, + 'done_reason': 'stop', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy-image', 'a sunset over mountains', width=1024, height=768, steps=20) + assert response['model'] == 'dummy-image' + assert response['image'] == PNG_BASE64 + assert response['done'] is True + + +def test_client_generate_image_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + # Progress updates + for i in range(1, 4): + yield ( + json.dumps( + { + 'model': 'dummy-image', + 'completed': i, + 'total': 3, + 'done': False, + } + ) + + '\n' + ) + # Final response with image + yield ( + json.dumps( + { + 'model': 'dummy-image', + 'image': PNG_BASE64, + 'done': True, + 'done_reason': 'stop', + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy-image', + 'prompt': 'a sunset over mountains', + 'stream': True, + 'width': 512, + 'height': 512, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy-image', 'a sunset over mountains', stream=True, width=512, height=512) + + parts = list(response) + # Check progress updates + assert parts[0]['completed'] == 1 + assert parts[0]['total'] == 3 + assert parts[0]['done'] is False + # Check final response + assert parts[-1]['image'] == PNG_BASE64 + assert parts[-1]['done'] is True + + +async def test_async_client_generate_image(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy-image', + 'prompt': 'a robot painting', + 'stream': False, + 'width': 1024, + 'height': 1024, + }, + ).respond_with_json( + { + 'model': 'dummy-image', + 'image': PNG_BASE64, + 'done': True, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy-image', 'a robot painting', width=1024, height=1024) + assert response['model'] == 'dummy-image' + assert response['image'] == PNG_BASE64 + + +def test_client_pull(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + response = client.pull('dummy') + assert response['status'] == 'success' + + +def test_client_pull_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'pulling manifest'}) + '\n' + yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' + yield json.dumps({'status': 'writing manifest'}) + '\n' + yield json.dumps({'status': 'removing any unused layers'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.pull('dummy', stream=True) + + it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) + for part in response: + assert part['status'] == next(it) + + +def test_client_push(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + response = client.push('dummy') + assert response['status'] == 'success' + + +def test_client_push_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'retrieving manifest'}) + '\n' + yield json.dumps({'status': 'pushing manifest'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.push('dummy', stream=True) + + it = iter(['retrieving manifest', 'pushing manifest', 'success']) + for part in response: + assert part['status'] == next(it) + + +@pytest.fixture +def userhomedir(): + with tempfile.TemporaryDirectory() as temp: + home = os.getenv('HOME', '') + os.environ['HOME'] = temp + yield Path(temp) + os.environ['HOME'] = home + + +def test_client_create_with_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) + assert response['status'] == 'success' + + +def test_client_create_with_parameters_roundtrip(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'quantize': 'q4_k_m', + 'from': 'mymodel', + 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + 'license': 'this is my license', + 'system': '\nUse\nmultiline\nstrings.\n', + 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = client.create( + 'dummy', + quantize='q4_k_m', + from_='mymodel', + adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + license='this is my license', + system='\nUse\nmultiline\nstrings.\n', + parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + stream=False, + ) + assert response['status'] == 'success' + + +def test_client_create_from_library(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'from': 'llama2', + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + + response = client.create('dummy', from_='llama2') + assert response['status'] == 'success' + + +def test_client_create_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +def test_client_create_blob_exists(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +def test_client_delete(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) + client = Client(httpserver.url_for('/api/delete')) + response = client.delete('dummy') + assert response['status'] == 'success' + + +def test_client_copy(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) + client = Client(httpserver.url_for('/api/copy')) + response = client.copy('dum', 'dummer') + assert response['status'] == 'success' + + +async def test_async_client_chat(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +async def test_async_client_chat_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['I ', "don't ", 'know.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': message, + }, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) + + it = iter(['I ', "don't ", 'know.']) + async for part in response: + assert part['message']['role'] == 'assistant' + assert part['message']['content'] == next(it) + + +async def test_async_client_chat_images(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [ + { + 'role': 'user', + 'content': 'Why is the sky blue?', + 'images': [PNG_BASE64], + }, + ], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?', 'images': [PNG_BYTES]}]) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +async def test_async_client_generate(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?') + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +async def test_async_client_generate_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['Because ', 'it ', 'is.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'response': message, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?', stream=True) + + it = iter(['Because ', 'it ', 'is.']) + async for part in response: + assert part['model'] == 'dummy' + assert part['response'] == next(it) + + +async def test_async_client_generate_images(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + 'images': [PNG_BASE64], + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as temp: + temp.write(PNG_BYTES) + temp.flush() + response = await client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +async def test_async_client_pull(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.pull('dummy') + assert response['status'] == 'success' + + +async def test_async_client_pull_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'pulling manifest'}) + '\n' + yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' + yield json.dumps({'status': 'writing manifest'}) + '\n' + yield json.dumps({'status': 'removing any unused layers'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.pull('dummy', stream=True) + + it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) + async for part in response: + assert part['status'] == next(it) + + +async def test_async_client_push(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.push('dummy') + assert response['status'] == 'success' + + +async def test_async_client_push_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'retrieving manifest'}) + '\n' + yield json.dumps({'status': 'pushing manifest'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.push('dummy', stream=True) + + it = iter(['retrieving manifest', 'pushing manifest', 'success']) + async for part in response: + assert part['status'] == next(it) + + +async def test_async_client_create_with_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = await client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) + assert response['status'] == 'success' + + +async def test_async_client_create_with_parameters_roundtrip(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'quantize': 'q4_k_m', + 'from': 'mymodel', + 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + 'license': 'this is my license', + 'system': '\nUse\nmultiline\nstrings.\n', + 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = await client.create( + 'dummy', + quantize='q4_k_m', + from_='mymodel', + adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + license='this is my license', + system='\nUse\nmultiline\nstrings.\n', + parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + stream=False, + ) + assert response['status'] == 'success' + + +async def test_async_client_create_from_library(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'from': 'llama2', + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + + response = await client.create('dummy', from_='llama2') + assert response['status'] == 'success' + + +async def test_async_client_create_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = await client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +async def test_async_client_create_blob_exists(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = await client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +async def test_async_client_delete(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) + client = AsyncClient(httpserver.url_for('/api/delete')) + response = await client.delete('dummy') + assert response['status'] == 'success' + + +async def test_async_client_copy(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) + client = AsyncClient(httpserver.url_for('/api/copy')) + response = await client.copy('dum', 'dummer') + assert response['status'] == 'success' + + +def test_headers(): + client = Client() + assert client._client.headers['content-type'] == 'application/json' + assert client._client.headers['accept'] == 'application/json' + assert client._client.headers['user-agent'].startswith('ollama-python/') + + client = Client( + headers={ + 'X-Custom': 'value', + 'Content-Type': 'text/plain', + } + ) + assert client._client.headers['x-custom'] == 'value' + assert client._client.headers['content-type'] == 'application/json' + + +def test_copy_tools(): + def func1(x: int) -> str: + """Simple function 1. + Args: + x (integer): A number + """ + + def func2(y: str) -> int: + """Simple function 2. + Args: + y (string): A string + """ + + # Test with list of functions + tools = list(_copy_tools([func1, func2])) + assert len(tools) == 2 + assert tools[0].function.name == 'func1' + assert tools[1].function.name == 'func2' + + # Test with empty input + assert list(_copy_tools()) == [] + assert list(_copy_tools(None)) == [] + assert list(_copy_tools([])) == [] + + # Test with mix of functions and tool dicts + tool_dict = { + 'type': 'function', + 'function': { + 'name': 'test', + 'description': 'Test function', + 'parameters': { + 'type': 'object', + 'properties': {'x': {'type': 'string', 'description': 'A string', 'enum': ['a', 'b', 'c']}, 'y': {'type': ['integer', 'number'], 'description': 'An integer'}}, + 'required': ['x'], + }, + }, + } + + tools = list(_copy_tools([func1, tool_dict])) + assert len(tools) == 2 + assert tools[0].function.name == 'func1' + assert tools[1].function.name == 'test' + + +def test_tool_validation(): + arbitrary_tool = {'type': 'custom_type', 'function': {'name': 'test'}} + tools = list(_copy_tools([arbitrary_tool])) + assert len(tools) == 1 + assert tools[0].type == 'custom_type' + assert tools[0].function.name == 'test' + + +def test_client_connection_error(): + client = Client('http://localhost:1234') + + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.generate('model', 'prompt') + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.show('model') + + +async def test_async_client_connection_error(): + client = AsyncClient('http://localhost:1234') + with pytest.raises(ConnectionError) as exc_info: + await client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) + assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + with pytest.raises(ConnectionError) as exc_info: + await client.generate('model', 'prompt') + assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + with pytest.raises(ConnectionError) as exc_info: + await client.show('model') + assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + + +def test_arbitrary_roles_accepted_in_message(): + _ = Message(role='somerandomrole', content="I'm ok with you adding any role message now!") + + +def _mock_request(*args: Any, **kwargs: Any) -> Response: + return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") + + +def test_arbitrary_roles_accepted_in_message_request(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(Client, '_request', _mock_request) + + client = Client() + + client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) + + +async def _mock_request_async(*args: Any, **kwargs: Any) -> Response: + return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") + + +async def test_arbitrary_roles_accepted_in_message_request_async(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(AsyncClient, '_request', _mock_request_async) + + client = AsyncClient() + + await client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) + + +def test_copy_messages_preserves_empty_string_content(): + from ollama._client import _copy_messages + + msgs = list( + _copy_messages( + [ + {'role': 'assistant', 'content': ''}, + {'role': 'tool', 'content': '', 'tool_name': 'web_search'}, + ] + ) + ) + assert msgs[0].content == '' + assert msgs[1].content == '' + assert msgs[1].tool_name == 'web_search' + + +def test_client_web_search_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + + client = Client() + + with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'): + client.web_search('test query') + + +def test_client_web_fetch_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + + client = Client() + + with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'): + client.web_fetch('https://example.com') + + +def _mock_request_web_search(self, cls, method, url, json=None, **kwargs): + assert method == 'POST' + assert url == 'https://ollama.com/api/web_search' + assert json is not None and 'query' in json and 'max_results' in json + return httpxResponse(status_code=200, content='{"results": {}, "success": true}') + + +def _mock_request_web_fetch(self, cls, method, url, json=None, **kwargs): + assert method == 'POST' + assert url == 'https://ollama.com/api/web_fetch' + assert json is not None and 'url' in json + return httpxResponse(status_code=200, content='{"results": {}, "success": true}') + + +def test_client_web_search_with_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') + monkeypatch.setattr(Client, '_request', _mock_request_web_search) + + client = Client() + client.web_search('what is ollama?', max_results=2) + + +def test_client_web_fetch_with_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') + monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) + + client = Client() + client.web_fetch('https://example.com') + + +def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + monkeypatch.setattr(Client, '_request', _mock_request_web_search) + + client = Client(headers={'Authorization': 'Bearer custom-token'}) + client.web_search('what is ollama?', max_results=1) + + +def test_client_web_fetch_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) + + client = Client(headers={'Authorization': 'Bearer custom-token'}) + client.web_fetch('https://example.com') + + +def test_client_bearer_header_from_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') + + client = Client() + assert client._client.headers['authorization'] == 'Bearer env-token' + + +def test_client_explicit_bearer_header_overrides_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') + monkeypatch.setattr(Client, '_request', _mock_request_web_search) + + client = Client(headers={'Authorization': 'Bearer explicit-token'}) + assert client._client.headers['authorization'] == 'Bearer explicit-token' + client.web_search('override check') + + +def test_client_close(): + client = Client() + client.close() + assert client._client.is_closed + + +@pytest.mark.anyio +async def test_async_client_close(): + client = AsyncClient() + await client.close() + assert client._client.is_closed + + +def test_client_context_manager(): + with Client() as client: + assert isinstance(client, Client) + assert not client._client.is_closed + + assert client._client.is_closed + + +@pytest.mark.anyio +async def test_async_client_context_manager(): + async with AsyncClient() as client: + assert isinstance(client, AsyncClient) + assert not client._client.is_closed + + assert client._client.is_closed + + +def test_generate_think_annotation_matches_chat(): + # The `think` parameter accepts bool or the 'low'/'medium'/'high' string levels. + # Client.generate must keep the same annotation as Client.chat and + # AsyncClient.generate so passing a string level does not raise a false type + # error (regression guard for the sync generate overloads/implementation). + expected = inspect.signature(Client.chat).parameters['think'].annotation + assert inspect.signature(Client.generate).parameters['think'].annotation == expected + assert inspect.signature(AsyncClient.generate).parameters['think'].annotation == expected From 0882c1655b38f87fbd1d6b56fb3405c243a6c7f5 Mon Sep 17 00:00:00 2001 From: r7mekmy4g67w6l Date: Thu, 6 Aug 2026 17:15:12 +0200 Subject: [PATCH 3/6] fix: preserve empty string message content in chat requests Signed-off-by: r7mekmy4g67w6l Normalize to LF so the PR shows the real 1-line change. --- ollama/_client.py | 2862 ++++++++++++++++++++++----------------------- 1 file changed, 1431 insertions(+), 1431 deletions(-) diff --git a/ollama/_client.py b/ollama/_client.py index 156d4e59..898854fd 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -1,1431 +1,1431 @@ -import contextlib -import ipaddress -import json -import os -import platform -import sys -import urllib.parse -from hashlib import sha256 -from os import PathLike -from pathlib import Path -from typing import ( - Any, - Callable, - Dict, - List, - Literal, - Mapping, - Optional, - Sequence, - Type, - TypeVar, - Union, - overload, -) - -import anyio -from pydantic.json_schema import JsonSchemaValue - -from ollama._utils import convert_function_to_tool - -if sys.version_info < (3, 9): - from typing import AsyncIterator, Iterator -else: - from collections.abc import AsyncIterator, Iterator - -from importlib import metadata - -try: - __version__ = metadata.version('ollama') -except metadata.PackageNotFoundError: - __version__ = '0.0.0' - -import httpx - -from ollama._types import ( - ChatRequest, - ChatResponse, - CopyRequest, - CreateRequest, - DeleteRequest, - EmbeddingsRequest, - EmbeddingsResponse, - EmbedRequest, - EmbedResponse, - GenerateRequest, - GenerateResponse, - Image, - ListResponse, - Message, - Options, - ProcessResponse, - ProgressResponse, - PullRequest, - PushRequest, - ResponseError, - ShowRequest, - ShowResponse, - StatusResponse, - Tool, - WebFetchRequest, - WebFetchResponse, - WebSearchRequest, - WebSearchResponse, -) - -T = TypeVar('T') - - -class BaseClient(contextlib.AbstractContextManager, contextlib.AbstractAsyncContextManager): - def __init__( - self, - client, - host: Optional[str] = None, - *, - follow_redirects: bool = True, - timeout: Any = None, - headers: Optional[Mapping[str, str]] = None, - **kwargs, - ) -> None: - """ - Creates a httpx client. Default parameters are the same as those defined in httpx - except for the following: - - `follow_redirects`: True - - `timeout`: None - `kwargs` are passed to the httpx client. - """ - - headers = { - k.lower(): v - for k, v in { - **(headers or {}), - 'Content-Type': 'application/json', - 'Accept': 'application/json', - 'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}', - }.items() - if v is not None - } - api_key = os.getenv('OLLAMA_API_KEY', None) - if not headers.get('authorization') and api_key: - headers['authorization'] = f'Bearer {api_key}' - - self._client = client( - base_url=_parse_host(host or os.getenv('OLLAMA_HOST')), - follow_redirects=follow_redirects, - timeout=timeout, - headers=headers, - **kwargs, - ) - - def __exit__(self, exc_type, exc_val, exc_tb): - self.close() - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.close() - - -CONNECTION_ERROR_MESSAGE = 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - - -class Client(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.Client, host, **kwargs) - - def close(self): - self._client.close() - - def _request_raw(self, *args, **kwargs): - try: - r = self._client.request(*args, **kwargs) - r.raise_for_status() - return r - except httpx.HTTPStatusError as e: - raise ResponseError(e.response.text, e.response.status_code) from None - except httpx.ConnectError: - raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None - - @overload - def _request( - self, - cls: Type[T], - *args, - stream: Literal[False] = False, - **kwargs, - ) -> T: ... - - @overload - def _request( - self, - cls: Type[T], - *args, - stream: Literal[True] = True, - **kwargs, - ) -> Iterator[T]: ... - - @overload - def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, Iterator[T]]: ... - - def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, Iterator[T]]: - if stream: - - def inner(): - with self._client.stream(*args, **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(): - part = json.loads(line) - if err := part.get('error'): - raise ResponseError(err) - yield cls(**part) - - return inner() - - return cls(**self._request_raw(*args, **kwargs).json()) - - @overload - def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> GenerateResponse: ... - - @overload - def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> Iterator[GenerateResponse]: ... - - def generate( - self, - model: str = '', - prompt: Optional[str] = None, - suffix: Optional[str] = None, - *, - system: Optional[str] = None, - template: Optional[str] = None, - context: Optional[Sequence[int]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: Optional[bool] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> Union[GenerateResponse, Iterator[GenerateResponse]]: - """ - 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. - """ - - return self._request( - GenerateResponse, - 'POST', - '/api/generate', - json=GenerateRequest( - model=model, - prompt=prompt, - suffix=suffix, - system=system, - template=template, - context=context, - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - raw=raw, - format=format, - images=list(_copy_images(images)) if images else None, - options=options, - keep_alive=keep_alive, - width=width, - height=height, - steps=steps, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> ChatResponse: ... - - @overload - def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> Iterator[ChatResponse]: ... - - def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> Union[ChatResponse, Iterator[ChatResponse]]: - """ - Create a chat response using the requested model. - - Args: - tools: - A JSON schema as a dict, an Ollama Tool or a Python Function. - Python functions need to follow Google style docstrings to be converted to an Ollama Tool. - For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings - stream: Whether to stream the response. - format: The format of the response. - - Example: - def add_two_numbers(a: int, b: int) -> int: - ''' - Add two numbers together. - - Args: - a: First number to add - b: Second number to add - - Returns: - int: The sum of a and b - ''' - return a + b - - client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) - - 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. - """ - return self._request( - ChatResponse, - 'POST', - '/api/chat', - json=ChatRequest( - model=model, - messages=list(_copy_messages(messages)), - tools=list(_copy_tools(tools)), - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - format=format, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - stream=stream, - ) - - def embed( - self, - model: str = '', - input: Union[str, Sequence[str]] = '', - truncate: Optional[bool] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - dimensions: Optional[int] = None, - ) -> EmbedResponse: - return self._request( - EmbedResponse, - 'POST', - '/api/embed', - json=EmbedRequest( - model=model, - input=input, - truncate=truncate, - options=options, - keep_alive=keep_alive, - dimensions=dimensions, - ).model_dump(exclude_none=True), - ) - - def embeddings( - self, - model: str = '', - prompt: Optional[str] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> EmbeddingsResponse: - """ - Deprecated in favor of `embed`. - """ - return self._request( - EmbeddingsResponse, - 'POST', - '/api/embeddings', - json=EmbeddingsRequest( - model=model, - prompt=prompt, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - ) - - @overload - def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> Iterator[ProgressResponse]: ... - - def pull( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return self._request( - ProgressResponse, - 'POST', - '/api/pull', - json=PullRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> Iterator[ProgressResponse]: ... - - def push( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return self._request( - ProgressResponse, - 'POST', - '/api/push', - json=PushRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[True] = True, - ) -> Iterator[ProgressResponse]: ... - - def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: bool = False, - ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return self._request( - ProgressResponse, - 'POST', - '/api/create', - json=CreateRequest( - model=model, - stream=stream, - quantize=quantize, - from_=from_, - files=files, - adapters=adapters, - license=license, - template=template, - system=system, - parameters=parameters, - messages=messages, - ).model_dump(exclude_none=True), - stream=stream, - ) - - def create_blob(self, path: Union[str, Path]) -> str: - sha256sum = sha256() - with open(path, 'rb') as r: - while True: - chunk = r.read(32 * 1024) - if not chunk: - break - sha256sum.update(chunk) - - digest = f'sha256:{sha256sum.hexdigest()}' - - with open(path, 'rb') as r: - self._request_raw('POST', f'/api/blobs/{digest}', content=r) - - return digest - - def list(self) -> ListResponse: - return self._request( - ListResponse, - 'GET', - '/api/tags', - ) - - def delete(self, model: str) -> StatusResponse: - r = self._request_raw( - 'DELETE', - '/api/delete', - json=DeleteRequest( - model=model, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - def copy(self, source: str, destination: str) -> StatusResponse: - r = self._request_raw( - 'POST', - '/api/copy', - json=CopyRequest( - source=source, - destination=destination, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - def show(self, model: str) -> ShowResponse: - return self._request( - ShowResponse, - 'POST', - '/api/show', - json=ShowRequest( - model=model, - ).model_dump(exclude_none=True), - ) - - def ps(self) -> ProcessResponse: - return self._request( - ProcessResponse, - 'GET', - '/api/ps', - ) - - def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: - """ - Performs a web search - - Args: - query: The query to search for - max_results: The maximum number of results to return (default: 3) - - Returns: - WebSearchResponse with the search results - Raises: - ValueError: If OLLAMA_API_KEY environment variable is not set - """ - if not self._client.headers.get('authorization', '').startswith('Bearer '): - raise ValueError('Authorization header with Bearer token is required for web search') - - return self._request( - WebSearchResponse, - 'POST', - 'https://ollama.com/api/web_search', - json=WebSearchRequest( - query=query, - max_results=max_results, - ).model_dump(exclude_none=True), - ) - - def web_fetch(self, url: str) -> WebFetchResponse: - """ - Fetches the content of a web page for the provided URL. - - Args: - url: The URL to fetch - - Returns: - WebFetchResponse with the fetched result - """ - if not self._client.headers.get('authorization', '').startswith('Bearer '): - raise ValueError('Authorization header with Bearer token is required for web fetch') - - return self._request( - WebFetchResponse, - 'POST', - 'https://ollama.com/api/web_fetch', - json=WebFetchRequest( - url=url, - ).model_dump(exclude_none=True), - ) - - -class AsyncClient(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.AsyncClient, host, **kwargs) - - async def close(self): - await self._client.aclose() - - async def _request_raw(self, *args, **kwargs): - try: - r = await self._client.request(*args, **kwargs) - r.raise_for_status() - return r - except httpx.HTTPStatusError as e: - raise ResponseError(e.response.text, e.response.status_code) from None - except httpx.ConnectError: - raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None - - @overload - async def _request( - self, - cls: Type[T], - *args, - stream: Literal[False] = False, - **kwargs, - ) -> T: ... - - @overload - async def _request( - self, - cls: Type[T], - *args, - stream: Literal[True] = True, - **kwargs, - ) -> AsyncIterator[T]: ... - - @overload - async def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, AsyncIterator[T]]: ... - - async def _request( - self, - cls: Type[T], - *args, - stream: bool = False, - **kwargs, - ) -> Union[T, AsyncIterator[T]]: - if stream: - - async def inner(): - async with self._client.stream(*args, **kwargs) as r: - try: - r.raise_for_status() - except httpx.HTTPStatusError as e: - await e.response.aread() - raise ResponseError(e.response.text, e.response.status_code) from None - - async for line in r.aiter_lines(): - part = json.loads(line) - if err := part.get('error'): - raise ResponseError(err) - yield cls(**part) - - return inner() - - return cls(**(await self._request_raw(*args, **kwargs)).json()) - - async def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: - """ - Performs a web search - - Args: - query: The query to search for - max_results: The maximum number of results to return (default: 3) - - Returns: - WebSearchResponse with the search results - """ - return await self._request( - WebSearchResponse, - 'POST', - 'https://ollama.com/api/web_search', - json=WebSearchRequest( - query=query, - max_results=max_results, - ).model_dump(exclude_none=True), - ) - - async def web_fetch(self, url: str) -> WebFetchResponse: - """ - Fetches the content of a web page for the provided URL. - - Args: - url: The URL to fetch - - Returns: - WebFetchResponse with the fetched result - """ - return await self._request( - WebFetchResponse, - 'POST', - 'https://ollama.com/api/web_fetch', - json=WebFetchRequest( - url=url, - ).model_dump(exclude_none=True), - ) - - @overload - async def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> GenerateResponse: ... - - @overload - async def generate( - self, - model: str = '', - prompt: str = '', - suffix: str = '', - *, - system: str = '', - template: str = '', - context: Optional[Sequence[int]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: bool = False, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> AsyncIterator[GenerateResponse]: ... - - async def generate( - self, - model: str = '', - prompt: Optional[str] = None, - suffix: Optional[str] = None, - *, - system: Optional[str] = None, - template: Optional[str] = None, - context: Optional[Sequence[int]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - raw: Optional[bool] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - images: Optional[Sequence[Union[str, bytes, Image]]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - width: Optional[int] = None, - height: Optional[int] = None, - steps: Optional[int] = None, - ) -> Union[GenerateResponse, AsyncIterator[GenerateResponse]]: - """ - 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. - """ - return await self._request( - GenerateResponse, - 'POST', - '/api/generate', - json=GenerateRequest( - model=model, - prompt=prompt, - suffix=suffix, - system=system, - template=template, - context=context, - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - raw=raw, - format=format, - images=list(_copy_images(images)) if images else None, - options=options, - keep_alive=keep_alive, - width=width, - height=height, - steps=steps, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - async def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[False] = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> ChatResponse: ... - - @overload - async def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: Literal[True] = True, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> AsyncIterator[ChatResponse]: ... - - async def chat( - self, - model: str = '', - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, - stream: bool = False, - think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, - logprobs: Optional[bool] = None, - top_logprobs: Optional[int] = None, - format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> Union[ChatResponse, AsyncIterator[ChatResponse]]: - """ - Create a chat response using the requested model. - - Args: - tools: - A JSON schema as a dict, an Ollama Tool or a Python Function. - Python functions need to follow Google style docstrings to be converted to an Ollama Tool. - For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings - stream: Whether to stream the response. - format: The format of the response. - - Example: - def add_two_numbers(a: int, b: int) -> int: - ''' - Add two numbers together. - - Args: - a: First number to add - b: Second number to add - - Returns: - int: The sum of a and b - ''' - return a + b - - await client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) - - 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. - """ - - return await self._request( - ChatResponse, - 'POST', - '/api/chat', - json=ChatRequest( - model=model, - messages=list(_copy_messages(messages)), - tools=list(_copy_tools(tools)), - stream=stream, - think=think, - logprobs=logprobs, - top_logprobs=top_logprobs, - format=format, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - stream=stream, - ) - - async def embed( - self, - model: str = '', - input: Union[str, Sequence[str]] = '', - truncate: Optional[bool] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - dimensions: Optional[int] = None, - ) -> EmbedResponse: - return await self._request( - EmbedResponse, - 'POST', - '/api/embed', - json=EmbedRequest( - model=model, - input=input, - truncate=truncate, - options=options, - keep_alive=keep_alive, - dimensions=dimensions, - ).model_dump(exclude_none=True), - ) - - async def embeddings( - self, - model: str = '', - prompt: Optional[str] = None, - options: Optional[Union[Mapping[str, Any], Options]] = None, - keep_alive: Optional[Union[float, str]] = None, - ) -> EmbeddingsResponse: - """ - Deprecated in favor of `embed`. - """ - return await self._request( - EmbeddingsResponse, - 'POST', - '/api/embeddings', - json=EmbeddingsRequest( - model=model, - prompt=prompt, - options=options, - keep_alive=keep_alive, - ).model_dump(exclude_none=True), - ) - - @overload - async def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - async def pull( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> AsyncIterator[ProgressResponse]: ... - - async def pull( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return await self._request( - ProgressResponse, - 'POST', - '/api/pull', - json=PullRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - async def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - async def push( - self, - model: str, - *, - insecure: bool = False, - stream: Literal[True] = True, - ) -> AsyncIterator[ProgressResponse]: ... - - async def push( - self, - model: str, - *, - insecure: bool = False, - stream: bool = False, - ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - return await self._request( - ProgressResponse, - 'POST', - '/api/push', - json=PushRequest( - model=model, - insecure=insecure, - stream=stream, - ).model_dump(exclude_none=True), - stream=stream, - ) - - @overload - async def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[False] = False, - ) -> ProgressResponse: ... - - @overload - async def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: Literal[True] = True, - ) -> AsyncIterator[ProgressResponse]: ... - - async def create( - self, - model: str, - quantize: Optional[str] = None, - from_: Optional[str] = None, - files: Optional[Dict[str, str]] = None, - adapters: Optional[Dict[str, str]] = None, - template: Optional[str] = None, - license: Optional[Union[str, List[str]]] = None, - system: Optional[str] = None, - parameters: Optional[Union[Mapping[str, Any], Options]] = None, - messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, - *, - stream: bool = False, - ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: - """ - Raises `ResponseError` if the request could not be fulfilled. - - Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. - """ - - return await self._request( - ProgressResponse, - 'POST', - '/api/create', - json=CreateRequest( - model=model, - stream=stream, - quantize=quantize, - from_=from_, - files=files, - adapters=adapters, - license=license, - template=template, - system=system, - parameters=parameters, - messages=messages, - ).model_dump(exclude_none=True), - stream=stream, - ) - - async def create_blob(self, path: Union[str, Path]) -> str: - sha256sum = sha256() - async with await anyio.open_file(path, 'rb') as r: - while True: - chunk = await r.read(32 * 1024) - if not chunk: - break - sha256sum.update(chunk) - - digest = f'sha256:{sha256sum.hexdigest()}' - - async def upload_bytes(): - async with await anyio.open_file(path, 'rb') as r: - while True: - chunk = await r.read(32 * 1024) - if not chunk: - break - yield chunk - - await self._request_raw('POST', f'/api/blobs/{digest}', content=upload_bytes()) - - return digest - - async def list(self) -> ListResponse: - return await self._request( - ListResponse, - 'GET', - '/api/tags', - ) - - async def delete(self, model: str) -> StatusResponse: - r = await self._request_raw( - 'DELETE', - '/api/delete', - json=DeleteRequest( - model=model, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - async def copy(self, source: str, destination: str) -> StatusResponse: - r = await self._request_raw( - 'POST', - '/api/copy', - json=CopyRequest( - source=source, - destination=destination, - ).model_dump(exclude_none=True), - ) - return StatusResponse( - status='success' if r.status_code == 200 else 'error', - ) - - async def show(self, model: str) -> ShowResponse: - return await self._request( - ShowResponse, - 'POST', - '/api/show', - json=ShowRequest( - model=model, - ).model_dump(exclude_none=True), - ) - - async def ps(self) -> ProcessResponse: - return await self._request( - ProcessResponse, - 'GET', - '/api/ps', - ) - - -def _copy_images(images: Optional[Sequence[Union[Image, Any]]]) -> Iterator[Image]: - for image in images or []: - yield image if isinstance(image, Image) else Image(value=image) - - -def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]]]) -> Iterator[Message]: - for message in messages or []: - # Keep empty strings (e.g. tool results with content='') — only drop None. - yield Message.model_validate( - {k: list(_copy_images(v)) if k == 'images' else v for k, v in dict(message).items() if v is not None}, - ) - - -def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None) -> Iterator[Tool]: - for unprocessed_tool in tools or []: - yield convert_function_to_tool(unprocessed_tool) if callable(unprocessed_tool) else Tool.model_validate(unprocessed_tool) - - -def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]: - if isinstance(s, (str, Path)): - try: - if (p := Path(s)).exists(): - return p - except Exception: - ... - return None - - -def _parse_host(host: Optional[str]) -> str: - """ - >>> _parse_host(None) - 'http://127.0.0.1:11434' - >>> _parse_host('') - 'http://127.0.0.1:11434' - >>> _parse_host('1.2.3.4') - 'http://1.2.3.4:11434' - >>> _parse_host(':56789') - 'http://127.0.0.1:56789' - >>> _parse_host('1.2.3.4:56789') - 'http://1.2.3.4:56789' - >>> _parse_host('http://1.2.3.4') - 'http://1.2.3.4:80' - >>> _parse_host('https://1.2.3.4') - 'https://1.2.3.4:443' - >>> _parse_host('https://1.2.3.4:56789') - 'https://1.2.3.4:56789' - >>> _parse_host('example.com') - 'http://example.com:11434' - >>> _parse_host('example.com:56789') - 'http://example.com:56789' - >>> _parse_host('http://example.com') - 'http://example.com:80' - >>> _parse_host('https://example.com') - 'https://example.com:443' - >>> _parse_host('https://example.com:56789') - 'https://example.com:56789' - >>> _parse_host('example.com/') - 'http://example.com:11434' - >>> _parse_host('example.com:56789/') - 'http://example.com:56789' - >>> _parse_host('example.com/path') - 'http://example.com:11434/path' - >>> _parse_host('example.com:56789/path') - 'http://example.com:56789/path' - >>> _parse_host('https://example.com:56789/path') - 'https://example.com:56789/path' - >>> _parse_host('example.com:56789/path/') - 'http://example.com:56789/path' - >>> _parse_host('[0001:002:003:0004::1]') - 'http://[0001:002:003:0004::1]:11434' - >>> _parse_host('[0001:002:003:0004::1]:56789') - 'http://[0001:002:003:0004::1]:56789' - >>> _parse_host('http://[0001:002:003:0004::1]') - 'http://[0001:002:003:0004::1]:80' - >>> _parse_host('https://[0001:002:003:0004::1]') - 'https://[0001:002:003:0004::1]:443' - >>> _parse_host('https://[0001:002:003:0004::1]:56789') - 'https://[0001:002:003:0004::1]:56789' - >>> _parse_host('[0001:002:003:0004::1]/') - 'http://[0001:002:003:0004::1]:11434' - >>> _parse_host('[0001:002:003:0004::1]:56789/') - 'http://[0001:002:003:0004::1]:56789' - >>> _parse_host('[0001:002:003:0004::1]/path') - 'http://[0001:002:003:0004::1]:11434/path' - >>> _parse_host('[0001:002:003:0004::1]:56789/path') - 'http://[0001:002:003:0004::1]:56789/path' - >>> _parse_host('https://[0001:002:003:0004::1]:56789/path') - 'https://[0001:002:003:0004::1]:56789/path' - >>> _parse_host('[0001:002:003:0004::1]:56789/path/') - 'http://[0001:002:003:0004::1]:56789/path' - """ - - host, port = host or '', 11434 - scheme, _, hostport = host.partition('://') - if not hostport: - scheme, hostport = 'http', host - elif scheme == 'http': - port = 80 - elif scheme == 'https': - port = 443 - - split = urllib.parse.urlsplit(f'{scheme}://{hostport}') - host = split.hostname or '127.0.0.1' - port = split.port or port - - try: - if isinstance(ipaddress.ip_address(host), ipaddress.IPv6Address): - # Fix missing square brackets for IPv6 from urlsplit - host = f'[{host}]' - except ValueError: - ... - - if path := split.path.strip('/'): - return f'{scheme}://{host}:{port}/{path}' - - return f'{scheme}://{host}:{port}' +import contextlib +import ipaddress +import json +import os +import platform +import sys +import urllib.parse +from hashlib import sha256 +from os import PathLike +from pathlib import Path +from typing import ( + Any, + Callable, + Dict, + List, + Literal, + Mapping, + Optional, + Sequence, + Type, + TypeVar, + Union, + overload, +) + +import anyio +from pydantic.json_schema import JsonSchemaValue + +from ollama._utils import convert_function_to_tool + +if sys.version_info < (3, 9): + from typing import AsyncIterator, Iterator +else: + from collections.abc import AsyncIterator, Iterator + +from importlib import metadata + +try: + __version__ = metadata.version('ollama') +except metadata.PackageNotFoundError: + __version__ = '0.0.0' + +import httpx + +from ollama._types import ( + ChatRequest, + ChatResponse, + CopyRequest, + CreateRequest, + DeleteRequest, + EmbeddingsRequest, + EmbeddingsResponse, + EmbedRequest, + EmbedResponse, + GenerateRequest, + GenerateResponse, + Image, + ListResponse, + Message, + Options, + ProcessResponse, + ProgressResponse, + PullRequest, + PushRequest, + ResponseError, + ShowRequest, + ShowResponse, + StatusResponse, + Tool, + WebFetchRequest, + WebFetchResponse, + WebSearchRequest, + WebSearchResponse, +) + +T = TypeVar('T') + + +class BaseClient(contextlib.AbstractContextManager, contextlib.AbstractAsyncContextManager): + def __init__( + self, + client, + host: Optional[str] = None, + *, + follow_redirects: bool = True, + timeout: Any = None, + headers: Optional[Mapping[str, str]] = None, + **kwargs, + ) -> None: + """ + Creates a httpx client. Default parameters are the same as those defined in httpx + except for the following: + - `follow_redirects`: True + - `timeout`: None + `kwargs` are passed to the httpx client. + """ + + headers = { + k.lower(): v + for k, v in { + **(headers or {}), + 'Content-Type': 'application/json', + 'Accept': 'application/json', + 'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}', + }.items() + if v is not None + } + api_key = os.getenv('OLLAMA_API_KEY', None) + if not headers.get('authorization') and api_key: + headers['authorization'] = f'Bearer {api_key}' + + self._client = client( + base_url=_parse_host(host or os.getenv('OLLAMA_HOST')), + follow_redirects=follow_redirects, + timeout=timeout, + headers=headers, + **kwargs, + ) + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.close() + + +CONNECTION_ERROR_MESSAGE = 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + + +class Client(BaseClient): + def __init__(self, host: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.Client, host, **kwargs) + + def close(self): + self._client.close() + + def _request_raw(self, *args, **kwargs): + try: + r = self._client.request(*args, **kwargs) + r.raise_for_status() + return r + except httpx.HTTPStatusError as e: + raise ResponseError(e.response.text, e.response.status_code) from None + except httpx.ConnectError: + raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None + + @overload + def _request( + self, + cls: Type[T], + *args, + stream: Literal[False] = False, + **kwargs, + ) -> T: ... + + @overload + def _request( + self, + cls: Type[T], + *args, + stream: Literal[True] = True, + **kwargs, + ) -> Iterator[T]: ... + + @overload + def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, Iterator[T]]: ... + + def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, Iterator[T]]: + if stream: + + def inner(): + with self._client.stream(*args, **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(): + part = json.loads(line) + if err := part.get('error'): + raise ResponseError(err) + yield cls(**part) + + return inner() + + return cls(**self._request_raw(*args, **kwargs).json()) + + @overload + def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> GenerateResponse: ... + + @overload + def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> Iterator[GenerateResponse]: ... + + def generate( + self, + model: str = '', + prompt: Optional[str] = None, + suffix: Optional[str] = None, + *, + system: Optional[str] = None, + template: Optional[str] = None, + context: Optional[Sequence[int]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: Optional[bool] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> Union[GenerateResponse, Iterator[GenerateResponse]]: + """ + 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. + """ + + return self._request( + GenerateResponse, + 'POST', + '/api/generate', + json=GenerateRequest( + model=model, + prompt=prompt, + suffix=suffix, + system=system, + template=template, + context=context, + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + raw=raw, + format=format, + images=list(_copy_images(images)) if images else None, + options=options, + keep_alive=keep_alive, + width=width, + height=height, + steps=steps, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> ChatResponse: ... + + @overload + def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> Iterator[ChatResponse]: ... + + def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> Union[ChatResponse, Iterator[ChatResponse]]: + """ + Create a chat response using the requested model. + + Args: + tools: + A JSON schema as a dict, an Ollama Tool or a Python Function. + Python functions need to follow Google style docstrings to be converted to an Ollama Tool. + For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings + stream: Whether to stream the response. + format: The format of the response. + + Example: + def add_two_numbers(a: int, b: int) -> int: + ''' + Add two numbers together. + + Args: + a: First number to add + b: Second number to add + + Returns: + int: The sum of a and b + ''' + return a + b + + client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) + + 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. + """ + return self._request( + ChatResponse, + 'POST', + '/api/chat', + json=ChatRequest( + model=model, + messages=list(_copy_messages(messages)), + tools=list(_copy_tools(tools)), + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + format=format, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + stream=stream, + ) + + def embed( + self, + model: str = '', + input: Union[str, Sequence[str]] = '', + truncate: Optional[bool] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + dimensions: Optional[int] = None, + ) -> EmbedResponse: + return self._request( + EmbedResponse, + 'POST', + '/api/embed', + json=EmbedRequest( + model=model, + input=input, + truncate=truncate, + options=options, + keep_alive=keep_alive, + dimensions=dimensions, + ).model_dump(exclude_none=True), + ) + + def embeddings( + self, + model: str = '', + prompt: Optional[str] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> EmbeddingsResponse: + """ + Deprecated in favor of `embed`. + """ + return self._request( + EmbeddingsResponse, + 'POST', + '/api/embeddings', + json=EmbeddingsRequest( + model=model, + prompt=prompt, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + ) + + @overload + def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> Iterator[ProgressResponse]: ... + + def pull( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return self._request( + ProgressResponse, + 'POST', + '/api/pull', + json=PullRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> Iterator[ProgressResponse]: ... + + def push( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return self._request( + ProgressResponse, + 'POST', + '/api/push', + json=PushRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[True] = True, + ) -> Iterator[ProgressResponse]: ... + + def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: bool = False, + ) -> Union[ProgressResponse, Iterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return self._request( + ProgressResponse, + 'POST', + '/api/create', + json=CreateRequest( + model=model, + stream=stream, + quantize=quantize, + from_=from_, + files=files, + adapters=adapters, + license=license, + template=template, + system=system, + parameters=parameters, + messages=messages, + ).model_dump(exclude_none=True), + stream=stream, + ) + + def create_blob(self, path: Union[str, Path]) -> str: + sha256sum = sha256() + with open(path, 'rb') as r: + while True: + chunk = r.read(32 * 1024) + if not chunk: + break + sha256sum.update(chunk) + + digest = f'sha256:{sha256sum.hexdigest()}' + + with open(path, 'rb') as r: + self._request_raw('POST', f'/api/blobs/{digest}', content=r) + + return digest + + def list(self) -> ListResponse: + return self._request( + ListResponse, + 'GET', + '/api/tags', + ) + + def delete(self, model: str) -> StatusResponse: + r = self._request_raw( + 'DELETE', + '/api/delete', + json=DeleteRequest( + model=model, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + def copy(self, source: str, destination: str) -> StatusResponse: + r = self._request_raw( + 'POST', + '/api/copy', + json=CopyRequest( + source=source, + destination=destination, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + def show(self, model: str) -> ShowResponse: + return self._request( + ShowResponse, + 'POST', + '/api/show', + json=ShowRequest( + model=model, + ).model_dump(exclude_none=True), + ) + + def ps(self) -> ProcessResponse: + return self._request( + ProcessResponse, + 'GET', + '/api/ps', + ) + + def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: + """ + Performs a web search + + Args: + query: The query to search for + max_results: The maximum number of results to return (default: 3) + + Returns: + WebSearchResponse with the search results + Raises: + ValueError: If OLLAMA_API_KEY environment variable is not set + """ + if not self._client.headers.get('authorization', '').startswith('Bearer '): + raise ValueError('Authorization header with Bearer token is required for web search') + + return self._request( + WebSearchResponse, + 'POST', + 'https://ollama.com/api/web_search', + json=WebSearchRequest( + query=query, + max_results=max_results, + ).model_dump(exclude_none=True), + ) + + def web_fetch(self, url: str) -> WebFetchResponse: + """ + Fetches the content of a web page for the provided URL. + + Args: + url: The URL to fetch + + Returns: + WebFetchResponse with the fetched result + """ + if not self._client.headers.get('authorization', '').startswith('Bearer '): + raise ValueError('Authorization header with Bearer token is required for web fetch') + + return self._request( + WebFetchResponse, + 'POST', + 'https://ollama.com/api/web_fetch', + json=WebFetchRequest( + url=url, + ).model_dump(exclude_none=True), + ) + + +class AsyncClient(BaseClient): + def __init__(self, host: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.AsyncClient, host, **kwargs) + + async def close(self): + await self._client.aclose() + + async def _request_raw(self, *args, **kwargs): + try: + r = await self._client.request(*args, **kwargs) + r.raise_for_status() + return r + except httpx.HTTPStatusError as e: + raise ResponseError(e.response.text, e.response.status_code) from None + except httpx.ConnectError: + raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None + + @overload + async def _request( + self, + cls: Type[T], + *args, + stream: Literal[False] = False, + **kwargs, + ) -> T: ... + + @overload + async def _request( + self, + cls: Type[T], + *args, + stream: Literal[True] = True, + **kwargs, + ) -> AsyncIterator[T]: ... + + @overload + async def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, AsyncIterator[T]]: ... + + async def _request( + self, + cls: Type[T], + *args, + stream: bool = False, + **kwargs, + ) -> Union[T, AsyncIterator[T]]: + if stream: + + async def inner(): + async with self._client.stream(*args, **kwargs) as r: + try: + r.raise_for_status() + except httpx.HTTPStatusError as e: + await e.response.aread() + raise ResponseError(e.response.text, e.response.status_code) from None + + async for line in r.aiter_lines(): + part = json.loads(line) + if err := part.get('error'): + raise ResponseError(err) + yield cls(**part) + + return inner() + + return cls(**(await self._request_raw(*args, **kwargs)).json()) + + async def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse: + """ + Performs a web search + + Args: + query: The query to search for + max_results: The maximum number of results to return (default: 3) + + Returns: + WebSearchResponse with the search results + """ + return await self._request( + WebSearchResponse, + 'POST', + 'https://ollama.com/api/web_search', + json=WebSearchRequest( + query=query, + max_results=max_results, + ).model_dump(exclude_none=True), + ) + + async def web_fetch(self, url: str) -> WebFetchResponse: + """ + Fetches the content of a web page for the provided URL. + + Args: + url: The URL to fetch + + Returns: + WebFetchResponse with the fetched result + """ + return await self._request( + WebFetchResponse, + 'POST', + 'https://ollama.com/api/web_fetch', + json=WebFetchRequest( + url=url, + ).model_dump(exclude_none=True), + ) + + @overload + async def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> GenerateResponse: ... + + @overload + async def generate( + self, + model: str = '', + prompt: str = '', + suffix: str = '', + *, + system: str = '', + template: str = '', + context: Optional[Sequence[int]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: bool = False, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> AsyncIterator[GenerateResponse]: ... + + async def generate( + self, + model: str = '', + prompt: Optional[str] = None, + suffix: Optional[str] = None, + *, + system: Optional[str] = None, + template: Optional[str] = None, + context: Optional[Sequence[int]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + raw: Optional[bool] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + images: Optional[Sequence[Union[str, bytes, Image]]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + width: Optional[int] = None, + height: Optional[int] = None, + steps: Optional[int] = None, + ) -> Union[GenerateResponse, AsyncIterator[GenerateResponse]]: + """ + 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. + """ + return await self._request( + GenerateResponse, + 'POST', + '/api/generate', + json=GenerateRequest( + model=model, + prompt=prompt, + suffix=suffix, + system=system, + template=template, + context=context, + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + raw=raw, + format=format, + images=list(_copy_images(images)) if images else None, + options=options, + keep_alive=keep_alive, + width=width, + height=height, + steps=steps, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + async def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[False] = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> ChatResponse: ... + + @overload + async def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: Literal[True] = True, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> AsyncIterator[ChatResponse]: ... + + async def chat( + self, + model: str = '', + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None, + stream: bool = False, + think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> Union[ChatResponse, AsyncIterator[ChatResponse]]: + """ + Create a chat response using the requested model. + + Args: + tools: + A JSON schema as a dict, an Ollama Tool or a Python Function. + Python functions need to follow Google style docstrings to be converted to an Ollama Tool. + For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings + stream: Whether to stream the response. + format: The format of the response. + + Example: + def add_two_numbers(a: int, b: int) -> int: + ''' + Add two numbers together. + + Args: + a: First number to add + b: Second number to add + + Returns: + int: The sum of a and b + ''' + return a + b + + await client.chat(model='llama3.2', tools=[add_two_numbers], messages=[...]) + + 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. + """ + + return await self._request( + ChatResponse, + 'POST', + '/api/chat', + json=ChatRequest( + model=model, + messages=list(_copy_messages(messages)), + tools=list(_copy_tools(tools)), + stream=stream, + think=think, + logprobs=logprobs, + top_logprobs=top_logprobs, + format=format, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + stream=stream, + ) + + async def embed( + self, + model: str = '', + input: Union[str, Sequence[str]] = '', + truncate: Optional[bool] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + dimensions: Optional[int] = None, + ) -> EmbedResponse: + return await self._request( + EmbedResponse, + 'POST', + '/api/embed', + json=EmbedRequest( + model=model, + input=input, + truncate=truncate, + options=options, + keep_alive=keep_alive, + dimensions=dimensions, + ).model_dump(exclude_none=True), + ) + + async def embeddings( + self, + model: str = '', + prompt: Optional[str] = None, + options: Optional[Union[Mapping[str, Any], Options]] = None, + keep_alive: Optional[Union[float, str]] = None, + ) -> EmbeddingsResponse: + """ + Deprecated in favor of `embed`. + """ + return await self._request( + EmbeddingsResponse, + 'POST', + '/api/embeddings', + json=EmbeddingsRequest( + model=model, + prompt=prompt, + options=options, + keep_alive=keep_alive, + ).model_dump(exclude_none=True), + ) + + @overload + async def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + async def pull( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> AsyncIterator[ProgressResponse]: ... + + async def pull( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return await self._request( + ProgressResponse, + 'POST', + '/api/pull', + json=PullRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + async def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + async def push( + self, + model: str, + *, + insecure: bool = False, + stream: Literal[True] = True, + ) -> AsyncIterator[ProgressResponse]: ... + + async def push( + self, + model: str, + *, + insecure: bool = False, + stream: bool = False, + ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + return await self._request( + ProgressResponse, + 'POST', + '/api/push', + json=PushRequest( + model=model, + insecure=insecure, + stream=stream, + ).model_dump(exclude_none=True), + stream=stream, + ) + + @overload + async def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[False] = False, + ) -> ProgressResponse: ... + + @overload + async def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: Literal[True] = True, + ) -> AsyncIterator[ProgressResponse]: ... + + async def create( + self, + model: str, + quantize: Optional[str] = None, + from_: Optional[str] = None, + files: Optional[Dict[str, str]] = None, + adapters: Optional[Dict[str, str]] = None, + template: Optional[str] = None, + license: Optional[Union[str, List[str]]] = None, + system: Optional[str] = None, + parameters: Optional[Union[Mapping[str, Any], Options]] = None, + messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None, + *, + stream: bool = False, + ) -> Union[ProgressResponse, AsyncIterator[ProgressResponse]]: + """ + Raises `ResponseError` if the request could not be fulfilled. + + Returns `ProgressResponse` if `stream` is `False`, otherwise returns a `ProgressResponse` generator. + """ + + return await self._request( + ProgressResponse, + 'POST', + '/api/create', + json=CreateRequest( + model=model, + stream=stream, + quantize=quantize, + from_=from_, + files=files, + adapters=adapters, + license=license, + template=template, + system=system, + parameters=parameters, + messages=messages, + ).model_dump(exclude_none=True), + stream=stream, + ) + + async def create_blob(self, path: Union[str, Path]) -> str: + sha256sum = sha256() + async with await anyio.open_file(path, 'rb') as r: + while True: + chunk = await r.read(32 * 1024) + if not chunk: + break + sha256sum.update(chunk) + + digest = f'sha256:{sha256sum.hexdigest()}' + + async def upload_bytes(): + async with await anyio.open_file(path, 'rb') as r: + while True: + chunk = await r.read(32 * 1024) + if not chunk: + break + yield chunk + + await self._request_raw('POST', f'/api/blobs/{digest}', content=upload_bytes()) + + return digest + + async def list(self) -> ListResponse: + return await self._request( + ListResponse, + 'GET', + '/api/tags', + ) + + async def delete(self, model: str) -> StatusResponse: + r = await self._request_raw( + 'DELETE', + '/api/delete', + json=DeleteRequest( + model=model, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + async def copy(self, source: str, destination: str) -> StatusResponse: + r = await self._request_raw( + 'POST', + '/api/copy', + json=CopyRequest( + source=source, + destination=destination, + ).model_dump(exclude_none=True), + ) + return StatusResponse( + status='success' if r.status_code == 200 else 'error', + ) + + async def show(self, model: str) -> ShowResponse: + return await self._request( + ShowResponse, + 'POST', + '/api/show', + json=ShowRequest( + model=model, + ).model_dump(exclude_none=True), + ) + + async def ps(self) -> ProcessResponse: + return await self._request( + ProcessResponse, + 'GET', + '/api/ps', + ) + + +def _copy_images(images: Optional[Sequence[Union[Image, Any]]]) -> Iterator[Image]: + for image in images or []: + yield image if isinstance(image, Image) else Image(value=image) + + +def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]]]) -> Iterator[Message]: + for message in messages or []: + # Keep empty strings (e.g. tool results with content='') — only drop None. + yield Message.model_validate( + {k: list(_copy_images(v)) if k == 'images' else v for k, v in dict(message).items() if v is not None}, + ) + + +def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None) -> Iterator[Tool]: + for unprocessed_tool in tools or []: + yield convert_function_to_tool(unprocessed_tool) if callable(unprocessed_tool) else Tool.model_validate(unprocessed_tool) + + +def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]: + if isinstance(s, (str, Path)): + try: + if (p := Path(s)).exists(): + return p + except Exception: + ... + return None + + +def _parse_host(host: Optional[str]) -> str: + """ + >>> _parse_host(None) + 'http://127.0.0.1:11434' + >>> _parse_host('') + 'http://127.0.0.1:11434' + >>> _parse_host('1.2.3.4') + 'http://1.2.3.4:11434' + >>> _parse_host(':56789') + 'http://127.0.0.1:56789' + >>> _parse_host('1.2.3.4:56789') + 'http://1.2.3.4:56789' + >>> _parse_host('http://1.2.3.4') + 'http://1.2.3.4:80' + >>> _parse_host('https://1.2.3.4') + 'https://1.2.3.4:443' + >>> _parse_host('https://1.2.3.4:56789') + 'https://1.2.3.4:56789' + >>> _parse_host('example.com') + 'http://example.com:11434' + >>> _parse_host('example.com:56789') + 'http://example.com:56789' + >>> _parse_host('http://example.com') + 'http://example.com:80' + >>> _parse_host('https://example.com') + 'https://example.com:443' + >>> _parse_host('https://example.com:56789') + 'https://example.com:56789' + >>> _parse_host('example.com/') + 'http://example.com:11434' + >>> _parse_host('example.com:56789/') + 'http://example.com:56789' + >>> _parse_host('example.com/path') + 'http://example.com:11434/path' + >>> _parse_host('example.com:56789/path') + 'http://example.com:56789/path' + >>> _parse_host('https://example.com:56789/path') + 'https://example.com:56789/path' + >>> _parse_host('example.com:56789/path/') + 'http://example.com:56789/path' + >>> _parse_host('[0001:002:003:0004::1]') + 'http://[0001:002:003:0004::1]:11434' + >>> _parse_host('[0001:002:003:0004::1]:56789') + 'http://[0001:002:003:0004::1]:56789' + >>> _parse_host('http://[0001:002:003:0004::1]') + 'http://[0001:002:003:0004::1]:80' + >>> _parse_host('https://[0001:002:003:0004::1]') + 'https://[0001:002:003:0004::1]:443' + >>> _parse_host('https://[0001:002:003:0004::1]:56789') + 'https://[0001:002:003:0004::1]:56789' + >>> _parse_host('[0001:002:003:0004::1]/') + 'http://[0001:002:003:0004::1]:11434' + >>> _parse_host('[0001:002:003:0004::1]:56789/') + 'http://[0001:002:003:0004::1]:56789' + >>> _parse_host('[0001:002:003:0004::1]/path') + 'http://[0001:002:003:0004::1]:11434/path' + >>> _parse_host('[0001:002:003:0004::1]:56789/path') + 'http://[0001:002:003:0004::1]:56789/path' + >>> _parse_host('https://[0001:002:003:0004::1]:56789/path') + 'https://[0001:002:003:0004::1]:56789/path' + >>> _parse_host('[0001:002:003:0004::1]:56789/path/') + 'http://[0001:002:003:0004::1]:56789/path' + """ + + host, port = host or '', 11434 + scheme, _, hostport = host.partition('://') + if not hostport: + scheme, hostport = 'http', host + elif scheme == 'http': + port = 80 + elif scheme == 'https': + port = 443 + + split = urllib.parse.urlsplit(f'{scheme}://{hostport}') + host = split.hostname or '127.0.0.1' + port = split.port or port + + try: + if isinstance(ipaddress.ip_address(host), ipaddress.IPv6Address): + # Fix missing square brackets for IPv6 from urlsplit + host = f'[{host}]' + except ValueError: + ... + + if path := split.path.strip('/'): + return f'{scheme}://{host}:{port}/{path}' + + return f'{scheme}://{host}:{port}' From 068efe0212af7e7d539332e13f487d2897d5b57b Mon Sep 17 00:00:00 2001 From: r7mekmy4g67w6l Date: Thu, 6 Aug 2026 17:15:19 +0200 Subject: [PATCH 4/6] fix: preserve empty string message content in chat requests Signed-off-by: r7mekmy4g67w6l Normalize to LF so the PR shows the real 1-line change. --- tests/test_client.py | 3030 +++++++++++++++++++++--------------------- 1 file changed, 1515 insertions(+), 1515 deletions(-) diff --git a/tests/test_client.py b/tests/test_client.py index ddc3c233..89170a60 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,1515 +1,1515 @@ -import base64 -import inspect -import json -import os -import re -import tempfile -from pathlib import Path -from typing import Any - -import pytest -from httpx import Response as httpxResponse -from pydantic import BaseModel -from pytest_httpserver import HTTPServer, URIPattern -from werkzeug.wrappers import Request, Response - -from ollama._client import CONNECTION_ERROR_MESSAGE, AsyncClient, Client, _copy_tools -from ollama._types import Image, Message - -PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC' -PNG_BYTES = base64.b64decode(PNG_BASE64) - -pytestmark = pytest.mark.anyio - - -@pytest.fixture -def anyio_backend(): - return 'asyncio' - - -class PrefixPattern(URIPattern): - def __init__(self, prefix: str): - self.prefix = prefix - - def match(self, uri): - return uri.startswith(self.prefix) - - -def test_client_chat(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -def test_client_chat_with_logprobs(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Hi'}], - 'tools': [], - 'stream': False, - 'logprobs': True, - 'top_logprobs': 3, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': 'Hello', - }, - 'logprobs': [ - { - 'token': 'Hello', - 'logprob': -0.1, - 'top_logprobs': [ - {'token': 'Hello', 'logprob': -0.1}, - {'token': 'Hi', 'logprob': -1.0}, - ], - } - ], - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Hi'}], logprobs=True, top_logprobs=3) - assert response['logprobs'][0]['token'] == 'Hello' - assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' - - -def test_client_chat_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['I ', "don't ", 'know.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': message, - }, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) - - it = iter(['I ', "don't ", 'know.']) - for part in response: - assert part['message']['role'] in 'assistant' - assert part['message']['content'] == next(it) - - -@pytest.mark.parametrize('message_format', ('dict', 'pydantic_model')) -@pytest.mark.parametrize('file_style', ('path', 'bytes')) -def test_client_chat_images(httpserver: HTTPServer, message_format: str, file_style: str, tmp_path): - from ollama._types import Image, Message - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [ - { - 'role': 'user', - 'content': 'Why is the sky blue?', - 'images': [PNG_BASE64], - }, - ], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = Client(httpserver.url_for('/')) - - if file_style == 'bytes': - image_content = PNG_BYTES - elif file_style == 'path': - image_path = tmp_path / 'transparent.png' - image_path.write_bytes(PNG_BYTES) - image_content = str(image_path) - - if message_format == 'pydantic_model': - messages = [Message(role='user', content='Why is the sky blue?', images=[Image(value=image_content)])] - elif message_format == 'dict': - messages = [{'role': 'user', 'content': 'Why is the sky blue?', 'images': [image_content]}] - else: - raise ValueError(f'Invalid message format: {message_format}') - - response = client.chat('dummy', messages=messages) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -def test_client_chat_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering"}', - }, - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' - - -def test_client_chat_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - }, - } - ) - - client = Client(httpserver.url_for('/')) - response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -async def test_async_client_chat_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering"}', - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' - - -async def test_async_client_chat_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -def test_client_generate(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?') - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -def test_client_generate_with_logprobs(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why', - 'stream': False, - 'logprobs': True, - 'top_logprobs': 2, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Hello', - 'logprobs': [ - { - 'token': 'Hello', - 'logprob': -0.2, - 'top_logprobs': [ - {'token': 'Hello', 'logprob': -0.2}, - {'token': 'Hi', 'logprob': -1.5}, - ], - } - ], - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why', logprobs=True, top_logprobs=2) - assert response['logprobs'][0]['token'] == 'Hello' - assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' - - -def test_client_generate_with_image_type(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'What is in this image?', - 'stream': False, - 'images': [PNG_BASE64], - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'A blue sky.', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'What is in this image?', images=[Image(value=PNG_BASE64)]) - assert response['model'] == 'dummy' - assert response['response'] == 'A blue sky.' - - -def test_client_generate_with_invalid_image(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'What is in this image?', - 'stream': False, - 'images': ['invalid_base64'], - }, - ).respond_with_json({'error': 'Invalid image data'}, status=400) - - client = Client(httpserver.url_for('/')) - with pytest.raises(ValueError): - client.generate('dummy', 'What is in this image?', images=[Image(value='invalid_base64')]) - - -def test_client_generate_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['Because ', 'it ', 'is.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'response': message, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?', stream=True) - - it = iter(['Because ', 'it ', 'is.']) - for part in response: - assert part['model'] == 'dummy' - assert part['response'] == next(it) - - -def test_client_generate_images(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - 'images': [PNG_BASE64], - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as temp: - temp.write(PNG_BYTES) - temp.flush() - response = client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -def test_client_generate_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering"}', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?', format='json') - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' - - -def test_client_generate_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -async def test_async_client_generate_format_json(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': 'json', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering"}', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?', format='json') - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' - - -async def test_async_client_generate_format_pydantic(httpserver: HTTPServer): - class ResponseFormat(BaseModel): - answer: str - confidence: float - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) - assert response['model'] == 'dummy' - assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' - - -def test_client_generate_image(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy-image', - 'prompt': 'a sunset over mountains', - 'stream': False, - 'width': 1024, - 'height': 768, - 'steps': 20, - }, - ).respond_with_json( - { - 'model': 'dummy-image', - 'image': PNG_BASE64, - 'done': True, - 'done_reason': 'stop', - } - ) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy-image', 'a sunset over mountains', width=1024, height=768, steps=20) - assert response['model'] == 'dummy-image' - assert response['image'] == PNG_BASE64 - assert response['done'] is True - - -def test_client_generate_image_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - # Progress updates - for i in range(1, 4): - yield ( - json.dumps( - { - 'model': 'dummy-image', - 'completed': i, - 'total': 3, - 'done': False, - } - ) - + '\n' - ) - # Final response with image - yield ( - json.dumps( - { - 'model': 'dummy-image', - 'image': PNG_BASE64, - 'done': True, - 'done_reason': 'stop', - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy-image', - 'prompt': 'a sunset over mountains', - 'stream': True, - 'width': 512, - 'height': 512, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.generate('dummy-image', 'a sunset over mountains', stream=True, width=512, height=512) - - parts = list(response) - # Check progress updates - assert parts[0]['completed'] == 1 - assert parts[0]['total'] == 3 - assert parts[0]['done'] is False - # Check final response - assert parts[-1]['image'] == PNG_BASE64 - assert parts[-1]['done'] is True - - -async def test_async_client_generate_image(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy-image', - 'prompt': 'a robot painting', - 'stream': False, - 'width': 1024, - 'height': 1024, - }, - ).respond_with_json( - { - 'model': 'dummy-image', - 'image': PNG_BASE64, - 'done': True, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy-image', 'a robot painting', width=1024, height=1024) - assert response['model'] == 'dummy-image' - assert response['image'] == PNG_BASE64 - - -def test_client_pull(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - response = client.pull('dummy') - assert response['status'] == 'success' - - -def test_client_pull_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'pulling manifest'}) + '\n' - yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' - yield json.dumps({'status': 'writing manifest'}) + '\n' - yield json.dumps({'status': 'removing any unused layers'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.pull('dummy', stream=True) - - it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) - for part in response: - assert part['status'] == next(it) - - -def test_client_push(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - response = client.push('dummy') - assert response['status'] == 'success' - - -def test_client_push_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'retrieving manifest'}) + '\n' - yield json.dumps({'status': 'pushing manifest'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = Client(httpserver.url_for('/')) - response = client.push('dummy', stream=True) - - it = iter(['retrieving manifest', 'pushing manifest', 'success']) - for part in response: - assert part['status'] == next(it) - - -@pytest.fixture -def userhomedir(): - with tempfile.TemporaryDirectory() as temp: - home = os.getenv('HOME', '') - os.environ['HOME'] = temp - yield Path(temp) - os.environ['HOME'] = home - - -def test_client_create_with_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) - assert response['status'] == 'success' - - -def test_client_create_with_parameters_roundtrip(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'quantize': 'q4_k_m', - 'from': 'mymodel', - 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - 'license': 'this is my license', - 'system': '\nUse\nmultiline\nstrings.\n', - 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = client.create( - 'dummy', - quantize='q4_k_m', - from_='mymodel', - adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - license='this is my license', - system='\nUse\nmultiline\nstrings.\n', - parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - stream=False, - ) - assert response['status'] == 'success' - - -def test_client_create_from_library(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'from': 'llama2', - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = Client(httpserver.url_for('/')) - - response = client.create('dummy', from_='llama2') - assert response['status'] == 'success' - - -def test_client_create_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -def test_client_create_blob_exists(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) - - client = Client(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -def test_client_delete(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) - client = Client(httpserver.url_for('/api/delete')) - response = client.delete('dummy') - assert response['status'] == 'success' - - -def test_client_copy(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) - client = Client(httpserver.url_for('/api/copy')) - response = client.copy('dum', 'dummer') - assert response['status'] == 'success' - - -async def test_async_client_chat(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -async def test_async_client_chat_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['I ', "don't ", 'know.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': message, - }, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], - 'tools': [], - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) - - it = iter(['I ', "don't ", 'know.']) - async for part in response: - assert part['message']['role'] == 'assistant' - assert part['message']['content'] == next(it) - - -async def test_async_client_chat_images(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/chat', - method='POST', - json={ - 'model': 'dummy', - 'messages': [ - { - 'role': 'user', - 'content': 'Why is the sky blue?', - 'images': [PNG_BASE64], - }, - ], - 'tools': [], - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'message': { - 'role': 'assistant', - 'content': "I don't know.", - }, - } - ) - - client = AsyncClient(httpserver.url_for('/')) - - response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?', 'images': [PNG_BYTES]}]) - assert response['model'] == 'dummy' - assert response['message']['role'] == 'assistant' - assert response['message']['content'] == "I don't know." - - -async def test_async_client_generate(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?') - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -async def test_async_client_generate_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - for message in ['Because ', 'it ', 'is.']: - yield ( - json.dumps( - { - 'model': 'dummy', - 'response': message, - } - ) - + '\n' - ) - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.generate('dummy', 'Why is the sky blue?', stream=True) - - it = iter(['Because ', 'it ', 'is.']) - async for part in response: - assert part['model'] == 'dummy' - assert part['response'] == next(it) - - -async def test_async_client_generate_images(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/generate', - method='POST', - json={ - 'model': 'dummy', - 'prompt': 'Why is the sky blue?', - 'stream': False, - 'images': [PNG_BASE64], - }, - ).respond_with_json( - { - 'model': 'dummy', - 'response': 'Because it is.', - } - ) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as temp: - temp.write(PNG_BYTES) - temp.flush() - response = await client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) - assert response['model'] == 'dummy' - assert response['response'] == 'Because it is.' - - -async def test_async_client_pull(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.pull('dummy') - assert response['status'] == 'success' - - -async def test_async_client_pull_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'pulling manifest'}) + '\n' - yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' - yield json.dumps({'status': 'writing manifest'}) + '\n' - yield json.dumps({'status': 'removing any unused layers'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/pull', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.pull('dummy', stream=True) - - it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) - async for part in response: - assert part['status'] == next(it) - - -async def test_async_client_push(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.push('dummy') - assert response['status'] == 'success' - - -async def test_async_client_push_stream(httpserver: HTTPServer): - def stream_handler(_: Request): - def generate(): - yield json.dumps({'status': 'retrieving manifest'}) + '\n' - yield json.dumps({'status': 'pushing manifest'}) + '\n' - yield json.dumps({'status': 'success'}) + '\n' - - return Response(generate()) - - httpserver.expect_ordered_request( - '/api/push', - method='POST', - json={ - 'model': 'dummy', - 'insecure': False, - 'stream': True, - }, - ).respond_with_handler(stream_handler) - - client = AsyncClient(httpserver.url_for('/')) - response = await client.push('dummy', stream=True) - - it = iter(['retrieving manifest', 'pushing manifest', 'success']) - async for part in response: - assert part['status'] == next(it) - - -async def test_async_client_create_with_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = await client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) - assert response['status'] == 'success' - - -async def test_async_client_create_with_parameters_roundtrip(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'quantize': 'q4_k_m', - 'from': 'mymodel', - 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - 'license': 'this is my license', - 'system': '\nUse\nmultiline\nstrings.\n', - 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile(): - response = await client.create( - 'dummy', - quantize='q4_k_m', - from_='mymodel', - adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, - template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', - license='this is my license', - system='\nUse\nmultiline\nstrings.\n', - parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, - messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], - stream=False, - ) - assert response['status'] == 'success' - - -async def test_async_client_create_from_library(httpserver: HTTPServer): - httpserver.expect_ordered_request( - '/api/create', - method='POST', - json={ - 'model': 'dummy', - 'from': 'llama2', - 'stream': False, - }, - ).respond_with_json({'status': 'success'}) - - client = AsyncClient(httpserver.url_for('/')) - - response = await client.create('dummy', from_='llama2') - assert response['status'] == 'success' - - -async def test_async_client_create_blob(httpserver: HTTPServer): - httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = await client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -async def test_async_client_create_blob_exists(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) - - client = AsyncClient(httpserver.url_for('/')) - - with tempfile.NamedTemporaryFile() as blob: - response = await client.create_blob(blob.name) - assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' - - -async def test_async_client_delete(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) - client = AsyncClient(httpserver.url_for('/api/delete')) - response = await client.delete('dummy') - assert response['status'] == 'success' - - -async def test_async_client_copy(httpserver: HTTPServer): - httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) - client = AsyncClient(httpserver.url_for('/api/copy')) - response = await client.copy('dum', 'dummer') - assert response['status'] == 'success' - - -def test_headers(): - client = Client() - assert client._client.headers['content-type'] == 'application/json' - assert client._client.headers['accept'] == 'application/json' - assert client._client.headers['user-agent'].startswith('ollama-python/') - - client = Client( - headers={ - 'X-Custom': 'value', - 'Content-Type': 'text/plain', - } - ) - assert client._client.headers['x-custom'] == 'value' - assert client._client.headers['content-type'] == 'application/json' - - -def test_copy_tools(): - def func1(x: int) -> str: - """Simple function 1. - Args: - x (integer): A number - """ - - def func2(y: str) -> int: - """Simple function 2. - Args: - y (string): A string - """ - - # Test with list of functions - tools = list(_copy_tools([func1, func2])) - assert len(tools) == 2 - assert tools[0].function.name == 'func1' - assert tools[1].function.name == 'func2' - - # Test with empty input - assert list(_copy_tools()) == [] - assert list(_copy_tools(None)) == [] - assert list(_copy_tools([])) == [] - - # Test with mix of functions and tool dicts - tool_dict = { - 'type': 'function', - 'function': { - 'name': 'test', - 'description': 'Test function', - 'parameters': { - 'type': 'object', - 'properties': {'x': {'type': 'string', 'description': 'A string', 'enum': ['a', 'b', 'c']}, 'y': {'type': ['integer', 'number'], 'description': 'An integer'}}, - 'required': ['x'], - }, - }, - } - - tools = list(_copy_tools([func1, tool_dict])) - assert len(tools) == 2 - assert tools[0].function.name == 'func1' - assert tools[1].function.name == 'test' - - -def test_tool_validation(): - arbitrary_tool = {'type': 'custom_type', 'function': {'name': 'test'}} - tools = list(_copy_tools([arbitrary_tool])) - assert len(tools) == 1 - assert tools[0].type == 'custom_type' - assert tools[0].function.name == 'test' - - -def test_client_connection_error(): - client = Client('http://localhost:1234') - - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.generate('model', 'prompt') - with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): - client.show('model') - - -async def test_async_client_connection_error(): - client = AsyncClient('http://localhost:1234') - with pytest.raises(ConnectionError) as exc_info: - await client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) - assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - with pytest.raises(ConnectionError) as exc_info: - await client.generate('model', 'prompt') - assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - with pytest.raises(ConnectionError) as exc_info: - await client.show('model') - assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' - - -def test_arbitrary_roles_accepted_in_message(): - _ = Message(role='somerandomrole', content="I'm ok with you adding any role message now!") - - -def _mock_request(*args: Any, **kwargs: Any) -> Response: - return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") - - -def test_arbitrary_roles_accepted_in_message_request(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(Client, '_request', _mock_request) - - client = Client() - - client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) - - -async def _mock_request_async(*args: Any, **kwargs: Any) -> Response: - return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") - - -async def test_arbitrary_roles_accepted_in_message_request_async(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(AsyncClient, '_request', _mock_request_async) - - client = AsyncClient() - - await client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) - - -def test_copy_messages_preserves_empty_string_content(): - from ollama._client import _copy_messages - - msgs = list( - _copy_messages( - [ - {'role': 'assistant', 'content': ''}, - {'role': 'tool', 'content': '', 'tool_name': 'web_search'}, - ] - ) - ) - assert msgs[0].content == '' - assert msgs[1].content == '' - assert msgs[1].tool_name == 'web_search' - - -def test_client_web_search_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - - client = Client() - - with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'): - client.web_search('test query') - - -def test_client_web_fetch_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - - client = Client() - - with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'): - client.web_fetch('https://example.com') - - -def _mock_request_web_search(self, cls, method, url, json=None, **kwargs): - assert method == 'POST' - assert url == 'https://ollama.com/api/web_search' - assert json is not None and 'query' in json and 'max_results' in json - return httpxResponse(status_code=200, content='{"results": {}, "success": true}') - - -def _mock_request_web_fetch(self, cls, method, url, json=None, **kwargs): - assert method == 'POST' - assert url == 'https://ollama.com/api/web_fetch' - assert json is not None and 'url' in json - return httpxResponse(status_code=200, content='{"results": {}, "success": true}') - - -def test_client_web_search_with_env_api_key(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') - monkeypatch.setattr(Client, '_request', _mock_request_web_search) - - client = Client() - client.web_search('what is ollama?', max_results=2) - - -def test_client_web_fetch_with_env_api_key(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') - monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) - - client = Client() - client.web_fetch('https://example.com') - - -def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - monkeypatch.setattr(Client, '_request', _mock_request_web_search) - - client = Client(headers={'Authorization': 'Bearer custom-token'}) - client.web_search('what is ollama?', max_results=1) - - -def test_client_web_fetch_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): - monkeypatch.delenv('OLLAMA_API_KEY', raising=False) - monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) - - client = Client(headers={'Authorization': 'Bearer custom-token'}) - client.web_fetch('https://example.com') - - -def test_client_bearer_header_from_env(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') - - client = Client() - assert client._client.headers['authorization'] == 'Bearer env-token' - - -def test_client_explicit_bearer_header_overrides_env(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') - monkeypatch.setattr(Client, '_request', _mock_request_web_search) - - client = Client(headers={'Authorization': 'Bearer explicit-token'}) - assert client._client.headers['authorization'] == 'Bearer explicit-token' - client.web_search('override check') - - -def test_client_close(): - client = Client() - client.close() - assert client._client.is_closed - - -@pytest.mark.anyio -async def test_async_client_close(): - client = AsyncClient() - await client.close() - assert client._client.is_closed - - -def test_client_context_manager(): - with Client() as client: - assert isinstance(client, Client) - assert not client._client.is_closed - - assert client._client.is_closed - - -@pytest.mark.anyio -async def test_async_client_context_manager(): - async with AsyncClient() as client: - assert isinstance(client, AsyncClient) - assert not client._client.is_closed - - assert client._client.is_closed - - -def test_generate_think_annotation_matches_chat(): - # The `think` parameter accepts bool or the 'low'/'medium'/'high' string levels. - # Client.generate must keep the same annotation as Client.chat and - # AsyncClient.generate so passing a string level does not raise a false type - # error (regression guard for the sync generate overloads/implementation). - expected = inspect.signature(Client.chat).parameters['think'].annotation - assert inspect.signature(Client.generate).parameters['think'].annotation == expected - assert inspect.signature(AsyncClient.generate).parameters['think'].annotation == expected +import base64 +import inspect +import json +import os +import re +import tempfile +from pathlib import Path +from typing import Any + +import pytest +from httpx import Response as httpxResponse +from pydantic import BaseModel +from pytest_httpserver import HTTPServer, URIPattern +from werkzeug.wrappers import Request, Response + +from ollama._client import CONNECTION_ERROR_MESSAGE, AsyncClient, Client, _copy_tools +from ollama._types import Image, Message + +PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC' +PNG_BYTES = base64.b64decode(PNG_BASE64) + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend(): + return 'asyncio' + + +class PrefixPattern(URIPattern): + def __init__(self, prefix: str): + self.prefix = prefix + + def match(self, uri): + return uri.startswith(self.prefix) + + +def test_client_chat(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +def test_client_chat_with_logprobs(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Hi'}], + 'tools': [], + 'stream': False, + 'logprobs': True, + 'top_logprobs': 3, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': 'Hello', + }, + 'logprobs': [ + { + 'token': 'Hello', + 'logprob': -0.1, + 'top_logprobs': [ + {'token': 'Hello', 'logprob': -0.1}, + {'token': 'Hi', 'logprob': -1.0}, + ], + } + ], + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Hi'}], logprobs=True, top_logprobs=3) + assert response['logprobs'][0]['token'] == 'Hello' + assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' + + +def test_client_chat_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['I ', "don't ", 'know.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': message, + }, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) + + it = iter(['I ', "don't ", 'know.']) + for part in response: + assert part['message']['role'] in 'assistant' + assert part['message']['content'] == next(it) + + +@pytest.mark.parametrize('message_format', ('dict', 'pydantic_model')) +@pytest.mark.parametrize('file_style', ('path', 'bytes')) +def test_client_chat_images(httpserver: HTTPServer, message_format: str, file_style: str, tmp_path): + from ollama._types import Image, Message + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [ + { + 'role': 'user', + 'content': 'Why is the sky blue?', + 'images': [PNG_BASE64], + }, + ], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = Client(httpserver.url_for('/')) + + if file_style == 'bytes': + image_content = PNG_BYTES + elif file_style == 'path': + image_path = tmp_path / 'transparent.png' + image_path.write_bytes(PNG_BYTES) + image_content = str(image_path) + + if message_format == 'pydantic_model': + messages = [Message(role='user', content='Why is the sky blue?', images=[Image(value=image_content)])] + elif message_format == 'dict': + messages = [{'role': 'user', 'content': 'Why is the sky blue?', 'images': [image_content]}] + else: + raise ValueError(f'Invalid message format: {message_format}') + + response = client.chat('dummy', messages=messages) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +def test_client_chat_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering"}', + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' + + +def test_client_chat_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +async def test_async_client_chat_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering"}', + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format='json') + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering"}' + + +async def test_async_client_chat_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +def test_client_generate(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?') + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +def test_client_generate_with_logprobs(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why', + 'stream': False, + 'logprobs': True, + 'top_logprobs': 2, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Hello', + 'logprobs': [ + { + 'token': 'Hello', + 'logprob': -0.2, + 'top_logprobs': [ + {'token': 'Hello', 'logprob': -0.2}, + {'token': 'Hi', 'logprob': -1.5}, + ], + } + ], + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why', logprobs=True, top_logprobs=2) + assert response['logprobs'][0]['token'] == 'Hello' + assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' + + +def test_client_generate_with_image_type(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'What is in this image?', + 'stream': False, + 'images': [PNG_BASE64], + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'A blue sky.', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'What is in this image?', images=[Image(value=PNG_BASE64)]) + assert response['model'] == 'dummy' + assert response['response'] == 'A blue sky.' + + +def test_client_generate_with_invalid_image(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'What is in this image?', + 'stream': False, + 'images': ['invalid_base64'], + }, + ).respond_with_json({'error': 'Invalid image data'}, status=400) + + client = Client(httpserver.url_for('/')) + with pytest.raises(ValueError): + client.generate('dummy', 'What is in this image?', images=[Image(value='invalid_base64')]) + + +def test_client_generate_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['Because ', 'it ', 'is.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'response': message, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?', stream=True) + + it = iter(['Because ', 'it ', 'is.']) + for part in response: + assert part['model'] == 'dummy' + assert part['response'] == next(it) + + +def test_client_generate_images(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + 'images': [PNG_BASE64], + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as temp: + temp.write(PNG_BYTES) + temp.flush() + response = client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +def test_client_generate_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering"}', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?', format='json') + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' + + +def test_client_generate_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +async def test_async_client_generate_format_json(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': 'json', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering"}', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?', format='json') + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering"}' + + +async def test_async_client_generate_format_pydantic(httpserver: HTTPServer): + class ResponseFormat(BaseModel): + answer: str + confidence: float + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'format': {'title': 'ResponseFormat', 'type': 'object', 'properties': {'answer': {'title': 'Answer', 'type': 'string'}, 'confidence': {'title': 'Confidence', 'type': 'number'}}, 'required': ['answer', 'confidence']}, + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?', format=ResponseFormat.model_json_schema()) + assert response['model'] == 'dummy' + assert response['response'] == '{"answer": "Because of Rayleigh scattering", "confidence": 0.95}' + + +def test_client_generate_image(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy-image', + 'prompt': 'a sunset over mountains', + 'stream': False, + 'width': 1024, + 'height': 768, + 'steps': 20, + }, + ).respond_with_json( + { + 'model': 'dummy-image', + 'image': PNG_BASE64, + 'done': True, + 'done_reason': 'stop', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy-image', 'a sunset over mountains', width=1024, height=768, steps=20) + assert response['model'] == 'dummy-image' + assert response['image'] == PNG_BASE64 + assert response['done'] is True + + +def test_client_generate_image_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + # Progress updates + for i in range(1, 4): + yield ( + json.dumps( + { + 'model': 'dummy-image', + 'completed': i, + 'total': 3, + 'done': False, + } + ) + + '\n' + ) + # Final response with image + yield ( + json.dumps( + { + 'model': 'dummy-image', + 'image': PNG_BASE64, + 'done': True, + 'done_reason': 'stop', + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy-image', + 'prompt': 'a sunset over mountains', + 'stream': True, + 'width': 512, + 'height': 512, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy-image', 'a sunset over mountains', stream=True, width=512, height=512) + + parts = list(response) + # Check progress updates + assert parts[0]['completed'] == 1 + assert parts[0]['total'] == 3 + assert parts[0]['done'] is False + # Check final response + assert parts[-1]['image'] == PNG_BASE64 + assert parts[-1]['done'] is True + + +async def test_async_client_generate_image(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy-image', + 'prompt': 'a robot painting', + 'stream': False, + 'width': 1024, + 'height': 1024, + }, + ).respond_with_json( + { + 'model': 'dummy-image', + 'image': PNG_BASE64, + 'done': True, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy-image', 'a robot painting', width=1024, height=1024) + assert response['model'] == 'dummy-image' + assert response['image'] == PNG_BASE64 + + +def test_client_pull(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + response = client.pull('dummy') + assert response['status'] == 'success' + + +def test_client_pull_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'pulling manifest'}) + '\n' + yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' + yield json.dumps({'status': 'writing manifest'}) + '\n' + yield json.dumps({'status': 'removing any unused layers'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.pull('dummy', stream=True) + + it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) + for part in response: + assert part['status'] == next(it) + + +def test_client_push(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + response = client.push('dummy') + assert response['status'] == 'success' + + +def test_client_push_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'retrieving manifest'}) + '\n' + yield json.dumps({'status': 'pushing manifest'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = Client(httpserver.url_for('/')) + response = client.push('dummy', stream=True) + + it = iter(['retrieving manifest', 'pushing manifest', 'success']) + for part in response: + assert part['status'] == next(it) + + +@pytest.fixture +def userhomedir(): + with tempfile.TemporaryDirectory() as temp: + home = os.getenv('HOME', '') + os.environ['HOME'] = temp + yield Path(temp) + os.environ['HOME'] = home + + +def test_client_create_with_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) + assert response['status'] == 'success' + + +def test_client_create_with_parameters_roundtrip(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'quantize': 'q4_k_m', + 'from': 'mymodel', + 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + 'license': 'this is my license', + 'system': '\nUse\nmultiline\nstrings.\n', + 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = client.create( + 'dummy', + quantize='q4_k_m', + from_='mymodel', + adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + license='this is my license', + system='\nUse\nmultiline\nstrings.\n', + parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + stream=False, + ) + assert response['status'] == 'success' + + +def test_client_create_from_library(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'from': 'llama2', + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = Client(httpserver.url_for('/')) + + response = client.create('dummy', from_='llama2') + assert response['status'] == 'success' + + +def test_client_create_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +def test_client_create_blob_exists(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) + + client = Client(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +def test_client_delete(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) + client = Client(httpserver.url_for('/api/delete')) + response = client.delete('dummy') + assert response['status'] == 'success' + + +def test_client_copy(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) + client = Client(httpserver.url_for('/api/copy')) + response = client.copy('dum', 'dummer') + assert response['status'] == 'success' + + +async def test_async_client_chat(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +async def test_async_client_chat_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['I ', "don't ", 'know.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': message, + }, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Why is the sky blue?'}], + 'tools': [], + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True) + + it = iter(['I ', "don't ", 'know.']) + async for part in response: + assert part['message']['role'] == 'assistant' + assert part['message']['content'] == next(it) + + +async def test_async_client_chat_images(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [ + { + 'role': 'user', + 'content': 'Why is the sky blue?', + 'images': [PNG_BASE64], + }, + ], + 'tools': [], + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': "I don't know.", + }, + } + ) + + client = AsyncClient(httpserver.url_for('/')) + + response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?', 'images': [PNG_BYTES]}]) + assert response['model'] == 'dummy' + assert response['message']['role'] == 'assistant' + assert response['message']['content'] == "I don't know." + + +async def test_async_client_generate(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?') + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +async def test_async_client_generate_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + for message in ['Because ', 'it ', 'is.']: + yield ( + json.dumps( + { + 'model': 'dummy', + 'response': message, + } + ) + + '\n' + ) + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.generate('dummy', 'Why is the sky blue?', stream=True) + + it = iter(['Because ', 'it ', 'is.']) + async for part in response: + assert part['model'] == 'dummy' + assert part['response'] == next(it) + + +async def test_async_client_generate_images(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + 'images': [PNG_BASE64], + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as temp: + temp.write(PNG_BYTES) + temp.flush() + response = await client.generate('dummy', 'Why is the sky blue?', images=[temp.name]) + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + +async def test_async_client_pull(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.pull('dummy') + assert response['status'] == 'success' + + +async def test_async_client_pull_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'pulling manifest'}) + '\n' + yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' + yield json.dumps({'status': 'writing manifest'}) + '\n' + yield json.dumps({'status': 'removing any unused layers'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/pull', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.pull('dummy', stream=True) + + it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) + async for part in response: + assert part['status'] == next(it) + + +async def test_async_client_push(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.push('dummy') + assert response['status'] == 'success' + + +async def test_async_client_push_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'retrieving manifest'}) + '\n' + yield json.dumps({'status': 'pushing manifest'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + + httpserver.expect_ordered_request( + '/api/push', + method='POST', + json={ + 'model': 'dummy', + 'insecure': False, + 'stream': True, + }, + ).respond_with_handler(stream_handler) + + client = AsyncClient(httpserver.url_for('/')) + response = await client.push('dummy', stream=True) + + it = iter(['retrieving manifest', 'pushing manifest', 'success']) + async for part in response: + assert part['status'] == next(it) + + +async def test_async_client_create_with_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'files': {'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = await client.create('dummy', files={'test.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}) + assert response['status'] == 'success' + + +async def test_async_client_create_with_parameters_roundtrip(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'quantize': 'q4_k_m', + 'from': 'mymodel', + 'adapters': {'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + 'template': '[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + 'license': 'this is my license', + 'system': '\nUse\nmultiline\nstrings.\n', + 'parameters': {'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + 'messages': [{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile(): + response = await client.create( + 'dummy', + quantize='q4_k_m', + from_='mymodel', + adapters={'someadapter.gguf': 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'}, + template='[INST] <>{{.System}}<>\n{{.Prompt}} [/INST]', + license='this is my license', + system='\nUse\nmultiline\nstrings.\n', + parameters={'stop': ['[INST]', '[/INST]', '<>', '<>'], 'pi': 3.14159}, + messages=[{'role': 'user', 'content': 'Hello there!'}, {'role': 'assistant', 'content': 'Hello there yourself!'}], + stream=False, + ) + assert response['status'] == 'success' + + +async def test_async_client_create_from_library(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/create', + method='POST', + json={ + 'model': 'dummy', + 'from': 'llama2', + 'stream': False, + }, + ).respond_with_json({'status': 'success'}) + + client = AsyncClient(httpserver.url_for('/')) + + response = await client.create('dummy', from_='llama2') + assert response['status'] == 'success' + + +async def test_async_client_create_blob(httpserver: HTTPServer): + httpserver.expect_ordered_request(re.compile('^/api/blobs/sha256[:-][0-9a-fA-F]{64}$'), method='POST').respond_with_response(Response(status=201)) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = await client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +async def test_async_client_create_blob_exists(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='POST').respond_with_response(Response(status=200)) + + client = AsyncClient(httpserver.url_for('/')) + + with tempfile.NamedTemporaryFile() as blob: + response = await client.create_blob(blob.name) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' + + +async def test_async_client_delete(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/delete'), method='DELETE').respond_with_response(Response(status=200)) + client = AsyncClient(httpserver.url_for('/api/delete')) + response = await client.delete('dummy') + assert response['status'] == 'success' + + +async def test_async_client_copy(httpserver: HTTPServer): + httpserver.expect_ordered_request(PrefixPattern('/api/copy'), method='POST').respond_with_response(Response(status=200)) + client = AsyncClient(httpserver.url_for('/api/copy')) + response = await client.copy('dum', 'dummer') + assert response['status'] == 'success' + + +def test_headers(): + client = Client() + assert client._client.headers['content-type'] == 'application/json' + assert client._client.headers['accept'] == 'application/json' + assert client._client.headers['user-agent'].startswith('ollama-python/') + + client = Client( + headers={ + 'X-Custom': 'value', + 'Content-Type': 'text/plain', + } + ) + assert client._client.headers['x-custom'] == 'value' + assert client._client.headers['content-type'] == 'application/json' + + +def test_copy_tools(): + def func1(x: int) -> str: + """Simple function 1. + Args: + x (integer): A number + """ + + def func2(y: str) -> int: + """Simple function 2. + Args: + y (string): A string + """ + + # Test with list of functions + tools = list(_copy_tools([func1, func2])) + assert len(tools) == 2 + assert tools[0].function.name == 'func1' + assert tools[1].function.name == 'func2' + + # Test with empty input + assert list(_copy_tools()) == [] + assert list(_copy_tools(None)) == [] + assert list(_copy_tools([])) == [] + + # Test with mix of functions and tool dicts + tool_dict = { + 'type': 'function', + 'function': { + 'name': 'test', + 'description': 'Test function', + 'parameters': { + 'type': 'object', + 'properties': {'x': {'type': 'string', 'description': 'A string', 'enum': ['a', 'b', 'c']}, 'y': {'type': ['integer', 'number'], 'description': 'An integer'}}, + 'required': ['x'], + }, + }, + } + + tools = list(_copy_tools([func1, tool_dict])) + assert len(tools) == 2 + assert tools[0].function.name == 'func1' + assert tools[1].function.name == 'test' + + +def test_tool_validation(): + arbitrary_tool = {'type': 'custom_type', 'function': {'name': 'test'}} + tools = list(_copy_tools([arbitrary_tool])) + assert len(tools) == 1 + assert tools[0].type == 'custom_type' + assert tools[0].function.name == 'test' + + +def test_client_connection_error(): + client = Client('http://localhost:1234') + + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.generate('model', 'prompt') + with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE): + client.show('model') + + +async def test_async_client_connection_error(): + client = AsyncClient('http://localhost:1234') + with pytest.raises(ConnectionError) as exc_info: + await client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}]) + assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + with pytest.raises(ConnectionError) as exc_info: + await client.generate('model', 'prompt') + assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + with pytest.raises(ConnectionError) as exc_info: + await client.show('model') + assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download' + + +def test_arbitrary_roles_accepted_in_message(): + _ = Message(role='somerandomrole', content="I'm ok with you adding any role message now!") + + +def _mock_request(*args: Any, **kwargs: Any) -> Response: + return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") + + +def test_arbitrary_roles_accepted_in_message_request(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(Client, '_request', _mock_request) + + client = Client() + + client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) + + +async def _mock_request_async(*args: Any, **kwargs: Any) -> Response: + return httpxResponse(status_code=200, content="{'response': 'Hello world!'}") + + +async def test_arbitrary_roles_accepted_in_message_request_async(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(AsyncClient, '_request', _mock_request_async) + + client = AsyncClient() + + await client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}]) + + +def test_copy_messages_preserves_empty_string_content(): + from ollama._client import _copy_messages + + msgs = list( + _copy_messages( + [ + {'role': 'assistant', 'content': ''}, + {'role': 'tool', 'content': '', 'tool_name': 'web_search'}, + ] + ) + ) + assert msgs[0].content == '' + assert msgs[1].content == '' + assert msgs[1].tool_name == 'web_search' + + +def test_client_web_search_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + + client = Client() + + with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'): + client.web_search('test query') + + +def test_client_web_fetch_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + + client = Client() + + with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'): + client.web_fetch('https://example.com') + + +def _mock_request_web_search(self, cls, method, url, json=None, **kwargs): + assert method == 'POST' + assert url == 'https://ollama.com/api/web_search' + assert json is not None and 'query' in json and 'max_results' in json + return httpxResponse(status_code=200, content='{"results": {}, "success": true}') + + +def _mock_request_web_fetch(self, cls, method, url, json=None, **kwargs): + assert method == 'POST' + assert url == 'https://ollama.com/api/web_fetch' + assert json is not None and 'url' in json + return httpxResponse(status_code=200, content='{"results": {}, "success": true}') + + +def test_client_web_search_with_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') + monkeypatch.setattr(Client, '_request', _mock_request_web_search) + + client = Client() + client.web_search('what is ollama?', max_results=2) + + +def test_client_web_fetch_with_env_api_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') + monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) + + client = Client() + client.web_fetch('https://example.com') + + +def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + monkeypatch.setattr(Client, '_request', _mock_request_web_search) + + client = Client(headers={'Authorization': 'Bearer custom-token'}) + client.web_search('what is ollama?', max_results=1) + + +def test_client_web_fetch_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv('OLLAMA_API_KEY', raising=False) + monkeypatch.setattr(Client, '_request', _mock_request_web_fetch) + + client = Client(headers={'Authorization': 'Bearer custom-token'}) + client.web_fetch('https://example.com') + + +def test_client_bearer_header_from_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') + + client = Client() + assert client._client.headers['authorization'] == 'Bearer env-token' + + +def test_client_explicit_bearer_header_overrides_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv('OLLAMA_API_KEY', 'env-token') + monkeypatch.setattr(Client, '_request', _mock_request_web_search) + + client = Client(headers={'Authorization': 'Bearer explicit-token'}) + assert client._client.headers['authorization'] == 'Bearer explicit-token' + client.web_search('override check') + + +def test_client_close(): + client = Client() + client.close() + assert client._client.is_closed + + +@pytest.mark.anyio +async def test_async_client_close(): + client = AsyncClient() + await client.close() + assert client._client.is_closed + + +def test_client_context_manager(): + with Client() as client: + assert isinstance(client, Client) + assert not client._client.is_closed + + assert client._client.is_closed + + +@pytest.mark.anyio +async def test_async_client_context_manager(): + async with AsyncClient() as client: + assert isinstance(client, AsyncClient) + assert not client._client.is_closed + + assert client._client.is_closed + + +def test_generate_think_annotation_matches_chat(): + # The `think` parameter accepts bool or the 'low'/'medium'/'high' string levels. + # Client.generate must keep the same annotation as Client.chat and + # AsyncClient.generate so passing a string level does not raise a false type + # error (regression guard for the sync generate overloads/implementation). + expected = inspect.signature(Client.chat).parameters['think'].annotation + assert inspect.signature(Client.generate).parameters['think'].annotation == expected + assert inspect.signature(AsyncClient.generate).parameters['think'].annotation == expected From e28b128f542b9d26e8fac22981e29382c1a483fe Mon Sep 17 00:00:00 2001 From: r7mekmy4g67w6l Date: Thu, 6 Aug 2026 17:17:46 +0200 Subject: [PATCH 5/6] fix: preserve empty string message content in chat requests Signed-off-by: r7mekmy4g67w6l Normalize to LF so the PR shows the real 1-line change. From 6ae0aa89c8de90523c8aa856a93873abdecdf80f Mon Sep 17 00:00:00 2001 From: r7mekmy4g67w6l Date: Thu, 6 Aug 2026 17:17:49 +0200 Subject: [PATCH 6/6] fix: preserve empty string message content in chat requests Signed-off-by: r7mekmy4g67w6l Normalize to LF so the PR shows the real 1-line change.