用FastAPI做本地大模型的API

本文讲述了如何使用FastAPI`langchain框架，包装本地大模型llama3.1，实现语言翻译功能的API。
相对于`Flask ，使用FastAPI做接口要简便得多。

`FastAPI`可以被看作是把 `Starlette`、`Pydantic`等框架粘合在一起的组合体。

- `FastAPI` 使用 `Pydantic` 进行数据验证，并使用 `Starlette` 作为工具，使其与 `Flask` 相比快得惊人，具有与 `Node` 或 `Go` 中的高速 Web APIs 相同的性能。
- `Starlette` + `Uvicorn` 提供异步请求能力，这是 `Flask` 所缺乏的。

安装依赖

pip install fastapi pydantic typing

定义翻译方法

下面的方法需要两个参数，其中：language是翻译目标语言，text是需要翻译的文本。

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_ollama.llms import OllamaLLM

# 翻译方法
def translate(language,text):
    # 1. 创建提示词模板
    system_template = "Translate the following into {language}:"
    prompt_template = ChatPromptTemplate.from_messages([
        ('system', system_template),
        ('user', '{text}')
    ])

    # 2. 创建本地大模型
    model = OllamaLLM(model="llama3.1")

    # 3. 创建解析器
    parser = StrOutputParser()

    # 4. 创建链
    chain = prompt_template | model | parser

    #5. 调用链
    result = chain.invoke({"language": language,"text":text})

    return result

定义接口输入和返回数据格式

从这里就可以看到代码比`Flask`简洁得多，这才是代码应有的样子😆

# 导入FastAPI和Pydantic
from fastapi import FastAPI
from pydantic import BaseModel, Field


# 定义一个Pydantic模型来校验输入的JSON数据
class query_model(BaseModel):
    lang: str = Field(min_length=2, max_length=20, description="语言名称" )
    text: str = Field(min_length=2, max_length=500, description="待翻译的文本" )

from enum import Enum
# 操作结果枚举
class code_enum(str,Enum):
    OK = 'ok'
    ERR = 'error'

# API返回的消息体
class response_model(BaseModel):
    code: code_enum = Field(description="操作结果" )
    desc: str = Field(description="具体内容" )

定义接口和路由

使用`@app`可以指定接口路由、返回的消息体格式，接口方法内部的注释可以被渲染生成playground。

# 创建一个FastAPI实例
app = FastAPI()

# 创建一个处理POST请求的端点
@app.post("/trans/", response_model=response_model)
async def translate_api(query: query_model):
    """
    翻译文本。

    参数:
    - Query: 翻译请求内容。

    返回:
    - Query: 测试
    """
    
    try:
        r = translate(query.lang.strip(),query.text.strip())
        return response_model(code=code_enum.OK,desc=r)
    except Exception as e:
        return response_model(code=code_enum.ERR,desc=str(e))

启动API

import uvicorn

if __name__ == '__main__':
    # 交互式API文档地址：
    # http://127.0.0.1:5001/docs/ 
    # http://127.0.0.1:5001/redoc/
    
    uvicorn.run(app, host="0.0.0.0", port=5001)