(agent) add fake streaming when the model does not support it

To have a better UX, we simulate a streamed return. The default
value is quite fast to not block the worker too long...
This commit is contained in:
Quentin BEY
2025-09-22 17:01:03 +02:00
parent 7e460b55a3
commit e04a050486
3 changed files with 45 additions and 7 deletions
+9 -2
View File
@@ -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,
@@ -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
+5
View File
@@ -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