多并发请求

import json
import aiohttp
import asyncio
import base64
import os
from typing import List, Dict, Any
from pathlib import Path
import glob


class AsyncExtractClient:
    def __init__(self, app_id: str, secret_code: str):
        self.app_id = app_id
        self.secret_code = secret_code

    async def extract(self, file_content: bytes, options: dict, request_body: dict) -> str:
        # 将文件内容转换为base64
        file_base64 = base64.b64encode(file_content).decode('utf-8')

        # 构建请求参数
        params = {}
        for key, value in options.items():
            params[key] = str(value)

        # 设置请求头
        headers = {
            "x-ti-app-id": self.app_id,
            "x-ti-secret-code": self.secret_code,
            "Content-Type": "application/json"
        }

        # 构建完整的请求体，包含base64文件数据
        full_request_body = {
            **request_body,
            "file": file_base64  # 添加base64编码的文件数据
        }

        # 发送异步请求
        async with aiohttp.ClientSession() as session:
            async with session.post(
                    "https://api.textin.com/ai/service/v2/entity_extraction",
                    params=params,
                    headers=headers,
                    json=full_request_body
            ) as response:
                # 检查响应状态
                response.raise_for_status()
                return await response.text()

    async def extract_single_file(self, file_path: str, output_dir: str, options: dict, request_body: dict) -> Dict[
        str, Any]:
        """处理单个文件的提取任务"""
        try:
            # 读取文件
            with open(file_path, "rb") as f:
                file_content = f.read()

            # 执行提取
            response = await self.extract(file_content, options, request_body)

            # 生成输出文件名
            file_name = Path(file_path).stem
            output_file = os.path.join(output_dir, f"{file_name}_result.json")

            # 保存结果到输出文件夹
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(response)

            return {
                "file_path": file_path,
                "output_file": output_file,
                "status": "success",
                "response": response
            }
        except Exception as e:
            return {
                "file_path": file_path,
                "status": "error",
                "error": str(e)
            }

    async def extract_folder(self, input_folder: str, output_folder: str, options: dict, request_body: dict,
                             file_extensions: List[str] = None, max_concurrent: int = 5) -> List[Dict[str, Any]]:
        """处理整个文件夹的文件"""

        # 确保输出文件夹存在
        os.makedirs(output_folder, exist_ok=True)

        # 如果没有指定文件扩展名，默认处理API支持的文档格式
        if file_extensions is None:
            file_extensions = [
                '*.png', '*.jpg', '*.jpeg', '*.pdf', '*.bmp', '*.tiff', '*.webp',
                '*.doc', '*.docx', '*.html', '*.mhtml', '*.xls', '*.xlsx', '*.csv',
                '*.ppt', '*.pptx', '*.txt', '*.ofd', '*.rtf'
            ]

        # 获取所有匹配的文件
        file_paths = []
        for ext in file_extensions:
            pattern = os.path.join(input_folder, ext)
            file_paths.extend(glob.glob(pattern))
            # 也搜索子文件夹
            file_paths.extend(glob.glob(os.path.join(input_folder, "**", ext), recursive=True))

        # 去重并排序
        file_paths = sorted(list(set(file_paths)))

        if not file_paths:
            print(f"在文件夹 {input_folder} 中没有找到匹配的文件")
            return []

        print(f"找到 {len(file_paths)} 个文件待处理")

        # 创建信号量限制并发数
        semaphore = asyncio.Semaphore(max_concurrent)

        async def extract_with_semaphore(file_path: str) -> Dict[str, Any]:
            async with semaphore:
                return await self.extract_single_file(file_path, output_folder, options, request_body)

        # 创建所有任务
        tasks = [extract_with_semaphore(file_path) for file_path in file_paths]

        # 并发执行所有任务
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # 处理异常结果
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                processed_results.append({
                    "file_path": file_paths[i],
                    "status": "error",
                    "error": str(result)
                })
            else:
                processed_results.append(result)

        return processed_results


async def main():
    # 创建客户端实例，需替换你的API Key
    client = AsyncExtractClient("你的x-ti-app-id", "你的x-ti-secret-code")

    # 指定输入和输出文件夹
    input_folder = "input_files"  # 输入文件夹路径
    output_folder = "output_results"  # 输出文件夹路径

    # 设置URL参数
    options = dict(
        parse_mode="scan"
    )

    # 设置请求体参数
    request_body = dict(
        prompt="请抽取这张小票中的实付金额、消费日期、店铺名称、订单号并以字段格式返回，请抽取货号、商品名称、数量、单价，并以表格格式返回"
    )

    # 可选：指定要处理的文件扩展名
    file_extensions = ['*.pdf', '*.jpg', '*.jpeg', '*.png']  # 可以根据需要调整

    try:
        # 处理整个文件夹
        results = await client.extract_folder(
            input_folder=input_folder,
            output_folder=output_folder,
            options=options,
            request_body=request_body,
            file_extensions=file_extensions,
            max_concurrent=5  # 可以根据需要调整并发数
        )

        # 统计处理结果
        success_count = sum(1 for r in results if r["status"] == "success")
        error_count = len(results) - success_count

        print(f"\n处理完成！")
        print(f"总文件数: {len(results)}")
        print(f"成功处理: {success_count}")
        print(f"处理失败: {error_count}")

        # 显示失败的文件
        if error_count > 0:
            print("\n失败的文件:")
            for result in results:
                if result["status"] == "error":
                    print(f"  - {result['file_path']}: {result['error']}")

        # 保存处理汇总信息
        summary_file = os.path.join(output_folder, "processing_summary.json")
        summary = {
            "input_folder": input_folder,
            "output_folder": output_folder,
            "total_files": len(results),
            "success_count": success_count,
            "error_count": error_count,
            "results": results
        }

        with open(summary_file, "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=2)

        print(f"\n处理汇总已保存到: {summary_file}")

    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    # 运行异步主函数
    asyncio.run(main())