Spaces:
Sleeping
Sleeping
| 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() |