使用Pydantic构建聊天完成API的数据模型-易微帮

使用Pydantic构建聊天完成API的数据模型

使用Pydantic构建聊天完成API的数据模型

使用Pydantic构建聊天完成API的数据模型

在构建大型语言模型(LLM)的API时,定义清晰的数据结构是至关重要的。本文将介绍如何使用Pydantic库来定义聊天完成API的数据模型,并提供相关的代码示例。

Pydantic简介

Pydantic是一个强大的数据验证和设置管理库,用于Python。它允许我们使用Python类型注解来定义数据模型,提供了自动的数据验证、序列化和文档生成功能。

基础概念

在我们的聊天完成API中,我们需要定义几个关键的数据结构:

模型信息
聊天消息
API请求和响应
工具和函数调用

让我们逐个看看这些结构的定义。

1. 模型信息

from pydantic import BaseModel, Field
import time
from typing import Optional, List

class ModelCard(BaseModel):
    id: str = ""
    object: str = "model"
    created: int = Field(default_factory=lambda: int(time.time()))
    owned_by: str = "owner"
    root: Optional[str] = None
    parent: Optional[str] = None
    permission: Optional[list] = None

class ModelList(BaseModel):
    object: str = "list"
    data: List[ModelCard] = ["glm-4"]

这里我们定义了ModelCard来存储模型的基本信息,以及ModelList来存储模型列表。

2. 聊天消息

from typing import Literal

class ChatMessage(BaseModel):
    role: Literal["user", "assistant", "system", "tool"]
    content: Optional[str] = None
    function_call: Optional[ChoiceDeltaToolCallFunction] = None
    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None

class DeltaMessage(BaseModel):
    role: Optional[Literal["user", "assistant", "system"]] = None
    content: Optional[str] = None
    function_call: Optional[ChoiceDeltaToolCallFunction] = None
    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None

ChatMessage定义了聊天消息的结构,包括角色、内容和可能的函数调用或工具调用。DeltaMessage用于流式响应中的增量消息。

3. API请求和响应

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.8
    top_p: Optional[float] = 0.8
    max_tokens: Optional[int] = None
    stream: Optional[bool] = False
    tools: Optional[Union[dict, List[dict]]] = None
    tool_choice: Optional[Union[str, dict]] = None
    repetition_penalty: Optional[float] = 1.1

class ChatCompletionResponse(BaseModel):
    model: str
    id: Optional[str] = Field(default_factory=lambda: generate_id('chatcmpl-', 29))
    object: Literal["chat.completion", "chat.completion.chunk"]
    choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]]
    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
    system_fingerprint: Optional[str] = Field(default_factory=lambda: generate_id('fp_', 9))
    usage: Optional[UsageInfo] = None

这些类定义了API的请求和响应结构。ChatCompletionRequest包含了请求的所有参数,而ChatCompletionResponse定义了API的响应格式。

4. 工具和函数调用

class FunctionCall(BaseModel):
    name: Optional[str] = None
    arguments: Optional[str] = None

class ChatCompletionMessageToolCall(BaseModel):
    index: Optional[int] = 0
    id: Optional[str] = None
    function: FunctionCall
    type: Optional[Literal["function"]] = 'function'

高级主题

除了基本的数据模型定义,还有一些高级主题值得探讨:

1. 自定义验证器

Pydantic允许我们定义自定义的验证器,以实现更复杂的验证逻辑。例如,我们可以为ChatCompletionRequest类添加一个验证器,以确保temperature和top_p的值在合理范围内:

from pydantic import validator

class ChatCompletionRequest(BaseModel):
    # ... 其他字段保持不变 ...

    @validator('temperature', 'top_p')
    def check_probability_range(cls, v):
        if v is not None and (v < 0 or v > 1):
            raise ValueError('Must be between 0 and 1')
        return v

2. 嵌套模型

我们的模型中已经使用了嵌套结构,例如ChatCompletionResponse包含了ChatCompletionResponseChoice。这种嵌套可以帮助我们组织复杂的数据结构。

3. 可选字段和默认值

注意到我们的模型中大量使用了Optional和默认值。这允许我们创建灵活的模型,可以处理不同情况下的数据。

4. 使用Literal进行枚举

我们使用Literal来限制某些字段的可能值,例如:

role: Literal["user", "assistant", "system", "tool"]

这确保了role字段只能是这四个值之一。

5. 使用Union处理多种可能的类型

在ChatCompletionResponse中,我们使用Union来处理可能的不同响应类型:

choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]]

这允许choices列表包含两种不同类型的选择。

实际应用示例

让我们看一个更完整的例子,展示如何在实际应用中使用这些模型:

import json
from pydantic import BaseModel, Field
from typing import List, Optional, Literal, Union
import time

# ... (在这里包含之前定义的所有模型) ...

def generate_id(prefix: str, length: int) -> str:
    import random
    import string
    return prefix + ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))

def process_chat_completion(request: ChatCompletionRequest) -> ChatCompletionResponse:
    # 这里应该是实际处理聊天完成请求的逻辑
    # 为了示例,我们只是创建一个模拟的响应
    response = ChatCompletionResponse(
        model=request.model,
        choices=[
            ChatCompletionResponseChoice(
                index=0,
                message=ChatMessage(
                    role="assistant",
                    content="This is a simulated response."
                ),
                finish_reason="stop"
            )
        ],
        usage=UsageInfo(prompt_tokens=10, completion_tokens=5, total_tokens=15)
    )
    return response

# 使用示例
if __name__ == "__main__":
    request = ChatCompletionRequest(
        model="gpt-3.5-turbo",
        messages=[
            ChatMessage(role="system", content="You are a helpful assistant."),
            ChatMessage(role="user", content="What's the weather like today?")
        ],
        temperature=0.7
    )

    response = process_chat_completion(request)

    print("Request:")
    print(json.dumps(request.dict(), indent=2))
    print("\nResponse:")
    print(json.dumps(response.dict(), indent=2))

这个例子展示了如何创建一个请求,处理它,并生成一个响应。在实际应用中,process_chat_completion函数将包含与语言模型交互的逻辑。

性能考虑

当处理大量请求时,Pydantic模型的验证可能会成为性能瓶颈。在这种情况下,可以考虑以下优化:

使用Config类禁用某些验证:

class Config:
    validate_assignment = False
    arbitrary_types_allowed = True

对于不需要验证的数据,可以使用construct()方法而不是常规的初始化:

response = ChatCompletionResponse.construct(**data)

如果确实需要最大性能,可以考虑使用其他更轻量级的方法,如简单的字典或自定义类。

结论

Pydantic为构建复杂的API数据模型提供了强大而灵活的工具。通过定义清晰的数据结构,我们可以提高代码的可读性、可维护性,并减少运行时错误。在构建聊天完成API时,使用Pydantic可以帮助我们处理复杂的请求和响应结构,同时提供自动的数据验证。

使用Pydantic构建聊天完成API的数据模型

使用Pydantic构建聊天完成API的数据模型