lingyue-zhibao/python-services/ner-service/app/services/deepseek_service.py

"""
DeepSeek API 服务（阿里云百炼平台）
用于调用 DeepSeek 模型进行 NER 提取
"""
import json
import re
import uuid
import httpx
from typing import List, Optional, Dict, Any
from loguru import logger

from ..config import settings
from ..models import EntityInfo, PositionInfo


class DeepSeekService:
    """DeepSeek API 服务"""
    
    def __init__(self):
        self.api_key = settings.deepseek_api_key
        self.base_url = settings.deepseek_base_url
        self.model = settings.deepseek_model
        self.timeout = settings.deepseek_timeout
        self.temperature = settings.deepseek_temperature
        self.max_tokens = settings.deepseek_max_tokens
        self.max_retries = settings.deepseek_max_retries
        self.chunk_size = settings.chunk_size
        self.chunk_overlap = settings.chunk_overlap
        
        logger.info(f"初始化 DeepSeek 服务: model={self.model}, base_url={self.base_url}")
    
    def _split_text(self, text: str) -> List[Dict[str, Any]]:
        """
        将长文本分割成多个块
        """
        if len(text) <= self.chunk_size:
            return [{"text": text, "start_pos": 0, "end_pos": len(text)}]
        
        chunks = []
        start = 0
        
        while start < len(text):
            end = min(start + self.chunk_size, len(text))
            
            # 尝试在句号、换行处分割
            if end < len(text):
                for sep in ['\n\n', '\n', '。', '；', '！', '？', '.']:
                    sep_pos = text.rfind(sep, start + self.chunk_size // 2, end)
                    if sep_pos > start:
                        end = sep_pos + len(sep)
                        break
            
            chunk_text = text[start:end]
            chunks.append({
                "text": chunk_text,
                "start_pos": start,
                "end_pos": end
            })
            
            start = end - self.chunk_overlap if end < len(text) else end
        
        logger.info(f"文本分割完成: 总长度={len(text)}, 分块数={len(chunks)}")
        return chunks
    
    def _build_ner_prompt(self, text: str, entity_types: Optional[List[str]] = None) -> str:
        """
        构建 NER 提取的 Prompt
        """
        types = entity_types or settings.entity_types
        
        prompt = f"""你是专业的命名实体识别(NER)专家。请从以下文本中提取有意义的命名实体。

## 实体类型定义

1. **PERSON** - 人名
   - 包括：姓名、职务+姓名（如"张总"、"李工程师"）
   - 不包括：职位本身（如"经理"、"主任"）

2. **ORG** - 机构/组织/部门
   - 包括：公司、政府机关、协会、院校、内部部门、委员会等
   - 例如："中国电建集团成都勘测设计研究院有限公司"、"国家能源局"、"安全质量环保部"、"人力资源部"、"安委会"
   - 不包括：泛指的简称（如"公司"、"集团"），除非是特定缩写（如"成都院"、"股份公司"）

3. **LOC** - 地点/位置
   - 包括：省市区县、具体地址、工程位置
   - 例如："四川省成都市"、"龙滩水电站"

4. **DATE** - 日期/时间
   - 包括：具体日期、时间段、年份
   - 例如："2024年7月13日"、"2019年版"
   - 不包括：单独的年份数字

5. **NUMBER** - 有意义的数值
   - 包括：带单位的数量、金额、评分、比例
   - 例如："93.33分"、"7000名"、"70多年"、"5个"
   - **不包括**：章节编号（如"1.1"、"第3章"）、纯序号（如"16"、"17"）、表格序号

6. **PROJECT** - 项目/工程名称
   - 包括：具体项目名称、工程名称
   - 例如："白鹤滩水电站工程"、"安全生产标准化建设项目"、"龙滩水电站"
   - **不包括**：文件编号（归类为 DOC_ID）

7. **METHOD** - 方法/标准/规范
   - 包括：技术标准、管理办法、评价标准的**完整名称**
   - 例如："电力建设企业安全生产标准化评价标准"、"安全生产风险抵押金管理办法"
   - **不包括**：编号（归类为 DOC_ID）

8. **DEVICE** - 设备/系统
   - 包括：设备名称、信息系统名称
   - 例如："OA系统"、"QHSE系统"、"造槽机"

9. **DOC_ID** - 文件编号/标准编号
   - 包括：文件编号、标准编号、规范编号、合同编号、项目编号、证书编号
   - 例如："AY-BZ-0092-2024-Z01"、"SQE.01C0213"、"中电建股〔2019〕122号"、"ZGDIDBOY-083"、"蓉设安质〔2024〕18号"
   - 特征：通常包含字母、数字、连字符、年份、括号等组合

10. **CERT** - 证书/资质/等级
    - 包括：**完整的**资质证书名称、等级认证、荣誉称号、专业资质
    - 例如："电力安全生产标准化一级企业证书"、"工程设计综合甲级资质"、"注册安全工程师证"、"电力、水利水电、市政公用工程施工总承包一级"
    - **不包括**：
      - 职务职称（归类为 TITLE）
      - 单独的行业名称（如单独的"电力"、"水利水电"不是资质，需要完整如"电力工程施工总承包一级"）
      - 资质的简称片段（如"甲级"本身不提取，需要完整如"工程设计综合甲级"）

11. **TITLE** - 职务/职称
    - 包括：行政职务、专业职称、岗位名称
    - 例如："董事长"、"总经理"、"安全总监"、"总工程师"、"首席专家"
    - 不包括：人名（归类为 PERSON）

12. **POLICY** - 政策法规
    - 包括：法律、法规、条例、办法的名称（不含编号）
    - 例如："安全生产法"、"企业安全生产费用提取和使用管理办法"、"劳动合同法"
    - 不包括：企业内部制度（归类为 METHOD）、文件编号（归类为 DOC_ID）

## 重要规则

1. **只提取有实际意义的实体**，忽略：
   - 章节编号（如"1"、"1.1"、"第一章"）
   - 表格序号
   - 单独的数字（如"16"、"17"）

2. **实体必须完整**，提取完整的名称而非片段：
   - 资质要完整：提取"工程设计综合甲级"而非"甲级"
   - 证书要完整：提取"注册安全工程师证"而非"安全工程师"
   - 机构要完整：提取"安全质量环保部"而非"环保部"

3. **不要拆分并列实体**：
   - 如"电力、水利水电、市政公用工程施工总承包一级"应作为一个完整的 CERT
   - 不要拆成"电力"、"水利水电"等单独的实体

4. **去除重复实体**，相同的实体只返回一次

5. **charStart和charEnd必须准确**，对应实体在原文中的字符位置（从0开始）

## 输出格式

直接输出JSON，格式如下：
{{"entities": [{{"name": "实体名称", "type": "实体类型", "charStart": 起始位置, "charEnd": 结束位置}}]}}

## 待处理文本

{text}

请直接输出JSON:"""
        return prompt
    
    async def _call_api(self, prompt: str) -> Optional[str]:
        """
        调用 DeepSeek API（OpenAI 兼容格式）
        """
        url = f"{self.base_url}/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.model,
            "messages": [
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
        }
        
        for attempt in range(self.max_retries):
            try:
                async with httpx.AsyncClient(timeout=self.timeout) as client:
                    response = await client.post(url, headers=headers, json=payload)
                    response.raise_for_status()
                    result = response.json()
                    
                    # OpenAI 格式响应
                    choices = result.get("choices", [])
                    if choices:
                        choice = choices[0]
                        message = choice.get("message", {})
                        content = message.get("content", "")
                        
                        # 检查是否因为 max_tokens 被截断
                        finish_reason = choice.get("finish_reason", "")
                        if finish_reason == "length":
                            logger.warning(f"API 响应被截断 (finish_reason=length), 考虑增加 max_tokens 或减小分块大小")
                        
                        return content
                    return None
                    
            except httpx.TimeoutException:
                logger.warning(f"DeepSeek API 请求超时 (尝试 {attempt + 1}/{self.max_retries})")
                if attempt == self.max_retries - 1:
                    logger.error(f"DeepSeek API 请求超时: timeout={self.timeout}s")
                    return None
            except httpx.HTTPStatusError as e:
                logger.error(f"DeepSeek API HTTP 错误: {e.response.status_code} - {e.response.text}")
                return None
            except Exception as e:
                logger.error(f"DeepSeek API 请求失败: {e}")
                if attempt == self.max_retries - 1:
                    return None
        
        return None
    
    def _parse_response(self, response: str, chunk_start_pos: int = 0) -> List[EntityInfo]:
        """
        解析 API 返回的 JSON 结果
        """
        entities = []
        
        try:
            # 移除 markdown code block 标记（支持多行模式）
            response = re.sub(r'```json\s*\n?', '', response, flags=re.IGNORECASE)
            response = re.sub(r'\n?```\s*$', '', response)
            response = re.sub(r'```\s*', '', response)
            response = response.strip()
            
            # 方法1：直接解析
            data = None
            try:
                data = json.loads(response)
            except json.JSONDecodeError as e:
                logger.debug(f"直接解析 JSON 失败: {e}")
            
            # 方法2：查找 JSON 对象（使用更宽松的正则）
            if not data or "entities" not in data:
                # 尝试匹配从 { 开始到最后一个 } 的内容
                json_match = re.search(r'\{[^{}]*"entities"\s*:\s*\[[\s\S]*?\]\s*\}', response)
                if json_match:
                    try:
                        data = json.loads(json_match.group())
                    except json.JSONDecodeError as e:
                        logger.debug(f"正则匹配 JSON 解析失败: {e}")
            
            # 方法3：尝试提取 entities 数组
            if not data or "entities" not in data:
                array_match = re.search(r'"entities"\s*:\s*(\[[\s\S]*\])', response)
                if array_match:
                    try:
                        entity_list = json.loads(array_match.group(1))
                        data = {"entities": entity_list}
                    except json.JSONDecodeError as e:
                        logger.debug(f"提取 entities 数组失败: {e}")
            
            # 方法4：处理被截断的 JSON，尝试逐个解析完整的实体对象
            if not data or "entities" not in data:
                logger.debug("尝试从截断的 JSON 中提取完整实体...")
                entity_pattern = r'\{\s*"name"\s*:\s*"([^"]+)"\s*,\s*"type"\s*:\s*"([^"]+)"\s*,\s*"charStart"\s*:\s*(\d+)\s*,\s*"charEnd"\s*:\s*(\d+)\s*\}'
                matches = re.findall(entity_pattern, response)
                if matches:
                    data = {"entities": []}
                    for match in matches:
                        data["entities"].append({
                            "name": match[0],
                            "type": match[1],
                            "charStart": int(match[2]),
                            "charEnd": int(match[3])
                        })
                    logger.info(f"从截断 JSON 中恢复了 {len(matches)} 个实体")
            
            if not data or "entities" not in data:
                logger.warning(f"未找到有效的 entities JSON, response_length={len(response)}, response_preview={response[:500]}...")
                return entities
            
            entity_list = data.get("entities", [])
            
            for item in entity_list:
                name = item.get("name", "").strip()
                entity_type = item.get("type", "").upper()
                char_start = item.get("charStart", 0)
                char_end = item.get("charEnd", 0)
                
                if not name or len(name) < 2:
                    continue
                
                # 校正位置
                adjusted_start = char_start + chunk_start_pos
                adjusted_end = char_end + chunk_start_pos
                
                entity = EntityInfo(
                    name=name,
                    type=entity_type,
                    value=name,
                    position=PositionInfo(
                        char_start=adjusted_start,
                        char_end=adjusted_end,
                        line=1
                    ),
                    confidence=0.95,  # DeepSeek 置信度较高
                    temp_id=str(uuid.uuid4())[:8]
                )
                entities.append(entity)
                
        except Exception as e:
            logger.error(f"解析响应失败: {e}")
        
        return entities
    
    async def extract_entities(
        self, 
        text: str, 
        entity_types: Optional[List[str]] = None
    ) -> List[EntityInfo]:
        """
        使用 DeepSeek API 提取实体
        """
        if not text or not text.strip():
            return []
        
        # 分割长文本
        chunks = self._split_text(text)
        
        all_entities = []
        seen_entities = set()
        
        for i, chunk in enumerate(chunks):
            logger.info(f"处理分块 {i+1}/{len(chunks)}: 长度={len(chunk['text'])}")
            
            prompt = self._build_ner_prompt(chunk["text"], entity_types)
            response = await self._call_api(prompt)
            
            if not response:
                logger.warning(f"分块 {i+1} API 返回为空")
                continue
            
            logger.debug(f"分块 {i+1} API 响应: {response[:500]}...")
            
            entities = self._parse_response(response, chunk["start_pos"])
            
            # 去重
            for entity in entities:
                entity_key = f"{entity.type}:{entity.name}"
                if entity_key not in seen_entities:
                    seen_entities.add(entity_key)
                    all_entities.append(entity)
            
            logger.info(f"分块 {i+1} 提取实体: {len(entities)} 个")
        
        logger.info(f"DeepSeek NER 提取完成: 总实体数={len(all_entities)}")
        return all_entities
    
    async def extract_entities_with_progress(
        self, 
        text: str, 
        entity_types: Optional[List[str]] = None
    ):
        """
        使用 DeepSeek API 提取实体（带进度生成器）
        
        Yields:
            SSE 事件字符串
        """
        import json
        
        async def sse_event(event: str, data: dict):
            return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
        
        if not text or not text.strip():
            yield await sse_event("complete", {"entities": [], "total_entities": 0})
            return
        
        # 分割长文本
        chunks = self._split_text(text)
        total_chunks = len(chunks)
        
        all_entities = []
        seen_entities = set()
        
        for i, chunk in enumerate(chunks):
            chunk_index = i + 1
            logger.info(f"处理分块 {chunk_index}/{total_chunks}: 长度={len(chunk['text'])}")
            
            # 发送进度事件
            yield await sse_event("progress", {
                "chunk_index": chunk_index,
                "total_chunks": total_chunks,
                "chunk_length": len(chunk['text']),
                "total_entities": len(all_entities),
                "progress_percent": int((chunk_index - 1) / total_chunks * 100),
                "message": f"正在处理第 {chunk_index}/{total_chunks} 个文本块..."
            })
            
            prompt = self._build_ner_prompt(chunk["text"], entity_types)
            response = await self._call_api(prompt)
            
            if not response:
                logger.warning(f"分块 {chunk_index} API 返回为空")
                continue
            
            logger.debug(f"分块 {chunk_index} API 响应: {response[:500]}...")
            
            entities = self._parse_response(response, chunk["start_pos"])
            
            # 去重并收集新实体
            new_entities = []
            for entity in entities:
                entity_key = f"{entity.type}:{entity.name}"
                if entity_key not in seen_entities:
                    seen_entities.add(entity_key)
                    all_entities.append(entity)
                    new_entities.append(entity)
            
            logger.info(f"分块 {chunk_index} 提取实体: {len(entities)} 个, 新增: {len(new_entities)} 个")
            
            # 发送分块完成事件
            yield await sse_event("chunk_complete", {
                "chunk_index": chunk_index,
                "total_chunks": total_chunks,
                "chunk_entities": len(entities),
                "new_entities": len(new_entities),
                "total_entities": len(all_entities),
                "progress_percent": int(chunk_index / total_chunks * 100)
            })
        
        logger.info(f"DeepSeek NER 提取完成: 总实体数={len(all_entities)}")
        
        # 发送实体数据事件（供调用方获取实体列表）
        yield await sse_event("entities_data", {
            "entities": [entity.model_dump(by_alias=True) for entity in all_entities],
            "total_entities": len(all_entities)
        })
    
    async def check_health(self) -> bool:
        """
        检查 DeepSeek API 是否可用
        """
        try:
            url = f"{self.base_url}/v1/models"
            headers = {
                "Authorization": f"Bearer {self.api_key}"
            }
            async with httpx.AsyncClient(timeout=10) as client:
                response = await client.get(url, headers=headers)
                return response.status_code == 200
        except Exception:
            return False


# 创建单例
deepseek_service = DeepSeekService()