| """Loopback OpenAI chat adapter. Native Spark tool calls are parsed, not repaired.""" |
| import argparse |
| import copy |
| import json |
| import re |
| import threading |
| import time |
| import uuid |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| from pathlib import Path |
|
|
| ROOT = Path(r'D:\webbrain-training\compass-tiny-v3-spark-x25-20260926') |
| MODEL_ID = 'webbrain-compass-tiny-v3-bf16' |
| PACKAGE = ROOT / 'release-compat-v3' |
|
|
|
|
| def parse_tools(text, tools): |
| """Only lossless native-format/schema decoding; errors stay visible.""" |
| definitions = {t['function']['name']: t['function'] for t in tools or []} |
| calls = [] |
| pattern = re.compile(r'<tool_call>(.*?)</tool_call>', re.S) |
| spans = list(pattern.finditer(text)) |
| if text.count('<tool_call>') != len(spans) or text.count('</tool_call>') != len(spans): |
| raise ValueError('Incomplete native tool call') |
| for span in spans: |
| body = span.group(1) |
| name = body.split('<arg_key>', 1)[0].strip() |
| if name not in definitions: |
| raise ValueError(f'Unknown model-emitted function: {name}') |
| tail = body[len(body.split('<arg_key>', 1)[0]):] |
| arg_pattern = re.compile(r'<arg_key>(.*?)</arg_key><arg_value>(.*?)</arg_value>', re.S) |
| args = {} |
| cursor = 0 |
| schema = definitions[name].get('parameters', {}) |
| for arg in arg_pattern.finditer(tail): |
| if tail[cursor:arg.start()].strip(): |
| raise ValueError('Malformed native argument sequence') |
| key, raw = arg.group(1), arg.group(2) |
| if key in args: |
| raise ValueError('Duplicate native argument key') |
| prop = schema.get('properties', {}).get(key, {}) |
| args[key] = raw if prop.get('type') == 'string' else json.loads(raw) |
| cursor = arg.end() |
| if tail[cursor:].strip() or any(k not in args for k in schema.get('required', [])): |
| raise ValueError('Malformed or incomplete native arguments') |
| calls.append({'id': 'call_' + uuid.uuid4().hex[:16], 'type': 'function', |
| 'function': {'name': name, 'arguments': json.dumps(args, ensure_ascii=False)}}) |
| return calls, pattern.sub('', text).strip() |
|
|
|
|
| def normalize_messages(messages): |
| messages = copy.deepcopy(messages) |
| for message in messages: |
| for tool in message.get('tool_calls', []): |
| arguments = tool['function']['arguments'] |
| if isinstance(arguments, str): |
| tool['function']['arguments'] = json.loads(arguments) |
| if not isinstance(tool['function']['arguments'], dict): |
| raise ValueError('Tool arguments must decode to an object') |
| return messages |
|
|
|
|
| class Bf16Backend: |
| def __init__(self): |
| import os |
| import sys |
| source = Path(__file__).parents[1] / 'spark_x25_17b_webbrain_compass_v3' |
| sys.path.insert(0, str(source)) |
| from train import gpu_guard |
| cfg = json.loads((source / 'train-config.json').read_text()) |
| gpu_guard(cfg) |
| import torch |
| import runtime |
| from transformers import AutoTokenizer |
| torch.set_num_threads(8) |
| self.torch = torch |
| self.tokenizer = AutoTokenizer.from_pretrained(PACKAGE, local_files_only=True) |
| self.model = runtime.load_native(str(PACKAGE), dtype=torch.bfloat16, device_map={'': 0}).eval() |
| self.lock = threading.Lock() |
| uploaded = json.loads((ROOT / 'upload-verified.json').read_text()) |
| self.info = {'id': MODEL_ID, 'precision': 'BF16', 'runtime': 'Transformers/PyTorch CUDA', |
| 'hardware': 'NVIDIA GeForce RTX 5090', 'revision': uploaded['revision'], |
| 'max_context_tokens': 16384, 'thinking_default': False, |
| 'onnx': False, 'webgpu': False, 'cloud_fallback': False} |
|
|
| def generate(self, request): |
| if request.get('model') != MODEL_ID: |
| raise ValueError(f'Use exact model id {MODEL_ID}') |
| tools = request.get('tools') |
| messages = normalize_messages(request['messages']) |
| options = request.get('chat_template_kwargs', {}) |
| if options.get('enable_thinking', False): |
| raise ValueError('This benchmark endpoint serves thinking-disabled Compass') |
| inputs = self.tokenizer.apply_chat_template(messages, tools=tools, enable_thinking=False, |
| add_generation_prompt=True, return_tensors='pt', return_dict=True).to('cuda') |
| tokens = int(inputs.input_ids.shape[1]) |
| maximum = int(request.get('max_completion_tokens', request.get('max_tokens', 512))) |
| if maximum < 1 or maximum > 4096 or tokens + maximum > self.info['max_context_tokens']: |
| raise ValueError('Requested input/output exceeds the deployed 16K context limit') |
| temperature = float(request.get('temperature', 0)) |
| if temperature < 0: |
| raise ValueError('Negative temperature') |
| kwargs = {'max_new_tokens': maximum, 'do_sample': temperature > 0, |
| 'pad_token_id': self.tokenizer.pad_token_id, 'use_cache': True} |
| if temperature > 0: |
| kwargs.update(temperature=temperature, top_p=float(request.get('top_p', 1.0))) |
| with self.lock, self.torch.no_grad(): |
| output = self.model.generate(**inputs, **kwargs) |
| tail = output[0, tokens:] |
| raw = self.tokenizer.decode(tail, skip_special_tokens=True) |
| finish = 'length' if len(tail) == maximum else 'stop' |
| stop = request.get('stop') or [] |
| if isinstance(stop, str): |
| stop = [stop] |
| for marker in stop: |
| if marker in raw: |
| raw = raw.split(marker, 1)[0] |
| finish = 'stop' |
| parser_error = None |
| try: |
| calls, content = parse_tools(raw, tools) |
| except (ValueError, KeyError, TypeError, json.JSONDecodeError) as exc: |
| calls, content, parser_error = [], raw, str(exc) |
| message = {'role': 'assistant', 'content': content or None} |
| if calls: |
| message['tool_calls'] = calls |
| finish = 'tool_calls' |
| result = {'id': 'chatcmpl-' + uuid.uuid4().hex, 'object': 'chat.completion', 'created': int(time.time()), |
| 'model': MODEL_ID, 'choices': [{'index': 0, 'message': message, 'finish_reason': finish}], |
| 'usage': {'prompt_tokens': tokens, 'completion_tokens': len(tail), 'total_tokens': tokens + len(tail)}, |
| 'x_runtime': self.info} |
| if parser_error: |
| result['x_parser_error'] = parser_error |
| return result |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--port', type=int, default=8560) |
| args = parser.parse_args() |
| backend = Bf16Backend() |
|
|
| class Handler(BaseHTTPRequestHandler): |
| def log_message(self, fmt, *args): |
| |
| print(json.dumps({'http': fmt % args, 'time': time.time()}), flush=True) |
|
|
| def send_json(self, status, data): |
| body = json.dumps(data, ensure_ascii=False).encode() |
| self.send_response(status) |
| self.send_header('Content-Type', 'application/json') |
| self.send_header('Content-Length', str(len(body))) |
| self.end_headers() |
| self.wfile.write(body) |
|
|
| def do_GET(self): |
| if self.path == '/v1/models': |
| self.send_json(200, {'object': 'list', 'data': [{**backend.info, 'object': 'model', 'owned_by': 'webbrain-one'}]}) |
| elif self.path == '/health': |
| self.send_json(200, {'ready': True, **backend.info}) |
| else: |
| self.send_json(404, {'error': 'Unknown route'}) |
|
|
| def do_POST(self): |
| if self.path != '/v1/chat/completions': |
| return self.send_json(404, {'error': 'Unknown route'}) |
| try: |
| length = int(self.headers.get('Content-Length', '0')) |
| if not 0 < length <= 4 * 1024 * 1024: |
| raise ValueError('Invalid body length') |
| request = json.loads(self.rfile.read(length)) |
| if request.get('stream', False): |
| raise ValueError('Use stream=false on this benchmark endpoint') |
| result = backend.generate(request) |
| self.send_json(200, result) |
| except (ValueError, KeyError, TypeError) as exc: |
| self.send_json(400, {'error': {'message': str(exc), 'type': 'invalid_request_error'}}) |
| except Exception as exc: |
| self.send_json(500, {'error': {'message': str(exc), 'type': 'runtime_error'}}) |
|
|
| print(json.dumps({'ready': True, 'baseUrl': f'http://127.0.0.1:{args.port}/v1', **backend.info}), flush=True) |
| ThreadingHTTPServer(('127.0.0.1', args.port), Handler).serve_forever() |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
