lingyue-zhibao/python-services/ner-service/app/models/request.py

"""
请求模型定义
"""
from pydantic import BaseModel, Field, ConfigDict
from typing import Optional, List


class PositionInfo(BaseModel):
    """位置信息"""
    model_config = ConfigDict(
        populate_by_name=True,
        # Pydantic v2: 序列化时默认使用 alias
        serialize_by_alias=True
    )
    
    char_start: int = Field(..., alias="charStart", description="字符起始位置")
    char_end: int = Field(..., alias="charEnd", description="字符结束位置")
    line: Optional[int] = Field(None, description="所在行号")
    page: Optional[int] = Field(None, description="所在页码")
    file_id: Optional[str] = Field(None, alias="fileId", description="文件ID")


class EntityInfo(BaseModel):
    """实体信息"""
    model_config = ConfigDict(
        populate_by_name=True,
        serialize_by_alias=True
    )
    
    name: str = Field(..., description="实体名称")
    type: str = Field(..., description="实体类型")
    value: Optional[str] = Field(None, description="实体值")
    position: Optional[PositionInfo] = Field(None, description="位置信息")
    context: Optional[str] = Field(None, description="上下文片段")
    confidence: Optional[float] = Field(None, description="置信度")
    temp_id: Optional[str] = Field(None, alias="tempId", description="临时ID")


class NerRequest(BaseModel):
    """NER 请求"""
    model_config = ConfigDict(populate_by_name=True)
    
    document_id: Optional[str] = Field(None, alias="documentId", description="文档ID")
    text: str = Field(..., description="待提取的文本内容")
    entity_types: Optional[List[str]] = Field(None, alias="entityTypes", description="指定提取的实体类型")
    extract_relations: bool = Field(True, alias="extractRelations", description="是否提取关系")
    user_id: Optional[str] = Field(None, alias="userId", description="用户ID")


class RelationRequest(BaseModel):
    """关系抽取请求"""
    model_config = ConfigDict(populate_by_name=True)
    
    document_id: Optional[str] = Field(None, alias="documentId", description="文档ID")
    text: str = Field(..., description="原始文本内容")
    entities: List[EntityInfo] = Field(..., description="已提取的实体列表")