【案例】Streamlit + LangChain 智能聊天机器人

# -*- coding: utf-8 -*-
"""
最小可跑：Streamlit 点单台 + LangChain 后厨
============================================
全文只有三件事：

  1. 建一条带记忆的链（编排层，不认识 Streamlit）
  2. 页面把会话 ID 和这一轮问题递进去
  3. 把返回的文本画成气泡

启动：streamlit run min_app.py
依赖：pip install streamlit langchain langchain-openai
"""

import os
import uuid

import streamlit as st
from langchain.chat_models import init_chat_model
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory

# ---------------------------------------------------------------- 后厨
_store = {}


def get_history(session_id: str) -> InMemoryChatMessageHistory:
    """一个会话 ID 对应一块独立的备忘板。"""
    if session_id not in _store:
        _store[session_id] = InMemoryChatMessageHistory()
    return _store[session_id]


@st.cache_resource(show_spinner=False)
def build_chain():
    """@st.cache_resource 让这条链只建一次，不随页面重跑反复创建。"""
    model = init_chat_model(
        os.environ.get("CHAT_MODEL", "openai:gpt-4o-mini"),
        temperature=0.3,
        api_key=os.environ.get("OPENAI_API_KEY"),
    )
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个中文智能助手，回答简洁、准确。"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ])
    chain = prompt | model | StrOutputParser()
    return RunnableWithMessageHistory(
        chain,
        get_history,
        input_messages_key="question",
        history_messages_key="history",
    )


# ---------------------------------------------------------------- 前厅
st.title("智能聊天机器人")

# 会话 ID 由服务端生成，存进 session_state，一个浏览器标签一份
if "sid" not in st.session_state:
    st.session_state["sid"] = str(uuid.uuid4())
if "bubbles" not in st.session_state:
    st.session_state["bubbles"] = []

for role, text in st.session_state["bubbles"]:
    with st.chat_message(role):
        st.markdown(text)

question = st.chat_input("说点什么")
if question:
    st.session_state["bubbles"].append(("user", question))
    with st.chat_message("user"):
        st.markdown(question)

    conversation = build_chain()
    cfg = {"configurable": {"session_id": st.session_state["sid"]}}
    with st.chat_message("assistant"):
        # 链的 stream 直接就是生成器，交给界面边收边渲染
        answer = st.write_stream(
            conversation.stream({"question": question}, config=cfg))
    st.session_state["bubbles"].append(("assistant", answer))

# -*- coding: utf-8 -*-
"""
编排层：chat_core.py —— 整个后厨都在这一个文件里
==================================================
它对外只暴露四个函数，界面层只认这四个：

  reply(session_id, question)          一次性拿完整答案
  reply_stream(session_id, question)   生成器，边算边吐
  history_of(session_id)               取这块备忘板上的消息，用来重放气泡
  reset(session_id)                    清空这块备忘板

这个文件不导入 streamlit，也不打印任何东西。
换 Flask、换命令行、换企业微信回调，都能原样搬走。
"""

import os
from typing import Iterator, List

from langchain.chat_models import init_chat_model
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory

from history_store import get_history
from trim_policy import trim_history

SYSTEM_PROMPT = os.environ.get(
    "SYSTEM_PROMPT",
    "你是一个中文智能助手。回答简洁、准确、有条理。\n"
    "只依据用户提供的信息和你确知的常识回答；不确定就直说不确定，不要编造。",
)

MODEL_NAME = os.environ.get("CHAT_MODEL", "openai:gpt-4o-mini")
TEMPERATURE = float(os.environ.get("CHAT_TEMPERATURE", "0.3"))


def _build_prompt() -> ChatPromptTemplate:
    """三段式：人设在最前、历史插槽在中间、本轮输入在最后。

    MessagesPlaceholder 就是「把备忘板整块贴进来」的那个插槽，
    顺序写错（历史放到本轮输入之后）会让模型把旧话当成最新指令。
    """
    return ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ])


def _build_model():
    """密钥与模型名全部走环境变量，源码里不出现任何凭据。"""
    return init_chat_model(
        MODEL_NAME,
        temperature=TEMPERATURE,
        api_key=os.environ.get("OPENAI_API_KEY"),
        timeout=float(os.environ.get("CHAT_TIMEOUT", "60")),
        max_retries=int(os.environ.get("CHAT_MAX_RETRIES", "2")),
    )


def _history_with_trim(session_id: str) -> BaseChatMessageHistory:
    """工厂函数：取出这个会话的备忘板，顺手按预算修剪一次。

    修剪放在这里而不是放在界面层，是因为「一块板子最多贴多少张便签」
    属于后厨的规矩，前厅不该知道。
    """
    history = get_history(session_id)
    trim_history(history)
    return history


_chain = _build_prompt() | _build_model() | StrOutputParser()

conversation = RunnableWithMessageHistory(
    _chain,
    _history_with_trim,
    input_messages_key="question",
    history_messages_key="history",
)


def _cfg(session_id: str) -> dict:
    return {"configurable": {"session_id": session_id}}


def reply(session_id: str, question: str) -> str:
    return conversation.invoke({"question": question}, config=_cfg(session_id))


def reply_stream(session_id: str, question: str) -> Iterator[str]:
    """把链的增量片段原样透出去；写历史由 RunnableWithMessageHistory 负责。"""
    for chunk in conversation.stream({"question": question},
                                     config=_cfg(session_id)):
        if chunk:
            yield chunk


def history_of(session_id: str) -> List[BaseMessage]:
    return list(get_history(session_id).messages)


def reset(session_id: str) -> None:
    get_history(session_id).clear()


if __name__ == "__main__":
    # 脱离界面自测：通了说明后厨没问题，之后页面出错就是前厅的事
    sid = "selftest"
    print(reply(sid, "我叫孙小空，请记住"))
    print(reply(sid, "我叫什么名字？"))
    print("这块备忘板上有 %d 条消息" % len(history_of(sid)))

# -*- coding: utf-8 -*-
"""
备忘板的容器：history_store.py
================================
会话历史要存在哪，是整个应用最容易被低估的一个决定。

  进程内字典   开发够用；重启即丢、多副本各存各的、内存只涨不降
  本地文件     单机部署可用，重启还在；并发写要加锁
  Redis / 库   多副本、要过期、要审计时的正解

换存储只要换掉 get_history 的实现，链上其余代码一个字不用改——
这正是把工厂函数单独拿出来的价值。
"""

import json
import os
import threading
from typing import Dict, List

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, messages_from_dict, messages_to_dict

BACKEND = os.environ.get("HISTORY_BACKEND", "memory")   # memory | file
DATA_DIR = os.environ.get("HISTORY_DIR", "./chat_history")

_lock = threading.Lock()
_mem: Dict[str, BaseChatMessageHistory] = {}


class FileChatMessageHistory(BaseChatMessageHistory):
    """把一个会话的消息落成一个 JSON 文件。

    只需要实现三个成员，LangChain 就认它：
      messages      属性，返回 List[BaseMessage]
      add_messages  追加若干条，传空列表要能安全返回
      clear         清空本会话，对应界面上的「新建对话」
    """

    def __init__(self, session_id: str, directory: str = DATA_DIR):
        # 会话 ID 直接当文件名会被 ../ 穿越，只留安全字符
        safe = "".join(c for c in session_id if c.isalnum() or c in "-_")[:64]
        os.makedirs(directory, exist_ok=True)
        self.path = os.path.join(directory, "%s.json" % (safe or "anonymous"))

    @property
    def messages(self) -> List[BaseMessage]:
        if not os.path.exists(self.path):
            return []
        try:
            with open(self.path, encoding="utf-8") as f:
                return messages_from_dict(json.load(f))
        except (ValueError, OSError):
            # 单个会话文件写坏，不应该让整个应用起不来
            return []

    def add_messages(self, messages: List[BaseMessage]) -> None:
        if not messages:
            return
        with _lock:
            merged = self.messages + list(messages)
            tmp = self.path + ".tmp"
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(messages_to_dict(merged), f, ensure_ascii=False)
            os.replace(tmp, self.path)     # 原子替换，进程被杀也不会留半个文件

    def clear(self) -> None:
        with _lock:
            if os.path.exists(self.path):
                os.remove(self.path)


def get_history(session_id: str) -> BaseChatMessageHistory:
    """工厂函数：给一个会话 ID，返回它专属的那块备忘板。"""
    if BACKEND == "file":
        return FileChatMessageHistory(session_id)

    from langchain_core.chat_history import InMemoryChatMessageHistory
    with _lock:
        if session_id not in _mem:
            _mem[session_id] = InMemoryChatMessageHistory()
        return _mem[session_id]


def known_sessions() -> List[str]:
    """运维用：当前有哪些会话。内存版看字典，文件版看目录。"""
    if BACKEND == "file":
        if not os.path.isdir(DATA_DIR):
            return []
        return sorted(f[:-5] for f in os.listdir(DATA_DIR)
                      if f.endswith(".json"))
    return sorted(_mem)


if __name__ == "__main__":
    from langchain_core.messages import AIMessage, HumanMessage

    h = get_history("demo-001")
    h.add_messages([HumanMessage(content="我叫孙小空"),
                    AIMessage(content="记住了")])
    print("demo-001：", len(h.messages), "条")
    print("demo-002：", len(get_history("demo-002").messages), "条")   # 0
    print("当前会话：", known_sessions())

# -*- coding: utf-8 -*-
"""
历史修剪与 Token 预算：trim_policy.py
=======================================
一块备忘板贴不下无限张便签。对话越长，每一轮送进模型的历史越长，
后果依次是：变慢 → 变贵 → 超出上下文窗口后被静默截断。

这里提供两把尺子，按需要选一把：

  按条数   trim_by_count   算得快，长短不均时不准
  按 Token trim_by_tokens  贵一点但准，生产建议用它

两把尺子都必须保证一件事：**system 那一条不能被裁掉**，
否则聊到几十轮后人设突然消失，表现为机器人「性格突变」。
"""

import os
from typing import List

from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, trim_messages

MAX_MESSAGES = int(os.environ.get("MAX_MESSAGES", "20"))
MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "3000"))
STRATEGY = os.environ.get("TRIM_STRATEGY", "count")     # count | tokens


def _approx_tokens(messages: List[BaseMessage]) -> int:
    """粗略估算：中文按 1 字 ≈ 1 token 上限估，宁可高估也不要低估。

    真实计费以模型侧 tokenizer 为准，这里只用来做预算控制。
    """
    total = 0
    for m in messages:
        text = m.content if isinstance(m.content, str) else str(m.content)
        total += len(text) + 4          # 每条消息本身的角色与分隔开销
    return total


def trim_by_count(messages: List[BaseMessage],
                  limit: int = MAX_MESSAGES) -> List[BaseMessage]:
    """保留最近 limit 条，并且从一条用户消息开始，避免以孤立回答开头。"""
    if len(messages) <= limit:
        return list(messages)
    kept = list(messages[-limit:])
    while kept and kept[0].type == "ai":
        kept.pop(0)
    return kept


def trim_by_tokens(messages: List[BaseMessage],
                   limit: int = MAX_TOKENS) -> List[BaseMessage]:
    """用 langchain_core 的 trim_messages 按预算从后往前保留。

    start_on="human" 保证裁剪后第一条是用户消息；
    include_system=True 保证 system 不参与裁剪。
    """
    return trim_messages(
        messages,
        max_tokens=limit,
        strategy="last",
        token_counter=_approx_tokens,
        start_on="human",
        include_system=True,
        allow_partial=False,
    )


def trim_history(history: BaseChatMessageHistory) -> int:
    """就地修剪一块备忘板，返回裁掉的条数。

    实现方式是「清空再写回保留部分」，因为 BaseChatMessageHistory
    的公共约定里只有 add_messages 和 clear，没有「删掉前 N 条」。
    """
    messages = list(history.messages)
    if not messages:
        return 0

    kept = (trim_by_tokens(messages) if STRATEGY == "tokens"
            else trim_by_count(messages))
    dropped = len(messages) - len(kept)
    if dropped > 0:
        history.clear()
        history.add_messages(kept)
    return dropped


if __name__ == "__main__":
    from langchain_core.messages import AIMessage, HumanMessage, SystemMessage

    demo = [SystemMessage(content="你是客服助手")]
    for i in range(1, 16):
        demo.append(HumanMessage(content="第 %d 个问题" % i))
        demo.append(AIMessage(content="第 %d 个回答" % i))

    by_count = trim_by_count(demo, 8)
    by_token = trim_by_tokens(demo, 120)
    print("原始 %d 条，估算 %d token" % (len(demo), _approx_tokens(demo)))
    print("按条数保留 %d 条，首条是 %s" % (len(by_count), by_count[0].type))
    print("按预算保留 %d 条，system 还在：%s"
          % (len(by_token), any(m.type == "system" for m in by_token)))

# -*- coding: utf-8 -*-
"""
会话标识：session_keys.py
==========================
一个线上的聊天应用永远不止一个人在用。谁看到哪块备忘板，
全靠一个 session_id 决定——**它是这个应用最敏感的一个字符串**。

三条规矩：

  1. 会话 ID 由服务端生成或校验，绝不直接采信前端传来的值
  2. 一个人可以有多段会话（新建对话 = 换一块空白备忘板）
  3. 会话 ID 不要拿用户手机号、邮箱这类可枚举的东西直接拼

这里给出两种生成方式：匿名随机 ID，和登录态派生 ID。
"""

import hashlib
import hmac
import os
import re
import uuid

SECRET = os.environ.get("SESSION_SECRET", "")
SAFE = re.compile(r"^[A-Za-z0-9_-]{8,64}$")


def new_anonymous_id() -> str:
    """匿名访客：随机一个，够长、不可猜。"""
    return "anon-" + uuid.uuid4().hex


def derive_id(user_id: str, conversation_id: str) -> str:
    """登录用户：由服务端已验证的身份派生，前端改不动。

    用 HMAC 而不是直接拼接，是为了让 ID 既稳定（同一个人同一段会话
    每次算出来都一样）又不可反推（拿到 ID 也猜不出用户是谁）。
    """
    if not SECRET:
        raise RuntimeError("未配置 SESSION_SECRET，拒绝派生会话 ID")
    raw = "%s|%s" % (user_id, conversation_id)
    digest = hmac.new(SECRET.encode("utf-8"), raw.encode("utf-8"),
                      hashlib.sha256).hexdigest()
    return "u-" + digest[:40]


def validate(session_id: str) -> str:
    """任何从外部进来的会话 ID 都要先过这一关。

    不合格就抛异常，而不是「兜底给一个默认值」——
    默认值意味着所有非法请求共用同一块备忘板，那是最糟的串话方式。
    """
    if not session_id or not SAFE.match(session_id):
        raise ValueError("非法的会话标识")
    return session_id


def fingerprint(session_id: str) -> str:
    """写日志时用这个短指纹，不要把完整会话 ID 落进日志文件。"""
    return hashlib.sha256(session_id.encode("utf-8")).hexdigest()[:8]


if __name__ == "__main__":
    a = new_anonymous_id()
    print("匿名会话：", a, "→ 校验通过：", validate(a) == a)
    print("日志指纹：", fingerprint(a))

    os.environ["SESSION_SECRET"] = "demo-secret-do-not-use-in-production"
    SECRET = os.environ["SESSION_SECRET"]
    s1 = derive_id("u1001", "c1")
    s2 = derive_id("u1001", "c2")
    s3 = derive_id("u2002", "c1")
    print("同人不同会话是否相同：", s1 == s2)      # False
    print("不同人同会话号是否相同：", s1 == s3)    # False
    print("同一入参两次是否稳定：", s1 == derive_id("u1001", "c1"))   # True

    for bad in ["", "../../etc/passwd", "abc"]:
        try:
            validate(bad)
        except ValueError as e:
            print("拒绝 %r：%s" % (bad, e))

# -*- coding: utf-8 -*-
"""
界面层：chat_app.py —— 只负责点单，不进后厨
==============================================
这个文件里没有一行 LangChain 代码。它只做四件事：

  1. 拿到这次会话的 session_id
  2. 把备忘板上的历史重放成气泡
  3. 收一句话，交给 chat_core
  4. 把返回的字流画出来

Streamlit 控件本身的机制（每次交互整脚本重跑、session_state 为什么必须用）
在「私有聊天机器人」那一讲已经讲透，这里只用结论。

启动：streamlit run chat_app.py
"""

import streamlit as st

import chat_core
import session_keys

st.set_page_config(page_title="智能聊天机器人", page_icon="💬")

# ① 会话标识：服务端生成一次，之后这一个浏览器标签一直用它
if "sid" not in st.session_state:
    st.session_state["sid"] = session_keys.new_anonymous_id()
sid = st.session_state["sid"]

# ② 侧边栏：把「当前连的是哪个模型、这块板上有多少条」摆出来
with st.sidebar:
    st.caption("会话指纹：%s" % session_keys.fingerprint(sid))
    st.caption("模型：%s" % chat_core.MODEL_NAME)
    st.caption("历史消息：%d 条" % len(chat_core.history_of(sid)))
    if st.button("新建对话"):
        chat_core.reset(sid)
        st.session_state["sid"] = session_keys.new_anonymous_id()
        st.rerun()

st.title("智能聊天机器人")

# ③ 重放历史：气泡的数据源是后厨的备忘板，界面自己不再存一份
#    两边各存一份迟早会对不上：界面上有、模型看不见，是最难查的那类 bug
ROLE = {"human": "user", "ai": "assistant"}
for message in chat_core.history_of(sid):
    role = ROLE.get(message.type)
    if role:
        with st.chat_message(role):
            st.markdown(message.content)

# ④ 收一句话，交给后厨
question = st.chat_input("输入你的问题")
if question:
    with st.chat_message("user"):
        st.markdown(question)

    with st.chat_message("assistant"):
        try:
            # write_stream 边收边画，返回值是拼好的完整文本
            st.write_stream(chat_core.reply_stream(sid, question))
        except Exception as exc:                      # noqa: BLE001
            # 后厨出故障时，页面要还在、错误要看得见，不能整页崩掉
            st.error("这次没能回答：%s" % exc)
            st.caption("稍后重试；若持续失败，检查模型服务与网络。")

# -*- coding: utf-8 -*-
"""
流式怎么一路从模型传到界面：stream_bridge.py
==============================================
链的 stream() 本身就是一个生成器，中间没有任何魔法：
模型每吐一小块，链上每一节把它加工一下往下传，最后到你手里。

这个文件演示三件事：

  1. 最朴素的消费：for chunk in chain.stream(...)
  2. 把生成器包一层，顺带统计首字延迟、片数、总字数
  3. 用 astream_events 观察「现在流的是链上哪一节」

界面侧只需要拿到一个生成器，Streamlit 的 st.write_stream 会边收边画，
并把拼好的完整文本作为返回值给你——所以「流式显示」和
「把完整回答写回历史」两件事一行就办完了。
"""

import os
import time
from typing import Iterable, Iterator

from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个中文助手，回答控制在 200 字以内。"),
    ("human", "{question}"),
])
model = init_chat_model(os.environ.get("CHAT_MODEL", "openai:gpt-4o-mini"),
                        temperature=0.3,
                        api_key=os.environ.get("OPENAI_API_KEY"))
chain = prompt | model | StrOutputParser()


def measured(chunks: Iterable[str]) -> Iterator[str]:
    """包一层做度量：首字延迟是聊天体验里最该盯的那个数字。

    注意它仍然是生成器——包装不能把流攒起来，否则就退化成一次性返回。
    """
    started = time.time()
    first_at = None
    pieces = 0
    chars = 0
    for chunk in chunks:
        if not chunk:
            continue                      # 末片常常是空的，跳过
        if first_at is None:
            first_at = time.time() - started
        pieces += 1
        chars += len(chunk)
        yield chunk
    total = time.time() - started
    print("首字 %.2fs | 共 %d 片 %d 字 | 总耗时 %.2fs"
          % (first_at or total, pieces, chars, total))


def stream_text(question: str) -> Iterator[str]:
    return measured(chain.stream({"question": question}))


async def watch_events(question: str) -> None:
    """想在界面上显示「正在检索…」「正在生成…」时用它。

    astream_events 给的是链上每一步的事件，而不只是最终文本：
    on_chat_model_stream 才是模型吐字，其余是节点的开始与结束。
    """
    async for ev in chain.astream_events({"question": question}):
        kind = ev["event"]
        if kind == "on_chat_model_stream":
            piece = ev["data"]["chunk"].content
            if piece:
                print(piece, end="", flush=True)
        elif kind in ("on_chain_start", "on_chain_end"):
            print("\n[%s] %s" % (kind, ev.get("name")))


if __name__ == "__main__":
    for piece in stream_text("用三句话解释什么是流式输出"):
        print(piece, end="", flush=True)
    print()

    # 事件版需要异步环境：
    # import asyncio; asyncio.run(watch_events("同上"))

# -*- coding: utf-8 -*-
"""
错误处理、重试与超时：resilience.py
=====================================
模型调用是一次跨网络的远程请求，它一定会失败，只是频率问题。
把失败分成三类，处理方式完全不同：

  可重试   429 限流、5xx、连接超时 —— 退避后再试，多半能成
  不可重试 401 密钥错、400 参数错  —— 再试一百次也一样，直接报错
  超时     用户已经等太久          —— 主动掐断，给一句人话

重试次数与超时时间最好交给模型客户端自己做（init_chat_model 的
max_retries / timeout），这里再包一层是为了：统一兜底文案、
把技术异常翻译成用户看得懂的话、给运维留一条可统计的记录。
"""

import os
import random
import time
from typing import Callable, Iterator

MAX_ATTEMPTS = int(os.environ.get("CHAT_MAX_ATTEMPTS", "3"))
BASE_DELAY = float(os.environ.get("CHAT_BASE_DELAY", "1.0"))
DEADLINE = float(os.environ.get("CHAT_DEADLINE", "45"))

RETRYABLE_HINTS = ("timeout", "timed out", "rate limit", "429",
                   "500", "502", "503", "504", "connection", "temporarily")

FRIENDLY = {
    "auth": "服务配置有问题，请联系管理员（凭据未通过校验）。",
    "rate": "当前访问的人有点多，请稍等几秒再试。",
    "timeout": "这次等待超时了，换个更短的问法试试。",
    "other": "这次没能回答，请稍后重试。",
}


def classify(exc: Exception) -> str:
    text = ("%s %s" % (type(exc).__name__, exc)).lower()
    if "401" in text or "api key" in text or "unauthor" in text:
        return "auth"
    if "429" in text or "rate limit" in text:
        return "rate"
    if "timeout" in text or "timed out" in text:
        return "timeout"
    return "other"


def is_retryable(exc: Exception) -> bool:
    text = ("%s %s" % (type(exc).__name__, exc)).lower()
    if classify(exc) == "auth":
        return False
    return any(h in text for h in RETRYABLE_HINTS)


def call_with_retry(fn: Callable[[], str]) -> str:
    """指数退避 + 抖动。抖动很关键：没有它，一批用户会在同一毫秒齐刷刷重试。"""
    started = time.time()
    last = None
    for attempt in range(1, MAX_ATTEMPTS + 1):
        if time.time() - started > DEADLINE:
            raise TimeoutError("超过整体时限 %.0fs" % DEADLINE)
        try:
            return fn()
        except Exception as exc:                      # noqa: BLE001
            last = exc
            if attempt == MAX_ATTEMPTS or not is_retryable(exc):
                break
            delay = BASE_DELAY * (2 ** (attempt - 1)) * (1 + random.random() * 0.3)
            time.sleep(delay)
    raise last


def safe_stream(make_stream: Callable[[], Iterator[str]]) -> Iterator[str]:
    """流式版本：**只在第一片之前重试**。

    已经吐了半句话再重试，用户会看到同一段话说两遍。
    所以首片之后出错只能收尾，把残句留在屏幕上并追加一句提示。
    """
    for attempt in range(1, MAX_ATTEMPTS + 1):
        produced = False
        try:
            for chunk in make_stream():
                produced = True
                yield chunk
            return
        except Exception as exc:                      # noqa: BLE001
            if produced:
                yield "\n\n（回答被中断：%s）" % FRIENDLY[classify(exc)]
                return
            if attempt == MAX_ATTEMPTS or not is_retryable(exc):
                yield FRIENDLY[classify(exc)]
                return
            time.sleep(BASE_DELAY * (2 ** (attempt - 1)))


if __name__ == "__main__":
    calls = {"n": 0}

    def flaky():
        calls["n"] += 1
        if calls["n"] < 3:
            raise RuntimeError("503 service temporarily unavailable")
        return "第 %d 次成功" % calls["n"]

    print(call_with_retry(flaky))

    def bad_key():
        raise RuntimeError("401 invalid api key")

    try:
        call_with_retry(bad_key)
    except Exception as e:                            # noqa: BLE001
        print("不重试，直接失败：", classify(e), "→", FRIENDLY[classify(e)])

# -*- coding: utf-8 -*-
"""
换厨师不换厨房：model_switch.py
=================================
编排层里唯一和「具体哪家模型」有关的，就是建模型那一行。
把它抽成一个函数，云端模型和本地 Ollama 就能共用同一套链、
同一套记忆、同一套界面——这正是分层的兑现时刻。

  云端     init_chat_model("openai:gpt-4o-mini")
  本地     init_chat_model("ollama:qwen2:7b")

init_chat_model 接受 "<provider>:<model>" 这种写法，
provider 对应装好的伙伴包（langchain-openai / langchain-ollama）。
本地模型不需要密钥，但需要 Ollama 服务已经起着。
"""

import os

from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

PRESETS = {
    "cloud": {
        "spec": os.environ.get("CLOUD_MODEL", "openai:gpt-4o-mini"),
        "needs_key": True,
        "note": "效果稳定、开箱即用；数据要出公网，按 token 计费",
    },
    "local": {
        "spec": os.environ.get("LOCAL_MODEL", "ollama:qwen2:7b"),
        "needs_key": False,
        "note": "数据不出内网、无调用费；要自备显卡，效果取决于模型大小",
    },
}


def make_model(profile: str = None, **overrides):
    """按档位建模型。除了这个函数，其余代码不该出现任何厂商名字。"""
    profile = profile or os.environ.get("MODEL_PROFILE", "cloud")
    if profile not in PRESETS:
        raise ValueError("未知档位 %r，可选：%s" % (profile, list(PRESETS)))

    preset = PRESETS[profile]
    kwargs = {
        "temperature": float(os.environ.get("CHAT_TEMPERATURE", "0.3")),
        "timeout": float(os.environ.get("CHAT_TIMEOUT", "60")),
    }
    if preset["needs_key"]:
        key = os.environ.get("OPENAI_API_KEY")
        if not key:
            raise RuntimeError("档位 %s 需要 OPENAI_API_KEY，未配置则改用 local"
                               % profile)
        kwargs["api_key"] = key
    else:
        # 本地模型服务地址也走环境变量，换机器不改源码
        kwargs["base_url"] = os.environ.get("OLLAMA_HOST",
                                            "http://127.0.0.1:11434")
    kwargs.update(overrides)
    return init_chat_model(preset["spec"], **kwargs)


def make_chain(profile: str = None):
    """链的结构与档位完全无关，换档位只换中间那一节。"""
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个中文助手，回答简洁。"),
        ("human", "{question}"),
    ])
    return prompt | make_model(profile) | StrOutputParser()


def available_profiles():
    """给界面用：哪些档位现在真的能选。缺密钥的档位不该出现在下拉框里。"""
    usable = []
    for name, preset in PRESETS.items():
        if preset["needs_key"] and not os.environ.get("OPENAI_API_KEY"):
            continue
        usable.append(name)
    return usable or ["local"]


if __name__ == "__main__":
    for name, preset in PRESETS.items():
        print("%-6s %-24s %s" % (name, preset["spec"], preset["note"]))
    print("当前可选档位：", available_profiles())

    # 同一个问题跑两个档位，比较回答风格与耗时；没配密钥的档位会跳过
    for name in available_profiles():
        try:
            print("[%s] %s" % (name, make_chain(name).invoke(
                {"question": "一句话介绍你自己"})))
        except Exception as exc:                      # noqa: BLE001
            print("[%s] 不可用：%s" % (name, exc))

# -*- coding: utf-8 -*-
"""
不开界面也能聊：cli_probe.py
==============================
同一个 chat_core，换一个前端。这份脚本的价值有三个：

  1. 证明「界面只是点单台」这句话不是口号——后厨一行没改
  2. 排障时分层定位：命令行能聊 → 界面的锅；也不能聊 → 后厨的锅
  3. 服务器上没浏览器时，这是最快的验收手段

内置几个斜杠命令，用来观察后厨的状态：
  /new    换一块空白备忘板
  /hist   打印这块板上的全部消息
  /stat   看条数与估算 token
  /quit   退出
"""

import sys

import chat_core
import session_keys
from trim_policy import _approx_tokens


def show_history(sid: str) -> None:
    role = {"human": "我", "ai": "助手", "system": "人设"}
    for i, m in enumerate(chat_core.history_of(sid), 1):
        text = m.content if isinstance(m.content, str) else str(m.content)
        head = text.replace("\n", " ")[:60]
        print("  %2d. %-4s %s" % (i, role.get(m.type, m.type), head))


def main() -> int:
    sid = session_keys.new_anonymous_id()
    print("会话指纹 %s ｜ 模型 %s ｜ /quit 退出"
          % (session_keys.fingerprint(sid), chat_core.MODEL_NAME))

    while True:
        try:
            question = input("\n我> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            return 0

        if not question:
            continue
        if question == "/quit":
            return 0
        if question == "/new":
            chat_core.reset(sid)
            sid = session_keys.new_anonymous_id()
            print("已换一块空白备忘板：%s" % session_keys.fingerprint(sid))
            continue
        if question == "/hist":
            show_history(sid)
            continue
        if question == "/stat":
            msgs = chat_core.history_of(sid)
            print("共 %d 条，估算 %d token" % (len(msgs), _approx_tokens(msgs)))
            continue

        print("助手> ", end="", flush=True)
        try:
            for piece in chat_core.reply_stream(sid, question):
                print(piece, end="", flush=True)
            print()
        except KeyboardInterrupt:
            # 这一轮答案不完整，把这一问一答从历史里撤掉，别让残句污染下一轮
            print("\n（已中断）")
            chat_core.reset(sid)
        except Exception as exc:                      # noqa: BLE001
            print("\n调用失败：%s" % exc)


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

# -*- coding: utf-8 -*-
"""
把后厨从「链」换成「店长」：agent_mode.py
==========================================
前面的应用是一条固定流水线：提示词 → 模型 → 解析器。
如果要让它自己决定「先查订单还是先查物流」，就把链换成 Agent。

对界面层来说，这个替换是无感的：它照样只调 reply / reply_stream。
后厨里换的东西有两处——

  记忆    RunnableWithMessageHistory → checkpointer + thread_id
  执行体  prompt | model | parser     → create_agent(...)

工具的协议层细节（模型怎么填参数、回执怎么贴回去）在 Function Call
那一讲；工具怎么写、description 怎么措辞在 Tools 与 Agent 实战那一讲。
这里只关心它怎么接进一个已有的聊天应用。
"""

import os
from typing import Iterator

from langchain.agents import create_agent
from langchain.tools import tool
from langgraph.checkpoint.memory import InMemorySaver

SYSTEM_PROMPT = (
    "你是电商平台的客服助手。需要订单信息时调用工具，"
    "不要凭空猜测订单号或物流状态。"
)

_ORDERS = {
    "12345": {"status": "已发货", "amount": 299.0},
    "12346": {"status": "待付款", "amount": 158.0},
}


@tool
def query_order(order_id: str) -> str:
    """按订单号查询订单状态与金额。order_id 是纯数字的订单编号。"""
    row = _ORDERS.get(order_id.strip())
    if not row:
        return "未查到订单 %s" % order_id
    return "订单 %s：%s，金额 %.2f 元" % (order_id, row["status"], row["amount"])


@tool
def refund_policy(keyword: str) -> str:
    """查询退换货政策。keyword 是要查的关键词，如「七天」「运费」。"""
    table = {
        "七天": "签收后七天内无理由退货，商品需不影响二次销售。",
        "运费": "质量问题运费由平台承担，非质量问题由买家承担。",
    }
    for k, v in table.items():
        if k in keyword:
            return v
    return "未找到相关条款，请转人工客服。"


# checkpointer 就是 Agent 这条路线上的备忘板；thread_id 相当于会话 ID
agent = create_agent(
    model=os.environ.get("CHAT_MODEL", "openai:gpt-4o-mini"),
    tools=[query_order, refund_policy],
    system_prompt=SYSTEM_PROMPT,
    checkpointer=InMemorySaver(),
)


def _cfg(session_id: str) -> dict:
    return {"configurable": {"thread_id": session_id}}


def reply(session_id: str, question: str) -> str:
    """和 chat_core.reply 同签名，界面层不用改一个字。"""
    result = agent.invoke(
        {"messages": [{"role": "user", "content": question}]},
        _cfg(session_id))
    return result["messages"][-1].content


def reply_stream(session_id: str, question: str) -> Iterator[str]:
    """流式：只把模型节点吐的字透给界面，工具执行过程不往界面上倒。

    v1 里模型节点名是 "model"（旧版叫 "agent"），流式过滤时别记错。
    """
    for chunk, meta in agent.stream(
            {"messages": [{"role": "user", "content": question}]},
            _cfg(session_id), stream_mode="messages"):
        if meta.get("langgraph_node") == "model" and chunk.content:
            yield chunk.content


if __name__ == "__main__":
    sid = "agent-demo"
    print(reply(sid, "帮我看下订单 12345 什么状态"))
    print(reply(sid, "那它多少钱来着？"))      # 考察 checkpointer 有没有生效
    print(reply(sid, "退货运费谁出？"))

# -*- coding: utf-8 -*-
"""
骨架模板：chatbot_skeleton.py —— 改 TODO 就是你自己的业务助手
================================================================
这一份把前面拆开讲的零件装在一起，并且刻意写成**单文件**，
方便直接复制到新项目里跑起来。文件内部仍然保持三层边界：

  配置区（TODO 全在这里） → 后厨（建链、记忆、修剪） → 前厅（Streamlit）

要拆成多文件时，把「后厨」整段剪进 chat_core.py 即可，前厅代码不用改。

启动：streamlit run chatbot_skeleton.py
依赖：pip install streamlit langchain langchain-openai
"""

import os
import uuid

import streamlit as st
from langchain.chat_models import init_chat_model
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory

# ================================ 配置区 ================================
# TODO 1：换成你的模型。云端 "openai:gpt-4o-mini"，本地 "ollama:qwen2:7b"
MODEL_SPEC = os.environ.get("CHAT_MODEL", "openai:gpt-4o-mini")

# TODO 2：写你的业务人设。务必写清边界，否则模型会一本正经地编
SYSTEM_PROMPT = (
    "你是某公司的智能助手。回答简洁、准确。\n"
    "只依据用户提供的信息和你确知的常识回答；不确定就说不确定，不要编造。"
)

# TODO 3：页面标题与开场提示
PAGE_TITLE = "智能助手"
INPUT_HINT = "输入你的问题"

# TODO 4：一块备忘板最多留多少条消息。太长会变慢、变贵，最后被静默截断
MAX_MESSAGES = int(os.environ.get("MAX_MESSAGES", "20"))

# TODO 5：需要持久化时，把 _store 换成文件或 Redis 实现
_store = {}

# ================================ 后厨 ==================================


def get_history(session_id: str) -> InMemoryChatMessageHistory:
    """工厂函数：一个会话 ID 一块备忘板，顺手修剪。"""
    if session_id not in _store:
        _store[session_id] = InMemoryChatMessageHistory()
    history = _store[session_id]

    messages = history.messages
    if len(messages) > MAX_MESSAGES:
        kept = list(messages[-MAX_MESSAGES:])
        while kept and kept[0].type == "ai":
            kept.pop(0)                 # 别让历史以一句孤立的回答开头
        history.clear()
        history.add_messages(kept)
    return history


@st.cache_resource(show_spinner=False)
def build_conversation():
    """只建一次。少了这个装饰器，页面每次重跑都会重新建链和客户端。"""
    model = init_chat_model(
        MODEL_SPEC,
        temperature=float(os.environ.get("CHAT_TEMPERATURE", "0.3")),
        api_key=os.environ.get("OPENAI_API_KEY"),
        timeout=float(os.environ.get("CHAT_TIMEOUT", "60")),
    )
    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ])
    chain = prompt | model | StrOutputParser()
    return RunnableWithMessageHistory(
        chain,
        get_history,
        input_messages_key="question",
        history_messages_key="history",
    )


# ================================ 前厅 ==================================
st.set_page_config(page_title=PAGE_TITLE, page_icon="💬")

if "sid" not in st.session_state:
    st.session_state["sid"] = uuid.uuid4().hex
sid = st.session_state["sid"]

with st.sidebar:
    st.caption("模型：%s" % MODEL_SPEC)
    st.caption("历史：%d 条（上限 %d）"
               % (len(get_history(sid).messages), MAX_MESSAGES))
    if st.button("新建对话"):
        get_history(sid).clear()
        st.session_state["sid"] = uuid.uuid4().hex
        st.rerun()

st.title(PAGE_TITLE)

ROLE = {"human": "user", "ai": "assistant"}
for message in get_history(sid).messages:
    role = ROLE.get(message.type)
    if role:
        with st.chat_message(role):
            st.markdown(message.content)

question = st.chat_input(INPUT_HINT)
if question:
    with st.chat_message("user"):
        st.markdown(question)
    with st.chat_message("assistant"):
        try:
            st.write_stream(build_conversation().stream(
                {"question": question},
                config={"configurable": {"session_id": sid}}))
        except Exception as exc:                      # noqa: BLE001
            st.error("这次没能回答：%s" % exc)

# -*- coding: utf-8 -*-
"""
交付前验收：smoke_test.py
===========================
手工点一遍页面，改一次代码就要再点一遍，早晚会偷懒。把验收写成脚本，
重点验的不是「有没有回答」，而是**记忆、隔离、流式、修剪**这四条链路。

退出码 0 表示全过，1 表示有失败，可以直接接进发版流程。
运行：python3 smoke_test.py
"""

import sys
import time

import chat_core
import session_keys
from trim_policy import trim_by_count

RESULTS = []


def check(name: str, ok: bool, detail: str = "") -> bool:
    RESULTS.append((name, ok, detail))
    print("%s %s %s" % ("[通过]" if ok else "[失败]", name, detail))
    return ok


def case_single_turn() -> bool:
    """1. 后厨能不能出菜。不通则后面的用例没有意义。"""
    text = chat_core.reply(session_keys.new_anonymous_id(), "回答一个字：好")
    return check("单轮调用", bool(text and text.strip()), "返回 %d 字" % len(text or ""))


def case_memory() -> bool:
    """2. 同一个会话 ID 里，第二轮能不能记得第一轮说的名字。"""
    sid = session_keys.new_anonymous_id()
    chat_core.reply(sid, "我叫孙小空，请记住这个名字")
    answer = chat_core.reply(sid, "我刚才说我叫什么？")
    return check("多轮记忆", "孙小空" in answer, answer.strip()[:40])


def case_isolation() -> bool:
    """3. 反向用例：换一个会话 ID 就不该记得。只验正向可能是模型蒙对的。"""
    sid_a, sid_b = (session_keys.new_anonymous_id() for _ in range(2))
    chat_core.reply(sid_a, "我叫孙小空，请记住这个名字")
    answer = chat_core.reply(sid_b, "我刚才说我叫什么？")
    return check("会话隔离", "孙小空" not in answer, answer.strip()[:40])


def case_stream() -> bool:
    """4. 流式要能收到多片；只收到一片说明流被谁缓冲住了。"""
    sid = session_keys.new_anonymous_id()
    started = time.time()
    first_at = None
    pieces = 0
    for chunk in chat_core.reply_stream(sid, "用三句话介绍中央厨房的分工"):
        if first_at is None:
            first_at = time.time() - started
        pieces += 1
    return check("流式分片", pieces > 1,
                 "%d 片，首字 %.2fs" % (pieces, first_at or -1))


def case_trim_keeps_system() -> bool:
    """5. 纯本地逻辑，不发请求：修剪之后人设必须还在。"""
    from langchain_core.messages import AIMessage, HumanMessage, SystemMessage

    msgs = [SystemMessage(content="你是客服助手")]
    for i in range(20):
        msgs += [HumanMessage(content="问%d" % i), AIMessage(content="答%d" % i)]
    kept = trim_by_count(msgs, 6)
    return check("修剪不裁人设", kept[0].type != "ai",
                 "保留 %d 条，首条 %s" % (len(kept), kept[0].type))


def main() -> int:
    if not case_single_turn():
        print("\n后厨不通，后续用例跳过。先查模型服务、密钥与网络。")
        return 1
    case_memory()
    case_isolation()
    case_stream()
    case_trim_keeps_system()

    failed = [n for n, ok, _ in RESULTS if not ok]
    print("\n共 %d 项，失败 %d 项" % (len(RESULTS), len(failed)))
    for name in failed:
        print("  · %s" % name)
    return 1 if failed else 0


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

# -*- coding: utf-8 -*-
"""
上线前自检：deploy_check.py
============================
「本机 streamlit run 打得开」离「同事随时能用、坏了有人知道」还差一整套。
这份脚本把交付清单写成可执行的检查项，缺什么一眼看得见。

它不发任何模型请求，只看配置与环境，所以可以随时跑、也能进 CI。
运行：python3 deploy_check.py
"""

import os
import shutil
import socket
import sys

CHECKS = []


def record(name, ok, detail, blocking=True):
    CHECKS.append((name, ok, detail, blocking))


def check_secrets():
    """密钥必须来自环境变量，源码里出现字面量是最常见的事故。"""
    key = os.environ.get("OPENAI_API_KEY", "")
    profile = os.environ.get("MODEL_PROFILE", "cloud")
    if profile == "local":
        record("模型凭据", True, "本地档位不需要密钥")
    else:
        record("模型凭据", bool(key),
               "已配置（长度 %d）" % len(key) if key else "未配置 OPENAI_API_KEY")


def check_session_secret():
    secret = os.environ.get("SESSION_SECRET", "")
    record("会话派生密钥", len(secret) >= 16,
           "已配置" if len(secret) >= 16 else "未配置或过短，登录态会话无法派生 ID")


def check_history_backend():
    backend = os.environ.get("HISTORY_BACKEND", "memory")
    if backend == "memory":
        record("历史存储", False,
               "进程内字典：重启即丢、多副本各存各的，仅适合开发", blocking=False)
    else:
        directory = os.environ.get("HISTORY_DIR", "./chat_history")
        writable = os.access(os.path.dirname(os.path.abspath(directory)) or ".",
                             os.W_OK)
        record("历史存储", writable, "%s（目录可写：%s）" % (backend, writable))


def check_trim():
    limit = int(os.environ.get("MAX_MESSAGES", "20"))
    record("历史上限", 2 <= limit <= 200,
           "MAX_MESSAGES=%d" % limit)


def check_timeout():
    timeout = float(os.environ.get("CHAT_TIMEOUT", "60"))
    record("超时设置", 5 <= timeout <= 180, "CHAT_TIMEOUT=%.0fs" % timeout)


def check_port(port=8501):
    """页面端口是否已被占用；同时提醒模型端口不应对外。"""
    with socket.socket() as s:
        s.settimeout(0.5)
        busy = s.connect_ex(("127.0.0.1", port)) == 0
    record("页面端口 %d" % port, True,
           "已有服务在监听" if busy else "空闲", blocking=False)


def check_process_manager():
    record("进程托管", shutil.which("systemctl") is not None,
           "有 systemctl，可写 service 常驻" if shutil.which("systemctl")
           else "未发现 systemctl，需另行托管进程", blocking=False)


def main():
    check_secrets()
    check_session_secret()
    check_history_backend()
    check_trim()
    check_timeout()
    check_port()
    check_process_manager()

    blocking_failed = 0
    for name, ok, detail, blocking in CHECKS:
        mark = "[通过]" if ok else ("[必修]" if blocking else "[提醒]")
        print("%s %-14s %s" % (mark, name, detail))
        if not ok and blocking:
            blocking_failed += 1

    print("\n还需人工确认：页面前置鉴权、模型端口只对网关开放、"
          "调用日志只记长度与短哈希。")
    print("必修项未通过 %d 条" % blocking_failed)
    return 1 if blocking_failed else 0


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