import os

from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 使用千问大模型进行对比
os.environ["DASHSCOPE_API_KEY"] = '通义千问API KEY'

# 文档加载器
loader = TextLoader(file_path='./agent.txt', encoding='utf-8')

# 加载文档
documents = loader.load()

# 文本分割器
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=100,
    chunk_overlap=0,
    length_function=len,
    is_separator_regex=False,
    separators=["\n\n", "\n"]
)

# 文档分割
docs = text_splitter.split_documents(documents)

# 调用大模型的向量库
embeddings = DashScopeEmbeddings(
    model='text-embedding-v4',
)
#
# # 构建本地向量库
# vector_db = Chroma.from_documents(
#     documents=docs,
#     embedding=embeddings,
#     persist_directory='data/chroma_storage',
# )
#
# 连接本地向量库
vector_db = Chroma(
    persist_directory='data/chroma_storage',
    embedding_function=embeddings,
)

retriever = vector_db.as_retriever()

# 大模型
llm = Tongyi(
    model='qwen-plus',
    temperature=0.1,
)

# 定义查询模板
prompt = ChatPromptTemplate.from_template(
    """请根据以下文档内容回答问题。
    文档内容: {context}
    问题: {question}
    如果文档内容中没有相关信息，请回答“抱歉，我无法回答这个问题”。
    """
)

# 构建 RAG 链
rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | llm
)

query = "请问：Agent有几种类型？分别是什么？"

answer = rag_chain.invoke(query)
print(answer)