Fastapi-Websocket sendet keine Fortschrittsaktualisierung, wenn die Nachricht gesendet wird ⇐ Python

[Guest]

Ich habe den folgenden Code, um das Ollama-Modell herunterzuladen und auf Anfrage den Fortschritt als Antwort zu aktualisieren.
Das Drucken des Protokolls funktioniert, aber websocket.send_text sendet kein Protokollobjekt oder ähnliches Objekt.
Sobald das Ollama das Ziehen des Modells abgeschlossen hat, werden alle diese Protokolle gesendet. Ich bin mir nicht sicher, wie ich das Problem beheben kann

Code: Select all

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()