api_test.py

import argparse
import time
from model import simple_llm
import cv2
import numpy as np
from mss import mss

def test_simple_chat(model="gpt-4o"):
    """测试基础的OpenAI API调用"""
    start_time = time.time()
    
    response = simple_llm(
        name=model,
        content="Hello!",
        sys="You are a helpful assistant."
    )
    
    end_time = time.time()
    
    print("回答内容:", response)
    print(f"耗时: {end_time - start_time:.2f}秒")

def test_screen_analysis(model="Qwen-VL-2.5"):
    """测试屏幕分析功能"""
    start_time = time.time()
    
    # 使用mss截取屏幕
    with mss() as sct:
        # 截取主屏幕
        monitor = {"top": 0, "left": 0, "width": 800, "height": 600}
        screenshot = sct.grab(monitor)
        
        # 转换为OpenCV格式
        img = np.array(screenshot)
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        
        # 调用模型分析屏幕
        prompt = "请分析这个屏幕截图并描述你看到的内容"
        response = simple_llm(
            name=model,
            content=prompt,
            img_input=img,
            sys="你是一个专业的屏幕分析助手，擅长分析和描述屏幕内容。"
        )
    
    end_time = time.time()
    
    print("分析结果:", response)
    print(f"耗时: {end_time - start_time:.2f}秒")

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='API测试工具')
    parser.add_argument('test_type', choices=['chat', 'screen'], 
                      help='选择测试类型: chat (基础聊天) 或 screen (屏幕分析)')
    parser.add_argument('--model', type=str,help='指定要使用的模型', default="gpt-4o-mini")
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 根据参数执行相应的测试
    if args.test_type == 'chat':
        print(f"执行基础聊天测试 (使用模型: {args.model})...")
        test_simple_chat(args.model)
    else:  # screen
        print(f"执行屏幕分析测试 (使用模型: {args.model})...")
        test_screen_analysis(args.model)

if __name__ == "__main__":
    main()
#python api_test.py chat --model gpt-4o