#利用假设提问检索技巧提升问答效果

prompt = """
   Generate a list of exactly 3 hypothetical questions that the below document could be used to answer:
   doc nanme: \n{doc_name}
   doc content: \n{doc}
   
   output format:
   ["腾讯2023年度的财务表现相比于2022年度有何变化？具体体现在哪些关键财务指标上?",
   "腾讯2023年度综合业绩报告中，未经重列和经重列的经营盈利、期内盈利以及本公司权益持有人应占盈利分别下降了多少百分比？",
   "2023年与2022年相比，董事会如何解释腾讯综合财务报表中出现的大幅盈利下滑，特别是基本每股盈利和摊薄每股盈利的下降幅度？"]
   """

# 用于将模型输出转换为列表
def ListOutputParser(input):
   import ast
   return ast.literal_eval(input.content)

from langchain_core.prompts import ChatPromptTemplate

# 定义任务 Chain
chain = (
   {"doc_name": lambda x: x.metadata["source"] ,"doc": lambda x: x.page_content}
   # Only asking for 3 hypothetical questions, but this could be adjusted
   | ChatPromptTemplate.from_template(prompt)
   | model | ListOutputParser
   )

hypothetical_questions = chain.invoke(data[0])

hypothetical_questions = chain.batch(data, {"max_concurrency": 5})

vectorstore = Chroma(
   collection_name="hypo-questions", embedding_function=embeddings
   )

from langchain.storage import InMemoryByteStore
# The storage layer for the parent documents
store = InMemoryByteStore()
id_key = "doc_id"

from langchain.retrievers.multi_vector import MultiVectorRetriever

retriever = MultiVectorRetriever(
vectorstore=vectorstore,
byte_store=store,
id_key=id_key,
)

import uuid
   
doc_ids = [str(uuid.uuid4()) for _ in data]

question_docs = []

for i, question_list in enumerate(hypothetical_questions):
   question_docs.extend(
      [Document(page_content=s, metadata={id_key: doc_ids[i]}) for s in question_list]
   )

retriever.vectorstore.add_documents(question_docs)

retriever.docstore.mset(list(zip(doc_ids, data)))

sub_docs = vectorstore.similarity_search("腾讯2023年收入比2022年高多少")

from langchain_core.prompts import PromptTemplate

template = """
   You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
   
   Context: {context}
   
   Question: {question}
   
   Answer:
   """

rag_prompt = PromptTemplate.from_template(template)

def format_docs(docs):
      return "\n\n".join(doc.page_content for doc in docs)

from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

rag_chain = (
   {"context": retriever | format_docs, "question": RunnablePassthrough()}
   | rag_prompt
   | model
   | StrOutputParser()
   )

question = "腾讯2023年收入比2022年高多少？"
print(f'\nQuestion: {question} \n Answer: {rag_chain.invoke(question)}')