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 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 | |
| def initialize_tools(): | |
| """ | |
| Initialize and select available tools (endpoints) from the configured hosts. | |
| Returns: | |
| tuple: A tuple containing three elements: | |
| - tool_setup (str): The endpoint or configuration for the main tool setup. | |
| - image_tool (str): The endpoint URL for image generation services. | |
| - audio_tool (str): The endpoint URL for audio generation services. | |
| Raises: | |
| Exception: If no available hosts are found or if any required tool endpoint is missing. | |
| Explanation: | |
| This function filters the list of hosts from 'auth' to find those not currently busy or whose busy period has expired. | |
| It randomly selects one available host, marks it as busy for one hour, | |
| and retrieves the required tool endpoints ('done', 'image', 'audio') from the selected host's configuration. | |
| If any of these endpoints are missing, it raises an exception. | |
| Finally, it returns the three tool endpoints as a tuple. | |
| """ | |
| # Get the current UTC time for busy status comparison | |
| now = datetime.utcnow() | |
| # Filter hosts that are either not marked busy or whose busy period has expired | |
| available = [ | |
| item for item in auth | |
| if item["jarvis"] not in busy or busy[item["jarvis"]] <= now | |
| ] | |
| # Raise an exception if no hosts are currently available | |
| if not available: | |
| raise Exception("No available hosts to initialize tools.") | |
| # Randomly select one host from the available list | |
| selected = random.choice(available) | |
| # Mark the selected host as busy for the next hour to prevent immediate reassignment | |
| mark(selected["jarvis"]) | |
| # Retrieve the tool endpoints from the selected host's configuration dictionary | |
| tool_setup = selected.get("done") # Main tool setup endpoint or configuration | |
| image_tool = selected.get("image") # Image generation service endpoint | |
| audio_tool = selected.get("audio") # Audio generation service endpoint | |
| # Verify that all required tool endpoints are present, raise exception if any is missing | |
| if not tool_setup or not image_tool or not audio_tool: | |
| raise Exception("Selected host is missing required tool endpoints.") | |
| # Return the three tool endpoints as a tuple | |
| return tool_setup, image_tool, audio_tool | |
| # Initialize the tools by selecting an available host and retrieving its endpoints | |
| tool_setup, image_tool, audio_tool = initialize_tools() |