# test_deepseek.py
import requests
import json
import time

def test_deepseek_model():
    """测试 DeepSeek 模型"""
    base_url = "http://localhost:11434"
    
    # 首先确保模型已下载
    print("检查模型列表...")
    try:
        response = requests.get(f"{base_url}/api/tags")
        models = response.json().get("models", [])
        print(f"已安装模型: {[model['name'] for model in models]}")
    except Exception as e:
        print(f"获取模型列表失败: {e}")
        return False
    
    # 测试模型响应
    test_prompt = "你好，请用中文介绍一下你自己，并说明你能在浏览器自动化中提供什么帮助？"
    
    try:
        print("测试 DeepSeek 模型响应...")
        response = requests.post(
            f"{base_url}/api/generate",
            json={
                "model": "deepseek-coder:6.7b",
                "prompt": test_prompt,
                "stream": False,
                "options": {
                    "temperature": 0.7,
                    "top_p": 0.9
                }
            },
            timeout=60
        )
        
        if response.status_code == 200:
            result = response.json()
            print("=" * 50)
            print("DeepSeek 模型测试成功!")
            print(f"回复: {result.get('response', '')}")
            print("=" * 50)
            return True
        else:
            print(f"模型响应失败: {response.status_code}")
            print(f"错误信息: {response.text}")
            return False
            
    except Exception as e:
        print(f"测试失败: {e}")
        return False

if __name__ == "__main__":
    # 等待服务启动
    print("等待 Ollama 服务启动...")
    time.sleep(10)
    
    success = test_deepseek_model()
    if success:
        print("DeepSeek 模型部署成功！")
    else:
        print("DeepSeek 模型部署失败，请检查。")