filename

main.py

@@ -148,17 +148,12 @@ import asyncio
 GENERATION_TIMEOUT_SEC = 60
 
 async def stream_generator(response, prompt):
-    buffer = ''  # Buffer to accumulate chunks
     async with async_timeout.timeout(GENERATION_TIMEOUT_SEC):
         try:
             async for chunk in response:
-                buffer += chunk
-                try:
-                    data = json.loads(buffer)  # Try to parse the accumulated buffer
-                    yield json.dumps({"prompt": prompt, "content": data})
-                    buffer = ''  # Clear the buffer after successful parsing
-                except json.JSONDecodeError:
-                    continue  # Continue accumulating data if JSON is incomplete
+                if isinstance(chunk, bytes):
+                    chunk = chunk.decode('utf-8')  # Convert bytes to str if needed
+                yield json.dumps({"prompt": prompt, "content": chunk})
         except asyncio.TimeoutError:
             raise HTTPException(status_code=504, detail="Stream timed out")