55 lines
2.9 KiB
Python
55 lines
2.9 KiB
Python
"""Small authenticated client for an external, already authenticated AI worker.
|
|
|
|
This module never launches a model, reads model credentials, or executes commands.
|
|
The worker calls claim(), then start() before invoking its trusted model adapter.
|
|
It renews the lease during execution and calls complete() with the real result.
|
|
A lost lease or ambiguous model exit requires uncertain(), never automatic replay.
|
|
"""
|
|
import json
|
|
from urllib.parse import urlsplit
|
|
from urllib.request import Request, build_opener, HTTPRedirectHandler, ProxyHandler
|
|
from .server import restricted_token
|
|
|
|
|
|
class NoRedirect(HTTPRedirectHandler):
|
|
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
|
return None
|
|
|
|
|
|
class Bridge:
|
|
def __init__(self, url, token_file, worker_id):
|
|
parsed = urlsplit(url)
|
|
if parsed.scheme not in ('http','https') or not parsed.hostname or parsed.username or parsed.password or parsed.query or parsed.fragment or parsed.path not in ('','/'):
|
|
raise ValueError('expected server origin')
|
|
self.url = url.rstrip('/')
|
|
self.token = restricted_token(token_file)
|
|
self.worker_id = worker_id
|
|
self.opener = build_opener(ProxyHandler({}), NoRedirect())
|
|
|
|
def request(self, path, data):
|
|
if not path.startswith('/api/') or '?' in path:
|
|
raise ValueError('invalid API path')
|
|
request = Request(self.url + path, data=json.dumps(data).encode('utf-8'), headers={
|
|
'Authorization':'Bearer ' + self.token, 'Content-Type':'application/json', 'X-CSRF-Token':'1'})
|
|
with self.opener.open(request, timeout=20) as response:
|
|
return json.load(response)
|
|
|
|
def claim(self, session_id, lease_seconds=300):
|
|
return self.request('/api/worker/claim', dict(worker_id=self.worker_id, session_id=session_id, lease_seconds=lease_seconds))['message']
|
|
|
|
def start(self, message):
|
|
return self.request('/api/worker/start', dict(worker_id=self.worker_id, session_id=message['session_id'],
|
|
message_id=message['id'], lease_id=message['lease_id']))
|
|
|
|
def renew(self, message, lease_seconds=300):
|
|
return self.request('/api/worker/renew', dict(worker_id=self.worker_id, session_id=message['session_id'],
|
|
message_id=message['id'], lease_id=message['lease_id'], lease_seconds=lease_seconds))
|
|
|
|
def uncertain(self, message):
|
|
return self.request('/api/worker/uncertain', dict(worker_id=self.worker_id, session_id=message['session_id'],
|
|
message_id=message['id'], lease_id=message['lease_id']))
|
|
|
|
def complete(self, message, text, thread_id=None, outcome='completed'):
|
|
return self.request('/api/worker/complete', dict(worker_id=self.worker_id, session_id=message['session_id'],
|
|
message_id=message['id'], lease_id=message['lease_id'], text=text, thread_id=thread_id, outcome=outcome))
|