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

"""
响应模型定义
"""
from pydantic import BaseModel, Field, ConfigDict
from typing import Optional, List
from .request import EntityInfo, PositionInfo


class RelationInfo(BaseModel):
    """关系信息"""
    model_config = ConfigDict(
        populate_by_name=True,
        serialize_by_alias=True
    )
    
    from_entity: str = Field(..., alias="fromEntity", description="源实体名称")
    from_entity_id: Optional[str] = Field(None, alias="fromEntityId", description="源实体临时ID")
    to_entity: str = Field(..., alias="toEntity", description="目标实体名称")
    to_entity_id: Optional[str] = Field(None, alias="toEntityId", description="目标实体临时ID")
    relation_type: str = Field(..., alias="relationType", description="关系类型")
    confidence: Optional[float] = Field(None, description="置信度")


class NerResponse(BaseModel):
    """NER 响应"""
    model_config = ConfigDict(
        populate_by_name=True,
        serialize_by_alias=True
    )
    
    document_id: Optional[str] = Field(None, alias="documentId", description="文档ID")
    entities: List[EntityInfo] = Field(default_factory=list, description="提取的实体列表")
    relations: List[RelationInfo] = Field(default_factory=list, description="提取的关系列表")
    processing_time: Optional[int] = Field(None, alias="processingTime", description="处理耗时(毫秒)")
    entity_count: int = Field(0, alias="entityCount", description="实体数量")
    relation_count: int = Field(0, alias="relationCount", description="关系数量")
    success: bool = Field(True, description="是否成功")
    error_message: Optional[str] = Field(None, alias="errorMessage", description="错误信息")

    @classmethod
    def success_response(cls, document_id: str, entities: List[EntityInfo], 
                         relations: List[RelationInfo], processing_time: int):
        return cls(
            document_id=document_id,
            entities=entities,
            relations=relations,
            processing_time=processing_time,
            entity_count=len(entities),
            relation_count=len(relations),
            success=True
        )
    
    @classmethod
    def error_response(cls, document_id: str, error_message: str):
        return cls(
            document_id=document_id,
            success=False,
            error_message=error_message
        )


class RelationResponse(BaseModel):
    """关系抽取响应"""
    model_config = ConfigDict(
        populate_by_name=True,
        serialize_by_alias=True
    )
    
    document_id: Optional[str] = Field(None, alias="documentId", description="文档ID")
    relations: List[RelationInfo] = Field(default_factory=list, description="提取的关系列表")
    processing_time: Optional[int] = Field(None, alias="processingTime", description="处理耗时(毫秒)")
    relation_count: int = Field(0, alias="relationCount", description="关系数量")
    success: bool = Field(True, description="是否成功")
    error_message: Optional[str] = Field(None, alias="errorMessage", description="错误信息")

    @classmethod
    def success_response(cls, document_id: str, relations: List[RelationInfo], processing_time: int):
        return cls(
            document_id=document_id,
            relations=relations,
            processing_time=processing_time,
            relation_count=len(relations),
            success=True
        )
    
    @classmethod
    def error_response(cls, document_id: str, error_message: str):
        return cls(
            document_id=document_id,
            success=False,
            error_message=error_message
        )


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str = "ok"
    version: str
    ner_model: str