[Python] fastapi websocket not sending progress update when message sent

I have the following piece of code to download the ollama model and on request and then update the progress in response

print log does work but websocket.send_text not sending log object or any object. one the ollama completes the model pulling then all these logs sent not sure how to resolve it

import asyncio

import ollama
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse

def print_logs(res):
try:
if res.total and res.completed:
percentage = res.completed / res.total * 100
log_message = f"{res.status} {percentage:.2f}%"
else:
log_message = f"{res.status}"
return log_message
except Exception as e:
raise TypeError(f"Error processing response: {res}") from e

async def download_model(model_name: str):
for res in ollama.Client().pull(model_name, stream=True):
log = print_logs(res)
yield log

app = FastAPI()

@app.websocket("/ws/download-model")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
try:
model_name = await websocket.receive_text()
print("model_name", model_name)
async for log in download_model(model_name):
# print("LOGa:", log)
print(log)
try:
await websocket.send_text(f"log {log}")
except Exception as e:
print("error ", e)
await websocket.send_text("Download complete.")
except Exception as e:
await websocket.send_text(f"Error: {str(e)}")
finally:
await websocket.close()

Continue reading...