Ollama API 与客户端接入

"""最小代码：不装任何第三方库，直接用 HTTP 调本地模型对话。

理解这 30 行，后面所有客户端（ChatBox、LangChain、你自己的后端）
都只是在这套请求上加包装。

    python3 min_chat.py
"""
import json
import os
import urllib.request

BASE = os.environ.get("OLLAMA_BASE", "http://127.0.0.1:11434")
MODEL = os.environ.get("OLLAMA_MODEL", "qwen2:1.5b")

# ① 请求体：model 用哪个模型，messages 是整段对话历史
payload = {
    "model": MODEL,
    "messages": [
        # system 定人设，放在最前面，整轮对话都受它约束
        {"role": "system", "content": "你是一个简洁的中文助手，回答不超过两句话。"},
        {"role": "user", "content": "什么是私有化部署？"},
    ],
    # stream 为 False：等模型全部生成完，一次性返回一个 JSON
    "stream": False,
    # options 里放生成参数，字段名和 /set parameter 里的一模一样
    "options": {"temperature": 0.3, "num_predict": 128},
}

# ② 发请求：路径固定是 /api/chat，方法是 POST，body 是 JSON
req = urllib.request.Request(
    BASE.rstrip("/") + "/api/chat",
    data=json.dumps(payload).encode("utf-8"),
    headers={"Content-Type": "application/json"},
)

with urllib.request.urlopen(req, timeout=120) as resp:
    result = json.loads(resp.read().decode("utf-8"))

# ③ 取答案：模型说的话在 message.content 里
print("模型回答：", result["message"]["content"])

# ④ 顺带看一眼性能，耗时字段单位是纳秒，除以 1e9 才是秒
total_s = result.get("total_duration", 0) / 1e9
eval_count = result.get("eval_count", 0)
eval_s = result.get("eval_duration", 1) / 1e9
print("总耗时 %.2f 秒，生成 %d 个 token，速度 %.1f token/s"
      % (total_s, eval_count, eval_count / eval_s if eval_s else 0))

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""api_probe.py —— Ollama HTTP 接口自检脚本

用途：接入之前先跑一遍，把「服务能不能用、模型在不在、流式通不通、
耗时统计对不对」一次性问清楚，避免带着环境问题去调业务代码。

用法：
    python3 api_probe.py                      # 用默认地址和模型
    OLLAMA_BASE=http://10.0.0.8:11434 python3 api_probe.py
    OLLAMA_MODEL=qwen2:7b python3 api_probe.py

说明：
  * 只读接口 + 一次极短的对话请求，不会改动服务端任何状态。
  * 令牌走环境变量 OLLAMA_GATEWAY_TOKEN，不接受命令行参数，
    避免凭据留在 shell history 和进程列表里。
  * 退出码 0 = 全部通过，1 = 有检查项失败，便于接进 CI 或巡检。
"""

import json
import os
import sys
import time
import urllib.error
import urllib.request

BASE = os.environ.get("OLLAMA_BASE", "http://127.0.0.1:11434").rstrip("/")
MODEL = os.environ.get("OLLAMA_MODEL", "qwen2:1.5b")
TOKEN = os.environ.get("OLLAMA_GATEWAY_TOKEN", "")
TIMEOUT = float(os.environ.get("OLLAMA_TIMEOUT", "30"))

OK = "[ OK ]"
NG = "[FAIL]"

results = []


def record(name, passed, detail=""):
    """记录一项检查结果，并立刻打印，方便边跑边看。"""
    results.append((name, passed))
    flag = OK if passed else NG
    line = "%s %s" % (flag, name)
    if detail:
        line += " —— " + detail
    print(line)


def headers():
    h = {"Content-Type": "application/json"}
    if TOKEN:
        h["Authorization"] = "Bearer " + TOKEN
    return h


def http_get(path):
    req = urllib.request.Request(BASE + path, headers=headers(), method="GET")
    with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
        return json.loads(resp.read().decode("utf-8"))


def http_post(path, payload, stream=False):
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(BASE + path, data=body,
                                 headers=headers(), method="POST")
    resp = urllib.request.urlopen(req, timeout=TIMEOUT)
    if stream:
        return resp
    with resp:
        return json.loads(resp.read().decode("utf-8"))


def check_version():
    """最轻的探活方式：能拿到版本号，说明服务在、网络通、鉴权过。"""
    try:
        data = http_get("/api/version")
        ver = data.get("version", "?")
        record("服务探活 GET /api/version", True, "版本 " + str(ver))
        return True
    except urllib.error.HTTPError as exc:
        hint = "鉴权失败，检查 OLLAMA_GATEWAY_TOKEN" if exc.code in (401, 403) else ""
        record("服务探活 GET /api/version", False, "HTTP %s %s" % (exc.code, hint))
    except Exception as exc:  # 连接被拒、超时、DNS 等
        record("服务探活 GET /api/version", False,
               "%s（先确认服务在跑、监听地址和防火墙）" % exc.__class__.__name__)
    return False


def check_model_present():
    """确认目标模型真的在磁盘上，避免把「模型没下载」误判成「接口不通」。"""
    try:
        data = http_get("/api/tags")
    except Exception as exc:
        record("模型列表 GET /api/tags", False, str(exc))
        return False
    names = [m.get("name", "") for m in data.get("models", [])]
    record("模型列表 GET /api/tags", True, "共 %d 个" % len(names))
    if MODEL in names:
        record("目标模型存在：" + MODEL, True)
        return True
    # 给出最接近的候选，多数情况是版本号写漏了
    near = [n for n in names if n.split(":")[0] == MODEL.split(":")[0]]
    detail = "本地没有；同名不同版本：%s" % (", ".join(near) if near else "无")
    record("目标模型存在：" + MODEL, False, detail)
    return False


def check_chat_nostream():
    """非流式：必须显式 stream=False，否则拿到的是多行流式响应。"""
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": "只回答两个字：收到"}],
        "stream": False,
        "options": {"temperature": 0},
    }
    started = time.time()
    try:
        data = http_post("/api/chat", payload)
    except Exception as exc:
        record("非流式对话 POST /api/chat", False, str(exc))
        return None
    text = data.get("message", {}).get("content", "")
    if not text:
        record("非流式对话 POST /api/chat", False, "响应里没有 message.content")
        return None
    record("非流式对话 POST /api/chat", True,
           "耗时 %.1fs，回答 %r" % (time.time() - started, text[:20]))
    return data


def check_stats(data):
    """耗时字段单位是纳秒；load_duration 很大说明模型刚被重新加载。"""
    if not data:
        return
    eval_count = data.get("eval_count") or 0
    eval_ns = data.get("eval_duration") or 0
    load_ns = data.get("load_duration") or 0
    if eval_count and eval_ns:
        speed = eval_count / (eval_ns / 1e9)
        record("生成速度统计", True, "%.1f token/s" % speed)
    else:
        record("生成速度统计", False, "响应里缺 eval_count / eval_duration")
    load_s = load_ns / 1e9
    if load_s > 1.0:
        record("模型常驻情况", True,
               "本次加载耗时 %.1fs，模型此前不在内存，可调大 OLLAMA_KEEP_ALIVE" % load_s)
    else:
        record("模型常驻情况", True, "模型已常驻，加载耗时 %.2fs" % load_s)


def check_chat_stream():
    """流式：一行一个 JSON，逐行解析；done 那片带统计、通常没有文本。"""
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": "从 1 数到 5，只输出数字"}],
        "stream": True,
        "options": {"temperature": 0},
    }
    started = time.time()
    first_chunk_at = None
    chunks = 0
    text = ""
    try:
        resp = http_post("/api/chat", payload, stream=True)
    except Exception as exc:
        record("流式对话 stream=true", False, str(exc))
        return
    with resp:
        for raw in resp:
            line = raw.decode("utf-8").strip()
            if not line:
                continue
            try:
                piece = json.loads(line)
            except json.JSONDecodeError:
                record("流式对话 stream=true", False,
                       "有一行不是合法 JSON，不要对整个响应做 json.loads")
                return
            seg = piece.get("message", {}).get("content", "")
            if seg:
                chunks += 1
                text += seg
                if first_chunk_at is None:
                    first_chunk_at = time.time() - started
            if piece.get("done"):
                break
    if chunks <= 1:
        record("流式对话 stream=true", False,
               "只收到 %d 片，检查网关是否开了缓冲（proxy_buffering）" % chunks)
        return
    record("流式对话 stream=true", True,
           "%d 片，首字 %.2fs，全文 %r" % (chunks, first_chunk_at or 0, text[:20]))


def check_delete_blocked():
    """网关应当拦掉模型管理类接口；直连时不拦是正常的，只提示。"""
    payload = {"model": "__probe_not_exist__"}
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(BASE + "/api/delete", data=body,
                                 headers=headers(), method="DELETE")
    try:
        urllib.request.urlopen(req, timeout=TIMEOUT)
        record("危险接口 DELETE /api/delete 被拦截", False,
               "未被拦截；若这是对外网关，必须按路径拦掉管理类接口")
    except urllib.error.HTTPError as exc:
        if exc.code in (401, 403, 404, 405):
            record("危险接口 DELETE /api/delete 被拦截", True, "HTTP %s" % exc.code)
        else:
            record("危险接口 DELETE /api/delete 被拦截", True,
                   "HTTP %s（非 2xx 即未放行）" % exc.code)
    except Exception as exc:
        record("危险接口 DELETE /api/delete 被拦截", True, exc.__class__.__name__)


def main():
    print("目标服务：%s" % BASE)
    print("目标模型：%s" % MODEL)
    print("鉴权令牌：%s" % ("已注入" if TOKEN else "未设置"))
    print("-" * 56)

    if not check_version():
        print("-" * 56)
        print("服务都没探通，后面的检查没有意义，先解决连通性。")
        return 1

    has_model = check_model_present()
    data = check_chat_nostream() if has_model else None
    check_stats(data)
    if has_model:
        check_chat_stream()
    check_delete_blocked()

    print("-" * 56)
    failed = [name for name, passed in results if not passed]
    print("通过 %d/%d" % (len(results) - len(failed), len(results)))
    if failed:
        for name in failed:
            print("  未通过：%s" % name)
        return 1
    print("全部通过，可以开始接业务代码。")
    return 0


if __name__ == "__main__":
    sys.exit(main())

"""完整案例：带上下文的流式命令行对话。

两个要点：
  1. stream 为 True 时，服务端按行返回一串 JSON，每行一小片文本，
     必须逐行解析、逐行打印，用户才会看到「一个字一个字冒出来」的效果。
  2. 模型服务端不保存任何历史。想让它记住上一句，就得把整段 messages
     每次都重新发过去——这正是私有化部署里最常被误解的一点。

    python3 stream_chat.py
"""
import json
import os
import sys
import urllib.request

BASE = os.environ.get("OLLAMA_BASE", "http://127.0.0.1:11434")
MODEL = os.environ.get("OLLAMA_MODEL", "qwen2:1.5b")

# 上下文窗口有限，历史无限增长会越来越慢、最后超出 num_ctx 被截断，
# 所以只保留最近 N 轮。system 那条要单独保住，不能被裁掉。
MAX_TURNS = 20

SYSTEM = {"role": "system", "content": "你是中文助手，回答简洁，不确定就直说。"}


def chat_stream(messages):
    """发一次流式请求，边收边打印，返回模型这一轮说的完整内容。"""
    payload = {
        "model": MODEL,
        "messages": messages,
        "stream": True,
        "options": {"temperature": 0.3},
    }
    req = urllib.request.Request(
        BASE.rstrip("/") + "/api/chat",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )

    pieces = []
    with urllib.request.urlopen(req, timeout=600) as resp:
        # 响应是「一行一个 JSON」，不是一个大 JSON 数组，不能整体 json.loads
        for raw in resp:
            line = raw.decode("utf-8").strip()
            if not line:
                continue
            chunk = json.loads(line)

            # 每一片的文本增量在 message.content 里
            part = chunk.get("message", {}).get("content", "")
            if part:
                pieces.append(part)
                sys.stdout.write(part)
                sys.stdout.flush()

            # done 为 True 的那一片是最后一片，带统计信息，本身通常没有文本
            if chunk.get("done"):
                dur = chunk.get("total_duration", 0) / 1e9
                cnt = chunk.get("eval_count", 0)
                print("\n[本轮 %d token，耗时 %.1f 秒]" % (cnt, dur))
    return "".join(pieces)


def main():
    print("私有模型对话（输入 exit 退出）  模型：%s" % MODEL)
    messages = [SYSTEM]

    while True:
        try:
            user = input("\n你：").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not user:
            continue
        if user.lower() in ("exit", "quit", "/bye"):
            break

        # 把用户这句追加进历史
        messages.append({"role": "user", "content": user})

        print("模型：", end="")
        answer = chat_stream(messages)

        # 关键一步：把模型的回答也追加回历史，下一轮它才「记得」自己说过什么
        messages.append({"role": "assistant", "content": answer})

        # 裁剪历史：保留 system + 最近 MAX_TURNS 条
        if len(messages) > MAX_TURNS + 1:
            messages = [SYSTEM] + messages[-MAX_TURNS:]

    print("已退出。")


if __name__ == "__main__":
    main()

"""OpenAI 兼容层：让已有的代码不改业务逻辑就切到私有模型。

Ollama 在 /v1 下实现了一部分 OpenAI 接口，所以用 openai 官方 SDK 时，
只需要换掉 base_url，剩下的调用写法一个字都不用动。

    pip install openai
    python3 openai_compat.py
"""
import os

from openai import OpenAI

# base_url 指向本地 Ollama 的 /v1，注意结尾这个斜杠要带上
BASE_URL = os.environ.get("OLLAMA_OPENAI_BASE", "http://127.0.0.1:11434/v1/")
MODEL = os.environ.get("OLLAMA_MODEL", "qwen2:1.5b")

# 本地 Ollama 不校验 key，但 SDK 要求这个参数必须有值，随便填一个占位串。
# 如果前面挂了带鉴权的网关，就把真实令牌放进环境变量再读出来，别写进源码。
API_KEY = os.environ.get("OLLAMA_API_KEY", "ollama")

client = OpenAI(base_url=BASE_URL, api_key=API_KEY)


def ask_once(question):
    """一次性返回：和调云端 OpenAI 的写法完全一致。"""
    resp = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": "你是中文助手，回答简洁。"},
            {"role": "user", "content": question},
        ],
        temperature=0.3,
    )
    return resp.choices[0].message.content


def ask_stream(question):
    """流式返回：逐片打印。"""
    stream = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": question}],
        stream=True,
    )
    for chunk in stream:
        # 最后一片的 delta.content 可能是 None，要判空
        part = chunk.choices[0].delta.content
        if part:
            print(part, end="", flush=True)
    print()


if __name__ == "__main__":
    print("一次性返回：")
    print(ask_once("用一句话说明 Ollama 是什么"))

    print("\n流式返回：")
    ask_stream("再用一句话说明为什么企业要私有化部署")

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""embed_search.py —— 用 /api/embed 做一个最小的语义检索

用途：演示「文本转向量 → 算相似度 → 找最接近的一条」这条链路，
把接口分组里「要答案」的第三个接口 /api/embed 用起来。

用法：
    python3 embed_search.py "公司年假怎么休"
    OLLAMA_EMBED_MODEL=bge-m3 python3 embed_search.py "报销流程"

说明：
  * 只用标准库，余弦相似度手写，不依赖 numpy。
  * 嵌入模型和对话模型是两类模型，不能混用：
    用 qwen2 这类对话模型去调 /api/embed，要么报错要么质量很差。
  * 老版本的接口路径是 /api/embeddings（复数），本脚本会在
    /api/embed 返回 404 时自动回退，并把实际用的路径打印出来。
"""

import json
import math
import os
import sys
import urllib.error
import urllib.request

BASE = os.environ.get("OLLAMA_BASE", "http://127.0.0.1:11434").rstrip("/")
EMBED_MODEL = os.environ.get("OLLAMA_EMBED_MODEL", "bge-m3")
TOKEN = os.environ.get("OLLAMA_GATEWAY_TOKEN", "")
TIMEOUT = float(os.environ.get("OLLAMA_TIMEOUT", "60"))

# 假装这是知识库里已经切好的片段。真实项目里它来自数据库或文件。
DOCS = [
    "年假按入职年限计算，满一年五天，此后每满一年增加一天，上限十五天。",
    "差旅报销需在回程后十个工作日内提交，超期需部门负责人签字说明。",
    "公司内网只允许通过堡垒机访问生产服务器，禁止在办公机上保存生产密钥。",
    "病假需提供二级以上医院证明，三天以内可由直属主管直接批准。",
    "试用期为三个月，期间双方可提前三日书面通知解除劳动合同。",
]


def headers():
    h = {"Content-Type": "application/json"}
    if TOKEN:
        h["Authorization"] = "Bearer " + TOKEN
    return h


def post(path, payload):
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(BASE + path, data=body,
                                 headers=headers(), method="POST")
    with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
        return json.loads(resp.read().decode("utf-8"))


def embed(texts):
    """返回每段文本对应的向量。优先新接口，404 时回退到旧接口。"""
    try:
        data = post("/api/embed", {"model": EMBED_MODEL, "input": texts})
        # 新接口统一返回 embeddings 列表
        vectors = data.get("embeddings")
        if vectors:
            return vectors, "/api/embed"
    except urllib.error.HTTPError as exc:
        if exc.code != 404:
            raise
    # 旧接口一次只能处理一段文本，字段名也不一样
    vectors = []
    for text in texts:
        data = post("/api/embeddings", {"model": EMBED_MODEL, "prompt": text})
        vectors.append(data["embedding"])
    return vectors, "/api/embeddings"


def cosine(a, b):
    """余弦相似度：两个向量夹角越小，值越接近 1。"""
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


def main():
    query = sys.argv[1] if len(sys.argv) > 1 else "年假有几天"
    print("嵌入模型：%s" % EMBED_MODEL)
    print("查询：%s" % query)

    try:
        vectors, used_path = embed(DOCS + [query])
    except urllib.error.HTTPError as exc:
        print("请求失败 HTTP %s" % exc.code)
        if exc.code == 404:
            print("两个接口路径都不通，先 GET /api/version 确认版本再对照文档。")
        return 1
    except Exception as exc:
        print("请求失败：%s（先确认嵌入模型已经 pull 下来）" % exc)
        return 1

    print("实际使用的接口：%s" % used_path)
    print("向量维度：%d" % len(vectors[0]))
    print("-" * 56)

    q_vec = vectors[-1]
    scored = []
    for doc, vec in zip(DOCS, vectors[:-1]):
        scored.append((cosine(q_vec, vec), doc))
    scored.sort(key=lambda item: item[0], reverse=True)

    for rank, (score, doc) in enumerate(scored, 1):
        mark = "  ←最接近" if rank == 1 else ""
        print("%d. %.4f  %s%s" % (rank, score, doc[:34], mark))

    print("-" * 56)
    print("把排第一的片段塞进 /api/chat 的 system 或 user 消息里，")
    print("就是最朴素的一套检索增强问答。")
    return 0


if __name__ == "__main__":
    sys.exit(main())

"""骨架模板：把私有模型封装成一个可以给同事直接 import 的客户端。

复制这个文件到你的项目里，按 TODO 改完就能用。
它把三件事固化下来：超时、重试、上下文裁剪——这三样不做，
上线后必然被「偶发超时」和「越聊越慢」找上门。

    from client_skeleton import OllamaClient
    bot = OllamaClient()
    print(bot.ask("你好"))
"""
import json
import os
import time
import urllib.error
import urllib.request


class OllamaClient:
    def __init__(self, base=None, model=None, system=None):
        # TODO: 换成你自己的服务地址；生产环境走网关域名，不要直连 11434
        self.base = (base or os.environ.get("OLLAMA_BASE", "http://127.0.0.1:11434")).rstrip("/")
        # TODO: 换成你们确定要用的模型全名，必须带冒号版本
        self.model = model or os.environ.get("OLLAMA_MODEL", "qwen2:1.5b")
        # TODO: 写你自己的人设
        self.system = system or "你是中文助手，回答简洁准确，不确定就直说。"
        # 令牌一律走环境变量，不准硬编码
        self.token = os.environ.get("OLLAMA_GATEWAY_TOKEN", "")

        # TODO: 按业务调整
        self.timeout = 120        # 单次请求超时（秒）；长文本生成要调大
        self.retries = 2          # 失败重试次数
        self.max_history = 20     # 保留最近多少条消息，防止上下文无限膨胀

        self.history = []

    # ---------- 内部：发一次请求 ----------
    def _post(self, path, payload):
        url = self.base + path
        headers = {"Content-Type": "application/json"}
        if self.token:
            headers["Authorization"] = "Bearer " + self.token
        body = json.dumps(payload).encode("utf-8")

        last = None
        for attempt in range(self.retries + 1):
            try:
                req = urllib.request.Request(url, data=body, headers=headers)
                with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                    return json.loads(resp.read().decode("utf-8"))
            except (urllib.error.URLError, TimeoutError) as exc:
                last = exc
                if attempt < self.retries:
                    # 退避重试：模型正在加载时第一次请求很容易超时
                    time.sleep(2 ** attempt)
        raise RuntimeError("调用 %s 失败：%s" % (url, last))

    # ---------- 对外：单轮问答，不带历史 ----------
    def ask(self, question, **options):
        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": self.system},
                {"role": "user", "content": question},
            ],
            "stream": False,
            # TODO: 确定性任务把 temperature 调到 0.1~0.3
            "options": {"temperature": 0.3, **options},
        }
        return self._post("/api/chat", payload)["message"]["content"]

    # ---------- 对外：多轮对话，自己维护历史 ----------
    def chat(self, question, **options):
        self.history.append({"role": "user", "content": question})
        payload = {
            "model": self.model,
            "messages": [{"role": "system", "content": self.system}] + self.history,
            "stream": False,
            "options": {"temperature": 0.3, **options},
        }
        answer = self._post("/api/chat", payload)["message"]["content"]

        # 模型的回答必须也存进历史，否则下一轮它不知道自己说过什么
        self.history.append({"role": "assistant", "content": answer})
        if len(self.history) > self.max_history:
            self.history = self.history[-self.max_history:]
        return answer

    def reset(self):
        """清空上下文，相当于对话里的 /clear。"""
        self.history = []


if __name__ == "__main__":
    bot = OllamaClient()
    print(bot.ask("用一句话介绍你自己"))

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""history_budget.py —— 上下文预算计算器

用途：回答「我的 num_ctx 够撑几轮对话」「历史该保留几条」这两个问题。
不去猜，按字数估 token，把预算算出来再定裁剪策略。

用法：
    python3 history_budget.py
    NUM_CTX=8192 AVG_USER_CHARS=120 python3 history_budget.py

估算口径（够用即可，不追求精确）：
  * 中文约 1.5 字/token，英文约 4 字符/token；本脚本按中文为主估。
  * 一轮对话 = 一条 user + 一条 assistant。
  * 生成也要占窗口，所以必须预留输出空间，不能把窗口全给历史。
"""

import os

NUM_CTX = int(os.environ.get("NUM_CTX", "4096"))
SYSTEM_CHARS = int(os.environ.get("SYSTEM_CHARS", "200"))
AVG_USER_CHARS = int(os.environ.get("AVG_USER_CHARS", "60"))
AVG_ASSIST_CHARS = int(os.environ.get("AVG_ASSIST_CHARS", "300"))
RESERVE_OUTPUT = int(os.environ.get("RESERVE_OUTPUT", "512"))

CHARS_PER_TOKEN = 1.5


def to_tokens(chars):
    """按中文口径把字数折成 token 数，向上取整。"""
    return int(chars / CHARS_PER_TOKEN + 0.999)


def main():
    sys_tokens = to_tokens(SYSTEM_CHARS)
    turn_tokens = to_tokens(AVG_USER_CHARS) + to_tokens(AVG_ASSIST_CHARS)
    usable = NUM_CTX - RESERVE_OUTPUT - sys_tokens

    print("上下文窗口 num_ctx      : %d token" % NUM_CTX)
    print("预留给生成的输出空间    : %d token" % RESERVE_OUTPUT)
    print("system 人设占用         : %d token（%d 字）" % (sys_tokens, SYSTEM_CHARS))
    print("可用于历史的额度        : %d token" % usable)
    print("单轮对话平均占用        : %d token（问 %d 字 + 答 %d 字）"
          % (turn_tokens, AVG_USER_CHARS, AVG_ASSIST_CHARS))
    print("-" * 56)

    if usable <= 0:
        print("窗口已被 system 和输出预留吃满，必须调大 num_ctx 或精简人设。")
        return 1

    max_turns = usable // turn_tokens
    print("理论可承载轮数          : %d 轮" % max_turns)

    # 留三成余量：真实对话里总有几轮特别长，卡着上限跑必然被截断。
    safe_turns = int(max_turns * 0.7)
    safe_msgs = max(safe_turns * 2, 2)
    print("建议保留最近            : %d 轮（即 messages 里 %d 条，不含 system）"
          % (safe_turns, safe_msgs))
    print("-" * 56)

    print("写进代码就是这一行：")
    print("    messages = [SYSTEM] + history[-%d:]" % safe_msgs)
    print()
    print("注意 system 要单独拼在前面，不能让它进入被裁剪的那一段。")
    print()

    print("不同窗口下的对照（按当前平均长度估）：")
    print("%-12s %-12s %-12s" % ("num_ctx", "可用额度", "建议保留条数"))
    for ctx in (2048, 4096, 8192, 16384, 32768):
        avail = ctx - RESERVE_OUTPUT - sys_tokens
        if avail <= 0:
            print("%-12d %-12s %-12s" % (ctx, "不足", "-"))
            continue
        turns = avail // turn_tokens
        msgs = max(int(turns * 0.7) * 2, 2)
        print("%-12d %-12d %-12d" % (ctx, avail, msgs))

    print("-" * 56)
    print("提醒：num_ctx 调大会同时抬高显存占用，")
    print("显存不够时模型会溢出到 CPU，速度断崖式下跌——")
    print("窗口不是越大越好，够用就行。")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())