Repository: LAION-AI/Open-Assistant
File: inference/server/oasst_inference_server/routes/chats.py (Python)

Code before:
        raise fastapi.HTTPException(status_code=400, detail="Only assistant messages can be streamed.")

    if message.has_finished:
        raise fastapi.HTTPException(status_code=204, detail=message.state)

    async def event_generator(chat_id: str, message_id: str, worker_compat_hash: str | None):
        redis_client = deps.make_redis_client()
        message_queue = queueing.message_queue(redis_client, message_id=message_id)
        work_queue = (
            queueing.work_queue(redis_client, worker_compat_hash=worker_compat_hash)
            if worker_compat_hash is not None
            else None
        )
        has_started = False
        try:
            while True:
                item = await message_queue.dequeue(timeout=settings.pending_event_interval)
                if item is None:
                    if not has_started:
                        if work_queue is None:
                            qpos, qlen = 0, 1
                        else:
                            # TODO: make more efficient, e.g. pipeline
                            qdeq = await work_queue.get_deq_counter()
                            qenq = await work_queue.get_enq_counter()
                            mpos = await queueing.get_pos_value(redis_client, message_id)
                            qpos = max(mpos - qdeq, 0)
                            qlen = max(qenq - qdeq, qpos)
                        yield {
                            "data": chat_schema.PendingResponseEvent(
                                queue_position=qpos,
                                queue_size=qlen,
                            ).json()
                        }
                    continue
                has_started = True

                _, response_packet_str = item
                response_packet = pydantic.parse_raw_as(inference.WorkerResponse, response_packet_str)

                if response_packet.response_type in ("error", "generated_text"):
                    logger.warning(
                        f"Received {response_packet.response_type=} response for {chat_id}. This should not happen."
                    )
                    break

                if response_packet.response_type == "internal_error":
                    yield {
                        "data": chat_schema.ErrorResponseEvent(
                            error=response_packet.error, message=response_packet.message
                        ).json(),

Review comment (bug):
could this be done in parallel like this? or is this a race condition?

```python
[qdeq, qenq, mpos] = await asyncio.gather(
  work_queue.get_deq_counter(),
  work_queue.get_enq_counter(),
  queueing.get_pos_value(redis_client, message_id)
)
```

Code after:
        message: models.DbMessage = await ucr.get_message_by_id(chat_id=chat_id, message_id=message_id)
    if message.role != "assistant":
        raise fastapi.HTTPException(status_code=400, detail="Only assistant messages can be streamed.")

    if message.has_finished:
        raise fastapi.HTTPException(status_code=204, detail=message.state)

    async def event_generator(chat_id: str, message_id: str, worker_compat_hash: str | None):
        redis_client = deps.make_redis_client()
        message_queue = queueing.message_queue(redis_client, message_id=message_id)
        work_queue = (
            queueing.work_queue(redis_client, worker_compat_hash=worker_compat_hash)
            if worker_compat_hash is not None
            else None
        )
        has_started = False
        try:
            while True:
                item = await message_queue.dequeue(timeout=settings.pending_event_interval)
                if item is None:
                    if not has_started:
                        if work_queue is None:
                            qpos, qlen = 0, 1
                        else:
                            # TODO: make more efficient, e.g. pipeline
                            [qdeq, qenq, mpos] = await asyncio.gather(
                                work_queue.get_deq_counter(),
                                work_queue.get_enq_counter(),
                                queueing.get_pos_value(redis_client, message_id),
                            )
                            qpos = max(mpos - qdeq, 0)
                            qlen = max(qenq - qdeq, qpos)
                        yield {
                            "data": chat_schema.PendingResponseEvent(
                                queue_position=qpos,
                                queue_size=qlen,
                            ).json()
                        }
                    continue
                has_started = True

                _, response_packet_str = item
                response_packet = pydantic.parse_raw_as(inference.WorkerResponse, response_packet_str)

                if response_packet.response_type in ("error", "generated_text"):
                    logger.warning(
                        f"Received {response_packet.response_type=} response for {chat_id}. This should not happen."
                    )
                    break

                if response_packet.response_type == "internal_error":