""" ============================================================================= CMU 17-630 · 前置讲 P2 共享模块：mock_server.py —— 本机的"模拟 LLM API"和它的 HTTP 客户端 ============================================================================= 【这个模块是做什么的】 云端 LLM 的 API 本质上是一个 HTTP 服务：你把一段 JSON（模型名、消息、参数）POST 到一个网址， 它回你一段 JSON（回答、停止原因、用量）。为了不联网、不花钱也能看清楚"线上传的到底是什么"， 这个模块在本机（127.0.0.1）起一个模拟服务器，接口的**形状**模仿 Anthropic 的 Messages API （POST /v1/messages、x-api-key 头、content 块、stop_reason、usage、SSE 流式事件）， 回答由 llm_api.SimLLM 生成。**它不是**任何真实服务的精确复刻，只是同一类东西的教学版。 MockLLMServer 服务器：校验 API key、解析 JSON、按"失败计划"故意返回 429 / 500、支持流式输出 HttpLLM 客户端：实现 llm_api.LLM 接口，可以和别的包装器（重试、缓存、记账）组合； stream() 演示怎样正确地接收流式输出（先按字节攒齐一行，再解码） raw_exchange 用最底层的 socket 发一个手写的 HTTP 请求，返回原始的请求和响应字节 【运行】python3 mock_server.py # 自检：起服务器、发几个请求、关服务器 【依赖】只用 Python 标准库；回答来自同目录的 llm_api.py。服务器只监听 127.0.0.1，**不对外开放、不联网**。 """ import codecs import json import os import socket import threading import time import urllib.error import urllib.request from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from llm_api import (LLM, AuthError, BadRequestError, LLMError, LLMRequest, LLMResponse, LLMTimeout, NotFoundError, RateLimitError, ServerError, SimLLM, Usage, count_request_tokens, count_tokens, tokenizer) KNOWN_MODELS = {"sim-small", "sim-large", "echo"} class MockLLMServer: """ 本机的模拟 LLM API 服务器。 【参数】 api_key ：正确的 API key；请求头 x-api-key 不等于它就返回 401 fail_plan ：失败计划，一个状态码列表，例如 [429, 429, 500]：前三个 /v1/messages 请求依次返回 这三种错误，之后正常。用来演示重试。 retry_after ：返回 429 时 retry-after 头的秒数 stream_delay：每生成一个 token 要多少秒（模拟模型一个 token 一个 token 地生成）。流式时每个 token 一生成就发出去；非流式时等全部生成完再一起发回 chunk_bytes ：流式输出时，服务器把要发的字节**故意**切成多大的块发送（模拟网络把数据切碎） response_delay：非流式请求在回答之前先等多少秒（模拟一个很慢的服务，用来演示客户端超时） 【用法】 with MockLLMServer(api_key="sk-test") as srv: srv.url → "http://127.0.0.1:端口号" 【Python 小知识】实现 __enter__ / __exit__ 的对象可以用 with 语句：进入时启动，离开时（哪怕出了异常）自动关闭。 ThreadingHTTPServer 每个请求开一个线程处理；port=0 表示让操作系统挑一个空闲端口。 """ def __init__(self, api_key="sk-test-0000", fail_plan=(), retry_after=1, stream_delay=0.0, chunk_bytes=0, response_delay=0.0, llm=None): self.api_key, self.fail_plan, self.retry_after = api_key, list(fail_plan), retry_after self.stream_delay, self.chunk_bytes, self.response_delay = stream_delay, chunk_bytes, response_delay self.completed = 0 # 服务器真正"处理完"（生成了回答）的请求数 self.llm = llm or SimLLM() self.requests_seen = 0 self.lock = threading.Lock() self.httpd = ThreadingHTTPServer(("127.0.0.1", 0), self._make_handler()) self.url = f"http://127.0.0.1:{self.httpd.server_address[1]}" self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) def __enter__(self): self.thread.start() return self def __exit__(self, *exc): self.httpd.shutdown() self.httpd.server_close() self.thread.join() def _make_handler(self): """生成处理请求的类（闭包里能访问服务器对象 srv）。""" srv = self class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def log_message(self, *args): # 不在终端里打印每个请求的访问日志 pass def _send_json(self, status, obj, headers=None): body = json.dumps(obj, ensure_ascii=False).encode("utf-8") self.send_response(status) self.send_header("content-type", "application/json") self.send_header("content-length", str(len(body))) for k, v in (headers or {}).items(): self.send_header(k, v) self.end_headers() self.wfile.write(body) def _error(self, status, etype, msg, headers=None): self._send_json(status, {"type": "error", "error": {"type": etype, "message": msg}}, headers) def do_GET(self): if self.path == "/health": self._send_json(200, {"ok": True}) else: self._error(404, "not_found_error", f"没有这个地址：{self.path}") def do_POST(self): length = int(self.headers.get("content-length", 0)) raw = self.rfile.read(length) if self.path != "/v1/messages": return self._error(404, "not_found_error", f"没有这个地址：{self.path}") if self.headers.get("x-api-key") != srv.api_key: return self._error(401, "authentication_error", "API key 缺失或无效") with srv.lock: srv.requests_seen += 1 planned = srv.fail_plan.pop(0) if srv.fail_plan else None if planned == 429: return self._error(429, "rate_limit_error", "请求太多，请稍后重试", {"retry-after": str(srv.retry_after)}) if planned: return self._error(planned, "api_error", "服务器内部错误") try: body = json.loads(raw.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as e: return self._error(400, "invalid_request_error", f"请求体不是合法的 JSON：{e}") if not isinstance(body, dict): # 第一版没有这一行：请求体是 "null" 时下面的 `f not in body` 直接让服务器线程崩溃。 # 合法的 JSON 不等于合法的请求——服务器必须检查类型。 return self._error(400, "invalid_request_error", "请求体必须是一个 JSON 对象") for f in ("model", "max_tokens", "messages"): if f not in body: return self._error(400, "invalid_request_error", f"缺少必填字段 {f}") if body["model"] not in KNOWN_MODELS: return self._error(404, "not_found_error", f"没有这个模型：{body['model']}") req = LLMRequest(messages=body["messages"], system=body.get("system", ""), model=body["model"], max_tokens=body["max_tokens"], temperature=body.get("temperature", 0.0), stop=tuple(body.get("stop_sequences", ())), seed=body.get("seed")) try: if body["model"] == "echo": # 原样复述最后一条用户消息（给流式输出的演示用） text = req.messages[-1]["content"] resp = LLMResponse(text=text, stop_reason="end_turn", model="echo", usage=Usage(count_request_tokens(req), count_tokens(text))) else: resp = srv.llm.complete(req) except BadRequestError as e: return self._error(400, "invalid_request_error", str(e)) with srv.lock: msg_id = f"msg_{srv.requests_seen:06d}" with srv.lock: srv.completed += 1 if body.get("stream"): return self._stream(msg_id, resp) if srv.response_delay: time.sleep(srv.response_delay) if srv.stream_delay: # 非流式也要"生成"同样多的 token，只是生成完了才一次性发回来（公平地比较两种方式） time.sleep(srv.stream_delay * len(tokenizer().encode(resp.text))) try: self._send_json(200, {"id": msg_id, "type": "message", "role": "assistant", "model": resp.model, "content": [{"type": "text", "text": resp.text}], "stop_reason": resp.stop_reason, "usage": {"input_tokens": resp.usage.input_tokens, "output_tokens": resp.usage.output_tokens}}) except (BrokenPipeError, ConnectionResetError): pass # 客户端已经超时断开：回答白生成了，没人收 def _stream(self, msg_id, resp): """ 以 SSE（server-sent events）的格式一块一块地发送回答。 【格式】每个事件是两行 + 一个空行： event: content_block_delta data: {"type": "content_block_delta", "delta": {"type": "text_delta", "text": "Vol"}} （空行） 【思路】把回答按 token 切成小片（每片一个 text_delta 事件），事件之间等 stream_delay 秒。 chunk_bytes > 0 时，再把要发送的**字节**随意切成固定大小的块——真实网络也会这样切， 一个汉字的 3 个字节可能分在两块里。正确的客户端必须能应付这种情况。 """ self.send_response(200) self.send_header("content-type", "text/event-stream") self.send_header("cache-control", "no-cache") self.send_header("connection", "close") self.end_headers() ids = tokenizer().encode(resp.text) events = [("message_start", {"type": "message_start", "message": {"id": msg_id, "model": resp.model}}), ("content_block_start", {"type": "content_block_start", "index": 0})] # 按 token 发送；但要保证每片文字都是完整的字符（不完整的字节先攒着，和下一个 token 一起发） dec = codecs.getincrementaldecoder("utf-8")() for i in ids: piece = dec.decode(tokenizer().vocab[i]) if piece: events.append(("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": piece}})) events += [("content_block_stop", {"type": "content_block_stop", "index": 0}), ("message_delta", {"type": "message_delta", "delta": {"stop_reason": resp.stop_reason}, "usage": {"input_tokens": resp.usage.input_tokens, "output_tokens": resp.usage.output_tokens}}), ("message_stop", {"type": "message_stop"})] for name, data in events: blob = f"event: {name}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n".encode("utf-8") step = srv.chunk_bytes or len(blob) for k in range(0, len(blob), step): self.wfile.write(blob[k:k + step]) self.wfile.flush() if srv.stream_delay and name == "content_block_delta": time.sleep(srv.stream_delay) self.close_connection = True return Handler # --------------------------------------------------------------------------- # 客户端 # --------------------------------------------------------------------------- def redact(secret): """把密钥打码成 "sk-te…00"（只留前 5 位和后 2 位），用于日志和报错信息。【边界】短于 8 位的全部打码。""" if not secret: return "<空>" return "*" * len(secret) if len(secret) < 8 else secret[:5] + "…" + secret[-2:] class HttpLLM(LLM): """ 通过 HTTP 调用 MockLLMServer 的客户端，实现 llm_api.LLM 接口。 【参数】base_url：服务器地址；api_key：API key（None 时从环境变量 MOCK_LLM_API_KEY 读取）； timeout：等待响应的最长秒数；name：模型名 【错误映射】HTTP 状态码 → llm_api 的错误类型：400 → BadRequestError，401 → AuthError，404 → NotFoundError， 429 → RateLimitError（带 retry-after），5xx → ServerError；连不上 → ServerError；超时 → LLMTimeout。 这样外面套的 RetryingLLM 不需要知道底下是 HTTP 还是别的什么。 【为什么从环境变量读 key】key 写在代码里，代码一提交到 git、一发给同事，key 就泄露了。 环境变量（或专门的密钥管理工具）把"代码"和"秘密"分开。 """ def __init__(self, base_url, api_key=None, timeout=10.0, name="sim-small"): self.base_url, self.timeout, self.name = base_url.rstrip("/"), timeout, name self.api_key = api_key if api_key is not None else os.environ.get("MOCK_LLM_API_KEY", "") def _body(self, req, stream=False): body = {"model": req.model, "max_tokens": req.max_tokens, "messages": req.messages} if req.system: body["system"] = req.system if req.temperature: body["temperature"] = req.temperature if req.stop: body["stop_sequences"] = list(req.stop) if req.seed is not None: body["seed"] = req.seed if stream: body["stream"] = True return json.dumps(body, ensure_ascii=False).encode("utf-8") def _open(self, req, stream=False): """发出 POST 请求，返回 urllib 的响应对象；HTTP 错误翻译成我们的错误类型。""" r = urllib.request.Request(self.base_url + "/v1/messages", data=self._body(req, stream), method="POST", headers={"content-type": "application/json", "x-api-key": self.api_key}) try: return urllib.request.urlopen(r, timeout=self.timeout) except urllib.error.HTTPError as e: try: msg = json.loads(e.read().decode("utf-8"))["error"]["message"] except Exception: # 错误体不是我们认识的格式：就用状态码 msg = f"HTTP {e.code}" if e.code == 400: raise BadRequestError(msg) from None if e.code == 401: raise AuthError(msg) from None if e.code == 404: raise NotFoundError(msg) from None if e.code == 429: ra = e.headers.get("retry-after") raise RateLimitError(msg, float(ra) if ra else None) from None if e.code >= 500: raise ServerError(f"{e.code}: {msg}") from None raise LLMError(f"{e.code}: {msg}") from None except socket.timeout: raise LLMTimeout(f"{self.timeout} 秒没有响应") from None except urllib.error.URLError as e: if isinstance(e.reason, socket.timeout): raise LLMTimeout(f"{self.timeout} 秒没有响应") from None raise ServerError(f"连不上服务器：{e.reason}") from None def complete(self, req): t0 = time.perf_counter() with self._open(req) as resp: data = json.loads(resp.read().decode("utf-8")) text = "".join(b["text"] for b in data["content"] if b["type"] == "text") u = data["usage"] return LLMResponse(text=text, stop_reason=data["stop_reason"], model=data["model"], usage=Usage(u["input_tokens"], u["output_tokens"]), latency_s=time.perf_counter() - t0) def stream(self, req, read_bytes=7, decode="buffered", on_text=None): """ 流式接收回答。 【参数】read_bytes：每次从网络上读多少字节（故意读得很小，让切碎的情况一定发生）； decode："buffered"（正确：先按字节攒到一个完整的事件 b"\\n\\n" 再解码）或 "naive"（错误：每读到一块字节就立刻按 UTF-8 解码显示）； on_text：每收到一段文字就调用一次的回调（可选），例如打印到屏幕 【返回】(完整的回答文字, 首个文字到达的时间, 总时间, 用户在屏幕上看到的文字) —— naive 模式下"屏幕上看到的文字"会出现 "�" 【思路（buffered）】SSE 事件以空行（b"\\n\\n"）分隔。按**字节**找分隔符：分隔符本身是 ASCII， 不会出现在多字节字符的中间，所以按它切出来的每一段都是完整的 UTF-8，可以放心解码。 """ t0 = time.perf_counter() ttft, text, shown = None, [], [] with self._open(req, stream=True) as resp: buf = b"" while True: chunk = resp.read(read_bytes) if not chunk: break if decode == "naive": shown.append(chunk.decode("utf-8", errors="replace")) buf += chunk while b"\n\n" in buf: event, buf = buf.split(b"\n\n", 1) lines = event.decode("utf-8").split("\n") data = json.loads(next(x[6:] for x in lines if x.startswith("data: "))) if data["type"] == "content_block_delta": piece = data["delta"]["text"] if ttft is None: ttft = time.perf_counter() - t0 text.append(piece) if decode == "buffered": shown.append(piece) if on_text: on_text(piece) full = "".join(text) if decode == "naive": raw_screen = "".join(shown) # 屏幕上显示的是整段 SSE 流（含 event: / data: 行）；只留下 text 字段里的部分方便对比 pieces = [p.split('"text": "', 1)[1].split('"', 1)[0] for p in raw_screen.split("\n\n") if '"text_delta"' in p and '"text": "' in p] screen = "".join(pieces) else: screen = "".join(shown) return full, ttft, time.perf_counter() - t0, screen def raw_exchange(base_url, body, headers): """ 用最底层的 socket 发一个手写的 HTTP/1.1 请求，返回 (请求的原始字节, 响应的原始字节)。 【思路】HTTP 只是"按约定格式写的文字"：一行请求行（方法 路径 版本），若干行头部（名字: 值）， 一个空行，然后是请求体。这里一个字节一个字节地把它拼出来发出去，让你看到 urllib、requests、 官方 SDK 在底下替你做了什么。 【边界】只处理带 content-length 的响应（MockLLMServer 的非流式响应都有）。 """ host, port = base_url.replace("http://", "").split(":") payload = json.dumps(body, ensure_ascii=False).encode("utf-8") head = [f"POST /v1/messages HTTP/1.1", f"host: {host}:{port}", f"content-length: {len(payload)}", "connection: close"] + [f"{k}: {v}" for k, v in headers.items()] request = ("\r\n".join(head) + "\r\n\r\n").encode("utf-8") + payload with socket.create_connection((host, int(port)), timeout=5) as s: s.sendall(request) chunks = [] while True: c = s.recv(4096) if not c: break chunks.append(c) return request, b"".join(chunks) if __name__ == "__main__": with MockLLMServer(api_key="sk-test-1234", fail_plan=[429]) as srv: llm = HttpLLM(srv.url, api_key="sk-test-1234") try: llm.chat("What is the capital of volel?") except RateLimitError as e: print(f" 第一次：429，retry_after = {e.retry_after}") r = llm.chat("What is the capital of volel?") print(f" 第二次：{r.text!r}（{r.stop_reason}，{r.usage}）") try: HttpLLM(srv.url, api_key="wrong").chat("hi") except AuthError as e: print(f" 错误的 key：AuthError（{e}）") full, ttft, total, screen = HttpLLM(srv.url, api_key="sk-test-1234", name="echo").stream( LLMRequest(messages=[{"role": "user", "content": "流式输出也要能处理中文。"}], model="echo")) print(f" 流式：{full!r}") print(f" key 打码：{redact('sk-test-1234')}")