# 本地
from langchain_community.chat_models import ChatOllama

# 创建一个名为llama2_chat的ChatOllama对象，使用llama2:13b-chat模型
llama2_chat = ChatOllama(model="llama2:13b-chat")

# 创建一个名为llama2_code的ChatOllama对象，使用codellama:7b-instruct模型
llama2_code = ChatOllama(model="codellama:7b-instruct")

# API
from langchain_community.llms import Replicate

# 设置REPLICATE_API_TOKEN
# REPLICATE_API_TOKEN = getpass()
# os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN

# 设置replicate_id
replicate_id = "meta/llama-2-13b-chat:f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d"

# 创建一个名为llama2_chat_replicate的Replicate对象，使用指定的模型和输入参数
llama2_chat_replicate = Replicate(
    model=replicate_id, input={"temperature": 0.01, "max_length": 500, "top_p": 1}
)

# 提示
from langchain_core.prompts import ChatPromptTemplate

# 根据 SQL 数据库的类型（如 MySQL、Microsoft SQL Server 等）更新模板
template = """根据下面的表结构，编写一个 SQL 查询来回答用户的问题：
{schema}

问题：{question}
SQL 查询："""
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "给定一个输入问题，将其转换为 SQL 查询。不需要前言。"),
        ("human", template),
    ]
)

# 链接到查询
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

sql_response = (
    RunnablePassthrough.assign(schema=get_schema)  # 从 get_schema 函数中获取表结构
    | prompt
    | llm.bind(stop=["\nSQLResult:"])  # 绑定到 llm，停止标记为 "\nSQLResult:"
    | StrOutputParser()  # 使用 StrOutputParser 解析输出
)

sql_response.invoke({"question": "Klay Thompson 在哪个球队？"})  # 调用 sql_response，传入问题参数

# 链式回答
template = """根据下面的表模式、问题、SQL查询和SQL响应，编写一个自然语言回答：
{schema}

问题: {question}
SQL查询: {query}
SQL响应: {response}"""
prompt_response = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "给定一个输入问题和SQL响应，将其转换为自然语言回答。没有前言。",
        ),
        ("human", template),
    ]
)

full_chain = (
    RunnablePassthrough.assign(query=sql_response)  # 将sql_response赋值给query
    | RunnablePassthrough.assign(
        schema=get_schema,  # 获取表模式
        response=lambda x: db.run(x["query"]),  # 运行SQL查询并将结果赋值给response
    )
    | prompt_response  # 使用prompt_response生成自然语言回答
    | llm  # 使用llm模型生成最终回答
)

full_chain.invoke({"question": "有多少个唯一的团队？"})  # 调用full_chain并传入问题参数

# 导入必要的模块
from langchain.memory import ConversationBufferMemory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# 定义模板字符串，用于生成对话提示
template = """给定一个输入问题，将其转换为 SQL 查询。不需要序文。基于下面的表结构，编写一个 SQL 查询，以回答用户的问题：
{schema}
"""
# 创建对话提示模板
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", template),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)

# 创建对话缓存内存
memory = ConversationBufferMemory(return_messages=True)

# 链式调用，将查询与内存连接起来
from langchain_core.runnables import RunnableLambda

sql_chain = (
    RunnablePassthrough.assign(
        schema=get_schema,
        history=RunnableLambda(lambda x: memory.load_memory_variables(x)["history"]),
    )
    | prompt
    | llm.bind(stop=["\nSQLResult:"])
    | StrOutputParser()
)

# 定义保存函数
def save(input_output):
    output = {"output": input_output.pop("output")}
    memory.save_context(input_output, output)
    return output["output"]

# 执行 SQL 查询并保存结果到内存
sql_response_memory = RunnablePassthrough.assign(output=sql_chain) | save
sql_response_memory.invoke({"question": "What team is Klay Thompson on?"})

# Chain to answer
# 定义模板字符串，包含表格模式、问题、SQL查询和SQL响应，用于生成自然语言响应
template = """Based on the table schema below, question, sql query, and sql response, write a natural language response:
{schema}

Question: {question}
SQL Query: {query}
SQL Response: {response}"""
# 从消息中创建聊天提示模板
prompt_response = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "Given an input question and SQL response, convert it to a natural language answer. No pre-amble.",
        ),
        ("human", template),
    ]
)

# 创建完整的处理链
full_chain = (
    RunnablePassthrough.assign(query=sql_response_memory)  # 将SQL响应分配给query
    | RunnablePassthrough.assign(  # 分配schema和response
        schema=get_schema,  # 获取表格模式
        response=lambda x: db.run(x["query"]),  # 运行SQL查询并获取响应
    )
    | prompt_response  # 使用聊天提示模板生成自然语言响应
    | llm  # 使用语言模型生成最终的自然语言响应
)

# 调用完整的处理链并传入问题
full_chain.invoke({"question": "What is his salary?"})