adds documentation to files

validators/database.py

@@ -6,11 +6,40 @@ from .streamer import ProcessedStreamResponse
 
 
 class LogDatabase:
+    """
+    A class to manage a log database stored as a JSONL (JSON Lines) file.
+
+    Attributes:
+        log_database_path (str): The path to the log database file.
+
+    Methods:
+        ensure_db_exists(file_path):
+            Ensures that the log database file exists. If it doesn't, an empty file is created.
+        
+        add_streams_to_db(stream_responses: ProcessedStreamResponse):
+            Asynchronously adds stream responses to the log database.
+        
+        append_dicts_to_file(file_path, dictionaries):
+            Asynchronously appends a list of dictionaries to the specified file.
+    """
+    
     def __init__(self, log_database_path: str):
+        """
+        Initializes the LogDatabase with the given log database file path.
+
+        Args:
+            log_database_path (str): The path to the log database file.
+        """
         self.log_database_path = log_database_path
         self.ensure_db_exists(log_database_path)
 
     def ensure_db_exists(self, file_path):
+        """
+        Ensures that the log database file exists. If it doesn't, creates an empty JSONL file.
+
+        Args:
+            file_path (str): The path to the log database file.
+        """
         if not os.path.exists(file_path):
             # Create an empty JSONL file
             with open(file_path, "w") as file:
@@ -21,6 +50,15 @@ class LogDatabase:
             bt.logging.info(f"File '{file_path}' already exists.")
 
     async def add_streams_to_db(self, stream_responses: ProcessedStreamResponse):
+        """
+        Asynchronously adds stream responses to the log database.
+
+        Args:
+            stream_responses (ProcessedStreamResponse): A list of processed stream responses to add to the log database.
+
+        Raises:
+            Exception: If an error occurs while adding streams to the database.
+        """
         bt.logging.info(f"Writing streams to the database...")
         try:
             stream_responses_dict = [
@@ -35,6 +73,13 @@ class LogDatabase:
             raise e
 
     async def append_dicts_to_file(self, file_path, dictionaries):
+        """
+        Asynchronously appends a list of dictionaries to the specified file.
+
+        Args:
+            file_path (str): The path to the file where dictionaries will be appended.
+            dictionaries (list): A list of dictionaries to append to the file.
+        """
         async with aiofiles.open(file_path, mode="a") as file:
             for dictionary in dictionaries:
                 await file.write(json.dumps(dictionary) + "\n")

validators/stream_manager.py

@@ -7,7 +7,24 @@ from aiohttp.web import Request
 
 
 class StreamManager:
+    """
+    A class to manage the processing of multiple asynchronous data streams and log their responses.
+
+    Attributes:
+        log_database (LogDatabase): The log database to store stream responses.
+
+    Methods:
+        process_streams(request, streams_responses, stream_uids):
+            Processes multiple asynchronous streams, logs their responses, and returns the selected stream response.
+    """
+    
     def __init__(self, log_database_path: str = "requests_db.jsonl"):
+        """
+        Initializes the StreamManager with the given log database file path.
+
+        Args:
+            log_database_path (str): The path to the log database file, defaults to "requests_db.jsonl".
+        """
         self.log_database = LogDatabase(log_database_path)
 
     async def process_streams(
@@ -16,6 +33,17 @@ class StreamManager:
         streams_responses: List[AsyncIterator],
         stream_uids: List[int],
     ):
+        """
+        Processes multiple asynchronous streams, logs their responses, and returns the selected stream response (stream from first non-empty chunk).
+
+        Args:
+            request (Request): The web request object.
+            streams_responses (List[AsyncIterator]): A list of asynchronous iterators representing the streams.
+            stream_uids (List[int]): A list of unique IDs for the streams.
+
+        Returns:
+            ProcessedStreamResponse: The response from the selected stream.
+        """
         lock = asyncio.Lock()
 
         streamers = [

validators/streamer.py

@@ -11,6 +11,18 @@ from prompting.protocol import StreamPromptingSynapse
 
 
 class StreamChunk(BaseModel):
+    """
+    A model representing a chunk of streaming data.
+
+    Attributes:
+        delta (str): The change in the stream.
+        finish_reason (Optional[str]): The reason for finishing the stream.
+        accumulated_chunks (List[str]): List of accumulated chunks.
+        accumulated_chunks_timings (List[float]): Timings for the accumulated chunks.
+        timestamp (str): The timestamp of the chunk.
+        sequence_number (int): The sequence number of the chunk.
+        selected_uid (int): The selected user ID.
+    """
     delta: str
     finish_reason: Optional[str]
     accumulated_chunks: List[str]
@@ -20,11 +32,29 @@ class StreamChunk(BaseModel):
     selected_uid: int
 
     def encode(self, encoding: str) -> bytes:
+        """
+        Encodes the StreamChunk instance to a JSON-formatted bytes object.
+
+        Args:
+            encoding (str): The encoding to use.
+
+        Returns:
+            bytes: The encoded JSON data.
+        """
         data = json.dumps(self.dict(), indent=4)
         return data.encode(encoding)
 
 
 class StreamError(BaseModel):
+    """
+    A model representing an error in the streaming data.
+
+    Attributes:
+        error (str): The error message.
+        timestamp (str): The timestamp of the error.
+        sequence_number (int): The sequence number at the time of error.
+        finish_reason (str): The reason for finishing the stream, defaults to "error".
+    """
     error: str
     timestamp: str
     sequence_number: int
@@ -39,6 +69,20 @@ ProcessedStreamResponse = Union[StreamChunk, StreamError]
 
 
 class AsyncResponseDataStreamer:
+    """
+    A class to manage asynchronous streaming of response data.
+
+    Attributes:
+        async_iterator (AsyncIterator): An asynchronous iterator for streaming data.
+        selected_uid (int): The selected user ID.
+        lock (asyncio.Lock): An asyncio lock to ensure exclusive access.
+        delay (float): Delay between processing chunks, defaults to 0.1 seconds.
+        accumulated_chunks (List[str]): List of accumulated chunks.
+        accumulated_chunks_timings (List[float]): Timings for the accumulated chunks.
+        finish_reason (str): The reason for finishing the stream.
+        sequence_number (int): The sequence number of the stream.
+        lock_acquired (bool): Flag indicating if the lock was acquired.
+    """
     def __init__(
         self,
         async_iterator: AsyncIterator,
@@ -59,6 +103,15 @@ class AsyncResponseDataStreamer:
     def ensure_response_is_created(
         self, initiated_response: web.StreamResponse
     ) -> web.StreamResponse:
+        """
+        Ensures that a StreamResponse is created if it does not already exist.
+
+        Args:
+            initiated_response (web.StreamResponse): The initiated response.
+
+        Returns:
+            web.StreamResponse: The ensured response.
+        """
         # Creates response if it was not created
         if initiated_response == None:
             initiated_response = web_response.StreamResponse(status=200, reason="OK")
@@ -74,6 +127,18 @@ class AsyncResponseDataStreamer:
         stream_chunk: StreamChunk,
         lock: asyncio.Lock,
     ) -> web.StreamResponse:
+        """
+        Writes a stream chunk to the response if the lock is acquired.
+
+        Args:
+            request (web.Request): The web request object.
+            initiated_response (web.StreamResponse): The initiated response.
+            stream_chunk (StreamChunk): The chunk of stream data to write.
+            lock (asyncio.Lock): The lock to ensure exclusive access.
+
+        Returns:
+            web.StreamResponse: The response with the written chunk.
+        """
         # 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()
@@ -93,6 +158,18 @@ class AsyncResponseDataStreamer:
         return initiated_response
 
     async def stream(self, request: web.Request) -> ProcessedStreamResponse:
+        """
+        Streams data from the async iterator and writes it to the response.
+
+        Args:
+            request (web.Request): The web request object.
+
+        Returns:
+            ProcessedStreamResponse: The final processed stream response.
+
+        Raises:
+            ValueError: If the stream does not return a valid synapse.
+        """
         try:
             start_time = time.time()
             client_response: web.Response = None