From 18b51e254a64f5d94566af16eb93588a654b3db0 Mon Sep 17 00:00:00 2001 From: Siddardha Date: Fri, 23 Oct 2026 18:05:13 +0530 Subject: [PATCH] fix: resolve windows tempfile locking, tool schema generation, and add client improvements - Fix #676: replace tempfile.NamedTemporaryFile with tmp_path in tests to prevent Windows file locking PermissionError - Fix #733: guard ChatRequest.serialize_model against None tool functions and parameters - Fix #735: filter *args and **kwargs in convert_function_to_tool and correctly map default parameter values - Preserve array items, enum, and in tool property generation - Preserve empty string content in _copy_messages - Add bearer token validation in AsyncClient.web_search and web_fetch to match Client - Add direct api_key support to BaseClient, Client, and AsyncClient - Add id and type fields to Message.ToolCall - Export public client methods in ollama.__all__ and add create_blob - Remove unused _as_path dead code in _client.py - Add CONTRIBUTING.md guide and new examples (web-fetch, context-manager) --- CONTRIBUTING.md | 117 +++++++++++++++++++++++++++++++ README.md | 18 ++++- examples/README.md | 7 ++ examples/context-manager.py | 44 ++++++++++++ examples/web-fetch.py | 37 ++++++++++ ollama/__init__.py | 17 +++++ ollama/_client.py | 34 ++++----- ollama/_types.py | 14 +++- ollama/_utils.py | 25 +++++-- tests/test_client.py | 98 ++++++++++++++++++-------- tests/test_type_serialization.py | 53 ++++++++++---- tests/test_utils.py | 49 +++++++++++++ 12 files changed, 441 insertions(+), 72 deletions(-) create mode 100644 CONTRIBUTING.md create mode 100644 examples/context-manager.py create mode 100644 examples/web-fetch.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 00000000..021cb0ce --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,117 @@ +# Contributing to Ollama Python + +Thank you for your interest in contributing to the official Ollama Python library! We welcome bug reports, bug fixes, feature proposals, documentation improvements, and example additions. + +--- + +## Prerequisites + +- **Python 3.8+** (Python 3.10+ recommended for development) +- **[uv](https://docs.astral.sh/uv/)** or **[hatch](https://hatch.pypa.io/)** +- A running instance of **[Ollama](https://ollama.com/download)** (optional for running unit/mock tests, required for real model execution) + +--- + +## Development Setup + +1. **Fork and clone the repository:** + + ```bash + git clone https://github.com//ollama-python.git + cd ollama-python + ``` + +2. **Set up virtual environment and install dependencies:** + + Using `uv`: + ```bash + uv sync --all-extras --all-groups + ``` + + Using `hatch`: + ```bash + hatch env create + ``` + +--- + +## Code Style and Quality + +We use [Ruff](https://docs.astral.sh/ruff/) to format and lint our codebase. + +- **Indentation**: 2 spaces (`indent-width = 2`, `indent-style = 'space'`) +- **Quotes**: Single quotes (`quote-style = 'single'`) +- **Line length**: 320 characters max + +### Formatting & Linting Commands + +Check formatting and linting: +```bash +uv run ruff format --check . +uv run ruff check . +``` + +Automatically fix format and lint issues: +```bash +uv run ruff format . +uv run ruff check --fix . +``` + +Or using `hatch`: +```bash +hatch fmt --check +``` + +--- + +## Running Tests + +All unit tests use mock HTTP servers with `pytest-httpserver` and doctests across module docstrings. + +Run the test suite with `uv`: +```bash +uv run pytest +``` + +Run tests across environments with `hatch`: +```bash +hatch test +``` + +Run a specific test file: +```bash +uv run pytest tests/test_client.py +uv run pytest tests/test_utils.py +uv run pytest tests/test_type_serialization.py +``` + +--- + +## Updating Lockfile and Requirements + +If dependencies in `pyproject.toml` are modified, make sure the lockfile and exported requirements remain in sync: + +```bash +uv lock +uv export > requirements.txt +``` + +Verify that requirements match: +```bash +git diff --exit-code requirements.txt +``` + +--- + +## Pull Request Guidelines + +1. **Branch naming**: Use descriptive branch names (e.g. `fix/named-tempfile-windows`, `feat/client-api-key-arg`). +2. **Commit messages**: Follow concise conventional commit style: + - `fix: ...` for bug fixes + - `feat: ...` for new features + - `docs: ...` for documentation updates + - `test: ...` for test additions + - `refactor: ...` for code refactoring +3. **Tests**: Add unit tests for any new functionality or bug fixes in `tests/`. +4. **Documentation & Examples**: Update `README.md` and `examples/README.md` if changing public APIs or adding new features. +5. **No regressions**: Ensure formatting (`ruff format`), linting (`ruff check`), and all tests pass before submitting. diff --git a/README.md b/README.md index 3038d7ce..2446546b 100644 --- a/README.md +++ b/README.md @@ -113,7 +113,8 @@ curl https://ollama.com/api/tags import os from ollama import Client -client = Client(host='https://ollama.com', headers={'Authorization': 'Bearer ' + os.environ.get('OLLAMA_API_KEY')}) +# You can pass api_key directly or rely on the OLLAMA_API_KEY environment variable +client = Client(host='https://ollama.com', api_key=os.environ.get('OLLAMA_API_KEY')) messages = [ { @@ -129,12 +130,16 @@ for part in client.chat('gpt-oss:120b', messages=messages, stream=True): ## Custom client A custom client can be created by instantiating `Client` or `AsyncClient` from `ollama`. -All extra keyword arguments are passed into the [`httpx.Client`](https://www.python-httpx.org/api/#client). +Clients accept `host`, `api_key`, `headers`, `timeout`, and any keyword arguments supported by [`httpx.Client`](https://www.python-httpx.org/api/#client). ```python from ollama import Client -client = Client(host='http://localhost:11434', headers={'x-some-header': 'some-value'}) +client = Client( + host='http://localhost:11434', + api_key='your_api_key', + headers={'x-some-header': 'some-value'}, +) response = client.chat( model='gemma4', messages=[ @@ -146,6 +151,13 @@ response = client.chat( ) ``` +Clients can also be used as context managers to ensure immediate resource and socket cleanup: + +```python +with Client() as client: + response = client.chat(model='gemma4', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}]) +``` + ## Async client The `AsyncClient` class is used to make asynchronous requests. It can be configured with the same fields as the `Client` class. diff --git a/examples/README.md b/examples/README.md index 66f58951..88d50c91 100644 --- a/examples/README.md +++ b/examples/README.md @@ -18,12 +18,14 @@ See [ollama/docs/api.md](https://github.com/ollama/ollama/blob/main/docs/api.md) - [async-chat.py](async-chat.py) - [chat-stream.py](chat-stream.py) - Streamed outputs - [chat-with-history.py](chat-with-history.py) - Chat with model and maintain history of the conversation +- [chat-logprobs.py](chat-logprobs.py) - Return log probabilities for chat completions ### Generate - Generate text with a model - [generate.py](generate.py) - [async-generate.py](async-generate.py) - [generate-stream.py](generate-stream.py) - Streamed outputs +- [generate-logprobs.py](generate-logprobs.py) - Return log probabilities for generation - [fill-in-middle.py](fill-in-middle.py) - Given a prefix and suffix, fill in the middle ### Tools/Function Calling - Call a function with a model @@ -46,6 +48,7 @@ export OLLAMA_API_KEY="your_api_key_here" ``` - [web-search.py](web-search.py) +- [web-fetch.py](web-fetch.py) - Fetch and extract content from a web page - [web-search-gpt-oss.py](web-search-gpt-oss.py) - Using browser research tools with gpt-oss #### MCP server @@ -131,3 +134,7 @@ Requirement: `pip install tqdm` ### System One - Score typed questions - [systemone.py](systemone.py) - Choice, boolean probability, and rubric score questions + +### Client Lifecycle & Context Managers - Manage connections and resources + +- [context-manager.py](context-manager.py) - Use sync and async clients with context managers for automatic resource cleanup diff --git a/examples/context-manager.py b/examples/context-manager.py new file mode 100644 index 00000000..92c9e6d7 --- /dev/null +++ b/examples/context-manager.py @@ -0,0 +1,44 @@ +# /// script +# requires-python = ">=3.8" +# dependencies = [ +# "ollama", +# ] +# /// +""" +Example demonstrating how to use Client and AsyncClient as context managers +to ensure proper cleanup and release of underlying HTTP connections. +""" + +import asyncio + +from ollama import AsyncClient, Client + + +def sync_example(): + print('=== Synchronous Client Context Manager ===') + # Entering the context manager initializes the client; + # exiting cleanly closes all httpx connection pools. + with Client() as client: + models = client.list() + print(f'Discovered {len(models.models)} installed models:') + for model in models.models[:3]: + print(f' - {model.model} (modified: {model.modified_at})') + + +async def async_example(): + print('\n=== Asynchronous Client Context Manager ===') + async with AsyncClient() as client: + response = await client.generate( + model='gemma4', + prompt='Why is the sky blue in one sentence?', + ) + print(f'Response: {response.response}') + + +if __name__ == '__main__': + sync_example() + # Run async example if an async event loop is available + try: + asyncio.run(async_example()) + except Exception as e: + print(f'Async example skipped or encountered error: {e}') diff --git a/examples/web-fetch.py b/examples/web-fetch.py new file mode 100644 index 00000000..f2cee657 --- /dev/null +++ b/examples/web-fetch.py @@ -0,0 +1,37 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "rich", +# "ollama", +# ] +# /// +import os +import sys + +from rich import print + +from ollama import Client + +api_key = os.getenv('OLLAMA_API_KEY') +if not api_key: + print('[bold red]Error:[/bold red] OLLAMA_API_KEY environment variable is required.') + print('Get your API key at https://ollama.com/settings/keys') + sys.exit(1) + +# You can pass api_key directly or rely on OLLAMA_API_KEY +client = Client(api_key=api_key) + +url = 'https://ollama.com' +print(f'Fetching content from: [bold cyan]{url}[/bold cyan]\n') + +# Fetch using the client +response = client.web_fetch(url) + +print(f'Title: [bold green]{response.title}[/bold green]') +if response.links: + print(f'Found {len(response.links)} links on the page:') + for link in response.links[:5]: + print(f' - {link}') + +print('\nContent Preview:') +print(response.content[:500] if response.content else 'No content returned') diff --git a/ollama/__init__.py b/ollama/__init__.py index 21b6f406..2e4d6902 100644 --- a/ollama/__init__.py +++ b/ollama/__init__.py @@ -63,6 +63,22 @@ 'Tool', 'WebFetchResponse', 'WebSearchResponse', + 'chat', + 'copy', + 'create', + 'create_blob', + 'delete', + 'embed', + 'embeddings', + 'generate', + 'list', + 'ps', + 'pull', + 'push', + 'show', + 'systemone', + 'web_fetch', + 'web_search', ] _client = Client() @@ -74,6 +90,7 @@ pull = _client.pull push = _client.push create = _client.create +create_blob = _client.create_blob delete = _client.delete list = _client.list copy = _client.copy diff --git a/ollama/_client.py b/ollama/_client.py index 3654e3c0..8648471c 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -6,7 +6,6 @@ import sys import urllib.parse from hashlib import sha256 -from os import PathLike from pathlib import Path from typing import ( Any, @@ -89,6 +88,7 @@ def __init__( follow_redirects: bool = True, timeout: Any = None, headers: Optional[Mapping[str, str]] = None, + api_key: Optional[str] = None, **kwargs, ) -> None: """ @@ -109,7 +109,7 @@ def __init__( }.items() if v is not None } - api_key = os.getenv('OLLAMA_API_KEY', None) + api_key = api_key or os.getenv('OLLAMA_API_KEY', None) if not headers.get('authorization') and api_key: headers['authorization'] = f'Bearer {api_key}' @@ -132,8 +132,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): class Client(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.Client, host, **kwargs) + def __init__(self, host: Optional[str] = None, *, api_key: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.Client, host, api_key=api_key, **kwargs) def close(self): self._client.close() @@ -752,8 +752,8 @@ def web_fetch(self, url: str) -> WebFetchResponse: class AsyncClient(BaseClient): - def __init__(self, host: Optional[str] = None, **kwargs) -> None: - super().__init__(httpx.AsyncClient, host, **kwargs) + def __init__(self, host: Optional[str] = None, *, api_key: Optional[str] = None, **kwargs) -> None: + super().__init__(httpx.AsyncClient, host, api_key=api_key, **kwargs) async def close(self): await self._client.aclose() @@ -832,7 +832,12 @@ async def web_search(self, query: str, max_results: int = 3) -> WebSearchRespons 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 await self._request( WebSearchResponse, 'POST', @@ -852,7 +857,12 @@ async def web_fetch(self, url: str) -> WebFetchResponse: Returns: WebFetchResponse with the fetched result + 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 fetch') + return await self._request( WebFetchResponse, 'POST', @@ -1379,7 +1389,7 @@ def _copy_images(images: Optional[Sequence[Union[Image, Any]]]) -> Iterator[Imag 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}, + {k: list(_copy_images(v)) if k == 'images' else v for k, v in dict(message).items() if v is not None}, ) @@ -1388,16 +1398,6 @@ def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable 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) diff --git a/ollama/_types.py b/ollama/_types.py index ad041013..0bee3b15 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -335,6 +335,12 @@ class ToolCall(SubscriptableBaseModel): Model tool calls. """ + id: Optional[str] = None + 'Identifier of the tool call.' + + type: Optional[str] = 'function' + 'Type of the tool call.' + class Function(SubscriptableBaseModel): """ Tool call function. @@ -390,8 +396,12 @@ def serialize_model(self, nxt): output = nxt(self) if output.get('tools'): for tool in output['tools']: - if 'function' in tool and 'parameters' in tool['function'] and 'defs' in tool['function']['parameters']: - tool['function']['parameters']['$defs'] = tool['function']['parameters'].pop('defs') + if isinstance(tool, dict): + fn = tool.get('function') + if isinstance(fn, dict): + params = fn.get('parameters') + if isinstance(params, dict) and 'defs' in params: + params['$defs'] = params.pop('defs') return output messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None diff --git a/ollama/_utils.py b/ollama/_utils.py index 15f1cc0c..a0153b5c 100644 --- a/ollama/_utils.py +++ b/ollama/_utils.py @@ -56,27 +56,38 @@ def _parse_docstring(doc_string: Union[str, None]) -> dict[str, str]: def convert_function_to_tool(func: Callable) -> Tool: doc_string_hash = str(hash(inspect.getdoc(func))) parsed_docstring = _parse_docstring(inspect.getdoc(func)) + sig = inspect.signature(func) + valid_params = {k: v for k, v in sig.parameters.items() if v.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)} + attrs = { + '__annotations__': {k: v.annotation if v.annotation != inspect._empty else str for k, v in valid_params.items()}, + '__signature__': sig.replace(parameters=list(valid_params.values())), + '__doc__': parsed_docstring[doc_string_hash], + **{k: v.default for k, v in valid_params.items() if v.default != inspect._empty}, + } schema = type( func.__name__, (pydantic.BaseModel,), - { - '__annotations__': {k: v.annotation if v.annotation != inspect._empty else str for k, v in inspect.signature(func).parameters.items()}, - '__signature__': inspect.signature(func), - '__doc__': parsed_docstring[doc_string_hash], - }, + attrs, ).model_json_schema() for k, v in schema.get('properties', {}).items(): # If type is missing, the default is string types = {t.get('type', 'string') for t in v.get('anyOf')} if 'anyOf' in v else {v.get('type', 'string')} if 'null' in types: - schema['required'].remove(k) + if 'required' in schema and k in schema['required']: + schema['required'].remove(k) types.discard('null') - schema['properties'][k] = { + prop = { 'description': parsed_docstring[k], 'type': ', '.join(types), } + if 'items' in v: + prop['items'] = v['items'] + if 'enum' in v: + prop['enum'] = v['enum'] + + schema['properties'][k] = prop tool = Tool( type='function', diff --git a/tests/test_client.py b/tests/test_client.py index c212a11d..49ee6753 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -442,7 +442,7 @@ def generate(): assert part['response'] == next(it) -def test_client_generate_images(httpserver: HTTPServer): +def test_client_generate_images(httpserver: HTTPServer, tmp_path): httpserver.expect_ordered_request( '/api/generate', method='POST', @@ -461,12 +461,11 @@ def test_client_generate_images(httpserver: HTTPServer): 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.' + image_path = tmp_path / 'temp.png' + image_path.write_bytes(PNG_BYTES) + response = client.generate('dummy', 'Why is the sky blue?', images=[str(image_path)]) + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' def test_client_generate_format_json(httpserver: HTTPServer): @@ -846,24 +845,26 @@ def test_client_create_from_library(httpserver: HTTPServer): assert response['status'] == 'success' -def test_client_create_blob(httpserver: HTTPServer): +def test_client_create_blob(httpserver: HTTPServer, tmp_path): 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' + blob = tmp_path / 'blob' + blob.write_bytes(b'') + response = client.create_blob(str(blob)) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' -def test_client_create_blob_exists(httpserver: HTTPServer): +def test_client_create_blob_exists(httpserver: HTTPServer, tmp_path): 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' + blob = tmp_path / 'blob' + blob.write_bytes(b'') + response = client.create_blob(str(blob)) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' def test_client_delete(httpserver: HTTPServer): @@ -1037,7 +1038,7 @@ def generate(): assert part['response'] == next(it) -async def test_async_client_generate_images(httpserver: HTTPServer): +async def test_async_client_generate_images(httpserver: HTTPServer, tmp_path): httpserver.expect_ordered_request( '/api/generate', method='POST', @@ -1056,12 +1057,11 @@ async def test_async_client_generate_images(httpserver: HTTPServer): 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.' + image_path = tmp_path / 'temp.png' + image_path.write_bytes(PNG_BYTES) + response = await client.generate('dummy', 'Why is the sky blue?', images=[str(image_path)]) + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' async def test_async_client_pull(httpserver: HTTPServer): @@ -1223,24 +1223,26 @@ async def test_async_client_create_from_library(httpserver: HTTPServer): assert response['status'] == 'success' -async def test_async_client_create_blob(httpserver: HTTPServer): +async def test_async_client_create_blob(httpserver: HTTPServer, tmp_path): 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' + blob = tmp_path / 'blob' + blob.write_bytes(b'') + response = await client.create_blob(str(blob)) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' -async def test_async_client_create_blob_exists(httpserver: HTTPServer): +async def test_async_client_create_blob_exists(httpserver: HTTPServer, tmp_path): 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' + blob = tmp_path / 'blob' + blob.write_bytes(b'') + response = await client.create_blob(str(blob)) + assert response == 'sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855' async def test_async_client_delete(httpserver: HTTPServer): @@ -1527,3 +1529,39 @@ def test_client_chat_with_think_level(httpserver: HTTPServer): assert response['model'] == 'qwen3.8:27b' assert response['message']['content'] == 'Hi there.' assert response['message']['thinking'] == 'Thinking deeply...' + + +def test_client_api_key(): + c = Client(api_key='my_secret_key') + assert c._client.headers.get('authorization') == 'Bearer my_secret_key' + + +def test_async_client_api_key(): + ac = AsyncClient(api_key='my_async_secret_key') + assert ac._client.headers.get('authorization') == 'Bearer my_async_secret_key' + + +async def test_async_client_web_search_requires_auth(): + client = AsyncClient() + with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'): + await client.web_search('test') + + +async def test_async_client_web_fetch_requires_auth(): + client = AsyncClient() + with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'): + await client.web_fetch('https://example.com') + + +def test_copy_messages_preserves_empty_content(): + from ollama._client import _copy_messages + + msgs = list(_copy_messages([{'role': 'assistant', 'content': ''}])) + assert msgs[0].content == '' + + +def test_init_create_blob_export(): + import ollama + + assert hasattr(ollama, 'create_blob') + assert 'create_blob' in ollama.__all__ diff --git a/tests/test_type_serialization.py b/tests/test_type_serialization.py index e1af3512..ab34cc56 100644 --- a/tests/test_type_serialization.py +++ b/tests/test_type_serialization.py @@ -1,4 +1,3 @@ -import tempfile from base64 import b64encode from pathlib import Path @@ -31,20 +30,18 @@ def test_image_serialization_plain_string(): assert img.model_dump() == 'not a path or base64' # Should return as-is -def test_image_serialization_path(): - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(b'test file content') - temp_file.flush() - img = Image(value=Path(temp_file.name)) - assert img.model_dump() == b64encode(b'test file content').decode() +def test_image_serialization_path(tmp_path): + temp_file = tmp_path / 'test_path.png' + temp_file.write_bytes(b'test file content') + img = Image(value=Path(temp_file)) + assert img.model_dump() == b64encode(b'test file content').decode() -def test_image_serialization_string_path(): - with tempfile.NamedTemporaryFile() as temp_file: - temp_file.write(b'test file content') - temp_file.flush() - img = Image(value=temp_file.name) - assert img.model_dump() == b64encode(b'test file content').decode() +def test_image_serialization_string_path(tmp_path): + temp_file = tmp_path / 'test_string_path.png' + temp_file.write_bytes(b'test file content') + img = Image(value=str(temp_file)) + assert img.model_dump() == b64encode(b'test file content').decode() with pytest.raises(ValueError): img = Image(value='some_path/that/does/not/exist.png') @@ -122,3 +119,33 @@ def test_think_boolean_serialization(): assert ChatRequest(model='test-model', think=True).model_dump(exclude_none=True)['think'] is True assert ChatRequest(model='test-model', think=False).model_dump(exclude_none=True)['think'] is False assert 'think' not in ChatRequest(model='test-model', think=None).model_dump(exclude_none=True) + + +def test_chat_request_with_none_function_or_parameters(): + from ollama._types import Tool + + req1 = ChatRequest(model='test', tools=[Tool(type='function', function=None)]) + dumped1 = req1.model_dump(exclude_none=True) + assert dumped1['model'] == 'test' + + req2 = ChatRequest(model='test', tools=[Tool(type='function', function=Tool.Function(name='test', parameters=None))]) + dumped2 = req2.model_dump(exclude_none=True) + assert dumped2['model'] == 'test' + + +def test_tool_call_id_and_type(): + from ollama._types import Message + + tool_call = Message.ToolCall( + id='call_abc123', + type='function', + function=Message.ToolCall.Function(name='my_func', arguments={'x': 1}), + ) + assert tool_call.id == 'call_abc123' + assert tool_call.type == 'function' + assert tool_call['id'] == 'call_abc123' + assert tool_call['type'] == 'function' + dumped = tool_call.model_dump(exclude_none=True) + assert dumped['id'] == 'call_abc123' + assert dumped['type'] == 'function' + assert dumped['function']['name'] == 'my_func' diff --git a/tests/test_utils.py b/tests/test_utils.py index cb9e0d4f..15ddb3ce 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -256,3 +256,52 @@ def func_with_parentheses_and_args(a: int, b: int): tool = convert_function_to_tool(func_with_parentheses_and_args).model_dump() assert tool['function']['parameters']['properties']['a']['description'] == 'First (:thing) number to add' assert tool['function']['parameters']['properties']['b']['description'] == 'Second number to add' + + +def test_function_with_varargs(): + def func_with_varargs(a: int, *args, b: str = 'default', **kwargs) -> str: + """ + A function with varargs and kwargs. + + Args: + a: The first number + b: The second string + """ + return f'{a}-{b}' + + tool = convert_function_to_tool(func_with_varargs).model_dump() + assert 'args' not in tool['function']['parameters']['properties'] + assert 'kwargs' not in tool['function']['parameters']['properties'] + assert tool['function']['parameters']['required'] == ['a'] + assert 'a' in tool['function']['parameters']['properties'] + assert 'b' in tool['function']['parameters']['properties'] + + +def test_function_with_array_items(): + def func_with_array(tags: List[str]): + """ + Function with array of strings. + + Args: + tags: List of tag strings + """ + + tool = convert_function_to_tool(func_with_array).model_dump() + assert tool['function']['parameters']['properties']['tags']['type'] == 'array' + assert tool['function']['parameters']['properties']['tags']['items'] == {'type': 'string'} + + +def test_function_with_literal_enum(): + from typing_extensions import Literal + + def func_with_enum(status: Literal['open', 'closed']): + """ + Function with literal choices. + + Args: + status: The status choice + """ + + tool = convert_function_to_tool(func_with_enum).model_dump() + assert tool['function']['parameters']['properties']['status']['type'] == 'string' + assert tool['function']['parameters']['properties']['status']['enum'] == ['open', 'closed']