MathApp/tools/loadtest/fake_upstream.py

"""
Fake GapGPT + Qdrant for load testing (embeddings, streaming chat, Qdrant points/query). Local only.

Timing profile is calibrated from the last 54 real assistant answers on the server
(Message.latency_ms, gpt-4.1-mini via GapGPT): median 5.2 s, p90 13.1 s, max 30 s.
Lognormal(median=FAKE_MEDIAN, sigma=FAKE_SIGMA), capped at FAKE_CAP.

FAKE_MODE=buffered  -> wait the whole latency, then emit all chunks (what GapGPT was
                       observed to do: first chunk arrives with the last).
FAKE_MODE=stream    -> first token after 20% of latency, rest spread over the remainder.
"""
import json
import math
import os
import random
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

MEDIAN = float(os.environ.get('FAKE_MEDIAN', '5.2'))
SIGMA = float(os.environ.get('FAKE_SIGMA', '0.715'))
CAP = float(os.environ.get('FAKE_CAP', '30'))
MODE = os.environ.get('FAKE_MODE', 'buffered')
EMB_DELAY = float(os.environ.get('FAKE_EMB_DELAY', '0.35'))
QDRANT_DELAY = float(os.environ.get('FAKE_QDRANT_DELAY', '0.08'))
PORT = int(os.environ.get('FAKE_PORT', '8901'))

ANSWER = ('برای حلِ این مسئله اول داده‌ها را می‌نویسیم. '
          'بعد قدم‌به‌قدم جلو می‌رویم و هر مرحله را وارسی می‌کنیم. ') * 14 + '✅ جواب: ۱۲'

active = 0
peak = 0
lock = threading.Lock()


class H(BaseHTTPRequestHandler):
    protocol_version = 'HTTP/1.1'

    def log_message(self, *a):
        pass

    def _json(self, obj):
        body = json.dumps(obj).encode()
        self.send_response(200)
        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):
        self._json({'active': active, 'peak': peak})

    def do_POST(self):
        global active, peak
        n = int(self.headers.get('Content-Length') or 0)
        req = json.loads(self.rfile.read(n) or b'{}')
        if '/points/query' in self.path:
            time.sleep(QDRANT_DELAY)
            pts = [{'id': i + 1, 'version': 1, 'score': 0.5 - i * 0.01,
                    'payload': {'lesson_title': 'درس نمونه', 'problem': 'مسئله‌ی نمونه ' * 30,
                                'generalSolution': 'روشِ کلی ' * 40, 'solution': 'راه‌حل ' * 60}}
                   for i in range(int(req.get('limit') or 4))]
            self._json({'result': {'points': pts}, 'status': 'ok', 'time': 0.001})
            return
        if self.path.endswith('/embeddings'):
            time.sleep(EMB_DELAY)
            vec = [random.uniform(-1, 1) for _ in range(1536)]
            norm = math.sqrt(sum(v * v for v in vec))
            self._json({'object': 'list', 'model': req.get('model'),
                        'data': [{'object': 'embedding', 'index': 0, 'embedding': [v / norm for v in vec]}],
                        'usage': {'prompt_tokens': 20, 'total_tokens': 20}})
            return
        with lock:
            active += 1
            peak = max(peak, active)
        try:
            latency = min(CAP, MEDIAN * math.exp(random.gauss(0, SIGMA)))
            words = ANSWER.split(' ')
            pieces = [' '.join(words[i:i + 5]) + ' ' for i in range(0, len(words), 5)]
            if not req.get('stream'):
                time.sleep(latency)
                self._json({'id': 'x', 'object': 'chat.completion', 'created': int(time.time()),
                            'model': req.get('model'),
                            'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': ANSWER},
                                         'finish_reason': 'stop'}],
                            'usage': {'prompt_tokens': 1400, 'completion_tokens': 400, 'total_tokens': 1800}})
                return
            self.send_response(200)
            self.send_header('Content-Type', 'text/event-stream')
            self.send_header('Transfer-Encoding', 'chunked')
            self.end_headers()

            def send(obj):
                data = f'data: {json.dumps(obj, ensure_ascii=False)}\n\n'.encode()
                self.wfile.write(f'{len(data):x}\r\n'.encode() + data + b'\r\n')
                self.wfile.flush()

            def chunk(text, finish=None):
                return {'id': 'x', 'object': 'chat.completion.chunk', 'created': int(time.time()),
                        'model': req.get('model'),
                        'choices': [{'index': 0, 'delta': {'content': text} if text else {},
                                     'finish_reason': finish}]}

            if MODE == 'buffered':
                time.sleep(latency)
                gap = 0.0
            else:
                time.sleep(latency * 0.2)
                gap = latency * 0.8 / len(pieces)
            for p in pieces:
                send(chunk(p))
                if gap:
                    time.sleep(gap)
            send(chunk('', 'stop'))
            send({'id': 'x', 'object': 'chat.completion.chunk', 'created': int(time.time()),
                  'model': req.get('model'), 'choices': [],
                  'usage': {'prompt_tokens': 1400, 'completion_tokens': 400, 'total_tokens': 1800}})
            last = b'data: [DONE]\n\n'
            self.wfile.write(f'{len(last):x}\r\n'.encode() + last + b'\r\n0\r\n\r\n')
            self.wfile.flush()
        finally:
            with lock:
                active -= 1


if __name__ == '__main__':
    ThreadingHTTPServer.daemon_threads = True
    ThreadingHTTPServer.request_queue_size = 256
    print(f'fake llm on :{PORT} mode={MODE} median={MEDIAN}', flush=True)
    ThreadingHTTPServer(('127.0.0.1', PORT), H).serve_forever()