fix: only separate streamed text that follows earlier text and a tool result

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Bug fixes
 
+* `.stream(content="text")` and `.chat()` now separate the text before and after a tool call with a blank line, as `.stream_async(content="text")` and `.chat_async()` already did. The synchronous versions used to run the two parts together. A response that starts with a tool call no longer starts with blank lines in the async versions either.
 * `ChatOllama()` can now be created without a `model`, so `ChatOllama().list_models()` works. Sending a message without a model still raises an error that lists the available models. (#166)
 * `ChatGoogle()` now reports `finish_reason == "content_filter"` for Gemini's `IMAGE_SAFETY`, `IMAGE_PROHIBITED_CONTENT`, `IMAGE_RECITATION` and `LANGUAGE` finish reasons, instead of passing the raw Gemini string through, and reports `"tool_use"` (rather than `"success"`) when a turn ends with a tool request, as `ChatOpenAI()` does. (#473)
 * The error raised for a tool request with invalid JSON arguments from `ChatOpenAI()` and `ChatOpenAICompletions()` (and the providers built on them) now shows the offending arguments, instead of the literal text `{args}`.

chatlas/_chat.py

@@ -155,6 +155,28 @@ def display_text(content: "Content") -> "Optional[str]":
     return None
 
 
+class _TextSeparator:
+    """
+    Decides when to put a blank line between streamed text from different rounds.
+
+    A round that ends in tool calls is followed by another round. When both have
+    text, the text is separated by a blank line. Text only gets separated from
+    text that was already emitted, so a response never starts with a blank line.
+    """
+
+    def __init__(self) -> None:
+        self.emitted_text = False
+        self.after_tool_results = False
+
+    def needed_before(self, chunk: object) -> bool:
+        if not isinstance(chunk, str) or not chunk:
+            return False
+        needed = self.emitted_text and self.after_tool_results
+        self.emitted_text = True
+        self.after_tool_results = False
+        return needed
+
+
 class Chat(Generic[SubmitInputArgsT, CompletionT]):
     """
     A chat object that can be used to interact with a language model.
@@ -2750,6 +2772,7 @@ def _chat_impl(
         agent_span = start_agent_span(self.provider, self._conversation_id)
         try:
             user_turn_result: UserTurn | None = user_turn
+            separator = _TextSeparator()
             while user_turn_result is not None:
                 for chunk in self._submit_turns(
                     user_turn_result,
@@ -2761,6 +2784,8 @@ def _chat_impl(
                     controller=controller,
                     _otel_parent=agent_span,
                 ):
+                    if separator.needed_before(chunk):
+                        yield "\n\n"
                     yield chunk
 
                 turn = self.get_last_turn(role="assistant")
@@ -2790,6 +2815,7 @@ def _chat_impl(
                             all_results.append(res)
 
                 if all_results:
+                    separator.after_tool_results = True
                     user_turn_result = UserTurn(all_results)
         except Exception as e:
             record_error(agent_span, e)
@@ -2837,6 +2863,7 @@ async def _chat_impl_async(
         agent_span = start_agent_span(self.provider, self._conversation_id)
         try:
             user_turn_result: UserTurn | None = user_turn
+            separator = _TextSeparator()
             while user_turn_result is not None:
                 turn_generator = self._submit_turns_async(
                     user_turn_result,
@@ -2850,6 +2877,8 @@ async def _chat_impl_async(
                 )
                 try:
                     async for chunk in turn_generator:
+                        if separator.needed_before(chunk):
+                            yield "\n\n"
                         yield chunk
                 finally:
                     await turn_generator.aclose()
@@ -2878,11 +2907,10 @@ async def _chat_impl_async(
                                 self._echo_content(res)
                             if content == "all":
                                 yield res
-                            else:
-                                yield "\n\n"
                             all_results.append(res)
 
                 if all_results:
+                    separator.after_tool_results = True
                     user_turn_result = UserTurn(all_results)
         except Exception as e:
             record_error(agent_span, e)

tests/test_stream_tool_separator.py

@@ -0,0 +1,88 @@
+import pytest
+from chatlas import Chat
+from chatlas._content import ContentText, ContentToolRequest
+from chatlas._turn import AssistantTurn
+
+from .test_provider_stream_contract import AccumulatingProvider, Chunk
+
+
+class TwoRoundProvider(AccumulatingProvider):
+    """Answers "Checking." with a tool request, then "Done." after the tool ran."""
+
+    def __init__(self, first_round_text: str = "Checking."):
+        super().__init__([])
+        self.first_round_text = first_round_text
+        self.round = 0
+
+    def _next(self):
+        self.round += 1
+        text = self.first_round_text if self.round == 1 else "Done."
+        return iter([Chunk(text)])
+
+    def chat_perform(self, *, stream, turns, tools, data_model, kwargs):
+        return self._next()
+
+    async def chat_perform_async(self, *, stream, turns, tools, data_model, kwargs):
+        it = self._next()
+
+        async def gen():
+            for c in it:
+                yield c
+
+        return gen()
+
+    def stream_turn(self, completion, has_data_model, turns=()):
+        contents = []
+        if completion["text"]:
+            contents.append(ContentText.model_construct(text=completion["text"]))
+        if self.round == 1:
+            contents.append(ContentToolRequest(id="c1", name="get_time", arguments={}))
+        return AssistantTurn(contents=contents, tokens=None, completion=None)
+
+
+def _chat(first_round_text: str = "Checking.") -> Chat:
+    chat = Chat(provider=TwoRoundProvider(first_round_text))
+
+    def get_time() -> str:
+        """Get the time."""
+        return "12:00"
+
+    chat.register_tool(get_time)
+    return chat
+
+
+def test_stream_text_separates_text_around_a_tool_call():
+    out = "".join(str(x) for x in _chat().stream("hi", content="text"))
+    assert out == "Checking.\n\nDone."
+
+
+@pytest.mark.asyncio
+async def test_stream_async_text_separates_text_around_a_tool_call():
+    stream = await _chat().stream_async("hi", content="text")
+    out = "".join([str(x) async for x in stream])
+    assert out == "Checking.\n\nDone."
+
+
+def test_stream_text_does_not_start_with_a_separator():
+    out = "".join(str(x) for x in _chat("").stream("hi", content="text"))
+    assert out == "Done."
+
+
+@pytest.mark.asyncio
+async def test_stream_async_text_does_not_start_with_a_separator():
+    stream = await _chat("").stream_async("hi", content="text")
+    out = "".join([str(x) async for x in stream])
+    assert out == "Done."
+
+
+def test_chat_text_separates_text_around_a_tool_call():
+    assert str(_chat().chat("hi", echo="none")) == "Checking.\n\nDone."
+    assert str(_chat("").chat("hi", echo="none")) == "Done."
+
+
+@pytest.mark.asyncio
+async def test_chat_async_text_separates_text_around_a_tool_call():
+    response = await _chat().chat_async("hi", echo="none")
+    assert await response.get_content() == "Checking.\n\nDone."
+    response = await _chat("").chat_async("hi", echo="none")
+    assert await response.get_content() == "Done."