信息提取Agent：结构化数据提取与整理

文档（PDF/Word/Excel/图片）
    ↓
[xParse Pipeline - Parse]
    └─ 解析文档，提取结构化元素（elements）
    ↓
聚合元素文本（elements[].text）
    ↓
[LangChain Agent]
    ├─ Tool 1: extract_invoice_info（提取发票信息）
    ├─ Tool 2: extract_medical_bill_info（提取医疗票据信息）
    ├─ Tool 3: extract_contract_info（提取合同信息）
    ├─ Tool 4: extract_resume_info（提取简历信息）
    ├─ Tool 5: extract_product_specs（提取产品规格）
    ├─ Tool 6: extract_api_info（提取API信息）
    └─ Tool 7: format_data（数据格式化）
    ↓
结构化数据（JSON/CSV）

python -m venv .venv && source .venv/bin/activate
pip install "xparse-client>=0.2.5" langchain langchain-community langchain-core \
            python-dotenv pandas
export XTI_APP_ID=your-app-id # 在 TextIn 官网注册获取
export XTI_SECRET_CODE=your-secret-code # 在 TextIn 官网注册获取
export DASHSCOPE_API_KEY=your-dashscope-api-key # 本教程使用通义千问大模型，也可以替换成其他大模型

from xparse_client import create_pipeline_from_config
import os
from dotenv import load_dotenv

load_dotenv()

EXTRACTION_PIPELINE_CONFIG = {
    "source": {
        "type": "local",
        "directory": "./extraction_documents",
        "pattern": ["*.pdf", "*.docx", "*.xlsx", "*.xls", "*.png", "*.jpg"]
    },
    "destination": {
        "type": "local",  # 使用本地存储，保存解析结果
        "output_dir": "./extraction_results"
    },
    "api_base_url": "https://api.textin.com/api/xparse",
    "api_headers": {
        "x-ti-app-id": os.getenv("XTI_APP_ID"),
        "x-ti-secret-code": os.getenv("XTI_SECRET_CODE")
    },
    "stages": [
        {
            "type": "parse",
            "config": {
                "provider": "textin"  # 使用TextIn解析引擎，对表格和列表识别效果好
            }
        }
    ]
}

def parse_document(file_path: str) -> list:
    """
    解析单个文档，返回 elements 列表
    
    Args:
        file_path: 文档路径
        
    Returns:
        list: elements 列表，每个元素包含 text、type、metadata 等字段
    """
    import os
    import json
    from xparse_client import Pipeline, LocalSource
    
    # 创建临时配置，使用包含该文件的目录作为source
    file_dir = os.path.dirname(os.path.abspath(file_path))
    file_name_pattern = os.path.basename(file_path)
    
    # 创建临时Pipeline配置
    temp_config = EXTRACTION_PIPELINE_CONFIG.copy()
    temp_config["source"] = {
        "type": "local",
        "directory": file_dir,
        "pattern": [file_name_pattern]  # 只处理指定的文件
    }
    
    # 创建Pipeline并运行（pipeline.run() 没有返回值，结果会保存到destination）
    pipeline = create_pipeline_from_config(temp_config)
    pipeline.run()
    
    # 从destination配置的输出目录读取解析结果
    output_dir = EXTRACTION_PIPELINE_CONFIG["destination"]["output_dir"]
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取文件名（不含路径和扩展名）
    file_name = os.path.splitext(os.path.basename(file_path))[0]
    result_file = os.path.join(output_dir, f"{file_name}.json")
    
    # 读取JSON文件
    if not os.path.exists(result_file):
        raise FileNotFoundError(
            f"解析结果文件不存在: {result_file}\n"
            f"请检查输出目录: {output_dir}\n"
            f"原始文件路径: {file_path}"
        )
    
    with open(result_file, 'r', encoding='utf-8') as f:
        elements = json.load(f)
    
    return elements

def aggregate_text_from_elements(elements: list) -> str:
    """
    聚合 elements 中的 text 字段，形成完整文档文本
    
    Args:
        elements: elements 列表
        
    Returns:
        str: 聚合后的完整文本
    """
    texts = []
    for element in elements:
        if isinstance(element, dict):
            text = element.get('text', '')
        else:
            text = getattr(element, 'text', '')
        
        if text and text.strip():
            texts.append(text.strip())
    
    return "\n\n".join(texts)

from langchain_core.tools import Tool
from langchain_community.chat_models import ChatTongyi
import os
import json

# 全局文档文本存储（实际应用中可以使用更持久化的存储）
_document_texts = {}

# 初始化 qwen-max 大模型
llm = ChatTongyi(
    model="qwen-max",
    dashscope_api_key=os.getenv("DASHSCOPE_API_KEY"),
    temperature=0,  # 使用较低温度以获得更确定性的输出
)

def set_document_text(file_path: str, text: str):
    """设置文档文本内容"""
    _document_texts[file_path] = text

def get_document_text(file_path: str = None) -> str:
    """获取文档文本内容"""
    if not _document_texts:
        return ""  # 如果没有加载任何文档，返回空字符串
    
    if file_path not in ("None", "none", None, "", "null"):
        return _document_texts.get(file_path, "")
    # 如果没有指定文件，返回第一个文档的文本
    return next(iter(_document_texts.values()), "")

def extract_invoice_info(file_path: str = None) -> str:
    """
    从发票中提取结构化信息（使用 qwen-max 大模型）
    
    提取内容包括：
    - 发票基本信息（发票代码、发票号码、开票日期）
    - 销售方信息（名称、纳税人识别号、地址电话、开户行及账号）
    - 购买方信息（名称、纳税人识别号、地址电话、开户行及账号）
    - 商品明细（名称、规格、单位、数量、单价、金额、税率、税额）
    - 金额信息（合计金额、合计税额、价税合计）
    - 其他信息（备注、收款人、复核人、开票人等）
    
    Args:
        file_path: 文档路径（可选），如果不提供则使用当前已加载的文档
    """
    # 获取文档文本
    context_text = get_document_text(file_path)
    
    # 如果没有文档文本，返回提示信息
    if not context_text:
        return "错误：未找到文档内容。请先使用 load_document() 方法加载文档，或提供文档路径。"
    
    # 构建 prompt，指导模型提取结构化信息
    prompt = f"""请从以下发票文本中提取结构化信息，并以 JSON 格式返回。

            要求提取的信息包括：
            1. 发票基本信息：invoice_code（发票代码）、invoice_number（发票号码）、date（开票日期）
            2. 销售方信息：name（名称）、tax_id（纳税人识别号）、address（地址电话）、bank_account（开户行及账号）
            3. 购买方信息：name（名称）、tax_id（纳税人识别号）、address（地址电话）、bank_account（开户行及账号）
            4. 商品明细（数组）：name、specification、unit、quantity、unit_price、amount、tax_rate、tax_amount
            5. 金额信息：total_amount（合计金额）、tax_amount（合计税额）、total_with_tax（价税合计）
            6. 其他信息：remark（备注）、payee（收款人）、reviewer（复核人）、drawer（开票人）

            请严格按照以下 JSON 格式返回，如果某个字段不存在，请使用空字符串 "" 或空对象 {{}} 或空数组 []：

            {{
            "invoice_info": {{"invoice_code": "", "invoice_number": "", "date": ""}},
            "seller": {{"name": "", "tax_id": "", "address": "", "bank_account": ""}},
            "buyer": {{"name": "", "tax_id": "", "address": "", "bank_account": ""}},
            "items": [{{"name": "", "specification": "", "unit": "", "quantity": "", "unit_price": "", "amount": "", "tax_rate": "", "tax_amount": ""}}],
            "amounts": {{"total_amount": "", "tax_amount": "", "total_with_tax": ""}},
            "other_info": {{"remark": "", "payee": "", "reviewer": "", "drawer": ""}}
            }}

            发票文本内容：
            {context_text}

            请只返回 JSON 格式的数据，不要包含任何其他解释或说明文字。"""

    # 调用大模型提取信息
    response_text = ""
    try:
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        
        # 尝试从响应中提取 JSON（可能包含 markdown 代码块）
        response_text = response.content.strip() if response.content else ""
        
        # 如果响应包含 markdown 代码块，提取其中的 JSON
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        elif "```" in response_text:
            json_start = response_text.find("```") + 3
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        
        # 解析 JSON
        invoice_data = json.loads(response_text)
        
        # 返回格式化的 JSON
        return json.dumps(invoice_data, ensure_ascii=False, indent=2)
        
    except json.JSONDecodeError as e:
        # 如果 JSON 解析失败，返回错误信息
        error_msg = f"JSON 解析失败：{str(e)}\n模型返回的原始内容：\n{response_text}"
        print(error_msg)
        return json.dumps({
            "error": "JSON 解析失败",
            "raw_response": response_text if response_text else "无响应",
            "error_detail": str(e)
        }, ensure_ascii=False, indent=2)
    except Exception as e:
        # 其他错误
        error_msg = f"提取信息时发生错误：{str(e)}"
        print(error_msg)
        if response_text:
            print(f"模型返回的原始内容：\n{response_text}")
        return json.dumps({
            "error": "提取信息失败",
            "error_detail": str(e),
            "raw_response": response_text if response_text else "无响应"
        }, ensure_ascii=False, indent=2)

def extract_medical_bill_info(file_path: str = None) -> str:
    """
    从医疗票据中提取结构化信息（使用 qwen-max 大模型）
    
    提取内容包括：
    - 患者信息（姓名、性别、年龄、身份证号、医保卡号）
    - 医疗机构信息（医院名称、科室、医生姓名）
    - 就诊信息（就诊日期、就诊类型、诊断结果）
    - 费用明细（项目名称、数量、单价、金额、医保类型）
    - 费用汇总（总费用、自费金额、医保支付、个人支付）
    - 其他信息（发票号码、结算方式等）
    
    Args:
        file_path: 文档路径（可选），如果不提供则使用当前已加载的文档
    """
    # 获取文档文本
    context_text = get_document_text(file_path)
    
    # 如果没有文档文本，返回提示信息
    if not context_text:
        return "错误：未找到文档内容。请先使用 load_document() 方法加载文档，或提供文档路径。"
    
    # 构建 prompt
    prompt = f"""请从以下医疗票据文本中提取结构化信息，包括患者信息、医疗机构信息、就诊信息、费用明细、费用汇总等，并以 JSON 格式返回。

            医疗票据文本内容：
            {context_text}

            请只返回 JSON 格式的数据，不要包含任何其他解释或说明文字。"""

    # 调用大模型提取信息
    response_text = ""
    try:
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        response_text = response.content.strip() if response.content else ""
        
        # 提取 JSON（处理 markdown 代码块）
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        elif "```" in response_text:
            json_start = response_text.find("```") + 3
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        
        medical_bill_data = json.loads(response_text)
        return json.dumps(medical_bill_data, ensure_ascii=False, indent=2)
        
    except json.JSONDecodeError as e:
        return json.dumps({
            "error": "JSON 解析失败",
            "raw_response": response_text if response_text else "无响应",
            "error_detail": str(e)
        }, ensure_ascii=False, indent=2)
    except Exception as e:
        return json.dumps({
            "error": "提取信息失败",
            "error_detail": str(e),
            "raw_response": response_text if response_text else "无响应"
        }, ensure_ascii=False, indent=2)

def extract_contract_info(file_path: str = None) -> str:
    """
    从合同中提取结构化信息（使用 qwen-max 大模型）
    
    提取内容包括：
    - 合同基本信息（合同编号、合同名称、签订日期、生效日期、到期日期）
    - 合同双方（甲方、乙方信息：名称、地址、法定代表人、联系方式）
    - 合同标的（标的物、服务内容、数量、金额）
    - 关键条款（付款方式、交付方式、违约责任、争议解决）
    - 金额信息（合同总金额、付款计划、保证金等）
    - 其他信息（签署地点、签署人、附件等）
    
    Args:
        file_path: 文档路径（可选），如果不提供则使用当前已加载的文档
    """
    # 获取文档文本
    context_text = get_document_text(file_path)
    
    # 如果没有文档文本，返回提示信息
    if not context_text:
        return "错误：未找到文档内容。请先使用 load_document() 方法加载文档，或提供文档路径。"
    
    # 构建 prompt
    prompt = f"""请从以下合同文本中提取结构化信息，包括合同基本信息、合同双方信息、合同标的、关键条款、金额信息等，并以 JSON 格式返回。

            合同文本内容：
            {context_text}

            请只返回 JSON 格式的数据，不要包含任何其他解释或说明文字。"""

    # 调用大模型提取信息
    response_text = ""
    try:
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        response_text = response.content.strip() if response.content else ""
        
        # 提取 JSON（处理 markdown 代码块）
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        elif "```" in response_text:
            json_start = response_text.find("```") + 3
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        
        contract_data = json.loads(response_text)
        return json.dumps(contract_data, ensure_ascii=False, indent=2)
        
    except json.JSONDecodeError as e:
        return json.dumps({
            "error": "JSON 解析失败",
            "raw_response": response_text if response_text else "无响应",
            "error_detail": str(e)
        }, ensure_ascii=False, indent=2)
    except Exception as e:
        return json.dumps({
            "error": "提取信息失败",
            "error_detail": str(e),
            "raw_response": response_text if response_text else "无响应"
        }, ensure_ascii=False, indent=2)

def extract_resume_info(file_path: str = None) -> str:
    """
    从简历中提取结构化信息（使用 qwen-max 大模型）
    
    提取内容包括：
    - 个人信息（姓名、性别、年龄、联系方式）
    - 教育经历（学校、专业、学历、时间）
    - 工作经历（公司、职位、时间、工作内容）
    - 技能（专业技能、语言能力、证书等）
    
    Args:
        file_path: 文档路径（可选），如果不提供则使用当前已加载的文档
    """
    # 获取文档文本
    context_text = get_document_text(file_path)
    
    # 如果没有文档文本，返回提示信息
    if not context_text:
        return "错误：未找到文档内容。请先使用 load_document() 方法加载文档，或提供文档路径。"
    
    # 构建 prompt
    prompt = f"""请从以下简历文本中提取结构化信息，包括个人信息、教育经历、工作经历、技能等，并以 JSON 格式返回。

            简历文本内容：
            {context_text}

            请只返回 JSON 格式的数据，不要包含任何其他解释或说明文字。"""

    # 调用大模型提取信息
    response_text = ""
    try:
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        response_text = response.content.strip() if response.content else ""
        
        # 提取 JSON（处理 markdown 代码块）
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        elif "```" in response_text:
            json_start = response_text.find("```") + 3
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        
        resume_data = json.loads(response_text)
        return json.dumps(resume_data, ensure_ascii=False, indent=2)
        
    except json.JSONDecodeError as e:
        return json.dumps({
            "error": "JSON 解析失败",
            "raw_response": response_text if response_text else "无响应",
            "error_detail": str(e)
        }, ensure_ascii=False, indent=2)
    except Exception as e:
        return json.dumps({
            "error": "提取信息失败",
            "error_detail": str(e),
            "raw_response": response_text if response_text else "无响应"
        }, ensure_ascii=False, indent=2)

def extract_product_specs(file_path: str = None) -> str:
    """
    从产品文档中提取产品规格和技术参数（使用 qwen-max 大模型）
    
    提取内容包括：
    - 产品名称和型号
    - 技术参数（尺寸、重量、性能指标等）
    - 功能特性
    - 价格信息
    
    Args:
        file_path: 文档路径（可选），如果不提供则使用当前已加载的文档
    """
    # 获取文档文本
    context_text = get_document_text(file_path)
    
    # 如果没有文档文本，返回提示信息
    if not context_text:
        return "错误：未找到文档内容。请先使用 load_document() 方法加载文档，或提供文档路径。"
    
    # 构建 prompt
    prompt = f"""请从以下产品文档文本中提取产品规格和技术参数，包括产品名称、型号、技术参数、功能特性、价格等，并以 JSON 格式返回。

            产品文档文本内容：
            {context_text}

            请只返回 JSON 格式的数据，不要包含任何其他解释或说明文字。"""

    # 调用大模型提取信息
    response_text = ""
    try:
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        response_text = response.content.strip() if response.content else ""
        
        # 提取 JSON（处理 markdown 代码块）
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        elif "```" in response_text:
            json_start = response_text.find("```") + 3
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        
        product_specs = json.loads(response_text)
        return json.dumps(product_specs, ensure_ascii=False, indent=2)
        
    except json.JSONDecodeError as e:
        return json.dumps({
            "error": "JSON 解析失败",
            "raw_response": response_text if response_text else "无响应",
            "error_detail": str(e)
        }, ensure_ascii=False, indent=2)
    except Exception as e:
        return json.dumps({
            "error": "提取信息失败",
            "error_detail": str(e),
            "raw_response": response_text if response_text else "无响应"
        }, ensure_ascii=False, indent=2)

def extract_api_info(file_path: str = None) -> str:
    """
    从技术文档中提取API接口信息（使用 qwen-max 大模型）
    
    提取内容包括：
    - API端点（URL路径）
    - 请求方法（GET、POST等）
    - 请求参数
    - 响应格式
    - 认证方式
    
    Args:
        file_path: 文档路径（可选），如果不提供则使用当前已加载的文档
    """
    # 获取文档文本
    context_text = get_document_text(file_path)
    
    # 如果没有文档文本，返回提示信息
    if not context_text:
        return "错误：未找到文档内容。请先使用 load_document() 方法加载文档，或提供文档路径。"
    
    # 构建 prompt
    prompt = f"""请从以下技术文档文本中提取API接口信息，包括API端点、请求方法、请求参数、响应格式、认证方式等，并以 JSON 格式返回。

                技术文档文本内容：
                {context_text}

                请只返回 JSON 格式的数据，不要包含任何其他解释或说明文字。"""

    # 调用大模型提取信息
    response_text = ""
    try:
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        response_text = response.content.strip() if response.content else ""
        
        # 提取 JSON（处理 markdown 代码块）
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        elif "```" in response_text:
            json_start = response_text.find("```") + 3
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        
        api_info = json.loads(response_text)
        return json.dumps(api_info, ensure_ascii=False, indent=2)
        
    except json.JSONDecodeError as e:
        return json.dumps({
            "error": "JSON 解析失败",
            "raw_response": response_text if response_text else "无响应",
            "error_detail": str(e)
        }, ensure_ascii=False, indent=2)
    except Exception as e:
        return json.dumps({
            "error": "提取信息失败",
            "error_detail": str(e),
            "raw_response": response_text if response_text else "无响应"
        }, ensure_ascii=False, indent=2)

import pandas as pd

def format_data(file_path: str = None) -> str:
    """
    将提取的数据格式化为标准格式（JSON、CSV等）
    
    支持格式：
    - JSON格式
    - CSV格式
    - 表格格式
    
    Args:
        file_path: 文档路径（可选），如果不提供则使用当前已加载的文档
    """
    # 获取文档文本
    context_text = get_document_text(file_path)
    
    # 如果没有文档文本，返回提示信息
    if not context_text:
        return "错误：未找到文档内容。请先使用 load_document() 方法加载文档，或提供文档路径。"
    
    # 使用大模型提取结构化数据
    prompt = f"""请从以下文本中提取所有键值对信息，并以 JSON 格式返回。

                要求：
                1. 提取所有形如"键：值"或"键:值"的键值对
                2. 返回格式为数组，每个元素包含 key、value、source（来源文件名）

                文本内容：
                {context_text}

                请返回 JSON 格式的数组，格式如下：
                [
                {{
                    "key": "键名",
                    "value": "值",
                    "source": "文件名"
                }}
                ]

                请只返回 JSON 格式的数据，不要包含任何其他解释或说明文字。"""

    try:
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        response_text = response.content.strip() if response.content else ""
        
        # 提取 JSON（处理 markdown 代码块）
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        elif "```" in response_text:
            json_start = response_text.find("```") + 3
            json_end = response_text.find("```", json_start)
            if json_end != -1:
                response_text = response_text[json_start:json_end].strip()
        
        data_list = json.loads(response_text)
        
        if not data_list:
            return "未找到可格式化的数据"
        
        # 格式化为JSON
        json_output = json.dumps(data_list, ensure_ascii=False, indent=2)
        
        # 格式化为CSV（使用pandas）
        try:
            df = pd.DataFrame(data_list)
            csv_output = df.to_csv(index=False)
        except:
            csv_output = "CSV格式化失败"
        
        return f"JSON格式：\n{json_output}\n\nCSV格式：\n{csv_output}"
        
    except Exception as e:
        return f"数据格式化失败：{str(e)}"

tools = [
    Tool(
        name="extract_invoice_info",
        description="从发票中提取结构化信息，包括发票基本信息、销售方/购买方信息、商品明细、金额信息等。如果已加载文档，可以直接调用无需参数；否则需要提供文档路径。",
        func=extract_invoice_info
    ),
    Tool(
        name="extract_medical_bill_info",
        description="从医疗票据中提取结构化信息，包括患者信息、医疗机构信息、就诊信息、费用明细、费用汇总等。如果已加载文档，可以直接调用无需参数；否则需要提供文档路径。",
        func=extract_medical_bill_info
    ),
    Tool(
        name="extract_contract_info",
        description="从合同中提取结构化信息，包括合同基本信息、合同双方信息、合同标的、关键条款、金额信息等。如果已加载文档，可以直接调用无需参数；否则需要提供文档路径。",
        func=extract_contract_info
    ),
    Tool(
        name="extract_resume_info",
        description="从简历中提取结构化信息，包括个人信息、教育经历、工作经历、技能等。如果已加载文档，可以直接调用无需参数；否则需要提供文档路径。",
        func=extract_resume_info
    ),
    Tool(
        name="extract_product_specs",
        description="从产品文档中提取产品规格和技术参数，包括产品名称、型号、技术参数、功能特性、价格等。如果已加载文档，可以直接调用无需参数；否则需要提供文档路径。",
        func=extract_product_specs
    ),
    Tool(
        name="extract_api_info",
        description="从技术文档中提取API接口信息，包括API端点、请求方法、请求参数、响应格式等。如果已加载文档，可以直接调用无需参数；否则需要提供文档路径。",
        func=extract_api_info
    ),
    Tool(
        name="format_data",
        description="将提取的数据格式化为标准格式（JSON、CSV等）。如果已加载文档，可以直接调用无需参数；否则需要提供文档路径。",
        func=format_data
    )
]

from langchain.agents import create_agent
from langchain_community.chat_models import ChatTongyi

llm = ChatTongyi(
    model="qwen-max",
    dashscope_api_key=os.getenv("DASHSCOPE_API_KEY"),
    temperature=0.2,  # 使用较低温度以获得更确定性的输出
)

agent = create_agent(
    model=llm,
    tools=tools,
    debug=True,
    system_prompt="""你是一个专业的信息提取助手。你的任务是帮助用户：
                    1. 从文档中提取结构化信息（发票、医疗票据、合同、简历、产品规格、API接口等）
                    2. 将提取的信息格式化为标准格式（JSON、CSV等）
                    3. 验证提取数据的完整性和准确性

                    在回答时，请：
                    - 提供结构化的提取结果
                    - 使用JSON或表格格式展示数据
                    - 如果数据不完整，说明缺失的部分
                    - 使用工具获取准确的信息，不要猜测
                    - 对于财务类文档（发票、医疗票据），确保金额和税务信息的准确性
                    - 对于合同文档，重点关注关键条款和风险点
                    """
)

agent = InformationExtractionAgent()

# 1. 加载文档
agent.load_document("./extraction_documents/invoice.pdf")

# 2. 提取信息
response = agent.query("从发票中提取发票代码、发票号码、销售方和购买方信息、商品明细和金额")
print(response)

# 加载医疗票据文档
agent.load_document("./extraction_documents/medical_bill.pdf")

# 提取信息
response = agent.query("从医疗票据中提取患者姓名、医院名称、诊断结果、总费用和医保支付金额")
print(response)

# 加载合同文档
agent.load_document("./extraction_documents/contract.pdf")

# 提取信息
response = agent.query("从合同中提取合同编号、甲方和乙方信息、合同金额、付款方式和违约责任")
print(response)

# 加载简历文档
agent.load_document("./extraction_documents/resume.pdf")

# 提取信息
response = agent.query("从简历中提取姓名、联系方式、教育经历和工作经历")
print(response)

# 加载产品文档
agent.load_document("./extraction_documents/product_spec.pdf")

# 提取信息
response = agent.query("从产品文档中提取产品名称、型号、技术参数和价格")
print(response)

# 加载技术文档
agent.load_document("./extraction_documents/api_docs.pdf")

# 提取信息
response = agent.query("从技术文档中提取所有API端点、请求方法和参数")
print(response)