fix(examples): record multi-tool assistant messages once

examples/multi-tool.py

@@ -56,14 +56,15 @@ def get_conditions(city: str) -> str:
   if chunk.message.content:
     print(chunk.message.content, end='', flush=True)
   if chunk.message.tool_calls:
+    # Record the assistant message once, before the results for its tool calls.
+    messages.append(chunk.message)
     for tool in chunk.message.tool_calls:
       if function_to_call := available_functions.get(tool.function.name):
         print('\nCalling function:', tool.function.name, 'with arguments:', tool.function.arguments)
         output = function_to_call(**tool.function.arguments)
         print('> Function output:', output, '\n')
 
-        # Add the assistant message and tool call result to the messages
-        messages.append(chunk.message)
+        # Add each tool result to the messages
         messages.append({'role': 'tool', 'content': str(output), 'tool_name': tool.function.name})
       else:
         print('Function', tool.function.name, 'not found')

tests/test_multi_tool_example.py

@@ -0,0 +1,48 @@
+import copy
+import runpy
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from ollama import ChatResponse
+
+
+@pytest.mark.parametrize('count', [1, 2])
+def test_multi_tool_example_records_one_assistant_for_multiple_tools(monkeypatch, capsys, count):
+  requests = []
+  tool_names = ['get_temperature', 'get_conditions'][:count]
+  tool_calls = [{'function': {'name': name, 'arguments': {'city': 'London'}}} for name in tool_names]
+  responses = [
+    [ChatResponse(message={'role': 'assistant', 'tool_calls': tool_calls})],
+    [ChatResponse(message={'role': 'assistant', 'content': 'The weather is mild.'})],
+  ]
+
+  def chat(*args, **kwargs):
+    requests.append(copy.deepcopy(kwargs['messages']))
+    return iter(responses.pop(0))
+
+  monkeypatch.setattr('ollama.Client', lambda: SimpleNamespace(chat=chat))
+
+  runpy.run_path(str(Path(__file__).parents[1] / 'examples' / 'multi-tool.py'))
+
+  assert len(requests) == 2
+  assert [message['role'] for message in requests[1]] == ['user', 'assistant', *(['tool'] * count)]
+  assert [message['tool_name'] for message in requests[1][2:]] == tool_names
+  assert len(requests[1][1]['tool_calls']) == count
+  assert 'The weather is mild.' in capsys.readouterr().out
+
+
+def test_multi_tool_example_without_tools_does_not_request_another_response(monkeypatch, capsys):
+  requests = []
+
+  def chat(*args, **kwargs):
+    requests.append(kwargs)
+    return iter([ChatResponse(message={'role': 'assistant', 'content': 'Hello.'})])
+
+  monkeypatch.setattr('ollama.Client', lambda: SimpleNamespace(chat=chat))
+
+  runpy.run_path(str(Path(__file__).parents[1] / 'examples' / 'multi-tool.py'))
+
+  assert len(requests) == 1
+  assert 'No tool calls returned' in capsys.readouterr().out