code/03-llm-apps/conversation.py

"""命令行多轮对话。

    python conversation.py                 正常聊天，输入空行退出
    python conversation.py --stateless     不保留历史，看看模型会不会"失忆"
    python conversation.py --max-messages 4 --summarize
                                           历史超过 4 条时，把旧的部分压缩成摘要

也可以用管道一次喂多行：printf "第一句\n第二句\n" | python conversation.py
"""
import argparse
import os
import sys

from openai import OpenAI

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.deepseek.com"),
)
MODEL = os.environ.get("LLM_MODEL", "deepseek-flash")
SYSTEM = {"role": "system", "content": "你是一个说话简短的编程助手，每次回答不超过三句话。"}
NO_THINKING = {"thinking": {"type": "disabled"}}

parser = argparse.ArgumentParser()
parser.add_argument("--stateless", action="store_true")
parser.add_argument("--max-messages", type=int, default=20)
parser.add_argument("--summarize", action="store_true")
args = parser.parse_args()

history = []  # 只放 user 和 assistant 消息，system 每次单独加在最前面


def summarize(messages):
    text = "\n".join(f"{m['role']}: {m['content']}" for m in messages)
    response = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": "把下面这段对话压缩成几句话的摘要，保留人名、项目、偏好等以后可能用到的事实：\n\n" + text}],
        extra_body=NO_THINKING,
    )
    return response.choices[0].message.content


def trim(history):
    if len(history) <= args.max_messages:
        return history
    # 保留最近的一半，更早的要么直接丢掉，要么压缩成一条摘要
    keep = history[-(args.max_messages // 2):]
    dropped = history[: len(history) - len(keep)]
    if not args.summarize:
        print(f"  [丢掉了最早的 {len(dropped)} 条消息]")
        return keep
    summary = summarize(dropped)
    print(f"  [把最早的 {len(dropped)} 条消息压缩成了摘要：{summary}]")
    return [{"role": "user", "content": f"（之前对话的摘要：{summary}）"},
            {"role": "assistant", "content": "好的，我记住了。"}] + keep


interactive = sys.stdin.isatty()
while True:
    try:
        user_text = input("你：" if interactive else "").strip()
    except EOFError:
        break
    if not user_text:
        break
    if not interactive:
        print(f"你：{user_text}")

    messages = [SYSTEM] + ([] if args.stateless else history) + [{"role": "user", "content": user_text}]
    response = client.chat.completions.create(model=MODEL, messages=messages, extra_body=NO_THINKING)
    answer = response.choices[0].message.content
    print(f"助手：{answer}  （本轮输入 {response.usage.prompt_tokens} 词元）")

    if not args.stateless:
        history += [{"role": "user", "content": user_text}, {"role": "assistant", "content": answer}]
        history = trim(history)