From 4cbaaad98dda95d0e9c47bbdc5b9ab8278ef9a78 Mon Sep 17 00:00:00 2001 From: emrgnt-cmplxty Date: Fri, 10 Jan 2025 11:59:14 -0800 Subject: [PATCH] cap streaming chunk size --- py/core/main/api/v3/retrieval_router.py | 12 ++++++++++-- py/core/main/services/management_service.py | 3 +-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/py/core/main/api/v3/retrieval_router.py b/py/core/main/api/v3/retrieval_router.py index b25375b19..f2dd3e040 100644 --- a/py/core/main/api/v3/retrieval_router.py +++ b/py/core/main/api/v3/retrieval_router.py @@ -440,7 +440,11 @@ async def rag_app( async def stream_generator(): try: async for chunk in response: - yield chunk + if len(chunk) > 1024: + for i in range(0, len(chunk), 1024): + yield chunk[i : i + 1024] + else: + yield chunk except GeneratorExit: # Clean up if needed, then return return @@ -667,7 +671,11 @@ async def agent_app( async def stream_generator(): try: async for chunk in response: - yield chunk + if len(chunk) > 1024: + for i in range(0, len(chunk), 1024): + yield chunk[i : i + 1024] + else: + yield chunk except GeneratorExit: # Clean up if needed, then return return diff --git a/py/core/main/services/management_service.py b/py/core/main/services/management_service.py index 5a4ab5487..5960a3a96 100644 --- a/py/core/main/services/management_service.py +++ b/py/core/main/services/management_service.py @@ -1068,7 +1068,6 @@ async def get_all_user_limits(self, user_id: UUID) -> dict[str, Any]: ) )["total_entries"] - storage_limits = { "chunks": { "limit": max_chunks, @@ -1084,7 +1083,7 @@ async def get_all_user_limits(self, user_id: UUID) -> dict[str, Any]: "limit": max_collections, "used": used_collections, "remaining": max_collections - used_collections, - } + }, } # 5) Return a structured response result = {