# # SPDX-FileCopyrightText: Hadad # SPDX-License-Identifier: Apache-2.0 # import random # Import the random module to enable functions for random selection and shuffling, which are essential for distributing workload evenly across multiple hosts and avoiding selection bias from datetime import datetime, timedelta # Import datetime and timedelta classes to handle precise time calculations, which are crucial for managing host availability windows and expiration of busy statuses from config import auth # Import the 'auth' configuration list, which contains dictionaries representing hosts with their respective credentials, endpoints, and tokens necessary for tool integration from src.utils.helper import busy, mark # Import 'busy', a dictionary tracking the busy state and expiration timestamps of hosts, and 'mark', a utility function to update this dictionary to indicate when a host has been assigned and for how long # Define a function to get available tools def initialize_tools(): """ This function is designed to robustly and perpetually search for an available host from a predefined list of hosts, ensuring that the selected host is fully configured with all required service endpoints and tokens before returning. It operates in an infinite loop to guarantee that it never returns a None value or raises an exception, thereby providing a reliable mechanism for host selection even in dynamic environments where host availability fluctuates frequently. The detailed workflow is as follows: - Obtain the current UTC timestamp to accurately evaluate the expiration of hosts' busy periods. - Filter the list of configured hosts to include only those that are either not currently marked as busy or whose busy period has elapsed, thus ensuring only eligible hosts are considered for assignment. - If no hosts meet the availability criteria, identify and remove expired busy entries from the tracking dictionary to free up hosts and retry immediately. - Randomly shuffle the filtered list of available hosts to prevent selection bias and promote equitable distribution of workload. - Iterate through the shuffled hosts, verifying that each host possesses all critical configuration elements: the main tool setup endpoint, image generation endpoint, audio generation endpoint, and Pollinations AI API token. - Upon finding a host with complete and valid configurations, mark it as busy for a fixed duration (one hour) to prevent immediate reassignment and potential conflicts. - Return a tuple containing the four key elements from the selected host, enabling downstream components to utilize these endpoints and tokens for their operations. - If no suitable host is found after checking all available hosts, perform a cleanup of expired busy statuses and continue the search without delay. This approach ensures continuous availability of a properly configured host, adapts dynamically to changing host states, and maintains system stability by preventing simultaneous conflicting assignments. Returns: tuple: A four-element tuple consisting of: - tool_setup (str): The endpoint or configuration string for the primary tool setup service. - image_tool (str): The URL endpoint for the image generation service. - audio_tool (str): The URL endpoint for the audio generation service. - poll_token (str): The API token string required for authenticating with Pollinations AI services. """ while True: # Start an endless loop that only breaks when a fully valid and available host is successfully selected and returned now = datetime.utcnow() # Capture the current coordinated universal time to accurately compare against host busy expiration timestamps # Create a filtered list of hosts that are eligible for assignment # Eligibility criteria include hosts not currently marked as busy or those whose busy period has expired, thereby allowing their reuse available = [ item for item in auth # Iterate over each host configuration dictionary in the 'auth' list imported from the configuration module if item["jarvis"] not in busy or busy[item["jarvis"]] <= now # Include the host if it is not present in the busy dictionary or if its busy timestamp is earlier than or equal to the current time, indicating availability ] # Check if the filtered list of available hosts is empty, which means all hosts are currently busy with unexpired busy periods if not available: # Identify all hosts in the busy dictionary whose busy periods have expired by comparing their timestamps to the current time expired_hosts = [host for host in busy if busy[host] <= now] # Generate a list of host identifiers whose busy status has expired and are thus candidates for cleanup for host in expired_hosts: # Iterate over each expired host to perform cleanup operations del busy[host] # Remove the host's busy status entry from the dictionary, effectively marking it as available again for assignment # After cleaning up expired busy statuses, continue the loop to reattempt host selection immediately, ensuring responsiveness and minimal delay in availability checks continue # Randomize the order of the available hosts list to prevent any selection bias and promote fair workload distribution among hosts random.shuffle(available) # Iterate through each host in the shuffled list to find one that has all the required endpoints and tokens properly configured for selected in available: # Extract the main tool setup endpoint from the selected host's configuration dictionary using the key 'done' tool_setup = selected.get("done") # Extract the image generation service endpoint URL using the key 'image' image_tool = selected.get("image") # Extract the audio generation service endpoint URL using the key 'audio' audio_tool = selected.get("audio") # Extract the Pollinations AI API token string using the key 'pol' poll_token = selected.get("pol") # Verify that all four critical configuration values are present and non-empty, ensuring the host is fully capable of handling the required services if tool_setup and image_tool and audio_tool and poll_token: # Mark the selected host as busy for the next hour by invoking the 'mark' function with the host's unique identifier # This prevents the same host from being assigned again immediately, allowing for load balancing and avoiding conflicts mark(selected["jarvis"]) # Return a tuple containing all four endpoints and the token, signaling successful host selection and configuration retrieval return tool_setup, image_tool, audio_tool, poll_token # If any required endpoint or token is missing, skip this host and continue checking the next one in the list # If none of the available hosts had all required endpoints and tokens, perform cleanup of expired busy statuses again before retrying expired_hosts = [host for host in busy if busy[host] <= now] # Identify hosts whose busy times have expired and are eligible to be freed up for host in expired_hosts: del busy[host] # Remove expired busy entries to free up hosts for reassignment # Continue the loop to retry host selection immediately, maintaining continuous operation without delay # Execute the host initialization process by calling the function to obtain an available host's endpoints and tokens tool_setup, image_tool, audio_tool, poll_token = initialize_tools() # This call blocks until a suitable host is found and returns the necessary configuration for downstream tool usage