From a0f5af0fdcc318e5e0101f0be58ba898bc5e20df Mon Sep 17 00:00:00 2001 From: Abi Raja Date: Fri, 8 Mar 2024 05:02:23 -0500 Subject: [PATCH] improve mock functionality --- backend/mock_llm.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/backend/mock_llm.py b/backend/mock_llm.py index 0c1060d..df0b8ed 100644 --- a/backend/mock_llm.py +++ b/backend/mock_llm.py @@ -4,6 +4,9 @@ from typing import Awaitable, Callable from custom_types import InputMode +STREAM_CHUNK_SIZE = 100 + + async def mock_completion( process_chunk: Callable[[str], Awaitable[None]], input_mode: InputMode ) -> str: @@ -13,8 +16,8 @@ async def mock_completion( else NO_IMAGES_NYTIMES_MOCK_CODE ) - for i in range(0, len(code_to_return), 100): - await process_chunk(code_to_return[i : i + 100]) + for i in range(0, len(code_to_return), STREAM_CHUNK_SIZE): + await process_chunk(code_to_return[i : i + STREAM_CHUNK_SIZE]) await asyncio.sleep(0.01) return code_to_return