Maas-Model-platform/backend/scripts/test_qwen_ocr.py

"""
测试通义千问OCR客户端（DashScope SDK版本）

用于验证OCR识别功能是否正常工作
"""
import asyncio
import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from app.services.qwen_ocr import QwenOCRClient
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()


async def test_ocr():
    """测试OCR识别"""
    
    # 获取API Key
    api_key = os.getenv('DASHSCOPE_API_KEY')
    if not api_key:
        print("❌ 错误：未找到 DASHSCOPE_API_KEY 环境变量")
        print("请在 .env 文件中配置：DASHSCOPE_API_KEY=your_api_key")
        return
    
    print("=" * 60)
    print("通义千问OCR测试（DashScope SDK）")
    print("=" * 60)
    
    # 创建客户端
    client = QwenOCRClient(api_key=api_key)
    print("✅ OCR客户端初始化成功")
    
    # 测试图片URL（使用官方文档中的示例图片）
    test_image_url = "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/ctdzex/biaozhun.jpg"
    
    print(f"\n📷 测试图片: {test_image_url}")
    print("\n" + "=" * 60)
    
    # 测试不同的任务类型
    test_cases = [
        ("advanced_recognition", None, "高精识别（内置任务）"),
        ("general_ocr", None, "通用文字识别（内置任务）"),
        ("table_parsing", None, "表格解析（内置任务）"),
        ("advanced_recognition", "请识别图片中的所有文字，按原文顺序输出", "高精识别（自定义提示词）"),
    ]
    
    for task_type, custom_prompt, description in test_cases:
        print(f"\n🔍 测试场景: {description}")
        print(f"   任务类型: {task_type}")
        if custom_prompt:
            print(f"   自定义提示词: {custom_prompt}")
        else:
            print(f"   使用内置任务: ocr_options={{\"task\": \"{task_type}\"}}")
        print("-" * 60)
        
        try:
            result = await client.recognize(
                image_url=test_image_url,
                task_type=task_type,
                custom_prompt=custom_prompt
            )
            
            text = result.get("text", "")
            input_tokens = result.get("input_tokens", 0)
            output_tokens = result.get("output_tokens", 0)
            
            print(f"✅ 识别成功")
            print(f"   输入Token: {input_tokens}")
            print(f"   输出Token: {output_tokens}")
            print(f"   文本长度: {len(text)} 字符")
            print(f"\n📝 识别结果:")
            print("-" * 60)
            # 只显示前500个字符
            if len(text) > 500:
                print(text[:500] + "\n... (已截断)")
            else:
                print(text)
            print("-" * 60)
            
        except Exception as e:
            print(f"❌ 识别失败: {str(e)}")
        
        print()
    
    print("=" * 60)
    print("测试完成")
    print("=" * 60)


if __name__ == "__main__":
    asyncio.run(test_ocr())