websocketfastapi

Sending websocket messages in FastAPI route


I am using FastAPI as backend and Next.js as Frontend. I am uploading many files while processing those files on backend side. Those files are pdf and word files.

And I am sending the progress to frontend using websocket connection.

I can see the logs are successfully printed on backend side.

0
0.06
0.12
0.18
0.24
0.3
0.36
0.42
0.48
0.54
0.6000000000000001
1.0

But on frontside I don't get any messages but after every processing has ended I got all of those websocket messages at one time.

I tried to log progress on fastapi backend and confirmed that server is working properly.

const ws = useRef<null | WebSocket>(null);

  useEffect(() => {
    ws.current = new WebSocket(`${WEBSOCKET_AI_SERVER}/ws`);

    ws.current.onmessage = (ev: MessageEvent) => {
      const data = JSON.parse(ev.data);
      switch (data.type) {
        case 'signin':
          setSocketId(data.id);
          break;
        case 'pleadings':
          setPleadingsProgress(data.percentage * 100);
          break;
        case 'bill':
          setBillProgress(data.percentage * 100);
        default:
          console.log(data);
      }
    }

    return () => {
      ws.current?.close();
    }
  }, [
    setSocketId,
    setPleadingsProgress
  ]);
@app.post("/uploadpleadingfiles")
async def upload_pleading_files(
    files: Annotated[list[UploadFile], File(description="Multiple files as UploadFile")],
    socketId: Annotated[str, Form()]
):
    websocket = manager.connection_pool.get(socketId)

    percentage = 0
    chunks = []
    chunk_text = ""
    file_count = len(files)

    # update progress
    await websocket.send_json({'type': 'pleadings', 'percentage': percentage})

    for file in files:
        # save file to local disk
        unique_filename = f"{uuid.uuid4()}_{file.filename}"
        file_location = os.path.join(UPLOAD_DIRECTORY, unique_filename)
        with open(file_location, "wb") as file_object:
            file_object.write(file.file.read())

        # update progress
        percentage += 0.6 * 0.5 / file_count
        await websocket.send_json({'type': 'pleadings', 'percentage': percentage})

        if unique_filename.endswith('.pdf'):
            text = extract_text_from_pdf(file_location)
            for sentence in text.sents:
                if len(chunk_text) + len(sentence.text) < CHUNK_SIZE:
                    chunk_text += sentence.text
                else:
                    chunks.append(chunk_text)
                    chunk_text = sentence.text

        elif unique_filename.endswith('.doc') or unique_filename.endswith('.docx'):
            document = Document()
            document.LoadFromFile(file_location)
            for index in range(document.Sections.Count):
                section = document.Sections.get_Item(index)
                for id in range(section.Paragraphs.Count):
                    paragraph = section.Paragraphs.get_Item(id)
                    if len(chunk_text) + len(paragraph.Text) < CHUNK_SIZE:
                        chunk_text += paragraph.Text
                    else:
                        chunks.append(chunk_text)
                        chunk_text = paragraph.Text

        os.unlink(file_location)

        # update progress
        percentage += 0.6 * 0.5 / file_count
        await websocket.send_json({'type': 'pleadings', 'percentage': percentage})

    if not len(chunk_text) == 0:
        chunks.append(chunk_text)

    # ...

    for index, chunk in enumerate(chunks):
        # ...

        # update progress
        percentage += 0.4 * 1 / chunk_count
        await websocket.send_json({'type': 'pleadings', 'percentage': percentage})

    # return the result

Solution

  • You should add these code lines so that you can confirm the websocket is connected correctly

    ws.current.onopen = () => {
       console.log('WebSocket connection established');
    };
    ws.current.onerror = (error) => {
        console.error('WebSocket error:', error);
    };
    ws.current.onclose = () => {
        console.log('WebSocket connection closed');
    };
    

    And then you can check the where error comes from exactly. To isolate the problem, test the WebSocket communication with a minimal setup. Create a simple FastAPI endpoint that sends periodic updates and a minimal React component to receive and display them. This will help you identify if the issue lies in the WebSocket handling or elsewhere. If the issues comes from websocket, then you cam simple add this line to FastAPI backend server.

    asyncio.sleep(0)
    

    In my case, it worked very correctly