90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
"""Per-user settings, re-read for new evaluations without restarting active chats."""
|
|
import base64
|
|
import copy
|
|
import ctypes
|
|
import json
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
DEFAULTS = json.loads(Path(__file__).with_name('jev-settings-defaults.json').read_text(encoding='utf-8'))
|
|
_condition = threading.Condition()
|
|
_active = 0
|
|
|
|
|
|
def home():
|
|
if '--home' in sys.argv:
|
|
return Path(sys.argv[sys.argv.index('--home') + 1])
|
|
return Path(os.environ.get('AITURK_CODEX_BRIDGE_HOME') or Path(__file__).parent)
|
|
|
|
|
|
def read():
|
|
try:
|
|
data = json.loads((home() / 'jev-settings.json').read_text(encoding='utf-8'))
|
|
except FileNotFoundError:
|
|
return copy.deepcopy(DEFAULTS), None
|
|
settings = data['settings']
|
|
# A corrupted file must not silently change routing or contact another host.
|
|
from urllib.parse import urlsplit
|
|
url = urlsplit(settings['endpoint'])
|
|
if url.scheme != 'https' or not url.netloc or url.username or url.password:
|
|
raise ValueError('Invalid Jev endpoint')
|
|
for name in ('ret', 'sart'):
|
|
n = settings['thresholds'][name]
|
|
if not isinstance(n, (int, float)) or not 0 <= n <= 1 or not settings['questions'][name]:
|
|
raise ValueError('Invalid Jev rule')
|
|
for name in ('timeoutSeconds', 'queueWaitSeconds'):
|
|
if not isinstance(settings[name], int) or not 1 <= settings[name] <= 3600:
|
|
raise ValueError('Invalid Jev timeout')
|
|
if not isinstance(settings['maxConcurrent'], int) or not 0 <= settings['maxConcurrent'] <= 512:
|
|
raise ValueError('Invalid Jev concurrency')
|
|
if not settings['fallbackModel'] or not isinstance(settings['rejectionRules'], list) or settings['failureAction'] not in ('pass', 'next'):
|
|
raise ValueError('Invalid Jev routing')
|
|
return settings, data.get('protectedKey')
|
|
|
|
|
|
def key():
|
|
_, encrypted = read()
|
|
if not encrypted:
|
|
return ''
|
|
if os.name != 'nt':
|
|
raise RuntimeError('Windows credential protection required')
|
|
from ctypes import wintypes
|
|
class Blob(ctypes.Structure):
|
|
_fields_ = [('cbData', wintypes.DWORD), ('pbData', ctypes.POINTER(ctypes.c_ubyte))]
|
|
raw = base64.b64decode(encrypted, validate=True)
|
|
buf = (ctypes.c_ubyte * len(raw)).from_buffer_copy(raw)
|
|
source, result = Blob(len(raw), buf), Blob()
|
|
decrypt = ctypes.windll.crypt32.CryptUnprotectData
|
|
decrypt.argtypes = [ctypes.POINTER(Blob), ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, wintypes.DWORD, ctypes.POINTER(Blob)]
|
|
decrypt.restype = wintypes.BOOL
|
|
if not decrypt(ctypes.byref(source), None, None, None, None, 0, ctypes.byref(result)):
|
|
raise ctypes.WinError()
|
|
try:
|
|
return ctypes.string_at(result.pbData, result.cbData).decode('utf-8')
|
|
finally:
|
|
ctypes.windll.kernel32.LocalFree.argtypes = [ctypes.c_void_p]
|
|
ctypes.windll.kernel32.LocalFree(result.pbData)
|
|
|
|
|
|
def acquire(settings):
|
|
global _active
|
|
deadline = time.monotonic() + settings['queueWaitSeconds']
|
|
with _condition:
|
|
while settings['maxConcurrent'] and _active >= settings['maxConcurrent']:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
return False
|
|
_condition.wait(remaining)
|
|
_active += 1
|
|
return True
|
|
|
|
|
|
def release():
|
|
global _active
|
|
with _condition:
|
|
_active -= 1
|
|
_condition.notify_all()
|