Ilyas KHIAT commited on
Commit
96ecc62
·
1 Parent(s): 9aa8f58
Files changed (1) hide show
  1. main.py +3 -8
main.py CHANGED
@@ -148,17 +148,12 @@ import asyncio
148
  GENERATION_TIMEOUT_SEC = 60
149
 
150
  async def stream_generator(response, prompt):
151
- buffer = '' # Buffer to accumulate chunks
152
  async with async_timeout.timeout(GENERATION_TIMEOUT_SEC):
153
  try:
154
  async for chunk in response:
155
- buffer += chunk
156
- try:
157
- data = json.loads(buffer) # Try to parse the accumulated buffer
158
- yield json.dumps({"prompt": prompt, "content": data})
159
- buffer = '' # Clear the buffer after successful parsing
160
- except json.JSONDecodeError:
161
- continue # Continue accumulating data if JSON is incomplete
162
  except asyncio.TimeoutError:
163
  raise HTTPException(status_code=504, detail="Stream timed out")
164
 
 
148
  GENERATION_TIMEOUT_SEC = 60
149
 
150
  async def stream_generator(response, prompt):
 
151
  async with async_timeout.timeout(GENERATION_TIMEOUT_SEC):
152
  try:
153
  async for chunk in response:
154
+ if isinstance(chunk, bytes):
155
+ chunk = chunk.decode('utf-8') # Convert bytes to str if needed
156
+ yield json.dumps({"prompt": prompt, "content": chunk})
 
 
 
 
157
  except asyncio.TimeoutError:
158
  raise HTTPException(status_code=504, detail="Stream timed out")
159