LCEL 与链的组合

# -*- coding: utf-8 -*-
"""
不用链 与 用链：同一件事的两种写法
==================================
同一个任务——「给定话题，讲一个简短笑话」——分别用
① 手动分步 invoke
② LCEL 管道组合
两种方式实现，输出完全一致。

环境：Python 3.10+
依赖：pip install langchain langchain-openai python-dotenv
"""

import os

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

load_dotenv()

# 密钥只从环境变量读，不写进源码
os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY", "")
os.environ["OPENAI_BASE_URL"] = os.environ.get("OPENAI_BASE_URL", "")

model = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)

prompt = PromptTemplate.from_template("给我讲一个关于 {topic} 话题的简短笑话")
parser = StrOutputParser()


# ---------------------------------------------------------------- 写法一
def run_step_by_step(topic: str) -> str:
    """手动分三步：自己端着中间结果从一个组件搬到下一个组件。"""
    prompt_value = prompt.invoke({"topic": topic})   # PromptValue
    message = model.invoke(prompt_value)             # AIMessage
    text = parser.invoke(message)                    # str
    return text


# ---------------------------------------------------------------- 写法二
# 管道符只负责「接线」：把上一个 Runnable 的输出接到下一个的输入。
# 这一行本身不调用模型，也不产生任何网络请求。
chain = prompt | model | parser


def run_with_chain(topic: str) -> str:
    """交给链：三个组件的调用与中间结果传递由 LCEL 统一处理。"""
    return chain.invoke({"topic": topic})


if __name__ == "__main__":
    topic = "冰淇淋"

    a = run_step_by_step(topic)
    print("【分步写法】", a)
    print("返回类型：", type(a))

    b = run_with_chain(topic)
    print("【管道写法】", b)
    print("返回类型：", type(b))

    # 看一眼链是由什么拼成的：RunnableSequence，里面按顺序装着三个组件
    print("链的类型：", type(chain).__name__)
    print("链上的步骤：", [type(s).__name__ for s in chain.steps])

    # 输出示例：
    # 链的类型： RunnableSequence
    # 链上的步骤： ['PromptTemplate', 'ChatOpenAI', 'StrOutputParser']

# -*- coding: utf-8 -*-
"""
Runnable 协议的五个统一入口
===========================
invoke / batch / stream 是同步三件套，ainvoke / astream 是对应的异步版本。
链上的每个组件，以及组合出来的整条链，用的都是这同一组方法名。

依赖：pip install langchain langchain-openai python-dotenv
"""

import asyncio

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini")
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一位擅长用一句话讲清概念的技术讲师"),
    ("human", "用一句话解释：{term}"),
])
chain = prompt | model | StrOutputParser()


def demo_invoke():
    """单输入单输出：最常用的一个。"""
    out = chain.invoke({"term": "管道符"})
    print("[invoke]", out)


def demo_batch():
    """批量：一次传一个列表，返回等长的结果列表，内部可并发。"""
    terms = [{"term": "Runnable"}, {"term": "LCEL"}, {"term": "顺序链"}]
    outs = chain.batch(terms, config={"max_concurrency": 3})
    for t, o in zip(terms, outs):
        print("[batch]", t["term"], "->", o[:40], "...")


def demo_stream():
    """流式：边生成边吐，首字延迟显著低于 invoke。"""
    print("[stream] ", end="", flush=True)
    for chunk in chain.stream({"term": "流式输出"}):
        # 因为链尾是 StrOutputParser，这里每个 chunk 已经是字符串片段
        print(chunk, end="", flush=True)
    print()


async def demo_async():
    """异步：ainvoke 与 astream，方法名只是多了一个 a 前缀。"""
    out = await chain.ainvoke({"term": "异步调用"})
    print("[ainvoke]", out)

    print("[astream] ", end="", flush=True)
    async for chunk in chain.astream({"term": "并发"}):
        print(chunk, end="", flush=True)
    print()


def demo_same_methods_everywhere():
    """同一组方法，在单个组件上也成立。"""
    print("[组件级 invoke] 提示词模板 ->", prompt.invoke({"term": "协议"}))
    print("[组件级 invoke] 解析器 ->", StrOutputParser().invoke("一段文本"))


if __name__ == "__main__":
    demo_invoke()
    demo_batch()
    demo_stream()
    demo_same_methods_everywhere()
    asyncio.run(demo_async())

# -*- coding: utf-8 -*-
"""
RunnablePassthrough 与 RunnableParallel
=======================================
链上每一步只吃上一步的输出。想让后面的步骤还能看见最初的输入，
就要靠这两个「接线件」把数据原样带下去、或者分叉再汇合。

依赖：pip install langchain langchain-openai python-dotenv
"""

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini")
parser = StrOutputParser()


# ------------------------------------------------ 1. 原样透传
# RunnablePassthrough 什么都不做，把输入原封不动交给下一步。
print(RunnablePassthrough().invoke({"product": "耳机"}))
# -> {'product': '耳机'}


# ------------------------------------------------ 2. 透传 + 补一个字段
# assign 在保留原有键的基础上，再算出一个新键。
price_prompt = ChatPromptTemplate.from_template(
    "请给出 {product} 的一个合理市场价，只回答数字，不要单位和多余文字"
)
add_price = RunnablePassthrough.assign(price=price_prompt | model | parser)

# 输入 {"product": "耳机"}，输出 {"product": "耳机", "price": "399"}
# 原来的 product 还在，后面的提示词才能同时用到商品名和价格。


# ------------------------------------------------ 3. 并行分叉
# RunnableParallel 把同一份输入同时喂给多个分支，结果汇成一个字典。
# 写成普通 dict 也可以，LCEL 会自动包装成 RunnableParallel。
selling_point = ChatPromptTemplate.from_template("用一句话说出 {product} 最大的卖点")
target_user = ChatPromptTemplate.from_template("{product} 最适合哪类人群？一句话")

fan_out = RunnableParallel(
    point=selling_point | model | parser,
    user=target_user | model | parser,
    product=RunnablePassthrough(),          # 把原始输入也留一份
)


# ------------------------------------------------ 4. 分叉之后再汇合
copy_prompt = ChatPromptTemplate.from_template(
    "为 {product} 写一条 50 字以内的促销文案。\n卖点：{point}\n目标人群：{user}"
)


def build_chain():
    """先补价格，再并行取卖点与人群，最后汇合成文案。"""
    return (
        add_price
        | RunnableParallel(
            product=lambda d: d["product"],
            price=lambda d: d["price"],
            point=selling_point | model | parser,
            user=target_user | model | parser,
        )
        | copy_prompt
        | model
        | parser
    )


if __name__ == "__main__":
    print("[并行分支]", fan_out.invoke({"product": "降噪耳机"}))

    chain = build_chain()
    print("[汇合结果]", chain.invoke({"product": "降噪耳机"}))

    # 并行分支是真的并发跑的：两个分支各自请求一次模型，总耗时接近单次而不是两次之和。

# -*- coding: utf-8 -*-
"""
RunnableLambda / .bind() / .with_config()
=========================================
三个把「普通 Python 世界」接进链里的开关：
  RunnableLambda   把任意函数变成链上的一个工位
  .bind()          给某个组件预先锁死一部分调用参数
  .with_config()   给某个组件或整条链挂运行期配置（标签、并发、超时）

依赖：pip install langchain langchain-openai python-dotenv
"""

import json
import re

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini")
parser = StrOutputParser()


# ---------------------------------------------------------- RunnableLambda
def clean_text(text: str) -> str:
    """去掉模型爱加的代码围栏和首尾空白。"""
    text = re.sub(r"^```(?:json|python)?|```$", "", text.strip(), flags=re.M)
    return text.strip()


def to_dict(text: str) -> dict:
    """把清洗后的字符串解析成字典；失败就原样兜住，不让整条链崩掉。"""
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return {"raw": text, "parsed": False}


extract_prompt = ChatPromptTemplate.from_template(
    "从下面这句话里抽取城市和日期，只输出 JSON，键为 city 和 date：\n{sentence}"
)

# 普通函数放进管道时，LCEL 会自动包成 RunnableLambda；
# 显式写出来可读性更好，也方便后面单独 .with_config() 起名字。
extract_chain = (
    extract_prompt
    | model
    | parser
    | RunnableLambda(clean_text).with_config({"run_name": "清洗围栏"})
    | RunnableLambda(to_dict).with_config({"run_name": "解析字典"})
)


# ---------------------------------------------------------------- .bind()
# bind 把参数「焊」在组件上，之后每次调用都自动带上，不用在调用处重复传。
strict_model = model.bind(temperature=0, max_tokens=256)

# 也可以绑定停止词：让模型生成到某个标记就停下
short_model = model.bind(stop=["\n\n"])


# ----------------------------------------------------------- .with_config()
# 配置挂在链上，对整条链生效；批量调用时用来限制并发。
traced_chain = extract_chain.with_config({
    "run_name": "信息抽取链",
    "tags": ["extract", "demo"],
    "max_concurrency": 4,
})


if __name__ == "__main__":
    sentence = "帮我看看 2025 年 12 月 10 日郑州的天气"
    print("[抽取结果]", extract_chain.invoke({"sentence": sentence}))

    # bind 之后的组件仍然是一个 Runnable，照样能进管道
    determin_chain = extract_prompt | strict_model | parser
    print("[temperature=0]", determin_chain.invoke({"sentence": sentence}))

    print("[带配置批量]", traced_chain.batch([
        {"sentence": "订一张 2026 年 1 月 3 日去上海的票"},
        {"sentence": "查一下明天北京冷不冷"},
    ]))

# -*- coding: utf-8 -*-
"""
最小可跑的一条链
================
提示词模板 | 大模型 | 输出解析器，三行装配，一行调用。
这是 LCEL 的最短路径，后面所有花样都是从这三步长出来的。

依赖：pip install langchain langchain-openai python-dotenv
环境变量：OPENAI_API_KEY
"""

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

load_dotenv()                                   # 从 .env 读 OPENAI_API_KEY

# ① 工位一：菜谱卡。花括号里的 topic 是待填的变量
prompt = ChatPromptTemplate.from_template("用一句话解释什么是 {topic}")

# ② 工位二：厨师。真正生成内容的是它
model = ChatOpenAI(model="gpt-4o-mini")

# ③ 工位三：摆盘。把 AIMessage 剥成纯字符串
parser = StrOutputParser()

# ④ 传送带：管道符把三个工位接成一条流水线。
#    这一行不发任何请求，只是把接线接好。
chain = prompt | model | parser

# ⑤ 开机：invoke 才真正让三个工位依次动起来
answer = chain.invoke({"topic": "管道符"})

print(answer)
print(type(answer))         # <class 'str'>

# -*- coding: utf-8 -*-
"""
单入单出的顺序链：剧名 → 大纲 → 剧评
=====================================
第一段的输出就是第二段的输入，中间不需要任何字段映射。
这是最简单的一类流水线，用管道符直接接起来即可。

依赖：pip install langchain langchain-openai python-dotenv
"""

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini")
parser = StrOutputParser()


# ------------------------------------------------ 第一段：给剧名写大纲
synopsis_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一名剧作家。给定剧本标题，写出一份分幕大纲。"),
    ("human", "剧本标题：{title}"),
])
synopsis_chain = synopsis_prompt | model | parser


# ------------------------------------------------ 第二段：给大纲写剧评
review_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是《纽约时报》的剧评家。读完剧情大纲后写一篇剧评。"),
    ("human", "剧情大纲：\n{synopsis}"),
])
review_chain = review_prompt | model | parser


# ------------------------------------------------ 串成一条
# 第一段吐出的是一个纯字符串，第二段的提示词要的是 {synopsis}，
# 中间用一个 dict 把字符串装进正确的键里。
full_chain = synopsis_chain | (lambda text: {"synopsis": text}) | review_chain


# ------------------------------------------------ 中途也想拿到大纲？
# 上面那条链只返回剧评，大纲被丢掉了。想两样都留下，用 assign：
from langchain_core.runnables import RunnablePassthrough  # noqa: E402

keep_both = (
    RunnablePassthrough.assign(synopsis=synopsis_chain)
    | RunnablePassthrough.assign(review=review_chain)
)
# 输入 {"title": "..."}，输出 {"title": ..., "synopsis": ..., "review": ...}


if __name__ == "__main__":
    title = "日落海滩上的悲剧"

    review = full_chain.invoke({"title": title})
    print("=" * 30, "剧评", "=" * 30)
    print(review)

    both = keep_both.invoke({"title": title})
    print("=" * 30, "两段都留下", "=" * 30)
    print("键：", list(both.keys()))
    print("大纲前 60 字：", both["synopsis"][:60])
    print("剧评前 60 字：", both["review"][:60])

    # 想看每一步的中间结果，不必加 verbose 参数，直接用 stream 观察事件：
    # for ev in keep_both.stream({"title": title}):
    #     print(ev)

# -*- coding: utf-8 -*-
"""
多入多出的顺序链：翻译 → 摘要 → 识别语种 → 评论
=================================================
四段处理，每段都要读到前面某一段的结果，而且四段的结果最后都要留下来。
做法：让链上流动的始终是一个 **字典**，每一段用 assign 往字典里加一个键。

依赖：pip install langchain langchain-openai python-dotenv
"""

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini")
parser = StrOutputParser()


def step(template: str):
    """把一段提示词包成「提示词 | 模型 | 解析器」的小链。"""
    return PromptTemplate.from_template(template) | model | parser


# 四段各自的小链。花括号里的变量名，就是它要从流动字典里取的键。
translate = step("把下面内容翻译成中文：\n\n{content}")
summarize = step("用一句话总结下面内容：\n\n{chinese_review}")
detect = step("下面内容是什么语言？只回答语言名称：\n\n{chinese_summary}")
comment = step(
    "请使用指定的语言对以下内容进行评论。\n\n"
    "内容：{chinese_summary}\n\n语言：{language}"
)


# ---------------------------------------------------------------- 串起来
# 每个 assign 都是「保留已有的键，再加一个新键」，所以字典越走越厚。
overall_chain = (
    RunnablePassthrough.assign(chinese_review=translate)
    | RunnablePassthrough.assign(chinese_summary=summarize)
    | RunnablePassthrough.assign(language=detect)
    | RunnablePassthrough.assign(comment=comment)
)


# ---------------------------------------------------------------- 只要部分输出
# 最后挂一个字典，挑出需要的键，等价于旧写法里的 output_variables。
pick_output = overall_chain | {
    "chinese_summary": lambda d: d["chinese_summary"],
    "comment": lambda d: d["comment"],
}


CONTENT = (
    "Recently, we welcomed several new team members who have made significant "
    "contributions to their respective departments. Furthermore, please remember "
    "that the open enrollment period for our employee benefits program is fast "
    "approaching."
)


if __name__ == "__main__":
    result = overall_chain.invoke({"content": CONTENT})

    print("最终字典里有这些键：", list(result.keys()))
    for k in ("chinese_review", "chinese_summary", "language", "comment"):
        print("-" * 60)
        print(k, "：", result[k])

    print("=" * 60)
    print("只取两个键：", pick_output.invoke({"content": CONTENT}).keys())

    # 注意：assign 之间是有先后依赖的，顺序不能随便调。
    # 把 language 那一步提到 chinese_summary 之前，会因为字典里还没有
    # chinese_summary 这个键而直接报 KeyError。

# -*- coding: utf-8 -*-
"""
路由链：按输入内容把请求分发到不同的子链
========================================
三个学科各有一条子链，外加一条兜底子链。
先让模型给问题分类，再用 RunnableBranch 按分类结果选一条路走下去。

依赖：pip install langchain langchain-openai python-dotenv
"""

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableBranch, RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini")
parser = StrOutputParser()


# ------------------------------------------------------------ 1. 分类器
# 只让模型吐一个词，方便下面做精确匹配。temperature=0 让分类结果稳定。
classify_prompt = ChatPromptTemplate.from_template(
    "下面的问题属于哪一类？只回答一个词：数学、物理、历史、其他。\n\n问题：{question}"
)
classifier = classify_prompt | model.bind(temperature=0) | parser


# ------------------------------------------------------------ 2. 各科子链
def expert(role: str, style: str):
    p = ChatPromptTemplate.from_messages([
        ("system", f"你是一位{role}。{style}"),
        ("human", "{question}"),
    ])
    return p | model | parser


math_chain = expert("数学教授", "先列出关键公式，再一步步推导，最后给出答案。")
physics_chain = expert("物理学教授", "先说清物理图像，再给出定量结论。")
history_chain = expert("历史学者", "先交代时代背景，再叙述事件与影响。")
default_chain = expert("耐心的通识讲师", "用平实的语言直接回答，不要装作自己是专家。")


# ------------------------------------------------------------ 3. 路由
# RunnableBranch 依次检查每个判断函数，命中就走对应的子链；
# 最后一个参数是兜底分支，所有条件都不满足时走它——这条一定要写。
branch = RunnableBranch(
    (lambda d: "数学" in d["topic"], math_chain),
    (lambda d: "物理" in d["topic"], physics_chain),
    (lambda d: "历史" in d["topic"], history_chain),
    default_chain,
)

# 先算出 topic 这个键，再把带着 topic 的整份字典交给 branch。
router_chain = RunnablePassthrough.assign(topic=classifier) | branch


# ------------------------------------------------------------ 4. 纯规则版路由
# 分类这一步不一定要花一次模型调用。关键词能判清楚时，直接写函数更快更省。
KEYWORDS = {
    "数学": ("方程", "求导", "积分", "概率", "几何"),
    "物理": ("加速度", "电场", "熵", "波长", "动量"),
    "历史": ("朝代", "战争", "条约", "王朝", "年间"),
}


def rule_classify(d: dict) -> dict:
    q = d["question"]
    topic = next((k for k, kws in KEYWORDS.items() if any(w in q for w in kws)), "其他")
    return {**d, "topic": topic}


fast_router = rule_classify | branch


if __name__ == "__main__":
    for q in ["请解释一下二次方程的求根公式", "唐朝的均田制是怎么回事", "今天适合出门吗"]:
        print("=" * 60)
        print("问题：", q)
        print("回答：", router_chain.invoke({"question": q})[:80], "...")

    print("=" * 60)
    print("[纯规则路由]", fast_router.invoke({"question": "帮我求导 x^2"})[:60], "...")

# -*- coding: utf-8 -*-
"""
流式输出：stream / astream / astream_events
===========================================
链上只要有一个组件不支持流式，流式就会在那里断掉，退化成「等它算完再一次性吐」。
这份代码演示三种粒度的流式读取，以及一个会把流式打断的典型写法。

依赖：pip install langchain langchain-openai python-dotenv
"""

import asyncio
import json
import time

from dotenv import load_dotenv
from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini")
prompt = ChatPromptTemplate.from_template("写一段 200 字左右的短文，主题是：{topic}")
chain = prompt | model | StrOutputParser()


# ------------------------------------------------------------ 1. 同步流式
def demo_stream():
    t0 = time.time()
    first = None
    for chunk in chain.stream({"topic": "清晨的地铁"}):
        if first is None:
            first = time.time() - t0
        print(chunk, end="", flush=True)
    print()
    print("首字延迟：%.2f 秒" % first)


# ------------------------------------------------------------ 2. 异步流式
async def demo_astream():
    async for chunk in chain.astream({"topic": "夜里的便利店"}):
        print(chunk, end="", flush=True)
    print()


# ------------------------------------------------------------ 3. 逐事件流式
async def demo_astream_events():
    """想知道「现在流的是链上哪一步」，用 astream_events。"""
    async for ev in chain.astream_events({"topic": "雨天的球场"}):
        kind = ev["event"]
        if kind == "on_chat_model_stream":
            print(ev["data"]["chunk"].content, end="", flush=True)
        elif kind in ("on_chain_start", "on_chain_end"):
            print("\n<%s %s>" % (kind, ev.get("name")))
    print()


# ------------------------------------------------------------ 4. 会打断流式的写法
def take_all(text: str) -> str:
    """普通函数必须拿到完整输入才能返回，链在这里必然攒齐再往下走。"""
    return text.upper()


blocked = chain | RunnableLambda(take_all)
# blocked.stream(...) 仍然能调用，但只会在最后一次性吐出一整块。


# 输出结构化数据时，换成支持增量解析的 JsonOutputParser，流式才有意义。
json_prompt = ChatPromptTemplate.from_template(
    "列出 3 个 {topic} 相关的关键词，输出 JSON 数组，键名为 keywords"
)
json_chain = json_prompt | model | JsonOutputParser()


async def demo_json_stream():
    async for partial in json_chain.astream({"topic": "向量数据库"}):
        # 每次拿到的是「到目前为止能解析出来的那部分」，字典会逐步长大
        print(json.dumps(partial, ensure_ascii=False))


if __name__ == "__main__":
    demo_stream()
    asyncio.run(demo_astream())
    asyncio.run(demo_astream_events())
    asyncio.run(demo_json_stream())

# -*- coding: utf-8 -*-
"""
SQL 查询链：自然语言 → SQL → 结果 → 人话
========================================
create_sql_query_chain 只负责「把问题翻译成一条 SQL」，它不执行 SQL。
执行由 SQLDatabase.run 完成，最后再用一段提示词把查询结果讲成人话。

依赖：pip install langchain langchain-classic langchain-community
      langchain-openai pymysql python-dotenv
"""

import os

from dotenv import load_dotenv
from langchain_classic.chains import create_sql_query_chain
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()

# 连接串里的账号口令一律走环境变量，不要写进源码，更不要提交进版本库
DB_USER = os.environ.get("DB_USER", "readonly")
DB_PASSWORD = os.environ.get("DB_PASSWORD", "")
DB_HOST = os.environ.get("DB_HOST", "127.0.0.1")
DB_PORT = os.environ.get("DB_PORT", "3306")
DB_NAME = os.environ.get("DB_NAME", "itcast")

db = SQLDatabase.from_uri(
    f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
)

print("数据库方言：", db.dialect)
print("可用数据表：", db.get_usable_table_names())

model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# 这一步产出的是一段 SQL 文本，不是查询结果
write_sql = create_sql_query_chain(llm=model, db=db)


def clean_sql(text: str) -> str:
    """模型有时会带上 `SQLQuery:` 前缀或 Markdown 围栏，先剥干净。"""
    text = text.strip()
    if text.lower().startswith("sqlquery:"):
        text = text.split(":", 1)[1]
    text = text.replace("```sql", "").replace("```", "")
    return text.strip().rstrip(";") + ";"


def guard(sql: str) -> str:
    """只放行只读查询，并强制加上行数上限。"""
    low = sql.lower().strip()
    if not low.startswith("select"):
        raise ValueError("只允许 SELECT 查询")
    if ";" in low[:-1]:
        raise ValueError("不允许多条语句")
    if " limit " not in low:
        sql = sql.rstrip(";") + " LIMIT 50;"
    return sql


answer_prompt = ChatPromptTemplate.from_template(
    "根据下面的信息，用中文回答用户的问题。\n\n"
    "问题：{question}\nSQL：{query}\n查询结果：{result}\n\n回答："
)

full_chain = (
    RunnablePassthrough.assign(query=write_sql | clean_sql | guard)
    | RunnablePassthrough.assign(result=lambda d: db.run(d["query"]))
    | answer_prompt
    | model
    | StrOutputParser()
)


if __name__ == "__main__":
    q = "employees 表里一共有多少个员工？"

    sql = guard(clean_sql(write_sql.invoke({"question": q})))
    print("模型写出的 SQL：", sql)
    print("直接执行的结果：", db.run(sql))

    print("整条链的回答：", full_chain.invoke({"question": q}))

    # 限制模型只看某几张表，可以显著降低写错表名的概率：
    # write_sql.invoke({"question": q, "table_names_to_use": ["employees"]})

# -*- coding: utf-8 -*-
"""
旧写法：LLMChain / SimpleSequentialChain / SequentialChain
==========================================================
这三个类现在住在 langchain-classic 包里，仍然可以安装、可以跑，
但它们已经不是推荐写法。放在这里是为了让你读懂线上还在跑的老代码。

安装：pip install langchain-classic langchain-openai python-dotenv
"""

from dotenv import load_dotenv
from langchain_classic.chains import LLMChain, SequentialChain, SimpleSequentialChain
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_openai import ChatOpenAI

load_dotenv()

llm = ChatOpenAI(model="gpt-4o-mini")


# ------------------------------------------------------- 1. 最基础的 LLMChain
prompt = PromptTemplate.from_template("桌上有 {number} 个苹果、四个桃子和 3 本书，一共有几个水果？")
llm_chain = LLMChain(llm=llm, prompt=prompt)

# 注意返回值：是一个字典，答案在 "text" 键里，而且原始输入也被塞了进去
# {'number': 2, 'text': '一共有 6 个水果。'}
print(llm_chain.invoke({"number": 2}))


# verbose=True 会把渲染后的提示词打印到终端，这是旧写法的调试手段
chat_template = ChatPromptTemplate.from_messages([
    ("system", "你是一位{area}领域具备丰富经验的高端技术人才"),
    ("human", "给我讲一个 {adjective} 笑话"),
])
verbose_chain = LLMChain(llm=llm, prompt=chat_template, verbose=True)
print(verbose_chain.invoke({"area": "互联网", "adjective": "上班的"}))


# --------------------------------------- 2. SimpleSequentialChain：单入单出
first = LLMChain(llm=llm, prompt=PromptTemplate.from_template(
    "我的邻居姓 {lastname}，他生了个儿子，给他儿子起个名字"))
second = LLMChain(llm=llm, prompt=PromptTemplate.from_template(
    "邻居的儿子名字叫 {child_name}，给他起一个小名"))

overall = SimpleSequentialChain(chains=[first, second], verbose=True)

# 调用时既不用 lastname 也不用 child_name，而是固定用 input；
# 返回值里的答案也固定在 output 键上——这两个键写死在源码里。
print(overall.invoke({"input": "王"}))


# ------------------------------------------- 3. SequentialChain：多入多出
translate = LLMChain(
    llm=llm,
    prompt=PromptTemplate.from_template("把下面内容翻译成中文：\n\n{content}"),
    output_key="chinese_review",
)
summarize = LLMChain(
    llm=llm,
    prompt=PromptTemplate.from_template("用一句话总结下面内容：\n\n{chinese_review}"),
    output_key="chinese_summary",
)

seq = SequentialChain(
    chains=[translate, summarize],
    input_variables=["content"],
    output_variables=["chinese_review", "chinese_summary"],
    verbose=True,
)

# 链与链之间靠 output_key / 模板变量名手动对齐，名字写错要等到运行时才报错
print(seq.invoke({"content": "Hello team, the enrollment period is approaching."}))

# -*- coding: utf-8 -*-
"""
把旧链改写成管道组合
====================
同一件事，上面是 langchain-classic 的旧写法，下面是当代的管道写法。
三组改写各自独立，可以单独运行任意一组。

依赖：pip install langchain langchain-classic langchain-openai python-dotenv
"""

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()

llm = ChatOpenAI(model="gpt-4o-mini")
parser = StrOutputParser()


# ========================================================== 一、LLMChain
def old_llm_chain():
    from langchain_classic.chains import LLMChain

    prompt = PromptTemplate.from_template("给我讲一个关于 {topic} 的笑话")
    chain = LLMChain(llm=llm, prompt=prompt)
    out = chain.invoke({"topic": "程序员"})
    return out["text"]          # 返回字典，答案要从 "text" 键里取


def new_llm_chain():
    prompt = PromptTemplate.from_template("给我讲一个关于 {topic} 的笑话")
    chain = prompt | llm | parser
    return chain.invoke({"topic": "程序员"})   # 直接就是字符串


# =============================================== 二、SimpleSequentialChain
def old_simple_sequential():
    from langchain_classic.chains import LLMChain, SimpleSequentialChain

    a = LLMChain(llm=llm, prompt=PromptTemplate.from_template(
        "我的邻居姓 {lastname}，给他刚出生的儿子起个名字"))
    b = LLMChain(llm=llm, prompt=PromptTemplate.from_template(
        "名字叫 {child_name}，再起一个小名"))
    chain = SimpleSequentialChain(chains=[a, b])
    return chain.invoke({"input": "王"})["output"]


def new_simple_sequential():
    a = PromptTemplate.from_template("我的邻居姓 {lastname}，给他刚出生的儿子起个名字") | llm | parser
    b = PromptTemplate.from_template("名字叫 {child_name}，再起一个小名") | llm | parser
    # 中间那个 lambda 就是旧写法帮你做的「把上一段的字符串塞进下一段的变量名」
    chain = a | (lambda name: {"child_name": name}) | b
    return chain.invoke({"lastname": "王"})


# ===================================================== 三、SequentialChain
def old_sequential():
    from langchain_classic.chains import LLMChain, SequentialChain

    one = LLMChain(llm=llm, output_key="chinese",
                   prompt=PromptTemplate.from_template("翻译成中文：\n{content}"))
    two = LLMChain(llm=llm, output_key="summary",
                   prompt=PromptTemplate.from_template("一句话总结：\n{chinese}"))
    chain = SequentialChain(chains=[one, two],
                            input_variables=["content"],
                            output_variables=["chinese", "summary"])
    return chain.invoke({"content": "Good morning, everyone."})


def new_sequential():
    one = PromptTemplate.from_template("翻译成中文：\n{content}") | llm | parser
    two = PromptTemplate.from_template("一句话总结：\n{chinese}") | llm | parser
    chain = (RunnablePassthrough.assign(chinese=one)
             | RunnablePassthrough.assign(summary=two))
    return chain.invoke({"content": "Good morning, everyone."})


if __name__ == "__main__":
    print("[旧] LLMChain          ->", old_llm_chain())
    print("[新] prompt | llm      ->", new_llm_chain())
    print("[旧] SimpleSequential  ->", old_simple_sequential())
    print("[新] a | 映射 | b       ->", new_simple_sequential())
    print("[旧] SequentialChain   ->", old_sequential())
    print("[新] assign 链         ->", new_sequential())

# -*- coding: utf-8 -*-
"""
本地模型上的同一条链
====================
换厂商只换最前面那一行模型实例，链的写法一个字不用改——
这正是 Runnable 统一协议带来的好处。

依赖：pip install langchain langchain-ollama
前置：本机已 ollama pull qwen3:8b，服务跑在 http://localhost:11434
"""

import os

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama

BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")

# 本地模型：langchain-ollama 是官方伙伴包，取代早期的社区版 Ollama 封装
model = ChatOllama(model="qwen3:8b", base_url=BASE_URL, temperature=0.7)

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一位取名很有讲究的长辈，只给出名字和一句解释。"),
    ("human", "我的邻居姓{lastname}，刚生了个儿子，帮忙起个名字"),
])

name_chain = prompt | model | StrOutputParser()


# ---------------------------------------------- 两段本地链串起来
nickname_prompt = ChatPromptTemplate.from_template(
    "大名叫「{full_name}」，帮忙起一个亲切的小名，只回答小名本身"
)
nickname_chain = nickname_prompt | model | StrOutputParser()

full_chain = name_chain | (lambda text: {"full_name": text}) | nickname_chain


# ---------------------------------------------- 换成云端模型只改这一行
# from langchain_openai import ChatOpenAI
# model = ChatOpenAI(model="gpt-4o-mini")
#
# 或者用统一入口，按字符串挑模型：
# from langchain.chat_models import init_chat_model
# model = init_chat_model("gpt-4o-mini", model_provider="openai")


if __name__ == "__main__":
    print("[大名]", name_chain.invoke({"lastname": "王"}))

    print("[小名]", end=" ", flush=True)
    for chunk in full_chain.stream({"lastname": "王"}):
        print(chunk, end="", flush=True)
    print()

    print("[批量]", name_chain.batch([
        {"lastname": "李"}, {"lastname": "赵"}, {"lastname": "钱"},
    ]))

# -*- coding: utf-8 -*-
"""
LCEL 链骨架（通用版）
=====================
一条可直接改造的流水线：装配 → 单次调用 → 批量 → 流式 → 异步。
改五处 TODO 就是你自己的链。

依赖：pip install langchain langchain-openai python-dotenv
环境变量：OPENAI_API_KEY（必填）、OPENAI_BASE_URL（走代理或第三方网关时填）
"""

import asyncio
import os

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable, RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()


# ---------------------------------------------------------------- TODO 1
# 换模型。密钥一律从环境变量读，不要写成字面量。
def build_model() -> ChatOpenAI:
    return ChatOpenAI(
        model=os.environ.get("CHAT_MODEL", "gpt-4o-mini"),
        temperature=0.3,
        timeout=60,
        max_retries=2,
    )


# ---------------------------------------------------------------- TODO 2
# 换提示词。花括号里的变量名，就是调用时字典要传的键。
SYSTEM = "你是一位严谨的技术文档作者，回答简洁、不说套话。"
HUMAN = "请用不超过 200 字说明：{question}"


def build_prompt() -> ChatPromptTemplate:
    return ChatPromptTemplate.from_messages([("system", SYSTEM), ("human", HUMAN)])


# ---------------------------------------------------------------- TODO 3
# 换后处理。不需要就把这一段从管道里去掉。
def postprocess(text: str) -> str:
    return text.strip()


# ---------------------------------------------------------------- 装配
def build_chain() -> Runnable:
    """提示词 | 模型 | 解析器 | 后处理。

    管道符不执行任何计算，它只负责把上一步的输出接到下一步的输入；
    真正干活的是链上的每一个组件。
    """
    return (
        build_prompt()
        | build_model()
        | StrOutputParser()
        | postprocess
    ).with_config({"run_name": "通用链"})          # TODO 4：起个便于追踪的名字


# ---------------------------------------------------------------- 想多留几个中间结果
def build_chain_keep_all() -> Runnable:
    """输入输出都是字典，每一步往字典里加一个键，中间结果全部保留。"""
    inner = build_prompt() | build_model() | StrOutputParser()
    return (
        RunnablePassthrough.assign(answer=inner)
        | RunnablePassthrough.assign(length=lambda d: len(d["answer"]))
    )


# ---------------------------------------------------------------- 四种调用方式
def run_once(chain: Runnable, question: str) -> str:
    return chain.invoke({"question": question})


def run_batch(chain: Runnable, questions: list) -> list:
    return chain.batch([{"question": q} for q in questions],
                       config={"max_concurrency": 5})


def run_stream(chain: Runnable, question: str) -> None:
    for chunk in chain.stream({"question": question}):
        print(chunk, end="", flush=True)
    print()


async def run_async(chain: Runnable, question: str) -> str:
    return await chain.ainvoke({"question": question})


if __name__ == "__main__":
    chain = build_chain()

    # TODO 5：换成你自己的问题
    q = "为什么链上的组件都要实现同一组方法？"

    print("[invoke]", run_once(chain, q))
    print("[batch ]", run_batch(chain, ["什么是 LCEL？", "什么是 Runnable？"]))

    print("[stream] ", end="")
    # 注意：后处理函数会打断流式，演示流式时用不带后处理的版本
    run_stream(build_prompt() | build_model() | StrOutputParser(), q)

    print("[ainvoke]", asyncio.run(run_async(chain, q)))
    print("[keep all]", build_chain_keep_all().invoke({"question": q}).keys())

# -*- coding: utf-8 -*-
"""
多入多出流水线骨架
==================
链上流动的始终是一个字典，每一段用 assign 往字典里加一个键，
所有中间结果都能保留到最后，再按需要挑几个键输出。

依赖：pip install langchain langchain-openai python-dotenv
"""

import os

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import Runnable, RunnableParallel, RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model=os.environ.get("CHAT_MODEL", "gpt-4o-mini"))
parser = StrOutputParser()


def step(template: str) -> Runnable:
    """一段处理 = 提示词 | 模型 | 解析器。"""
    return PromptTemplate.from_template(template) | model | parser


# ---------------------------------------------------------------- TODO 1
# 逐段定义。花括号里的变量名必须是此前字典里已经存在的键，
# 否则会在运行时报 KeyError。
STEPS = [
    ("draft", "根据下面的要点写一段 150 字的初稿：\n\n{outline}"),
    ("polish", "润色下面这段文字，保持原意，语气更专业：\n\n{draft}"),
    ("title", "给下面这段文字起一个不超过 12 字的标题：\n\n{polish}"),
]


# ---------------------------------------------------------------- TODO 2
# 需要并行跑的段落放这里：同一份输入同时喂给多个分支，互不等待。
PARALLEL_STEPS = {
    "keywords": step("列出下面文字的 3 个关键词，用顿号分隔：\n\n{polish}"),
    "summary": step("用一句话总结下面文字：\n\n{polish}"),
}


# ---------------------------------------------------------------- TODO 3
# 最终要交出去的键。想全都要就把这一行设成 None。
OUTPUT_KEYS = ["title", "polish", "keywords", "summary"]


def build_pipeline() -> Runnable:
    chain: Runnable = RunnablePassthrough()

    for key, template in STEPS:
        chain = chain | RunnablePassthrough.assign(**{key: step(template)})

    if PARALLEL_STEPS:
        chain = chain | RunnablePassthrough.assign(**PARALLEL_STEPS)

    if OUTPUT_KEYS:
        chain = chain | RunnableParallel(
            **{k: (lambda d, k=k: d[k]) for k in OUTPUT_KEYS}
        )
    return chain.with_config({"run_name": "多段流水线"})


if __name__ == "__main__":
    pipeline = build_pipeline()

    # TODO 4：换成你自己的初始输入，键名要和第一段模板里的变量对上
    result = pipeline.invoke({"outline": "为什么链上的每个组件都要实现同一组方法"})

    for k, v in result.items():
        print("-" * 60)
        print(k, "：", v)

    # 批量跑多份输入，并限制并发，避免打爆上游限流
    batch = pipeline.batch(
        [{"outline": "什么是路由链"}, {"outline": "流式输出为什么会断"}],
        config={"max_concurrency": 2},
    )
    print("批量结果条数：", len(batch))

# -*- coding: utf-8 -*-
"""
路由链骨架
==========
按输入把请求分发到不同的子链，并且一定带兜底分支。
注册表写法：加一路只需要往 ROUTES 里加一项，不用改 if/else。

依赖：pip install langchain langchain-openai python-dotenv
"""

import os

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable, RunnableBranch, RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()

model = ChatOpenAI(model=os.environ.get("CHAT_MODEL", "gpt-4o-mini"))
parser = StrOutputParser()


def make_branch(system_prompt: str) -> Runnable:
    """一路子链 = 一段人设提示词 + 模型 + 解析器。"""
    p = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", "{question}"),
    ])
    return p | model | parser


# ---------------------------------------------------------------- TODO 1
# 路由表：键是分类名，值是这一路的子链。加一路就在这里加一项。
ROUTES = {
    "售前": make_branch("你是售前顾问，重点讲清产品能力边界与适用场景，不承诺没有的功能。"),
    "售后": make_branch("你是售后工程师，先复述问题，再给出排查步骤，最后给出联系方式。"),
    "账单": make_branch("你是账单专员，只回答计费与发票相关的问题，金额一律以系统记录为准。"),
}

# ---------------------------------------------------------------- TODO 2
# 兜底分支：所有分类都没命中时走它。这一路一定要写。
FALLBACK = make_branch("你是客服接线员，礼貌说明这个问题需要转接人工，并请用户补充信息。")


# ---------------------------------------------------------------- TODO 3
# 分类器。规则能判清楚就别花一次模型调用；判不清楚再交给模型。
KEYWORDS = {
    "售前": ("能不能", "支持吗", "多少钱", "试用"),
    "售后": ("报错", "打不开", "崩溃", "退货", "维修"),
    "账单": ("发票", "扣费", "账单", "续费"),
}

classify_prompt = ChatPromptTemplate.from_template(
    "下面的问题属于哪一类？只回答一个词：" + "、".join(ROUTES) + "、其他。\n\n问题：{question}"
)
model_classifier = classify_prompt | model.bind(temperature=0) | parser


def classify(d: dict) -> str:
    q = d["question"]
    for name, words in KEYWORDS.items():
        if any(w in q for w in words):
            return name
    return model_classifier.invoke(d)


# ---------------------------------------------------------------- 装配
def build_router() -> Runnable:
    """先算出 route 键，再按它选一条子链。"""
    branch = RunnableBranch(
        *[(lambda d, k=k: k in d["route"], chain) for k, chain in ROUTES.items()],
        FALLBACK,
    )
    return RunnablePassthrough.assign(route=classify) | branch


if __name__ == "__main__":
    router = build_router()

    # TODO 4：换成你自己的问题
    for q in ["你们支持私有化部署吗", "登录页一直报错 500", "上个月的发票怎么开", "在吗"]:
        print("=" * 60)
        print("问题：", q)
        print("回答：", router.invoke({"question": q})[:80], "...")