Error trying to stream Assistant API with websocket

I am trying to implement an assistant into my personal website, but I am having some trouble passing the stream from the back end to the front end using a websocket. I seem to be successfully passing the input message from the front end to the back end, but in the process of trying to create a stream I keep running into this error:

ERROR:    Exception in ASGI application
Traceback (most recent call last):
  File "/Users/nickhenderson/Projects/react-website/run.py", line 90, in chat_with_assistant
    with client.beta.threads.runs.stream(
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/openai/lib/streaming/_assistants.py", line 447, in __enter__
    self.__event_handler._init(self.__stream)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/openai/lib/streaming/_assistants.py", line 59, in _init
    if self.__stream:
       ^^^^^^^^^^^^^
AttributeError: 'EventHandler' object has no attribute '_AssistantEventHandler__stream'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/_exception_handler.py", line 42, in wrapped_app
    await app(scope, receive, sender)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/routing.py", line 93, in app
    await func(session)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/fastapi/routing.py", line 383, in app
    await dependant.call(**solved_result.values)
  File "/Users/nickhenderson/Projects/react-website/run.py", line 99, in chat_with_assistant
    raise HTTPException(status_code=500, detail=str(e))
fastapi.exceptions.HTTPException: 500: 'EventHandler' object has no attribute '_AssistantEventHandler__stream'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/uvicorn/protocols/websockets/websockets_impl.py", line 243, in run_asgi
    result = await self.app(self.scope, self.asgi_receive, self.asgi_send)  # type: ignore[func-returns-value]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/uvicorn/middleware/proxy_headers.py", line 60, in __call__
    return await self.app(scope, receive, send)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/fastapi/applications.py", line 1054, in __call__
    await super().__call__(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/applications.py", line 113, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/middleware/errors.py", line 152, in __call__
    await self.app(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/middleware/cors.py", line 77, in __call__
    await self.app(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/middleware/exceptions.py", line 62, in __call__
    await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/_exception_handler.py", line 53, in wrapped_app
    raise exc
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/_exception_handler.py", line 42, in wrapped_app
    await app(scope, receive, sender)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/routing.py", line 715, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/routing.py", line 735, in app
    await route.handle(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/routing.py", line 362, in handle
    await self.app(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/routing.py", line 95, in app
    await wrap_app_handling_exceptions(app, session)(scope, receive, send)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/_exception_handler.py", line 63, in wrapped_app
    await response(scope, receive, sender)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/responses.py", line 148, in __call__
    await send(
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/_exception_handler.py", line 39, in sender
    await send(message)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/starlette/_exception_handler.py", line 39, in sender
    await send(message)
  File "/Users/nickhenderson/Projects/react-website/venv/lib/python3.11/site-packages/uvicorn/protocols/websockets/websockets_impl.py", line 342, in asgi_send
    raise RuntimeError(msg % message_type)
RuntimeError: Expected ASGI message 'websocket.send' or 'websocket.close', but got 'websocket.http.response.start'.

Does anyone have any idea why this might be happening or any advice on how to resolve? This is my first time implementing a websocket. I am just trying to get the streaming effect on the front end. Here is my back end implementation:

class MessageRequest(BaseModel):
    message: str
    num_msg: int

class EventHandler(AssistantEventHandler):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket

    @override
    async def on_text_created(self, text) -> None:
        await self.websocket.send_text("\nassistant > ")

    @override
    async def on_text_delta(self, delta, snapshot):
        await self.websocket.send_text(delta.value)

    async def on_tool_call_created(self, tool_call):
        await self.websocket.send_text(f"\nassistant > {tool_call.type}\n")

    async def on_tool_call_delta(self, delta, snapshot):
        if delta.type == 'code_interpreter':
            if delta.code_interpreter.input:
                await self.websocket.send_text(delta.code_interpreter.input)
            if delta.code_interpreter.outputs:
                await self.websocket.send_text("\n\noutput >")
                for output in delta.code_interpreter.outputs:
                    if output.type == "logs":
                        await self.websocket.send_text(f"\n{output.logs}")

@app.get("/")
def read_root():
    return {"message": "Hello, FastAPI!"}


@app.websocket("/chat")
async def chat_with_assistant(websocket: WebSocket):
    await websocket.accept()
    
    try:
        while True:
            data = await websocket.receive_json()
            user_input = data['message']
            num_msg = data['num_msg']

            global t_id

            
            if num_msg == 0:
                thread = client.beta.threads.create()
                t_id = thread.id

            message = client.beta.threads.messages.create(
                thread_id=t_id,
                role="user",
                content= user_input
            )

            event_handler = EventHandler(websocket)


            with client.beta.threads.runs.stream(
                thread_id=t_id,
                assistant_id=assistant_id,
                instructions="Please address the user as Jane Doe. The user has a premium account.",
                event_handler=event_handler,
            ) as stream:
                await stream.until_done()

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

and here is the front end implementation:

export function Chat() {
  const [messages, setMessages] = useState<Message[]>([])
  const [input, setInput] = useState('')
  const [numMessages, setNumMessages] = useState(0)
  const [isConnected, setIsConnected] = useState(false)
  const textareaRef = useRef<HTMLTextAreaElement>(null)
  const scrollAreaRef = useRef<HTMLDivElement>(null)
  const socketRef = useRef<WebSocket | null>(null)

  useEffect(() => {
    socketRef.current = new WebSocket('ws://localhost:8000/chat')


    socketRef.current.onopen = () => {
      console.log('WebSocket connection established')
      setIsConnected(true)

    }

    socketRef.current.onmessage = (event) => {
      const newMessage = { text: event.data, isUser: false }
      setMessages((prev) => [...prev, newMessage])
    }

    socketRef.current.onclose = () => {
      console.log('WebSocket connection closed')
      setIsConnected(false)
    }

    socketRef.current.onerror = (error) => {
      console.error('WebSocket error:', error)
      }

    return () => {
      if (socketRef.current) {
        socketRef.current.close()
      }
    }
  }, [])

  const handleSend = () => {

    if (input.trim() && isConnected && socketRef.current) {

      const userMessage = { text: input, isUser: true }
      setMessages((prev) => [...prev, userMessage])

  
      socketRef.current.send(JSON.stringify({
        message: input,
        num_msg: numMessages
      }))

      setInput('')
      setNumMessages((prev) => prev + 1)

    }
  }

Again I am very new to this. Do I even need to use web sockets? Is there a simpler way for me to display the streamed assistant output in my UI?

The conventional wisdom is to not use websockets unless you have a real need for them, as they can be a bit arcane.

But hey why not right?

I think you need to describe the symptoms of your problem a bit more. It’s hard to guess what might be going wrong from the code alone.

I would try adding logging to the event handler methods so you can tell what’s executing and what’s not. That may give you a clue where to focus next.

I wish I could be of more help. Good Luck!