Spaces:
Running
Running
| # | |
| # SPDX-FileCopyrightText: Hadad <[email protected]> | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| import random # Import random module to enable random selection from a list | |
| from datetime import datetime # Import datetime class to work with current UTC time | |
| from typing import Dict, List # Import type hints for dictionaries and lists (not explicitly used here but imported) | |
| from config import auth # Import authentication configuration, likely a list of host dictionaries with credentials | |
| from src.utils.helper import busy, mark # Import 'busy' dictionary and 'mark' function to track and update host busy status | |
| # Dictionary to map session IDs to their assigned host information | |
| mapping = {} | |
| def get_host(session_id: str): | |
| """ | |
| Retrieve or assign a host for the given session ID. | |
| Args: | |
| session_id (str): A unique identifier for the current session. | |
| Returns: | |
| dict: The selected host dictionary from the auth configuration. | |
| Raises: | |
| Exception: If no available hosts are found to assign. | |
| Explanation: | |
| This function manages host assignment per session. If the session ID already has a host assigned, | |
| it returns that host immediately. Otherwise, it filters the list of hosts from 'auth' to find those | |
| that are currently not busy or whose busy period has expired (based on the 'busy' dictionary). | |
| From the available hosts, it randomly selects one, records the assignment in 'mapping', | |
| marks the selected host as busy for one hour, and returns the selected host. | |
| """ | |
| # Check if the session ID already has an assigned host in the mapping dictionary | |
| if session_id in mapping: | |
| # Return the previously assigned host for this session | |
| return mapping[session_id] | |
| # Get the current UTC time to compare against busy timestamps | |
| now = datetime.utcnow() | |
| # Filter hosts from auth that are either not marked busy or whose busy period has expired | |
| connect = [ | |
| h for h in auth | |
| if h["jarvis"] not in busy or busy[h["jarvis"]] <= now | |
| ] | |
| # If no hosts are available after filtering, raise an exception to indicate no hosts can be assigned | |
| if not connect: | |
| raise Exception("No available hosts to assign.") | |
| # Randomly select one host from the list of available hosts | |
| selected = random.choice(connect) | |
| # Map the session ID to the selected host for future reference | |
| mapping[session_id] = selected | |
| # Mark the selected host as busy for the next hour to prevent immediate reassignment | |
| mark(selected["jarvis"]) | |
| # Return the selected host dictionary | |
| return selected |