diff --git a/examples/gpt-oss-tools-stream.py b/examples/gpt-oss-tools-stream.py index 97be624e..44d3ca47 100644 --- a/examples/gpt-oss-tools-stream.py +++ b/examples/gpt-oss-tools-stream.py @@ -70,6 +70,7 @@ def get_weather_conditions(city: str) -> str: tool_calls.extend(chunk.message.tool_calls) if chunk.message.content: + content += chunk.message.content if not (chunk.message.thinking or chunk.message.thinking == '') and final: print('\n\n' + '=' * 10) print('Final result: ') diff --git a/tests/test_examples.py b/tests/test_examples.py new file mode 100644 index 00000000..96796fa7 --- /dev/null +++ b/tests/test_examples.py @@ -0,0 +1,35 @@ +import copy +import runpy +import sys +from pathlib import Path +from types import SimpleNamespace + +from ollama import ChatResponse + + +def test_gpt_oss_stream_preserves_content_before_tool_result(monkeypatch): + requests = [] + responses = [ + [ + ChatResponse(message={'role': 'assistant', 'content': 'Checking '}), + ChatResponse(message={'role': 'assistant', 'content': 'London.', 'tool_calls': [{'function': {'name': 'get_weather', 'arguments': {'city': 'London'}}}]}), + ], + [ChatResponse(message={'role': 'assistant', 'content': 'The weather is mild.'})], + ] + + def chat(**kwargs): + requests.append(copy.deepcopy(kwargs['messages'])) + return iter(responses.pop(0)) + + monkeypatch.setattr('ollama.Client', lambda: SimpleNamespace(chat=chat)) + monkeypatch.setitem(sys.modules, 'rich', SimpleNamespace(print=lambda *args, **kwargs: None)) + + runpy.run_path(str(Path(__file__).parents[1] / 'examples' / 'gpt-oss-tools-stream.py')) + + assert len(requests) == 2 + assistant, tool = requests[1][1:] + assert assistant['role'] == 'assistant' + assert assistant['content'] == 'Checking London.' + assert len(assistant['tool_calls']) == 1 + assert tool['role'] == 'tool' + assert tool['tool_name'] == 'get_weather'