fix async streaming (#12187)

llama-index-core/llama_index/core/agent/react/step.py

@@ -615,6 +615,8 @@ class ReActAgentWorker(BaseAgentWorker):
                 )
             )
             # wait until response writing is done
+            agent_response._ensure_async_setup()
+
             await agent_response._is_function_false_event.wait()
 
         return self._get_task_step_response(agent_response, step, is_done)

llama-index-core/llama_index/core/chat_engine/types.py

@@ -113,7 +113,6 @@ class StreamingAgentChatResponse:
         self._is_function_not_none_thread_event.set()
 
     def aput_in_queue(self, delta: Optional[str]) -> None:
-        self._ensure_async_setup()
         self._aqueue.put_nowait(delta)
         self._new_item_event.set()
 
@@ -167,6 +166,8 @@ class StreamingAgentChatResponse:
         memory: BaseMemory,
         on_stream_end_fn: Optional[callable] = None,
     ) -> None:
+        self._ensure_async_setup()
+
         if self.achat_stream is None:
             raise ValueError(
                 "achat_stream is None. Cannot asynchronously write to "

llama-index-integrations/agent/llama-index-agent-openai-legacy/llama_index/agent/openai_legacy/openai_agent.py

@@ -248,7 +248,9 @@ class BaseOpenAIAgent(BaseAgent):
             chat_stream_response.awrite_response_to_history(self.memory)
         )
         # wait until openAI functions stop executing
+        chat_stream_response._ensure_async_setup()
         await chat_stream_response._is_function_false_event.wait()
+
         # return response stream
         return chat_stream_response
 

llama-index-integrations/agent/llama-index-agent-openai-legacy/pyproject.toml

@@ -28,7 +28,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-agent-openai-legacy"
 readme = "README.md"
-version = "0.1.2"
+version = "0.1.3"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"

llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py

@@ -312,8 +312,11 @@ class OpenAIAgentWorker(BaseAgentWorker):
                 on_stream_end_fn=partial(self.finalize_task, task),
             )
         )
+        chat_stream_response._ensure_async_setup()
+
         # wait until openAI functions stop executing
         await chat_stream_response._is_function_false_event.wait()
+
         # return response stream
         return chat_stream_response
 

llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml

@@ -28,7 +28,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-agent-openai"
 readme = "README.md"
-version = "0.1.6"
+version = "0.1.7"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"