fromtypingimportList,OptionalfrompydanticimportBaseModelclassChatMessage(BaseModel):role:strcontent:strclassChatCompletionRequest(BaseModel):model:str="mock-gpt-model"messages:List[ChatMessage]max_tokens:Optional[int]=512temperature:Optional[float]=0.1stream:Optional[bool]=False

PyDantic 模型代表了来自客户端的请求，旨在复制 API 参考。为了简洁起见，这个模型没有实现完整的规范，而是实现了使其正常工作的基本部分。如果你缺少一个属于API 规范的参数（比如top_p），你只需将其添加到模型中。

ChatCompletionRequest模型化了 OpenAI 在请求中使用的参数。聊天 API 规范要求指定一个ChatMessage列表（类似于聊天历史，通常由客户端负责保持并在每次请求时反馈）。每个聊天消息有一个role属性（通常是system、assistant或user）和一个content属性，包含实际的消息文本。

接下来，我们将编写 FastAPI 聊天完成端点：

importtimefromfastapiimportFastAPI app=FastAPI(title="OpenAI-compatible API")@app.post("/chat/completions")asyncdefchat_completions(request:ChatCompletionRequest):ifrequest.messagesandrequest.messages[0].role=='user':resp_content="As a mock AI Assitant, I can only echo your last message:"+request.messages[-1].contentelse:resp_content="As a mock AI Assitant, I can only echo your last message, but there were no messages!"return{"id":"1337","object":"chat.completion","created":time.time(),"model":request.model,"choices":[{"message":ChatMessage(role="assistant",content=resp_content)}]}

就这么简单。

测试我们的实现

假设这两个代码块位于名为main.py的文件中，我们将在选择的环境中安装两个 Python 库（最好创建一个新的环境）：pip install fastapi openai，并通过终端启动服务器：

uvicorn main:app

使用另一个终端（或在后台启动服务器），我们将打开一个 Python 控制台，并复制粘贴以下代码，这些代码直接来自OpenAI 的 Python 客户端参考：

fromopenaiimportOpenAI# init client and connect to localhost serverclient=OpenAI(api_key="fake-api-key",base_url="http://localhost:8000"# change the default port if needed)# call APIchat_completion=client.chat.completions.create(messages=[{"role":"user","content":"Say this is a test",}],model="gpt-1337-turbo-pro-max",)# print the top "choice"print(chat_completion.choices[0].message.content)

如果你一切操作正确，服务器的响应应该会正确打印。检查chat_completion对象也很有价值，可以确认所有相关属性是否与我们服务器发送的一致。你应该能看到类似下面的内容：

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/b41ac9b9749709c7180ac0bf9d59a281.png

由作者编写的代码，使用Carbon格式化

升级——支持流式传输

由于 LLM 生成通常较慢（计算开销大），将生成的内容流式传输回客户端是值得的，这样用户可以在生成过程中看到响应，而不必等到它完全生成完毕。如果你记得，我们给ChatCompletionRequest添加了一个布尔型的stream属性——这让客户端可以请求将数据流式传输回去，而不是一次性发送。

这使得事情变得稍微复杂了一些。我们将创建一个生成器函数来包装我们的模拟响应（在实际场景中，我们将需要一个连接到 LLM 生成的生成器）

importasyncioimportjsonasyncdef_resp_async_generator(text_resp:str):# let's pretend every word is a token and return it over timetokens=text_resp.split(" ")fori,tokeninenumerate(tokens):chunk={"id":i,"object":"chat.completion.chunk","created":time.time(),"model":"blah","choices":[{"delta":{"content":token+" "}}],}yieldf"data:{json.dumps(chunk)}\n\n"awaitasyncio.sleep(1)yield"data: [DONE]\n\n"

现在，我们将修改原始的端点，当stream==True时返回一个 StreamingResponse

importtimefromstarlette.responsesimportStreamingResponse app=FastAPI(title="OpenAI-compatible API")@app.post("/chat/completions")asyncdefchat_completions(request:ChatCompletionRequest):ifrequest.messages:resp_content="As a mock AI Assitant, I can only echo your last message:"+request.messages[-1].contentelse:resp_content="As a mock AI Assitant, I can only echo your last message, but there wasn't one!"ifrequest.stream:returnStreamingResponse(_resp_async_generator(resp_content),media_type="application/x-ndjson")return{"id":"1337","object":"chat.completion","created":time.time(),"model":request.model,"choices":[{"message":ChatMessage(role="assistant",content=resp_content)}]}

测试流式传输实现

重启 uvicorn 服务器后，我们将打开 Python 控制台并输入这段代码（同样是来自 OpenAI 的文档）

fromopenaiimportOpenAI# init client and connect to localhost serverclient=OpenAI(api_key="fake-api-key",base_url="http://localhost:8000"# change the default port if needed)stream=client.chat.completions.create(model="mock-gpt-model",messages=[{"role":"user","content":"Say this is a test"}],stream=True,)forchunkinstream:print(chunk.choices[0].delta.contentor"")