validators/streamer.py

import json
import time
import asyncio
import traceback
import bittensor as bt
from pydantic import BaseModel
from datetime import datetime
from typing import AsyncIterator, Optional, List, Union
from aiohttp import web, web_response
from prompting.protocol import StreamPromptingSynapse


class StreamChunk(BaseModel):
    delta: str
    finish_reason: Optional[str]
    accumulated_chunks: List[str]
    accumulated_chunks_timings: List[float]
    timestamp: str
    sequence_number: int
    selected_uid: int

    def encode(self, encoding: str) -> bytes:
        data = json.dumps(self.dict(), indent=4)
        return data.encode(encoding)


class StreamError(BaseModel):
    error: str
    timestamp: str
    sequence_number: int
    finish_reason: str = "error"

    def encode(self, encoding: str) -> bytes:
        data = json.dumps(self.dict(), indent=4)
        return data.encode(encoding)


ProcessedStreamResponse = Union[StreamChunk, StreamError]


class AsyncResponseDataStreamer:
    def __init__(
        self,
        async_iterator: AsyncIterator,
        selected_uid: int,
        lock: asyncio.Lock,
        delay: float = 0.1,
    ):
        self.async_iterator = async_iterator
        self.delay = delay
        self.selected_uid = selected_uid
        self.accumulated_chunks: List[str] = []
        self.accumulated_chunks_timings: List[float] = []
        self.finish_reason: str = None
        self.sequence_number: int = 0
        self.lock = lock
        self.lock_acquired = False

    def ensure_response_is_created(
        self, initiated_response: web.StreamResponse
    ) -> web.StreamResponse:
        # Creates response if it was not created
        if initiated_response == None:
            initiated_response = web_response.StreamResponse(status=200, reason="OK")
            initiated_response.headers["Content-Type"] = "application/json"
            return initiated_response

        return initiated_response

    async def write_to_stream(
        self,
        request: web.Request,
        initiated_response: web.StreamResponse,
        stream_chunk: StreamChunk,
        lock: asyncio.Lock,
    ) -> web.StreamResponse:
        # Try to acquire the lock and sets the lock_acquired flag. Only the stream that acquires the lock should write to the response
        if lock.locked() == False:
            self.lock_acquired = await lock.acquire()

        if initiated_response == None and self.lock_acquired:
            initiated_response = self.ensure_response_is_created(initiated_response)
            # Prepare and send the headers
            await initiated_response.prepare(request)

        if self.lock_acquired:
            await initiated_response.write(stream_chunk.encode("utf-8"))
        else:
            bt.logging.info(
                f"Stream of uid {stream_chunk.selected_uid} was not the first to return, skipping..."
            )

        return initiated_response

    async def stream(self, request: web.Request) -> ProcessedStreamResponse:
        # response = web_response.StreamResponse(status=200, reason="OK")
        # response.headers["Content-Type"] = "application/json"
        # await response.prepare(request)  # Prepare and send the headers

        try:
            start_time = time.time()
            client_response: web.Response = None
            final_response: ProcessedStreamResponse

            async for chunk in self.async_iterator:
                if isinstance(chunk, str):
                    # Chunks are currently returned in string arrays, so we need to concatenate them
                    concatenated_chunks = "".join(chunk)
                    self.accumulated_chunks.append(concatenated_chunks)
                    self.accumulated_chunks_timings.append(time.time() - start_time)
                    # Gets new response state
                    self.sequence_number += 1
                    new_response_state = self._create_chunk_response(
                        concatenated_chunks
                    )
                    # Writes the new response state to the response
                    client_response = await self.write_to_stream(
                        request, client_response, new_response_state, self.lock
                    )
                    # await response.write(new_response_state.encode('utf-8'))

            if chunk is not None and isinstance(chunk, StreamPromptingSynapse):
                if len(self.accumulated_chunks) == 0:
                    self.accumulated_chunks.append(chunk.completion)
                    self.accumulated_chunks_timings.append(time.time() - start_time)

                self.finish_reason = "completed"
                self.sequence_number += 1
                # Assuming the last chunk holds the last value yielded which should be a synapse with the completion filled
                synapse = chunk
                final_response = self._create_chunk_response(synapse.completion)

                if synapse.completion:
                    client_response = await self.write_to_stream(
                        request, client_response, final_response, self.lock
                    )
            else:
                raise ValueError("Stream did not return a valid synapse.")

        except Exception as e:
            bt.logging.error(
                f"Encountered an error while processing stream for uid {self.selected_uid} get_stream_response:\n{traceback.format_exc()}"
            )
            error_response = self._create_error_response(str(e))
            final_response = error_response

            # Only the stream that acquires the lock should write the error response
            if self.lock_acquired:
                self.ensure_response_is_created(client_response)
                client_response.set_status(500, reason="Internal error")
                client_response.write(error_response.encode("utf-8"))
        finally:
            # Only the stream that acquires the lock should close the response
            if self.lock_acquired:
                self.ensure_response_is_created(client_response)
                # Ensure to close the response properly
                await client_response.write_eof()

            return final_response

    def _create_chunk_response(self, chunk: str) -> StreamChunk:
        """
        Creates a StreamChunk object with the current state.

        :param chunk: List of strings representing the current chunk.
        :return: StreamChunk object.
        """
        return StreamChunk(
            delta=chunk,
            finish_reason=self.finish_reason,
            accumulated_chunks=self.accumulated_chunks,
            accumulated_chunks_timings=self.accumulated_chunks_timings,
            timestamp=self._current_timestamp(),
            sequence_number=self.sequence_number,
            selected_uid=self.selected_uid,
        )

    def _create_error_response(self, error_message: str) -> StreamError:
        """
        Creates a StreamError object.

        :param error_message: Error message to include in the StreamError.
        :return: StreamError object.
        """
        return StreamError(
            error=error_message,
            timestamp=self._current_timestamp(),
            sequence_number=self.sequence_number,
        )

    def _current_timestamp(self) -> str:
        """
        Returns the current timestamp in ISO format.

        :return: Current timestamp as a string.
        """
        return datetime.utcnow().isoformat()