from langchain_core.messages import HumanMessage

from _004_langgraph_more_nodes.agent_state import AgentState
from common.llm import my_llm
import json


def neo4j_answer_generate_node(state: AgentState) -> AgentState:
    print("图谱答案节点执行开始了")
    # 读转写结果而不是 state["input"]：回答要针对补全后的完整问题，
    # 否则模型看到的是"那第二呢？"，拿着查询结果也不知道在答什么
    user_input = state.get("input_semantic_trans") or ""
    cypher_results = state.get("cypher_results", [])

    # 把 cypher_results 转成字符串，方便喂给大模型
    cypher_results_str = json.dumps(cypher_results, ensure_ascii=False, indent=2)

    prompt = f"""
    你是一个中医知识图谱问答助手。
    用户提出了问题：{user_input}

    我已经在 Neo4j 图数据库中执行了查询，查询结果如下：
    {cypher_results_str}

    请你根据这些查询结果，用简洁、清晰、自然的中文回答用户的问题。
    如果查询结果无法回答用户的问题，请如实告知用户没有找到相关答案。
    """
    # print(prompt)
    response = my_llm.invoke([HumanMessage(content=prompt)])

    # 保存结果
    state["neo4j_answer"] = response.content.strip()
    state["output"] = response.content.strip()
    print("完成进行neo4j输入大模型的回答")

    return state

if __name__ == '__main__':
    result=neo4j_answer_generate_node(
        {
            "input_semantic_trans": "请给我推荐一个Herb",
            "cypher_results": [
                [
                    {
                        "herb": "人参",
                        "effect": "解热镇痛"
                    }
                ]
            ]
        }
    )

    print(result)