# python环境 3.12.4
# streamlit 1.39.0
# toml

pip install langchain
pip install streamlit

import streamlit as st

# 设置 st 的标题和布局
st.set_page_config(page_title="RAG测试问答", layout="wide")

# 设置应用标题
st.title("RAG测试问答")

# 支持上传 txt 文件
upload_file = st.sidebar.file_uploader(label="上传文件", type=["txt"])

if not upload_file:
    st.info("请上传 txt 文件")
    st.stop()

import streamlit as st
import tempfile
import os
from langchain.memory import ConversationBufferMemory # 会话记录到内存
from langchain_community.chat_message_histories import StreamlitChatMessageHistory #  Streamlit聊天记录存储
from langchain_community.document_loaders import TextLoader # 文本加载器
from langchain_ollama.embeddings import OllamaEmbeddings # Ollama Eembeddings 语言模型
from langchain_chroma import Chroma # Chroma 向量数据库
from langchain_text_splitters import RecursiveCharacterTextSplitter # 文本分割器

# 设置 st 的标题和布局
st.set_page_config(page_title="RAG测试问答", layout="wide")

# 设置应用标题
st.title("RAG测试问答")

# 支持上传 txt 文件
upload_file = st.sidebar.file_uploader(label="上传文件", type=["txt"])

if not upload_file:
    st.info("请上传 txt 文件")
    st.stop()

# step1 实现知识库生成
@st.cache_resource(ttl="1h")
def get_knowledge_base(uploaded_file):
    # 读取上传的文档
    docs = []
    # 将 uploaded_file 存到 /tmp
    temp_dir = tempfile.TemporaryDirectory(dir=r"/tmp")
    tempfilepath = os.path.join(temp_dir.name, uploaded_file.name)
    with open(tempfilepath, "wb") as f:
        f.write(uploaded_file.getvalue())
    # 使用 TextLoader 加载文档
    docs = TextLoader(tempfilepath, encoding="utf-8").load()
    # 使用 RecursiveCharacterTextSplitter 分割文档
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = splitter.split_documents(docs)
    # 使用 OllamaEmbeddings 生成文档向量
    embeddings = OllamaEmbeddings(base_url="http://127.0.0.1:11434", model="shaw/dmeta-embedding-zh")
    # 使用 Chroma 向量数据库存储文档向量
    chroma_db = Chroma.from_documents(splits, embeddings)
    # 创建文档检索器 约等于 知识库
    retriever = chroma_db.as_retriever()
    return retriever
    
retriever = get_knowledge_base(upload_file) 

# step2 初始化聊天消息界面
# 如果用户输入"清空聊天记录"，则重新初始化界面
if "messages" not in st.session_state or st.sidebar.button("清空聊天记录"):
    st.session_state["messages"] = [{
        "role": "你好",
        "content": "我是测试 RAG 问答小助手"
    }]

# 显示历史聊天记录
for msg in st.session_state["messages"]:
    st.chat_message(msg["role"], msg["content"])  # 

# 创建历史聊天记录
msgs = StreamlitChatMessageHistory()
# 创建对话缓存
memory = ConversationBufferMemory(
    chat_memory=msgs,
    return_messages=True,
    memory_key="chat_history",
    output_key="out"
)
# 创建 UI 输入框
user_query = st.chat_input(placeholder="请输入要测试的问题")

# step3 创建检索 agent 
from langchain.tools.retriever import create_retriever_tool

# step3-1 用于创建文档检索的工具
tool = create_retriever_tool(
    retriever=retriever,
    name="文档检索",
    description="根据输入的关键词，检索相关文档",
)
tools = [tool]

# step3-2 创建 LLM对话模型
# 创建指令 Prompt
instruction = """你是一个设计用于查询文档回答问题的代理
您可以使用文档检索工具，并基于检索内容来回答问题。
可能你不查询文档就知道答案，但是仍然要去查询文档来获得答案。
如果从文档找不到任何信息和答案来回答问题，则需要返回“非常抱歉，这个问题暂时没有录入到知识库中。”作为答案。
"""

base_template = """
{instruction}

TOOLS:
----------

你可以使用以下工具：

{tools}

使用工具中，你可以参考这样子：

ZWJ```
思考：我是否需要使用工具？ 是的
动作：我需要使用工具：[{tool_names}]
动作：输入:{input}
动作执行后： 返回动作执行后的结果
ZWJ```

当你需要返回一个答案，且这个答案不需要使用工具时，你可以参考这样子：

ZWJ```
思考：我是否需要使用工具？ 不是
答案： [你的答案]
ZWJ```

开始！

上一次历史对话内容如下：

{chat_history}

新的问题是：{input}
{agent_scratchpad}"""

# agent_scratchpad 是 agent 的 scratchpad，用于存储 agent 的状态

# 基础模板
base_prompt = PromptTemplate.from_template(base_template)

# 填充基础模板
prompt = base_prompt.partial(instruction=instruction)

# 创建 LLM 模型
llm = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:7b")


# step3-3 创建 agent
agent = create_react_agent(llm=llm, prompt=prompt, tools=tools)

agent_excutor = AgentExecutor(
    agent=agent,
    tools=tools,
    memory=memory,
    verbose=True,
    handle_parsing_errors="从知识库没找到对应内容或者答案"
)

# step5 用户输入查询与返回
if  user_query:
    # 添加到 session 历史记录
    st.session_state["messages"].append({"role": "user", "content": user_query})
    # 显示用户信息
    st.chat_message("user").write(user_query)
    
    with st.chat_message("assistant"):
        # 创建回调
        callback = StreamlitCallbackHandler(st.container())
        # 将 agent执行过程 显示在 streamlit中，如：思考、选择工具、执行查询等等
        config = {"callbacks": [callback]}
        # agent 执行
        response = agent_excutor.invoke({"input": user_query}, config=config)
        # 保存agent 执行结果到聊天记录 
        st.session_state["messages"].append({"role": "assistant", "content": response["output"]})
        # 显示在 streamlit中
        st.write(response["output"])

from fastapi import FastAPI
from fastapi.responses import RedirectResponse
from langserve import add_routes

app = FastAPI()

# 这里是定义路由和对应路由实现的方法
@app.get("/")
async def redirect_root_to_docs():
    return RedirectResponse("/docs")


# 这里我们可以添加路由
# add_routes(app, NotImplemented)

if __name__ == "__main__":
    import uvicorn
    # 这里通过 uvicorn启动服务，端口为8000
    uvicorn.run(app, host="0.0.0.0", port=8000)

from fastapi import FastAPI
from fastapi.responses import RedirectResponse
from langserve import add_routes
from langchain_ollama.llms import OllamaLLM

deepseek = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:32b")

app = FastAPI(
    title="Qborfy个人 LangServer",
    version="0.1.0",
    description="Qborfy个人 LangServer，学习测试使用",
)


@app.get("/")
async def redirect_root_to_docs():
    return RedirectResponse("/docs")

# 添加 deepseek路由
add_routes(app, deepseek, path="/deepseek")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)

# 实现langchain调用远端调用 LangServer
from langchain.schema.runnable import RunnableMap
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langserve import RemoteRunnable

llm = RemoteRunnable("http://127.0.0.1:8000/deepseek")

template = "你是世界级的 AI 技术专家, {input}"
# 这里我们使用一个简单的模板
prompt = PromptTemplate(
    input_variables=["input"],
    template=template
)
# 创建一个简单的链式调用
chain = prompt | llm | StrOutputParser()

# 执行链式调用
for chunk in llm.stream("海洋是什么颜色"):
    print(chunk, end="|", flush=True)

import os

# 设置LangSimth 环境变量
os.environ["LANGSMITH_TRACING"] = "true"
os.environ["LANGSMITH_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGSMITH_API_KEY"] = "<LANG_SIMTH_KEY>"
os.environ["LANGSMITH_PROJECT"] = "mylangserver"

from langchain.globals import set_verbose
# 全局 verbose配置开关
set_verbose(True)

# 针对部分链接调用设置详细日志开关

llm = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:32b", verbose=True)

from langchain.globals import set_debug
# 全局 debug配置开关
set_debug(True)

# 实现langchain调用 ollama 大模型
from langchain_ollama.llms import OllamaLLM

llm = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:32b")

chunks = []
# llm.stream 会返回一个流
for chunk in llm.stream("海洋是什么颜色"):
    chunks.append(chunk)
    print(chunk, end="|", flush=True)

# 实现langchain调用 ollama 大模型
from langchain_ollama.llms import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import asyncio

llm = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:32b")
propmt = ChatPromptTemplate.from_template("给我讲一个关于{input}的笑话")
parser = StrOutputParser()
chain = propmt | llm | parser

# 异步调用需要定义 async 方法
async def async_stream():
    async for chunk in chain.astream("公鸡"):
        print(chunk, end="|", flush=True)

# 调用的话需要通过  asyncio.run() 方法       
asyncio.run(async_stream())

# 实现langchain调用 ollama 大模型
from langchain_ollama.llms import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
import asyncio

llm = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:32b")
propmt = ChatPromptTemplate.from_template(
    """
    以 JSON格式返回{x}的人口列表
    使用一个`省份`作为字段列表返回
    每个省份都应有有字段`省份名`+`人口`字段                                  
    """)
parser = JsonOutputParser() # 保证每次输出都是json格式
chain = propmt | llm | parser

# 异步调用需要定义 async 方法
async def async_stream():
    async for chunk in chain.astream("广东省、福建省、广西省"):
        # 可以看到每次 chunk都是一个完整的 json 格式
        print(chunk, end="\n", flush=True)

# 调用的话需要通过  asyncio.run() 方法       
asyncio.run(async_stream())

# 实现langchain调用 ollama 大模型
from langchain_ollama.llms import OllamaLLM
import asyncio

llm = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:32b")
async def async_stream():
    async for event in llm.astream_events("你好", version="v2"):
        print(event)
# 调用的话需要通过  asyncio.run() 方法       
asyncio.run(async_stream())

from langchain_core.prompts import ChatPromptTemplate

# 通过一个消息数组创建聊天消息模板
# 数组每一个元素代表一条消息
# 第一个参数是消息角色  system 代表系统消息， 第二个参数代表消息内容
# 消息角色 system 代表系统消息
# 消息角色 human 代表系统消息代表人类
# 消息角色 ai 代表LLM大模型返回的消息内容
# {xxx} 定义 模板参数，如下定义两个模板参数  name代表人工智能名字 user_input 代表用户输入的文本 
chat_template = ChatPromptTemplate.from_messages(
    [
        ("system", "你是人工智能助手， 你的名字是{name}"), 
        ("human", "你好"), 
        ("ai", "你好，我是人工智能助手{name}，很高兴为您服务"), 
        ("human", "{user_input}"),
    ]
)

# 通过模板参数格式化模板内容
message = chat_template.format_messages(name="小爱同学", user_input="你的名字叫什么？")

print(message)

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

chat_template = ChatPromptTemplate.from_messages(
    [
        ("system", "你是人工智能助手"), 
        MessagesPlaceholder("msgs")
    ]
)
# 这里我们可以之前的定义的消息列表放在一起
msgs=[SystemMessage(content='你的名字是小爱同学'), HumanMessage(content='你好'), AIMessage(content='你好，我是人工智能助手，很高兴为您服务')]

print(chat_template.invoke(msgs))

from langchain.prompts.few_shot import FewShotPromptTemplate
from langchain.prompts import PromptTemplate

examples = [
    {
        "question": "谁的寿命更长， 穆罕默德二世还是爱因斯坦？",
        "answer": 
        """
        爱因斯坦活了 76 岁。
        穆罕默德二世活了 89 岁。
        因此，穆罕默德二世比爱因斯坦活得更长。
        """
    },
    {
        "question": "目前电影票房第一名是谁？",
        "answer":
        """
        《阿凡达》的票房是 27.9 亿美元。
        《复仇者联盟 4：终局之战》的票房是 27.8 亿美元。
        因此，《阿凡达》的票房更高。
        """
    },
]
# 这里我们创建一个简单示例模板，用来分析解析示例的 question 和 answer
example_prompt = PromptTemplate(input_variables=["question", "answer"], template="问题：{question}\\n答案：{answer}")

# 用 FewShotPromptTemplate 可以根据模板+示例去生成一个拥有示例集的提示语模板
prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    suffix="问题：{input}",
    input_variables=["input"],
)

print(prompt.format(input="谁的寿命更长， 穆罕默德二世还是爱因斯坦？"))

# ExapleSelector 示例筛选器
from langchain.prompts.example_selector import SemanticSimilarityExampleSelector
from langchain_chroma import Chroma
from langchain_ollama.embeddings import OllamaEmbeddings
from langchain.prompts import FewShotPromptTemplate, PromptTemplate

# 引用 shaw/dmeta-embedding-zh 模型做为嵌入模型，其对中文支持度更加友好
ollama_emb = OllamaEmbeddings(
    base_url="http://127.0.0.1:11434", 
    model="shaw/dmeta-embedding-zh",
)

examples = [
    {
        "question": "谁的寿命更长， 穆罕默德二世还是爱因斯坦？",
        "answer": 
        """
        爱因斯坦活了 76 岁。
        穆罕默德二世活了 89 岁。
        因此，穆罕默德二世比爱因斯坦活得更长。
        """
    },
    {
        "question": "目前电影票房第一名是谁？",
        "answer":
        """
        《阿凡达》的票房是 27.9 亿美元。
        《复仇者联盟 4：终局之战》的票房是 27.8 亿美元。
        因此，《阿凡达》的票房更高。
        """
    },
    {
        "question": "深圳第一高楼是哪个？",
        "answer":
        """
        深圳平安国际金融中心（平安中心）的楼高是 593米。
        深圳京基100的楼高是441.8米。
        所以深圳第一高楼是平安国际金融中心。
        """
    },
]

example_selector = SemanticSimilarityExampleSelector.from_examples(
    # 这里是示例集合
    examples=examples,
    # 用户生成嵌入的嵌入类，用于衡量语义的相似度
    embeddings=ollama_emb,
    # 用于计算相似性的向量存储库，这里使用的是 Chroma， 一个保存在内容的向量存储库
    vectorstore_cls=Chroma(),
    # 选择前 k 个最相似的示例 这里设置为 1
    k=1,
)

question = "穆罕默德二世？"
# 选择最相似的示例
selected_examples = example_selector.select_examples({"question": question})

# 接下来结合 FewShotPromptTemplate 我们就可以得到更加准备且少量的示例 PromptTemplate
example_prompt = PromptTemplate(input_variables=["question", "answer"], template="问题：{question}\\n答案：{answer}")

prompt = FewShotPromptTemplate(
    # 这里调整完最相似的示例集合
    examples=selected_examples, 
    example_prompt=example_prompt,
    suffix="问题：{input}",
    input_variables=["input"],
)

print(prompt.format(input=question))

# 实现langchain调用 ollama 大模型
from langchain_ollama.llms import OllamaLLM
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

llm = OllamaLLM(base_url="http://127.0.0.1:11434", model="deepseek-r1:32b")

template = "你是世界级的 AI 技术专家, {input}"
# 这里我们使用一个简单的模板
prompt = PromptTemplate(
    input_variables=["input"],
    template=template
)
# 创建一个简单的链式调用
chain = prompt | llm | StrOutputParser()

# 执行链式调用
response = chain.invoke({
    "input":"请写一篇关于 AI 的文章，字数不大于 100"
})
print(response)