proposa-ai/src/step4_writing/company_info_extractor.py

"""
公司信息提取器

从通用标书生成资料中提取结构化的公司基本信息。
使用 LLM 一次性提取，结果缓存。
"""

from __future__ import annotations

import logging
from typing import Dict, List, Optional

from llm_client import LLMClient
from models import CompanyInfo, ParsedDocument

logger = logging.getLogger(__name__)


def extract_company_info(
    general_materials: List[ParsedDocument],
) -> CompanyInfo:
    """从通用标书资料中提取结构化公司信息

    优先从"企业基本情况简介"类 MD 文件中提取。

    Args:
        general_materials: Step 1 加载的通用资料文档列表

    Returns:
        CompanyInfo — 结构化公司信息
    """
    # 1. 优先找公司简介类文档
    intro_text = ""
    for doc in general_materials:
        if not doc.content:
            continue
        name = doc.file_name.lower()
        if any(kw in name for kw in ["企业基本情况", "公司简介", "企业介绍", "公司概况"]):
            intro_text = doc.content
            break

    # 2. 降级：合并所有短文档
    if not intro_text:
        combined = []
        for doc in general_materials:
            if doc.content:
                combined.append(f"--- {doc.file_name} ---\n{doc.content[:3000]}")
        intro_text = "\n\n".join(combined)

    if not intro_text:
        logger.warning("未找到公司信息资料，返回空 CompanyInfo")
        return CompanyInfo()

    # 3. LLM 提取
    llm = LLMClient()

    prompt = """请从以下公司资料中提取公司基本信息，以 JSON 格式输出。
如果某项信息未找到，请填空字符串""。

输出格式：
{
  "full_name": "公司全称",
  "short_name": "公司简称",
  "address": "注册地址或办公地址",
  "phone": "联系电话",
  "fax": "传真",
  "zip_code": "邮编",
  "bank_name": "开户银行",
  "bank_account": "银行账号",
  "legal_rep": "法定代表人姓名",
  "legal_rep_id": "法定代表人身份证号",
  "registration_capital": "注册资本",
  "unified_code": "统一社会信用代码",
  "website": "公司官网"
}

注意：
- 公司全称必须与营业执照一致，不要删减或添加文字
- 地址必须完整，包含省市（直辖）区
- 电话格式保留原样"""

    try:
        result = llm.extract_json(
            system_prompt="你是企业信息提取专家。请从文本中准确提取公司基本信息。",
            user_prompt=f"{prompt}\n\n---\n{intro_text[:12000]}",
            max_tokens=32768,
        )
        if isinstance(result, dict):
            # LLM 可能返回 null（而非空字符串）：str(None) 会生成 "None"
            # 污染占位符替换结果，统一转回空字符串
            def _safe_str(v) -> str:
                if v is None:
                    return ""
                s = str(v).strip()
                return "" if s.lower() in ("none", "null") else s

            info = CompanyInfo(
                full_name=_safe_str(result.get("full_name")),
                short_name=_safe_str(result.get("short_name")),
                address=_safe_str(result.get("address")),
                phone=_safe_str(result.get("phone")),
                fax=_safe_str(result.get("fax")),
                zip_code=_safe_str(result.get("zip_code")),
                bank_name=_safe_str(result.get("bank_name")),
                bank_account=_safe_str(result.get("bank_account")),
                legal_rep=_safe_str(result.get("legal_rep")),
                legal_rep_id=_safe_str(result.get("legal_rep_id")),
                registration_capital=_safe_str(result.get("registration_capital")),
                unified_code=_safe_str(result.get("unified_code")),
                website=_safe_str(result.get("website")),
            )
            logger.info(f"公司信息提取完成: {info.full_name}")
            return info
        else:
            logger.warning(f"公司信息提取返回非 dict: {type(result)}")
            return CompanyInfo()

    except Exception as e:
        logger.warning(f"公司信息提取失败: {e}")
        # 降级：正则提取
        return _fallback_extract(intro_text)


def _fallback_extract(text: str) -> CompanyInfo:
    """降级方案：正则表达式提取公司名"""
    import re
    info = CompanyInfo()

    # 尝试匹配公司名（中文名+有限公司）
    m = re.search(r"((?:上海|北京|深圳|广州|杭州|南京|成都|武汉|重庆|天津)[一-龥（）\(\)]{2,20}(?:有限公司|股份有限公司))", text)
    if m:
        info.full_name = m.group(1)

    # 地址
    m = re.search(r"(?:地址|注册地址|办公地址)[：:]\s*(.+?)(?:[\n，。]|$)", text)
    if m:
        info.address = m.group(1).strip()

    # 电话
    m = re.search(r"(?:电话|联系电话|Tel)[：:]\s*([\d\-]+)", text)
    if m:
        info.phone = m.group(1).strip()

    return info


__all__ = ["extract_company_info"]