File size: 4,715 Bytes
f99ad65
 
 
 
 
d17e7ef
 
 
 
f99ad65
d17e7ef
f99ad65
 
 
d17e7ef
 
 
f99ad65
 
d17e7ef
 
 
 
 
 
 
f99ad65
 
 
d17e7ef
 
 
f99ad65
 
 
 
 
d17e7ef
 
 
 
 
 
f99ad65
 
d17e7ef
 
f99ad65
d17e7ef
 
f99ad65
 
 
d17e7ef
 
 
 
 
 
 
 
 
f99ad65
d17e7ef
 
 
 
f99ad65
d17e7ef
f99ad65
d17e7ef
 
 
 
 
 
f99ad65
 
 
d17e7ef
 
 
 
 
 
 
 
 
 
 
f99ad65
d17e7ef
f99ad65
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
#
# SPDX-FileCopyrightText: Hadad <[email protected]>
# SPDX-License-Identifier: Apache-2.0
#

import asyncio  # Import the asyncio library to handle asynchronous operations and events
import requests  # Import the requests library for HTTP requests and session management
import uuid  # Import the uuid library to generate unique identifiers
import threading  # Import threading to run background timers for delayed operations

from config import LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS  # Import configuration variables that track marked provider keys and their failure attempts

class SessionWithID(requests.Session):
    """
    Custom session class extending requests.Session to add unique session identification 
    and asynchronous cancellation control. This allows tracking individual user sessions 
    and managing cancellation of ongoing HTTP requests asynchronously.
    """
    def __init__(self):
        super().__init__()  # Initialize the base requests.Session class
        self.session_id = str(uuid.uuid4())  
        # Generate and assign a unique string ID for this session instance to identify it uniquely
        self.stop_event = asyncio.Event()    
        # Create an asyncio Event object used to signal when the session should stop or cancel operations
        self.cancel_token = {"cancelled": False}  
        # Dictionary flag to indicate if the current session's operations have been cancelled

def create_session():
    """
    Factory function to create and return a new SessionWithID instance.
    This should be called whenever a new user session starts or a chat session is reset,
    ensuring each session has its own unique ID and cancellation controls.
    """
    return SessionWithID()

def ensure_stop_event(sess):
    """
    Utility function to verify that a given session object has the required asynchronous 
    control attributes: stop_event and cancel_token. If they are missing (e.g., when restoring 
    sessions from storage), this function adds them to maintain consistent session behavior.
    
    Parameters:
    - sess: The session object to check and update.
    """
    if not hasattr(sess, "stop_event"):
        sess.stop_event = asyncio.Event()  
        # Add an asyncio Event to signal stop requests if missing
    if not hasattr(sess, "cancel_token"):
        sess.cancel_token = {"cancelled": False}  
        # Add a cancellation flag dictionary if missing

def marked_item(item, marked, attempts):
    """
    Mark a provider key or host as temporarily problematic after repeated failures to prevent 
    using unreliable providers continuously. This function adds the item to a 'marked' set 
    and increments its failure attempt count. If the failure count reaches 3 or more, a timer 
    is started to automatically unmark the item after 5 minutes (300 seconds), allowing retries.
    
    Parameters:
    - item: The provider key or host identifier to mark as problematic.
    - marked: A set containing currently marked items.
    - attempts: A dictionary tracking the number of failure attempts per item.
    """
    marked.add(item)  
    # Add the item to the set of marked problematic providers
    attempts[item] = attempts.get(item, 0) + 1  
    # Increment the failure attempt count for this item, initializing if necessary
    if attempts[item] >= 3:
        # If the item has failed 3 or more times, schedule removal from marked after 5 minutes
        def remove():
            marked.discard(item)  
            # Remove the item from the marked set to allow retrying
            attempts.pop(item, None)  
            # Remove the attempt count entry for this item to reset its failure state
        threading.Timer(300, remove).start()  
        # Start a background timer that will call remove() after 300 seconds (5 minutes)

def get_model_key(display, MODEL_MAPPING, DEFAULT_MODEL_KEY):
    """
    Translate a human-readable model display name into its internal model key identifier.
    Searches the MODEL_MAPPING dictionary for the key whose value matches the display name.
    Returns the DEFAULT_MODEL_KEY if no matching display name is found.
    
    Parameters:
    - display: The display name of the model as a string.
    - MODEL_MAPPING: Dictionary mapping internal model keys to display names.
    - DEFAULT_MODEL_KEY: The fallback model key to return if no match is found.
    
    Returns:
    - The internal model key string corresponding to the display name.
    """
    # Iterate through the MODEL_MAPPING dictionary items and return the key where the value matches the display name
    return next((k for k, v in MODEL_MAPPING.items() if v == display), DEFAULT_MODEL_KEY)