class GraphState(TypedDict):
    keys: Dict[str, any]

def retrieve(state):
'''
从state中取出问题 => 调用检索器获取相关文档
把获取的文档添加到state中，并返回新的状态即可
'''
    state_dict = state["keys"]
    question = state_dict["question"]
    documents = retriever.get_relevant_documents(question)
    return {"keys": {"documents": documents, "question": question}}

def grade_documents(state):

    #取出state中的问题和检索的文档
    print("---CHECK RELEVANCE---")
    state_dict = state["keys"]
    question = state_dict["question"]
    documents = state_dict["documents"]

    #此处省略定义llm_with_tool/parser_tool工具，可参考langchain文档

    #Prompt提示模版
    prompt = PromptTemplate(
        template="""您是一个评分人员，评估检索到的文档与用户问题 
                    的相关性。以下是检索到的文档：
        {context}
        以下是用户的问题：{question}
        如果文档包含与用户问题相关的关键词或语义意义，请将其评为相关。给出一个yes或no来表明文档是否与问题相关。""",
        input_variables=["context", "question"],
    )

    #构建一个评估的Chain
    chain = prompt | llm_with_tool | parser_tool

    #对文档做评估
    filtered_docs = []
    search = "No" 
    for d in documents:
        score = chain.invoke({"question": question, "context": d.page_content})
        grade = score[0].binary_score

        #如果相关，则添加到关联文档
        if grade == "yes":
            print("---GRADE: DOCUMENT RELEVANT---")
            filtered_docs.append(d)

        #如果不相关，则要求进行web搜索
        else:
            print("---GRADE: DOCUMENT NOT RELEVANT---")
            search = "Yes"  # Perform web search
            continue

    #将过滤后的文档、是否需要web搜索保存到state返回
    return {
        "keys": {
            "documents": filtered_docs,
            "question": question,
            "run_web_search": search,
        }
    }

def generate(state):

    print("---GENERATE---")
    state_dict = state["keys"]
    question = state_dict["question"]
    documents = state_dict["documents"]

    # 获取一个现成的Prompt（借助langchain hub）
    prompt = hub.pull("rlm/rag-prompt")

    # 模型
    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, streaming=True)

    #Chain：提示词=>llm=>输出解析
    rag_chain = prompt | llm | StrOutputParser()

    #调用
    generation = rag_chain.invoke({"context": documents, "question": question})

    #结果放入state，返回
    return {
        "keys": {"documents": documents, "question": question, "generation": generation}
    }

def transform_query(state):、

    print("---TRANSFORM QUERY---")
    state_dict = state["keys"]
    question = state_dict["question"]
    documents = state_dict["documents"]

    #提示模版
    prompt = PromptTemplate(
        template="""你需要生成针对检索优化的问题。请根据输入内容，尝试推理其中的语义意图/含义。这是初始问题：
               {question}
                 请提出一个改进的问题：""",
                 input_variables=["question"],
    )

    # 模型
    model = ChatOpenAI(temperature=0, model="gpt-3.5-turbo-1106", streaming=True)

    # 链
    chain = prompt | model | StrOutputParser()

    # 调用获得改进的问题
    better_question = chain.invoke({"question": question})

    # 放入state返回
    return {"keys": {"documents": documents, "question": better_question}}

def web_search(state):
    print("---WEB SEARCH---")
    state_dict = state["keys"]
    question = state_dict["question"]
    documents = state_dict["documents"]

    #bing搜索
    tool = BingSearchAPIWrapper(k=3)
    docs = tool.run(question)
    web_results = "\n".join([d for d in docs])
    web_results = Document(page_content=web_results)

    #搜索结果添加到已经有的documents后面返回
    documents.append(web_results)

    return {"keys": {"documents": documents, "question": question}}

def decide_to_generate(state):
    print("---DECIDE TO GENERATE---")
    state_dict = state["keys"]
    question = state_dict["question"]
    filtered_documents = state_dict["documents"]

    #取出评估节点放入的评估结果：是否需要web搜索
    search = state_dict["run_web_search"]

    #如果需要，返回下一个节点：transform_query
    if search == "Yes":
        print("---DECISION: TRANSFORM QUERY and RUN WEB SEARCH---")
        return "transform_query"

    #如果不需要，返回下一个节点：generate
    else:
        # We have relevant documents, so generate answer
        print("---DECISION: GENERATE---")
        return "generate"

#定一个一个Graph
workflow = StateGraph(GraphState)

#添加上面准备的Node
workflow.add_node("retrieve", retrieve) # retrieve
workflow.add_node("grade_documents", grade_documents) # grade documents
workflow.add_node("generate", generate) # generatae
workflow.add_node("transform_query", transform_query) # transform_query
workflow.add_node("web_search", web_search) # web search

#特殊边：入口
workflow.set_entry_point("retrieve")

#普通边：检索完成后，进入文档评估
workflow.add_edge("retrieve", "grade_documents")

#条件边：借助辅助方法决定进入哪一个节点
workflow.add_conditional_edges(
    "grade_documents",
    decide_to_generate,
    {
        "transform_query": "transform_query",
        "generate": "generate",
    },
)

#普通边：改写问题后总是去搜索
workflow.add_edge("transform_query", "web_search")

#普通边：搜索后，生成最后答案
workflow.add_edge("web_search", "generate")

#特殊边：生成答案后结束
workflow.add_edge("generate", END)

#编译成应用，完成！
app = workflow.compile()

inputs = {"keys": {"question": "百度的comate是什么？有哪些特点？"}}
for output in app.stream(inputs):

    #output中保存了每一步节点完成信息，key是节点名，value是state
    for key, value in output.items():
        pprint.pprint(f"Node '{key}':")
        pprint.pprint("\n---\n")

# 打印最后的State中的generation，就是最终答案
pprint.pprint(value["keys"]["generation"])