Add separator when LLM starts streaming

src/neuroagent/agents/simple_chat_agent.py

@@ -73,12 +73,9 @@ async def astream(
             streamed_response = self.agent.astream_events(
                 {"messages": query}, version="v2", config=config
             )
+            is_streaming = False
             async for event in streamed_response:
                 kind = event["event"]
-
-                # newline everytime model starts streaming.
-                if kind == "on_chat_model_start":
-                    yield "\n\n"
                 # check for the model stream.
                 if kind == "on_chat_model_stream":
                     # check if we are calling the tools.
@@ -95,6 +92,9 @@ async def astream(
 
                     content = data_chunk.content
                     if content:
+                        if not is_streaming:
+                            yield "\n<begin_llm_response>\n"
+                        is_streaming = True
                         yield content
             yield "\n"
 

tests/agents/test_simple_chat_agent.py

@@ -84,8 +84,8 @@ async def test_astream(fake_llm_with_tools, httpx_mock):
 
         msg_list = "".join([el async for el in response])
         assert (
-            msg_list == "\n\n\nCalling tool : get-morpho-tool with arguments :"
-            ' {"brain_region_id":"http://api.brain-map.org/api/v2/data/Structure/549"}\n\nGreat'
+            msg_list == "\nCalling tool : get-morpho-tool with arguments :"
+            ' {"brain_region_id":"http://api.brain-map.org/api/v2/data/Structure/549"}\n<begin_llm_response>\nGreat'
             " answer\n"
         )