# Runnable 四大核心类型
# 原文示例: AI 数据分析流水线

from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableParallel
from langchain_core.prompts import ChatPromptTemplate

print("1. RunnablePassthrough — 原样传递输入")
print("   用途: 保存原始用户输入供后续步骤使用")
print()
print("2. RunnableLambda — 把任意 Python 函数接入管道")
print("   用途: 数据库查询、文件读写、复杂计算")
print()
print("3. RunnableParallel (Map) — 数据路由器")
print("   用途: 从输入提取多个字段, 分别处理")
print()
print("4. 管道符 | — 把 Runnable 连接起来")
print("   等价于: f(g(h(x))), 但更直观")

# 从原文提取的完整流水线模式
# (模拟, 不需要真实 LangChain 安装)

pipeline_steps = {
    "step1_prompt": {
        "name": "SQL生成",
        "input": "用户问题",
        "tool": "LLM + SQL Prompt",
        "output": "SQL 查询语句"
    },
    "step2_exec": {
        "name": "数据获取",
        "input": "SQL 查询语句",
        "tool": "RunnableLambda(数据库查询)",
        "output": "DataFrame"
    },
    "step3_plot": {
        "name": "可视化",
        "input": "DataFrame + 用户问题",
        "tool": "LLM + Plotly Prompt",
        "output": "Plotly 绘图代码"
    },
    "step4_render": {
        "name": "渲染",
        "input": "Plotly 代码",
        "tool": "RunnableLambda(exec)",
        "output": "图表图片"
    }
}

print("=== AI 数据分析流水线 ===")
for step, info in pipeline_steps.items():
    print(f"  {info['name']:8s} | 输入: {info['input']:15s} → 工具: {info['tool']:25s} → 输出: {info['output']}")

# 用本地 LLM 模拟整个流水线的工作流
import httpx

def run_ai_pipeline(user_question):
    """模拟 Runnable 流水线"""
    print(f"用户输入: {user_question}")
    print()
    
    # Step 1: Prompt 填充 (类似 RunnableParallel)
    context = {
        "user_question": user_question,
        "table_schema": "sales(region, product, amount, date)"
    }
    print(f"[Step 1] 注入上下文: {context}")
    
    # Step 2: LLM 生成 SQL (类似 PromptTemplate | LLM)
    print(f"[Step 2] LLM 生成 SQL... (模拟)")
    sql = f"SELECT region, SUM(amount) FROM sales GROUP BY region"
    print(f"  SQL: {sql}")
    
    # Step 3: 执行查询 (类似 RunnableLambda)
    print(f"[Step 3] 执行查询... (模拟)")
    result = {"东北": 100, "华北": 200, "华东": 350}
    
    # Step 4: LLM 生成可视化 (类似 RunnablePassthrough 传递数据)
    print(f"[Step 4] LLM 生成 Plotly 代码... (模拟)")
    plot_code = """
import plotly.express as px
fig = px.bar(x=['东北','华北','华东'], y=[100,200,350])
fig.show()
"""
    print(f"  输出: 柱状图 (各区域销售额)")
    print()
    print("=== 流水线完成 ===")
    return result

run_ai_pipeline("各区域的销售额是多少？")

# 关键设计模式: Runnable 作为函数组合
# 原文核心代码结构解构

print("原文中的核心模式:")
print()
print("# 管道操作符 | 等价于函数组合")
print("runnable1 | runnable2 | runnable3")
print("  = f(g(h(x))), 输出自动传递")
print()
print("# RunnablePassthrough 保存原始输入")
print("RunnablePassthrough() | prompt | llm")
print("  # 原始问题存为 'my_question' 键")
print()
print("# RunnableLambda 接入业务代码")
print("RunnableLambda(lambda x: pd.DataFrame(x).head(3))")
print("  # 把数据库查询结果的前3行传给LLM")
print()
print("# RunnableParallel 数据路由")
print("{'question': itemgetter('user_input'), 'data': execute_sql}")
print("  # 同时处理多个字段, 结果合并到字典")