From 9d8af88b1cadd377aad929777d987ae74ebbf387 Mon Sep 17 00:00:00 2001
From: Logan <logan.markewich@live.com>
Date: Fri, 22 Mar 2024 18:06:11 -0600
Subject: [PATCH] fix async streaming (#12187)
---
llama-index-core/llama_index/core/agent/react/step.py | 2 ++
llama-index-core/llama_index/core/chat_engine/types.py | 3 ++-
.../llama_index/agent/openai_legacy/openai_agent.py | 2 ++
.../agent/llama-index-agent-openai-legacy/pyproject.toml | 2 +-
.../llama-index-agent-openai/llama_index/agent/openai/step.py | 3 +++
.../agent/llama-index-agent-openai/pyproject.toml | 2 +-
6 files changed, 11 insertions(+), 3 deletions(-)
diff --git a/llama-index-core/llama_index/core/agent/react/step.py b/llama-index-core/llama_index/core/agent/react/step.py
index 7253d19b5..272ea226d 100644
--- a/llama-index-core/llama_index/core/agent/react/step.py
+++ b/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)
diff --git a/llama-index-core/llama_index/core/chat_engine/types.py b/llama-index-core/llama_index/core/chat_engine/types.py
index a4a0e0903..929156526 100644
--- a/llama-index-core/llama_index/core/chat_engine/types.py
+++ b/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 "
diff --git a/llama-index-integrations/agent/llama-index-agent-openai-legacy/llama_index/agent/openai_legacy/openai_agent.py b/llama-index-integrations/agent/llama-index-agent-openai-legacy/llama_index/agent/openai_legacy/openai_agent.py
index 22cda7e5e..a807f8b04 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai-legacy/llama_index/agent/openai_legacy/openai_agent.py
+++ b/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
diff --git a/llama-index-integrations/agent/llama-index-agent-openai-legacy/pyproject.toml b/llama-index-integrations/agent/llama-index-agent-openai-legacy/pyproject.toml
index 3ab491664..f53dadfc6 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai-legacy/pyproject.toml
+++ b/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"
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
index 3a3f7b22e..c9892bd47 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
+++ b/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
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml b/llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml
index a11750db8..a62b0b71a 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/pyproject.toml
+++ b/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"
--
GitLab