import json, sys, time, urllib.request

PORT = sys.argv[1] if len(sys.argv) > 1 else "8080"
MODEL = "Qwen3.8-27B-Uncensored-HauhauCS-Aggressive-Q4_K_P.gguf"

def call(messages, max_tokens=300, tools=None):
    body = {"model": MODEL, "messages": messages, "max_tokens": max_tokens, "stream": True}
    if tools:
        body["tools"] = tools
    req = urllib.request.Request(f"http://localhost:{PORT}/v1/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"})
    t0 = time.time(); ttf = None; n_gen = 0; n_reason = 0; toolcall = None
    with urllib.request.urlopen(req) as r:
        for line in r:
            line = line.decode().strip()
            if not line.startswith("data: ") or line[6:] == "[DONE]":
                continue
            d = json.loads(line[6:])
            delta = d["choices"][0].get("delta", {})
            c = delta.get("content"); rc = delta.get("reasoning_content")
            if rc: n_reason += 1
            if c: n_gen += 1
            if delta.get("tool_calls"):
                toolcall = delta["tool_calls"][0].get("function", {})
            if (c or rc) and ttf is None:
                ttf = time.time() - t0
    tot = time.time() - t0
    n = n_gen + n_reason
    return dict(ttft=ttf, total=tot, out=n, gen=n_gen, reason=n_reason,
                speed=n / max(tot - ttf, 0.01), toolcall=toolcall)

long_ctx = "背景资料。" + ("人工智能是计算机科学的一个重要分支。" * 700)
TOOLSCHEMA = [{"type": "function", "function": {"name": "get_weather", "parameters": {
    "type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]

r = call([{"role": "user", "content": "你好，1+1等于几？简短回答。"}], max_tokens=100)
print(f"[short]     ttft={r['ttft']:.2f}s total={r['total']:.2f}s out={r['out']} (reason={r['reason']}) {r['speed']:.1f} t/s")

r = call([{"role": "user", "content": long_ctx + " 请总结上面的资料，50字以内。"}], max_tokens=150)
print(f"[16k ctx]   ttft={r['ttft']:.2f}s total={r['total']:.2f}s out={r['out']} (reason={r['reason']}) {r['speed']:.1f} t/s")

r = call([{"role": "user", "content": long_ctx + " 再总结一次，30字以内。"}], max_tokens=150)
print(f"[16k cached] ttft={r['ttft']:.2f}s total={r['total']:.2f}s out={r['out']} (reason={r['reason']}) {r['speed']:.1f} t/s")

r = call([{"role": "user", "content": "请计算 27*43 + 15 等于多少？只给最终数字。"}], max_tokens=200)
print(f"[math]      out={r['out']} toolcall={r['toolcall']}")

r = call([{"role": "system", "content": "You are a helpful assistant with tools. Use tools when appropriate."},
          {"role": "user", "content": "What is the weather in Beijing right now?"}], max_tokens=200, tools=TOOLSCHEMA)
print(f"[toolcall]  out={r['out']} tool={json.dumps(r['toolcall'], ensure_ascii=False)}")