types: preserve tool call function indexes

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.'
 

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.

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():