lingyue-zhibao/python-services/ner-service/app/routers/relation.py

"""
关系抽取路由
"""
from fastapi import APIRouter, HTTPException
from loguru import logger
import time

from ..models import RelationRequest, RelationResponse
from ..services.relation_service import relation_service

router = APIRouter()


@router.post("/relations", response_model=RelationResponse)
async def extract_relations(request: RelationRequest):
    """
    从已提取的实体中抽取关系
    """
    start_time = time.time()
    
    try:
        logger.info(f"开始提取关系: document_id={request.document_id}, "
                   f"entity_count={len(request.entities)}")
        
        if not request.entities or len(request.entities) < 2:
            return RelationResponse.success_response(
                document_id=request.document_id,
                relations=[],
                processing_time=0
            )
        
        # 调用关系抽取服务
        relations = await relation_service.extract_relations(
            text=request.text,
            entities=request.entities
        )
        
        processing_time = int((time.time() - start_time) * 1000)
        
        logger.info(f"关系提取完成: document_id={request.document_id}, "
                   f"relation_count={len(relations)}, processing_time={processing_time}ms")
        
        return RelationResponse.success_response(
            document_id=request.document_id,
            relations=relations,
            processing_time=processing_time
        )
        
    except Exception as e:
        logger.error(f"关系提取失败: document_id={request.document_id}, error={str(e)}")
        return RelationResponse.error_response(
            document_id=request.document_id,
            error_message=str(e)
        )