←返回首页

使用 Neo4j 分析代码结构并辅助 LLM 代码生成

2025年7月5日·数据库
Neo4j 对代码结构和关系进行建模，为大型语言模型（LLM）提供项目特定的上下文，从而指导其生成更准确、更相关的代码片段

问题描述

大型语言模型（LLM）能够根据自然语言提示生成代码片段，但它们无法直接访问大型代码库的具体结构和依赖关系，导致生成的代码虽然语法正确，却可能与实际项目上下文无关或存在冗余。

为提升生成代码的相关性和正确性，需要将项目特定的上下文知识融入生成过程。

解决方案概述

本项目采用 Neo4j 原生图数据库，将代码库表示为知识图谱，捕捉文件、类、函数等实体及其相互关系，从而支持复杂的代码结构和依赖查询。

从 Neo4j 检索到的丰富上下文可作为 LLM 提示的一部分，指导生成更准确、具备项目语境的代码片段。

代码图模型

图谱模式建模了关键元素及其关系：

  • (:File {name})：源文件节点
  • (:Class {name})：类节点
  • (:Function {name})：函数/方法节点
  • (:Module {name})：模块/包节点

关系类型：

  • (:File)-[:CONTAINS]->(:Class)
  • (:Class)-[:DEFINES]->(:Function)
  • (:Function)-[:CALLS]->(:Function)（函数调用）
  • (:Function)-[:IMPORTS]->(:Module)

该模式支持查找调用层级、依赖关系、模块使用情况等丰富查询。

工具与库

  • Neo4j Community Edition：本地图数据库
  • Python ast 模块：解析 Python 源文件为抽象语法树（AST）
  • Neo4j Python Driver（neo4j 包）：与 Neo4j 数据库交互
  • LLM API 或本地模型：如 OpenAI GPT、HuggingFace Transformers

实现细节

1. 使用 ast 解析 Python 代码

Python 标准库的 ast 模块可将 Python 文件解析为树结构，便于提取函数、类定义，以及函数调用和导入信息。

import ast

def extract_functions_from_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        source = f.read()
    tree = ast.parse(source)
    functions = [node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)]
    return functions

2. 写入 Neo4j 图数据库

通过官方 Python 驱动和 Cypher 查询，将函数及其关系插入 Neo4j。

from neo4j import GraphDatabase

driver = GraphDatabase.driver("bolt://localhost:7687", auth=("neo4j", "password"))

def add_function(tx, function_name, file_name):
    tx.run("""
        MERGE (f:Function {name: $function_name})
        MERGE (file:File {name: $file_name})
        MERGE (file)-[:CONTAINS]->(f)
    """, function_name=function_name, file_name=file_name)

with driver.session() as session:
    session.write_transaction(add_function, "save_user", "user_service.py")

3. 构建函数调用关系图

遍历 AST 时，检测函数体内的调用，并以 CALLS 关系表示：

class FunctionCallVisitor(ast.NodeVisitor):
    def __init__(self):
        self.calls = []

    def visit_Call(self, node):
        if isinstance(node.func, ast.Name):
            self.calls.append(node.func.id)
        self.generic_visit(node)

def get_function_calls(func_node):
    visitor = FunctionCallVisitor()
    visitor.visit(func_node)
    return visitor.calls

将调用关系记录到 Neo4j：

def add_call_relation(tx, caller, callee):
    tx.run("""
        MATCH (caller:Function {name: $caller}), (callee:Function {name: $callee})
        MERGE (caller)-[:CALLS]->(callee)
    """, caller=caller, callee=callee)

4. 查询 Neo4j 获取上下文

示例 Cypher 查询，查找所有调用某函数的函数：

MATCH (caller:Function)-[:CALLS]->(callee:Function {name: 'save_user'})
RETURN caller.name

通过 Python 发起查询，丰富 LLM 提示：

def get_callers(session, function_name):
    result = session.run("""
        MATCH (caller:Function)-[:CALLS]->(callee:Function {name: $function_name})
        RETURN caller.name AS caller_name
    """, function_name=function_name)
    return [record["caller_name"] for record in result]

5. LLM 提示工程

在调用 LLM 生成代码前，先检索相关图谱上下文（如已有相关函数或模块），并将其融入提示：

context = "项目使用 pandas 进行数据处理。已有函数：save_user, load_user。"

prompt = f"""
鉴于项目背景：
{context}

请编写一个 Python 函数，将用户数据导出为 CSV 文件。
"""

generated_code =
#Neo4j