TxtAI项目中Graph RAG功能增强：属性同步与类型化图查询实践

2025-05-22 21:28:43作者：柯茵沙

引言

在知识图谱与自然语言处理结合的领域，TxtAI作为一个强大的开源项目，提供了丰富的语义搜索和图分析能力。本文将深入探讨如何增强TxtAI中的Graph RAG功能，特别是针对属性同步和类型化图查询这两个关键需求的技术实现方案。

属性同步的核心挑战

在TxtAI的当前版本中，图组件仅同步文本数据，这限制了图查询的表达能力。例如，当我们需要查询特定类型的节点（如"person"或"skill"）时，系统无法识别这些类型属性，导致查询结果不完整或为空。

这一限制在实际应用中尤为明显。考虑以下场景：一个包含人员节点（A1、B1）和技能节点（S1、S2）的知识图谱，我们需要找出所有掌握特定技能的人员。由于类型属性的缺失，基于类型的查询变得困难。

技术实现方案

基础属性同步实现

最直接的解决方案是扩展TxtAI的Graph类，利用NetworkX的内置属性处理功能：

class EnhancedGraph(Graph):
    def sync_attributes(self, attributes):
        nx.set_node_attributes(self.graph, attributes)
    
    def get_node_attributes(self, attribute):
        return nx.get_node_attributes(self.graph, attribute)

这种方法简单直接，通过NetworkX提供的set_node_attributes和get_node_attributes函数，实现了节点属性的同步和查询功能。

类型化图查询增强

为了支持更复杂的类型化查询，我们可以进一步扩展功能：

class TypedGraph(EnhancedGraph):
    def add_typed_node(self, node, node_type, **attr):
        attr['type'] = node_type
        self.graph.add_node(node, **attr)
    
    def get_nodes_by_type(self, node_type):
        return [n for n, data in self.graph.nodes(data=True) 
                if data.get('type') == node_type]

这种实现允许我们为节点添加类型属性，并基于类型进行查询，解决了最初提出的类型查询问题。

高级图查询能力

Cypher查询集成

对于需要更复杂查询的场景，可以集成Cypher查询语言的支持：

from grandcypher import GrandCypher

class CypherEnabledGraph(TypedGraph):
    def __init__(self):
        super().__init__()
        self.cypher = GrandCypher(self.graph)
    
    def find_connected_nodes(self, start_node, rel_type, end_node_type):
        query = f"""
        MATCH (start {{id: '{start_node}'}})-[r:{rel_type}]->(end {{type: '{end_node_type}'}})
        RETURN end
        """
        return self.cypher.run(query)

这种实现使得我们可以执行类似"查找所有掌握Python技能的人员"这样的复杂查询。

RDF/OWL集成方案

对于需要更强语义能力的场景，可以考虑集成RDF和OWL：

from owlready2 import *

class OntologyGraph(CypherEnabledGraph):
    def __init__(self):
        super().__init__()
        self.onto = get_ontology("http://example.org/onto.owl")
        
    def add_typed_node(self, node, node_type, **attr):
        with self.onto:
            if not self.onto[node_type]:
                types.new_class(node_type, (Thing,))
            new_individual = self.onto[node_type](node)
            for key, value in attr.items():
                setattr(new_individual, key, value)
        super().add_typed_node(node, node_type, **attr)

这种方案不仅支持类型化查询，还能利用OWL的推理能力发现隐含关系。

RAG功能增强

将属性信息整合到RAG流程中可以显著提升生成质量：

class AttributeAwareRAG:
    def __init__(self, graph):
        self.graph = graph
        self.tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        self.model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
    
    def get_node_embedding(self, node):
        attrs = self.graph.graph.nodes[node]
        text = f"Node: {node}, Type: {attrs.get('type','Unknown')}, " + \
               ", ".join(f"{k}:{v}" for k,v in attrs.items() if k != 'type')
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True)
        with torch.no_grad():
            outputs = self.model(**inputs)
        return outputs.last_hidden_state.mean(dim=1).squeeze().numpy()

这种方法在生成嵌入时考虑了节点类型和其他属性，使得检索到的上下文更加精确。

实际应用示例

让我们看一个完整的应用示例：

# 创建增强图实例
graph = OntologyGraph()

# 添加节点和边
graph.add_typed_node("A1", "Person", name="Alice", age=30)
graph.add_typed_node("B1", "Person", name="Bob", age=35)
graph.add_typed_node("S1", "Skill", name="Python")
graph.add_typed_edge("A1", "S1", "hasSkill")

# 执行类型查询
persons = graph.get_nodes_by_type("Person")
print(f"所有人员节点: {persons}")

# 执行Cypher查询
python_experts = graph.find_connected_nodes("S1", "hasSkill", "Person")
print(f"掌握Python的人员: {python_experts}")

# 使用增强RAG
rag = AttributeAwareRAG(graph)
response = rag.generate_response("谁懂Python?")
print(response)