基于知识图谱的RAG

导入所需库

1import streamlit as st
2from LLM import get_entities, generate_response
3from Neo4j import Neo4jHandler

设置应用标题并初始化会话状态以存储聊天记录和用户输入

1# Streamlit 应用标题
2st.title("基于知识图谱的聊天：查询你的数据集")
3
4# 初始化会话状态以存储聊天历史和用户输入
5if "messages" not in st.session_state:
6    st.session_state["messages"] = []
7
8if "user_input" not in st.session_state:
9    st.session_state["user_input"] = ""

创建 Neo4jHandler 对象，使用凭据连接数据库；该类的具体实现将在后面的 Neo4j.py 文件中展示。

1neo4j_handler = Neo4jHandler(connection_uri, username, password)

获取用户输入

1user_input = st.text_input("提出与你的数据集相关的问题：", value=st.session_state["user_input"])

从提交查询到生成回复的处理步骤

1if st.button("发送") and user_input:
2  # 步骤 1：从用户查询中提取实体
3  extracted_entities = get_entities(user_input)
4  print("查询中的实体：", extracted_entities)
5  
6  # 步骤 2：从 Neo4j 知识图谱中检索相关的关系
7  related_relationships_tuple_list = neo4j_handler.get_entities_and_relationships(extracted_entities)
8  print("从知识图谱中提取的上下文：", related_relationships_tuple_list)
9  
10  # 步骤 3：使用知识图谱上下文增强查询
11  messages = [
12    {"role": "system", "content": f"请使用以下知识图谱上下文（以关系元组列表形式提供：(实体1, 关系, 实体2)）来回答用户问题。\n{related_relationships_tuple_list}"},
13    {"role": "user", "content": user_input},
14  ]
15  
16  # 步骤 4：使用增强后的查询调用 GPT-4 生成回复
17  bot_response = generate_response(messages)
18  
19  # 步骤 5：将用户消息和机器人回复添加到会话状态的聊天历史中
20  st.session_state["messages"].append({"role": "user", "content": user_input})
21  st.session_state["messages"].append({"role": "assistant", "content": bot_response})
22
23  # 清空会话状态中的输入框内容
24  st.session_state["user_input"] = ""

使用不同颜色区分用户和机器人的消息样式

1def display_message(message, is_user):
2  if is_user:
3    st.markdown(f"""
4    <div style="background-color:#DCF8C6;padding:10px;border-radius:10px;margin-bottom:10px;">
5    {message}
6    </div>
7    """, unsafe_allow_html=True)
8  else:
9    st.markdown(f"""
10    <div style="background-color:#E6E6FA;padding:10px;border-radius:10px;margin-bottom:10px;">
11    {message}
12    </div>
13    """, unsafe_allow_html=True)

显示聊天消息

1if st.session_state["messages"]:
2  for msg in st.session_state["messages"]:
3    if msg["role"] == "user":
4      display_message(msg["content"], is_user=True)
5    elif msg["role"] == "assistant":
6      display_message(msg["content"], is_user=False)

关闭 Neo4j 驱动连接

1neo4j_handler.close()

从查询中提取实体

1def get_entities(query):
2  content = extract_entities(query)
3  entities = parse_llm_response_content(content)
4  return entities

生成对查询的回复

1def generate_response(messages):
2  try:
3    response = client.chat.completions.create(
4    model="gpt-4",
5    messages=messages,
6    )
7    response_content = response.choices[0].message.content
8    return response_content
9  except Exception as e:
10    return f"错误：{str(e)}"

连接数据库并实现检索流程

1from neo4j import GraphDatabase
2
3class Neo4jHandler:
4  def __init__(self, connection_uri, username, password):
5    self.driver = GraphDatabase.driver(connection_uri, auth=(username, password))
6  
7  def close(self):
8    self.driver.close()
9  
10  def get_entities_and_relationships(self, entities):
11    with self.driver.session() as session:
12      results = []
13      for entity in entities:
14        result = session.read_transaction(self._retrieve_relationships, entity)
15        results.extend(result)
16      return results
17
18  @staticmethod
19  def _retrieve_relationships(tx, entity_name):
20    query = """
21    MATCH (e1:Entity {name: $entity_name})-[r:RELATION]->(e2:Entity)
22    RETURN e1.name AS entity1, r.type AS relationship, e2.name AS entity2
23    """
24    result = tx.run(query, entity_name=entity_name)
25    return [(record["entity1"], record["relationship"], record["entity2"]) for record in result]