From 05d685fa59cd67a0af87545a26a8b5f49a5a9f76 Mon Sep 17 00:00:00 2001 From: Anton Dubovik Date: Wed, 26 Jun 2024 15:26:33 +0100 Subject: [PATCH] fix: index fields are moved in non-streaming mode --- aidial_sdk/utils/streaming.py | 4 ++-- tests/examples/test_echo.py | 1 - 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/aidial_sdk/utils/streaming.py b/aidial_sdk/utils/streaming.py index 327fcdc..02d9c34 100644 --- a/aidial_sdk/utils/streaming.py +++ b/aidial_sdk/utils/streaming.py @@ -2,7 +2,7 @@ from typing import Any, AsyncGenerator, Dict from aidial_sdk.utils.logging import log_debug -from aidial_sdk.utils.merge_chunks import merge +from aidial_sdk.utils.merge_chunks import cleanup_indices, merge DONE_MARKER = "[DONE]" @@ -15,7 +15,7 @@ async def merge_chunks( response = merge(response, chunk) for choice in response["choices"]: - choice["message"] = choice["delta"] + choice["message"] = cleanup_indices(choice["delta"]) del choice["delta"] return response diff --git a/tests/examples/test_echo.py b/tests/examples/test_echo.py index a80dcbc..bddfba6 100644 --- a/tests/examples/test_echo.py +++ b/tests/examples/test_echo.py @@ -34,5 +34,4 @@ def test_app(): assert response_content == content response_attachment = response_message["custom_content"]["attachments"][0] - del response_attachment["index"] assert response_attachment == attachment