使用FastAPI和langchain做本地大模型的API

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

FastAPI 使用 Pydantic 进行数据验证，并使用 Starlette 作为工具，使其与 Flask 相比快得惊人，具有与 Node 或 Go 中的高速 Web APIs 相同的性能。

Starlette + Uvicorn 提供异步请求能力，这是 Flask 所缺乏的。”

本文讲述了如何使用FastAPI和langchain框架，包装本地大模型llama3.1，实现语言翻译功能的API。  

相对于Flask，使用FastAPI做接口要简便得多。

安装依赖

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)

验证API

显然，上述API的地址为：http://127.0.0.1:5001/trans。可以使用多种方法验证API。

1. 使用第三方工具

下图使用ApiFox来验证接口。

2. 使用flasgger生成的API

使用浏览器打开地址：http://127.0.0.1:5001/docs/，依图示对接口进行测试。

文章转自微信公众号@AI很有趣