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

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


class PositionInfo(BaseModel):
    """位置信息"""
    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 Config:
        populate_by_name = True
        # 序列化时使用 alias (charStart, charEnd)
        by_alias = True
        # 确保嵌套对象也使用 alias
        json_encoders = {}


class EntityInfo(BaseModel):
    """实体信息"""
    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 Config:
        populate_by_name = True
        by_alias = True


class NerRequest(BaseModel):
    """NER 请求"""
    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 Config:
        populate_by_name = True


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