"""私有嵌入接口：Cloudreve 与 DSH 共享单模型、推理锁和请求上限。""" import hmac import ipaddress from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import json import os from pathlib import Path import threading from engine import Engine TOKEN = Path(os.environ['AUTH_FILE']).read_text().strip() ENGINE = Engine(os.environ.get('MODEL_DIR', '/models')) LOCK = threading.Lock() ADMISSION = threading.BoundedSemaphore(8) METRICS = {'requests': 0, 'texts': 0, 'windows': 0, 'errors': 0} MODEL_NAME = 'bge-small-zh-v1.5-rknn' DSH_MODEL_NAME = MODEL_NAME + '-symmetric-v1' class Handler(BaseHTTPRequestHandler): protocol_version = 'HTTP/1.0' def setup(self): super().setup() self.connection.settimeout(10 if self.server.dsh_compatible else 60) def log_message(self, *args): # 不记录正文、密钥、文件名和向量。 pass def respond(self, status, data): body = json.dumps(data, ensure_ascii=False).encode() self.send_response(status) self.send_header('Content-Type', 'application/json; charset=utf-8') self.send_header('Content-Length', str(len(body))) self.end_headers() self.wfile.write(body) def authenticated(self): return hmac.compare_digest(self.headers.get('Authorization', ''), 'Bearer ' + self.server.token) def do_GET(self): if self.path == '/health': self.respond(200, {'status': 'ok', 'backend': 'rknn', 'dimensions': 512}) elif self.path == '/metrics' and self.authenticated(): self.respond(200, METRICS.copy()) else: self.respond(404, {'error': '接口不存在'}) def do_POST(self): if self.path != '/v1/embeddings': return self.respond(404, {'error': '接口不存在'}) if not self.authenticated(): return self.respond(401, {'error': '需要认证'}) try: length = int(self.headers.get('Content-Length', '0')) if length <= 0 or length > 1024 * 1024: return self.respond(413, {'error': '请求大小超限'}) payload = json.loads(self.rfile.read(length)) if not isinstance(payload, dict): raise ValueError('请求必须为对象') inputs = payload.get('input') if isinstance(inputs, str): inputs = [inputs] if not isinstance(inputs, list) or not 1 <= len(inputs) <= int(os.environ.get('MAX_BATCH', '10')): raise ValueError('批量数量超限') if any(not isinstance(text, str) or not text.strip() or len(text) > int(os.environ.get('MAX_CHARS', '16384')) for text in inputs): raise ValueError('文本为空或长度超限') kind = payload.get('input_type', 'auto') if kind not in ['auto', 'query', 'document']: raise ValueError('输入类型无效') if self.server.dsh_compatible: if payload.get('model') != DSH_MODEL_NAME or payload.get('dimensions', 512) != 512: raise ValueError('模型或维度不匹配') # 原版插件不区分类型：查询与记忆统一使用无指令的对称编码，禁止猜测正文。 kind = 'document' if any(int(path.read_text()) >= 75000 for path in Path('/sys/class/thermal').glob('thermal_zone*/temp')): return self.respond(503, {'error': '温度保护，请稍后重试'}) if not LOCK.acquire(timeout=5 if self.server.dsh_compatible else 15): return self.respond(429, {'error': '推理队列繁忙'}) try: data = [] total = 0 for i, text in enumerate(inputs): vector, tokens, windows = ENGINE.embed(text, kind) data.append({'object': 'embedding', 'index': i, 'embedding': vector}) total += tokens METRICS['windows'] += windows METRICS['requests'] += 1 METRICS['texts'] += len(inputs) self.respond(200, {'object': 'list', 'model': DSH_MODEL_NAME if self.server.dsh_compatible else MODEL_NAME, 'data': data, 'usage': {'total_tokens': total}}) finally: LOCK.release() except (ValueError, json.JSONDecodeError): METRICS['errors'] += 1 self.respond(400, {'error': '输入格式、类型、模型、维度或窗口限制不满足'}) except (BrokenPipeError, ConnectionResetError, TimeoutError): METRICS['errors'] += 1 except Exception: METRICS['errors'] += 1 self.respond(500, {'error': '本地推理失败'}) class Server(ThreadingHTTPServer): daemon_threads = True request_queue_size = 8 def __init__(self, address, token, dsh_compatible=False, allowed_networks=None): self.token = token self.dsh_compatible = dsh_compatible self.allowed_networks = allowed_networks super().__init__(address, Handler) def verify_request(self, request, address): if self.allowed_networks is None: return True remote = ipaddress.ip_address(address[0]) return any(remote in network for network in self.allowed_networks) def process_request(self, request, address): if not ADMISSION.acquire(blocking=False): self.shutdown_request(request) return try: super().process_request(request, address) except Exception: ADMISSION.release() raise def process_request_thread(self, request, address): try: super().process_request_thread(request, address) finally: ADMISSION.release() def main(): servers = [] try: servers.append(Server((os.environ.get('BIND_ADDRESS', '127.0.0.1'), int(os.environ.get('BIND_PORT', '8090'))), TOKEN)) dsh_address = os.environ.get('DSH_BIND_ADDRESS') if dsh_address: address = ipaddress.ip_address(dsh_address) # 新接口只允许具体的RFC1918内网地址，拒绝通配／公网／回环监听。 private_ranges = [ipaddress.ip_network(value) for value in ['10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16']] if not any(address in network for network in private_ranges): raise ValueError('DSH监听必须使用指定IPv4内网地址') allowed = [ipaddress.ip_network(value.strip()) for value in os.environ['DSH_ALLOWED_NETWORKS'].split(',') if value.strip()] if not allowed or any(network.prefixlen != 32 or not any(network.subnet_of(private) for private in private_ranges) for network in allowed): raise ValueError('DSH来源白名单必须为具体IPv4内网主机') token = Path(os.environ['DSH_AUTH_FILE']).read_text().strip() if not token or hmac.compare_digest(token, TOKEN): raise ValueError('DSH必须使用独立非空认证密钥') servers.append(Server((dsh_address, int(os.environ['DSH_BIND_PORT'])), token, True, allowed)) for server in servers[1:]: threading.Thread(target=server.serve_forever, kwargs={'poll_interval': 1}, daemon=True).start() print('本地NPU嵌入接口已启动；指定监听、独立认证、共享推理。', flush=True) servers[0].serve_forever(poll_interval=1) finally: for server in servers: server.server_close() ENGINE.release() if __name__ == '__main__': main()