From 0dd8ea2c4c3c2ddd8f8434187e62e05fccc994ed Mon Sep 17 00:00:00 2001 From: er-s-an <92594860+er-s-an@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:44:22 +0800 Subject: [PATCH] types: preserve tool call function indexes --- ollama/_types.py | 3 +++ tests/test_client.py | 37 ++++++++++++++++++++++++++++++++ tests/test_type_serialization.py | 25 ++++++++++++++++++++- 3 files changed, 64 insertions(+), 1 deletion(-) diff --git a/ollama/_types.py b/ollama/_types.py index cf7264b5..95a199f0 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -340,6 +340,9 @@ class Function(SubscriptableBaseModel): Tool call function. """ + index: Optional[int] = None + 'Index of the tool call.' + name: str 'Name of the function.' diff --git a/tests/test_client.py b/tests/test_client.py index c212a11d..4746cc46 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1273,6 +1273,43 @@ def test_headers(): assert client._client.headers['content-type'] == 'application/json' +@pytest.fixture +def tool_call_index_server(httpserver: HTTPServer, stream: bool): + calls = [{'function': {'index': index, 'name': 'lookup', 'arguments': {'city': city}}} for index, city in enumerate(['London', 'Paris'])] + request = httpserver.expect_request('/api/chat', method='POST') + if stream: + chunks = [{'message': {'role': 'assistant', 'tool_calls': [call]}, 'done': False} for call in calls] + chunks.append({'message': {'role': 'assistant', 'content': ''}, 'done': True}) + request.respond_with_data(''.join(json.dumps(chunk) + '\n' for chunk in chunks), content_type='application/x-ndjson') + else: + request.respond_with_json({'message': {'role': 'assistant', 'tool_calls': calls}, 'done': True}) + return httpserver + + +@pytest.mark.parametrize('stream', [False, True]) +def test_client_chat_tool_call_function_index(tool_call_index_server: HTTPServer, stream: bool): + with Client(tool_call_index_server.url_for('/')) as client: + response = client.chat('model', stream=stream) + responses = list(response) if stream else [response] + + calls = [call for response in responses for call in response.message.tool_calls or []] + assert [call.function.index for call in calls] == [0, 1] + assert [call.function.name for call in calls] == ['lookup', 'lookup'] + assert [call.function.arguments for call in calls] == [{'city': 'London'}, {'city': 'Paris'}] + + +@pytest.mark.parametrize('stream', [False, True]) +async def test_async_client_chat_tool_call_function_index(tool_call_index_server: HTTPServer, stream: bool): + async with AsyncClient(tool_call_index_server.url_for('/')) as client: + response = await client.chat('model', stream=stream) + responses = [part async for part in response] if stream else [response] + + calls = [call for response in responses for call in response.message.tool_calls or []] + assert [call.function.index for call in calls] == [0, 1] + assert [call.function.name for call in calls] == ['lookup', 'lookup'] + assert [call.function.arguments for call in calls] == [{'city': 'London'}, {'city': 'Paris'}] + + def test_copy_tools(): def func1(x: int) -> str: """Simple function 1. diff --git a/tests/test_type_serialization.py b/tests/test_type_serialization.py index e1af3512..38f3b57f 100644 --- a/tests/test_type_serialization.py +++ b/tests/test_type_serialization.py @@ -4,7 +4,30 @@ import pytest -from ollama._types import ChatRequest, CreateRequest, GenerateRequest, Image +from ollama._types import ChatRequest, ChatResponse, CreateRequest, GenerateRequest, Image + + +@pytest.mark.parametrize('index', [0, 2]) +def test_chat_response_preserves_tool_call_function_index(index): + message = {'role': 'assistant', 'tool_calls': [{'function': {'index': index, 'name': 'lookup', 'arguments': {'city': 'London'}}}]} + response = ChatResponse(message=message) + function = response.message.tool_calls[0].function + + assert function.index == index + assert function['index'] == index + assert response.model_dump(exclude_none=True)['message'] == message + + request = ChatRequest(model='test-model', messages=[response.message]) + assert request.model_dump(exclude_none=True)['messages'] == [message] + assert ChatRequest.model_validate_json(request.model_dump_json(exclude_none=True)).model_dump(exclude_none=True)['messages'] == [message] + + +def test_chat_response_without_tool_call_function_index(): + message = {'role': 'assistant', 'tool_calls': [{'function': {'name': 'lookup', 'arguments': {}}}]} + response = ChatResponse(message=message) + + assert getattr(response.message.tool_calls[0].function, 'index', None) is None + assert response.model_dump(exclude_none=True)['message'] == message def test_image_serialization_bytes():