BJQA/DeepSeekAPI/findMax.py

import json
import requests
import string
import random
import time

# 随机字符串生成器
def generate_random_text(length):
    return ''.join(random.choices(string.ascii_letters + string.digits + " ", k=length))

# 测试 prompt 长度上限
def test_prompt_limit(base_prompt, model, url, step=1000, max_len=1000000):
    current_length = 0
    success_length = 0

    while current_length <= max_len:
        random_text = generate_random_text(current_length)
        test_prompt = base_prompt + "\n附加内容：" + random_text

        payload = {
            "model": model,
            "prompt": test_prompt,
            "stream": False
        }

        try:
            print(f"🔎 正在测试长度: {len(test_prompt)} 字符 ... ", end="")
            response = requests.post(url, json=payload, timeout=30)

            if response.status_code == 200:
                result = response.json()
                output = result.get("response", "")
                print(f"✅ 成功，模型响应 {len(output)} 字符")
                success_length = current_length
                current_length += step
                time.sleep(0.5)  # 防止接口过载
            else:
                print(f"❌ 失败，状态码: {response.status_code}")
                break

        except Exception as e:
            print(f"❌ 异常: {str(e)}")
            break

    print("\n✅ 测试完成")
    print(f"模型可接受的最大 prompt 追加内容约为：{success_length} 字符（约 {success_length // 4} tokens）")

# === 你的基础 prompt（可从文件加载，也可直接写入） ===
base_prompt = "你是一个AI助手，以下是我想给你的内容："

# === 调用测试 ===
OLLAMA_URL = "http://192.168.1.207:11434/api/generate"
MODEL_NAME = "deepseek-r1:8b"  # 替换成你实际加载的模型名

test_prompt_limit(
    base_prompt=base_prompt,
    model=MODEL_NAME,
    url=OLLAMA_URL,
    step=10000,      # 每次增加 1 万字符
    max_len=500000   # 最大 50 万字符（可根据模型能力增大）
)