跳至内容

将 Instructor 与 SQLModel 集成

SQLModel 是一个旨在使用 Python 对象与 SQL 数据库交互的 Python 库。

SQLModel 基于 PydanticSQLAlchemy,由也开发了 FastAPItiangolo 创建。

因此,你可以期待这些库之间的无缝集成,减少代码重复并改善你的开发者体验。

示例:将 Instructor 的响应直接添加到你的数据库

定义模型

首先,我们将定义一个模型,它将作为数据库的表以及 Instructor 输出的结构。

模型定义

你需要让你的模型继承 SQLModelinstructor.OpenAISchema,以便它们能与 SQLModel 一起工作。

from typing import Optional
from uuid import UUID, uuid4
from pydantic.json_schema import SkipJsonSchema
from sqlmodel import Field, SQLModel
import instructor


class Hero(SQLModel, instructor.OpenAISchema, table=True):
    id: SkipJsonSchema[UUID] = Field(default_factory=lambda: uuid4(), primary_key=True)
    name: str
    secret_name: str
    age: Optional[int] = None

使用 SkipJsonSchema 的重要性

注意在 id 字段上使用了 SkipJsonSchema

这将从发送到 LLM 的 JSON schema 中移除该字段,因此它不会尝试生成 UUID。

当 Instructor 解包响应并加载到 Hero 模型时,它将使用 default_factory 自动生成一个 UUID。

这种方法在 LLM 生成过程中节省了 token,更重要的是可以避免 LLM 生成不正确 UUID 格式时可能发生的错误。

发送到 LLM 的最终 JSON schema 将看起来像

{
    "properties": {
        "name": {
            "title": "Name",
            "type": "string"
        },
        "secret_name": {
            "title": "Secret Name",
            "type": "string"
        },
        "age": {
            "anyOf": [
                {
                    "type": "integer"
                },
                {
                    "type": "null"
                }
            ],
            "default": null,
            "title": "Age"
        }
    },
    "required": [
        "name",
        "secret_name"
    ],
    "title": "Hero",
    "type": "object"
}

生成一条记录

create_hero 函数将向 OpenAI 查询一条 Hero 记录

import instructor
from openai import OpenAI


client = instructor.from_openai(OpenAI())


def create_hero() -> Hero:
    return client.chat.completions.create(
        model="gpt-3.5-turbo",
        response_model=Hero,
        messages=[
            {"role": "user", "content": "Make a new superhero"},
        ],
    )

将响应插入数据库

engine = create_engine("sqlite:///database.db")
SQLModel.metadata.create_all(engine)

hero = create_hero()

# The Raw Response from the LLM will not have an id due to the SkipJsonSchema
print(hero._raw_response.choices[0].message.content)
#> {'name': 'Superman', 'secret_name': 'Clark Kent', 'age': 30}

# The model_dump() method will include the generated id as it has been loaded as a Hero object
print(hero.model_dump())
#> {'name': 'Superman', 'secret_name': 'Clark Kent', 'age': 30, 'id': UUID('1234-5678-...')}

with Session(engine) as session:
    session.add(hero)
    session.commit()

Image of hero record in the database

瞧!你现在可以将相同的模型用于数据库和 Instructor,使它们无缝协作!还可以查阅FastAPI 指南,了解如何在 API 中使用这些模型。