diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 0850952..3033ac4 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -102,6 +102,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes # Feature flags self._is_document_upload_enabled = is_feature_enabled(self.user, "document_upload") self._is_web_search_enabled = is_feature_enabled(self.user, "web_search") + self._fake_streaming_delay = settings.FAKE_STREAMING_DELAY @property def _stop_cache_key(self): @@ -490,14 +491,20 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes # A user prompt node => The user has provided input pass - elif Agent.is_model_request_node(node): + elif Agent.is_model_request_node(node): # pylint: disable=too-many-nested-blocks # A model request node => agent is asking the model to generate a response if not self._support_streaming: result = await node.run(run.ctx) logger.debug("node.run result: %s", result) for part in result.model_response.parts: if isinstance(part, TextPart): - yield events_v4.TextPart(text=part.content) + if self._fake_streaming_delay: + for i in range(0, len(part.content), 4): + await self._agent_stop_streaming() + yield events_v4.TextPart(text=part.content[i : i + 4]) + time.sleep(self._fake_streaming_delay) + else: + yield events_v4.TextPart(text=part.content) elif isinstance(part, ToolCallPart): yield events_v4.ToolCallPart( tool_call_id=part.tool_call_id, diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation.py b/src/backend/chat/tests/views/chat/conversations/test_conversation.py index ed1b0af..e94cc3c 100644 --- a/src/backend/chat/tests/views/chat/conversations/test_conversation.py +++ b/src/backend/chat/tests/views/chat/conversations/test_conversation.py @@ -1098,16 +1098,19 @@ def test_post_conversation_model_selection_new(api_client, mock_openai_stream, s @freeze_time("2025-07-25T10:36:35.297675Z") +@pytest.mark.parametrize("stream_delay", [None, 0.0001]) @respx.mock def test_post_conversation_data_protocol_no_stream( api_client, mock_openai_no_stream, mock_uuid4, settings, + stream_delay, ): """ Test posting messages to a conversation using the 'data' protocol without streaming. """ + settings.FAKE_STREAMING_DELAY = stream_delay settings.LLM_CONFIGURATIONS = { "default-model": LLModel( hrid="model-1", @@ -1149,12 +1152,35 @@ def test_post_conversation_data_protocol_no_stream( assert response.get("x-vercel-ai-data-stream") == "v1" assert response.streaming - # Wait for the streaming content to be fully received + # Wait for the content to be fully received response_content = b"".join(response.streaming_content).decode("utf-8") - assert response_content == ( - '0:"The sky appears blue due to a phenomenon called Rayleigh scattering."\n' - 'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":135}}\n' - ) + + if stream_delay: + assert response_content == ( + '0:"The "\n' + '0:"sky "\n' + '0:"appe"\n' + '0:"ars "\n' + '0:"blue"\n' + '0:" due"\n' + '0:" to "\n' + '0:"a ph"\n' + '0:"enom"\n' + '0:"enon"\n' + '0:" cal"\n' + '0:"led "\n' + '0:"Rayl"\n' + '0:"eigh"\n' + '0:" sca"\n' + '0:"tter"\n' + '0:"ing."\n' + 'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":135}}\n' + ) + else: + assert response_content == ( + '0:"The sky appears blue due to a phenomenon called Rayleigh scattering."\n' + 'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":135}}\n' + ) assert mock_openai_no_stream.called diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 6b7d789..446eda1 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -471,6 +471,11 @@ class Base(Configuration): LLM_ROUTING_MODEL_HRID = values.Value( "default-routing-model", environ_name="LLM_ROUTING_MODEL_HRID", environ_prefix=None ) + FAKE_STREAMING_DELAY = values.FloatValue( + default=0.0025, + environ_name="FAKE_STREAMING_DELAY", + environ_prefix=None, # 25ms + ) # These settings are default values used in the default LLM_CONFIGURATIONS # They allow a deployment with only one model without a specific configuration file