826 lines
33 KiB
Python
826 lines
33 KiB
Python
|
"""
|
|||
|
|
graph_builder.py - 知识图谱构建的核心流程控制器
|
|||
|
|
包含完整的构建流程和交互逻辑
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
import json
|
|||
|
|
try:
|
|||
|
|
import faiss
|
|||
|
|
except ImportError: # pragma: no cover - depends on deployment environment
|
|||
|
|
faiss = None
|
|||
|
|
try:
|
|||
|
|
import numpy as np
|
|||
|
|
except ImportError: # pragma: no cover - depends on deployment environment
|
|||
|
|
np = None
|
|||
|
|
import time
|
|||
|
|
from datetime import datetime
|
|||
|
|
from pathlib import Path
|
|||
|
|
import logging
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
from collections import defaultdict
|
|||
|
|
|
|||
|
|
from models import GraphNode, GraphEdge, NodeType, EdgeType
|
|||
|
|
from code_parser import (
|
|||
|
|
parser, CPP_EXTENSIONS, IGNORE_DIRS,
|
|||
|
|
safe_decode, extract_comment_before, extract_code_snippet,
|
|||
|
|
extract_full_function_name, collect_call_expressions,
|
|||
|
|
extract_class_name, extract_member_function_name, extract_base_classes
|
|||
|
|
)
|
|||
|
|
from llm_processor import generate_code_logic, generate_code_summary, get_qwen_embedding
|
|||
|
|
from config import PROJECT_ROOT, MAX_CODE_LENGTH, QWEN_EMBEDDING_MODEL
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# 全局图数据结构
|
|||
|
|
FUNCTION_CALL_GRAPH = defaultdict(list)
|
|||
|
|
FILE_DEPENDENCIES = defaultdict(list)
|
|||
|
|
CALLED_BY_GRAPH = defaultdict(list)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _as_node_type(value):
|
|||
|
|
return value.value if hasattr(value, "value") else value
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _normalize_output_dir(output_dir):
|
|||
|
|
if output_dir is None:
|
|||
|
|
return Path(".")
|
|||
|
|
path = Path(output_dir)
|
|||
|
|
path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
return path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _qualified_name(node):
|
|||
|
|
if node.id.startswith("Function:"):
|
|||
|
|
return node.id[len("Function:"):]
|
|||
|
|
return node.name
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _has_real_embedding(embedding):
|
|||
|
|
if not embedding:
|
|||
|
|
return False
|
|||
|
|
try:
|
|||
|
|
return any(abs(float(value)) > 1e-12 for value in embedding)
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _function_metadata(node, embedding_dim=None):
|
|||
|
|
raw = node.raw_attributes or {}
|
|||
|
|
embedding_available = raw.get("embedding_available")
|
|||
|
|
if embedding_available is None:
|
|||
|
|
embedding_available = _has_real_embedding(node.embedding)
|
|||
|
|
return {
|
|||
|
|
"node_id": node.id,
|
|||
|
|
"name": node.name,
|
|||
|
|
"function_name": node.name,
|
|||
|
|
"qualified_name": raw.get("qualified_name") or _qualified_name(node),
|
|||
|
|
"file": node.file_path,
|
|||
|
|
"file_path": node.file_path,
|
|||
|
|
"start_line": node.start_line,
|
|||
|
|
"end_line": node.end_line,
|
|||
|
|
"signature": node.signature,
|
|||
|
|
"summary": node.summary or "",
|
|||
|
|
"logic_flow": node.logic_flow or "",
|
|||
|
|
"code_snippet": raw.get("code_snippet", ""),
|
|||
|
|
"calls": raw.get("calls", []),
|
|||
|
|
"called_by": raw.get("called_by", []),
|
|||
|
|
"includes": raw.get("includes", []),
|
|||
|
|
"embedding_model": QWEN_EMBEDDING_MODEL,
|
|||
|
|
"embedding_dim": embedding_dim or (len(node.embedding) if node.embedding else 0),
|
|||
|
|
"embedding_available": bool(embedding_available),
|
|||
|
|
"embedding_error": raw.get("embedding_error", ""),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SymbolResolver:
|
|||
|
|
"""符号解析器：将代码符号映射为图谱节点ID"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self._symbol_table = {}
|
|||
|
|
self._file_to_entities = defaultdict(list)
|
|||
|
|
|
|||
|
|
def register_entity(self, node: GraphNode):
|
|||
|
|
"""注册新发现的实体到符号表"""
|
|||
|
|
self._symbol_table[node.name] = node.id
|
|||
|
|
if node.file_path:
|
|||
|
|
self._file_to_entities[node.file_path].append(node.id)
|
|||
|
|
|
|||
|
|
def resolve_name(self, name: str, current_context: dict) -> str:
|
|||
|
|
"""
|
|||
|
|
解析代码中的名称到节点ID
|
|||
|
|
:param name: 代码中的名称
|
|||
|
|
:param current_context: 当前上下文（文件、类等）
|
|||
|
|
:return: 节点ID或None
|
|||
|
|
"""
|
|||
|
|
if name in self._symbol_table:
|
|||
|
|
return self._symbol_table[name]
|
|||
|
|
|
|||
|
|
for qualified_name, node_id in self._symbol_table.items():
|
|||
|
|
if qualified_name.endswith('::' + name) or qualified_name == name:
|
|||
|
|
return node_id
|
|||
|
|
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def process_single_file(file_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map):
|
|||
|
|
"""处理单个C/C++文件，提取代码实体和结构关系"""
|
|||
|
|
logger.info(f"处理单个文件: {file_path}")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
with open(file_path, 'rb') as f:
|
|||
|
|
source_bytes = f.read()
|
|||
|
|
source_lines = source_bytes.split(b'\n')
|
|||
|
|
tree = parser.parse(source_bytes)
|
|||
|
|
|
|||
|
|
# 1. 创建文件节点
|
|||
|
|
file_node_id = f"File:{file_path}"
|
|||
|
|
file_node = GraphNode(
|
|||
|
|
id=file_node_id,
|
|||
|
|
type=NodeType.FILE,
|
|||
|
|
name=Path(file_path).name,
|
|||
|
|
file_path=file_path
|
|||
|
|
)
|
|||
|
|
all_nodes.append(file_node)
|
|||
|
|
symbol_resolver.register_entity(file_node)
|
|||
|
|
|
|||
|
|
# 2. 收集头文件依赖
|
|||
|
|
includes = []
|
|||
|
|
for node in tree.root_node.named_children:
|
|||
|
|
if node.type == "preproc_include":
|
|||
|
|
path_node = node.child_by_field_name("path")
|
|||
|
|
if path_node:
|
|||
|
|
include_text = path_node.text.decode('utf-8').strip('"<>')
|
|||
|
|
includes.append(include_text)
|
|||
|
|
FILE_DEPENDENCIES[file_path] = includes
|
|||
|
|
|
|||
|
|
# 3. 遍历文件内的顶级声明
|
|||
|
|
for node in tree.root_node.named_children:
|
|||
|
|
# 3.1 处理函数定义
|
|||
|
|
if node.type == "function_definition":
|
|||
|
|
func_name = extract_full_function_name(node)
|
|||
|
|
if func_name == "<unknown>":
|
|||
|
|
logger.warning(f" 跳过无法解析函数名的节点: {node.text[:50]}...")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 创建函数节点
|
|||
|
|
func_node_id = f"Function:{func_name}"
|
|||
|
|
func_node = GraphNode(
|
|||
|
|
id=func_node_id,
|
|||
|
|
type=NodeType.FUNCTION,
|
|||
|
|
name=func_name,
|
|||
|
|
signature=node.text.decode('utf-8', errors='ignore').split('{')[0].strip() + ' {...}',
|
|||
|
|
file_path=file_path,
|
|||
|
|
start_line=node.start_point[0] + 1,
|
|||
|
|
end_line=node.end_point[0] + 1,
|
|||
|
|
raw_attributes={'ast_node_type': 'function_definition'}
|
|||
|
|
)
|
|||
|
|
all_nodes.append(func_node)
|
|||
|
|
symbol_resolver.register_entity(func_node)
|
|||
|
|
|
|||
|
|
# 存储原始信息
|
|||
|
|
comment = extract_comment_before(node, source_lines)
|
|||
|
|
code_snippet = extract_code_snippet(tree, node, source_lines)
|
|||
|
|
function_raw_info_map[func_node_id] = {
|
|||
|
|
'ast_node': node,
|
|||
|
|
'comment': comment,
|
|||
|
|
'code_snippet': code_snippet,
|
|||
|
|
'source_lines': source_lines,
|
|||
|
|
'tree': tree
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 添加 FILE -> CONTAINS -> FUNCTION 边
|
|||
|
|
all_edges.append(GraphEdge(
|
|||
|
|
source_id=file_node_id,
|
|||
|
|
target_id=func_node_id,
|
|||
|
|
type=EdgeType.CONTAINS
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
# 收集函数调用关系
|
|||
|
|
callees = collect_call_expressions(node)
|
|||
|
|
FUNCTION_CALL_GRAPH[func_name] = callees
|
|||
|
|
|
|||
|
|
# 3.2 处理类/结构体定义
|
|||
|
|
elif node.type in ["class_specifier", "struct_specifier"]:
|
|||
|
|
class_name = extract_class_name(node)
|
|||
|
|
if not class_name or class_name == "<unknown>":
|
|||
|
|
logger.debug(f" 跳过无法解析类名的节点")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 创建类节点
|
|||
|
|
class_node_id = f"Class:{class_name}"
|
|||
|
|
class_node = GraphNode(
|
|||
|
|
id=class_node_id,
|
|||
|
|
type=NodeType.CLASS,
|
|||
|
|
name=class_name,
|
|||
|
|
signature=node.text.decode('utf-8', errors='ignore').split('{')[0].strip() + ' {...}',
|
|||
|
|
file_path=file_path,
|
|||
|
|
start_line=node.start_point[0] + 1,
|
|||
|
|
end_line=node.end_point[0] + 1,
|
|||
|
|
raw_attributes={
|
|||
|
|
'ast_type': node.type,
|
|||
|
|
'is_struct': node.type == "struct_specifier"
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
all_nodes.append(class_node)
|
|||
|
|
symbol_resolver.register_entity(class_node)
|
|||
|
|
|
|||
|
|
# 添加 FILE -> CONTAINS -> CLASS 边
|
|||
|
|
all_edges.append(GraphEdge(
|
|||
|
|
source_id=file_node_id,
|
|||
|
|
target_id=class_node_id,
|
|||
|
|
type=EdgeType.CONTAINS
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
# 处理基类/继承关系
|
|||
|
|
base_clause = node.child_by_field_name("base_clause")
|
|||
|
|
if base_clause:
|
|||
|
|
base_classes = extract_base_classes(base_clause)
|
|||
|
|
for base_class_info in base_classes:
|
|||
|
|
base_name = base_class_info.get('name')
|
|||
|
|
access_type = base_class_info.get('access', 'public')
|
|||
|
|
edge_type = EdgeType.EXTENDS if access_type in ['public', 'protected',
|
|||
|
|
'private'] else EdgeType.IMPLEMENTS
|
|||
|
|
|
|||
|
|
if base_name:
|
|||
|
|
base_id = symbol_resolver.resolve_name(
|
|||
|
|
base_name,
|
|||
|
|
{'current_file': file_path, 'current_class': class_name}
|
|||
|
|
)
|
|||
|
|
if base_id:
|
|||
|
|
all_edges.append(GraphEdge(
|
|||
|
|
source_id=class_node_id,
|
|||
|
|
target_id=base_id,
|
|||
|
|
type=edge_type,
|
|||
|
|
properties={'access': access_type}
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
# 处理类内方法
|
|||
|
|
for member in node.named_children:
|
|||
|
|
if member.type == "function_definition":
|
|||
|
|
member_func_name = extract_member_function_name(member, class_name)
|
|||
|
|
if not member_func_name or member_func_name == "<unknown>":
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
member_node_id = f"Function:{class_name}::{member_func_name}"
|
|||
|
|
existing_member_node = next((n for n in all_nodes if n.id == member_node_id), None)
|
|||
|
|
|
|||
|
|
if not existing_member_node:
|
|||
|
|
member_node = GraphNode(
|
|||
|
|
id=member_node_id,
|
|||
|
|
type=NodeType.FUNCTION,
|
|||
|
|
name=member_func_name,
|
|||
|
|
signature=member.text.decode('utf-8', errors='ignore').split('{')[0].strip() + ' {...}',
|
|||
|
|
file_path=file_path,
|
|||
|
|
start_line=member.start_point[0] + 1,
|
|||
|
|
end_line=member.end_point[0] + 1,
|
|||
|
|
raw_attributes={
|
|||
|
|
'is_member': True,
|
|||
|
|
'parent_class': class_name,
|
|||
|
|
'ast_node_type': 'member_function_definition'
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
all_nodes.append(member_node)
|
|||
|
|
symbol_resolver.register_entity(member_node)
|
|||
|
|
|
|||
|
|
# 存储原始信息
|
|||
|
|
member_comment = extract_comment_before(member, source_lines)
|
|||
|
|
member_code = extract_code_snippet(tree, member, source_lines)
|
|||
|
|
function_raw_info_map[member_node_id] = {
|
|||
|
|
'ast_node': member,
|
|||
|
|
'comment': member_comment,
|
|||
|
|
'code_snippet': member_code,
|
|||
|
|
'source_lines': source_lines,
|
|||
|
|
'tree': tree
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 收集成员函数调用
|
|||
|
|
member_callees = collect_call_expressions(member)
|
|||
|
|
FUNCTION_CALL_GRAPH[f"{class_name}::{member_func_name}"] = member_callees
|
|||
|
|
|
|||
|
|
# 添加类包含边
|
|||
|
|
all_edges.append(GraphEdge(
|
|||
|
|
source_id=class_node_id,
|
|||
|
|
target_id=member_node_id,
|
|||
|
|
type=EdgeType.CONTAINS,
|
|||
|
|
properties={'member_type': 'method'}
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"解析文件 {file_path} 失败: {e}", exc_info=True)
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
|
|||
|
|
def process_project_files(project_root, all_nodes, all_edges, symbol_resolver, function_raw_info_map):
|
|||
|
|
"""处理整个项目目录"""
|
|||
|
|
logger.info(f"处理项目目录: {project_root}")
|
|||
|
|
|
|||
|
|
for root, dirs, files in os.walk(project_root):
|
|||
|
|
# 跳过忽略目录
|
|||
|
|
dirs[:] = [d for d in dirs if d not in IGNORE_DIRS]
|
|||
|
|
|
|||
|
|
for file in files:
|
|||
|
|
file_path_obj = Path(file)
|
|||
|
|
if file_path_obj.suffix.lower() not in CPP_EXTENSIONS:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
full_file_path = os.path.join(root, file)
|
|||
|
|
process_single_file(full_file_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)
|
|||
|
|
|
|||
|
|
return len(files)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def semantic_enhancement(all_nodes, all_edges, function_raw_info_map):
|
|||
|
|
"""阶段三：语义增强，为函数节点生成摘要与向量嵌入"""
|
|||
|
|
logger.info("阶段三：语义增强，为函数节点生成摘要与向量嵌入...")
|
|||
|
|
|
|||
|
|
# 1. 过滤出需要处理的函数节点
|
|||
|
|
function_nodes = [n for n in all_nodes if n.type == NodeType.FUNCTION]
|
|||
|
|
nodes_with_raw_info = []
|
|||
|
|
nodes_without_raw_info = []
|
|||
|
|
|
|||
|
|
for func_node in function_nodes:
|
|||
|
|
if func_node.id in function_raw_info_map:
|
|||
|
|
nodes_with_raw_info.append(func_node)
|
|||
|
|
else:
|
|||
|
|
nodes_without_raw_info.append(func_node)
|
|||
|
|
|
|||
|
|
logger.info(f"有原始信息的节点: {len(nodes_with_raw_info)} 个")
|
|||
|
|
if nodes_without_raw_info:
|
|||
|
|
logger.warning(f"无原始信息的节点: {len(nodes_without_raw_info)} 个")
|
|||
|
|
|
|||
|
|
# 2. 处理有原始信息的节点
|
|||
|
|
if nodes_with_raw_info:
|
|||
|
|
print(f"\n开始为 {len(nodes_with_raw_info)} 个函数生成摘要和向量...")
|
|||
|
|
print("这可能需要一些时间，请耐心等待...")
|
|||
|
|
|
|||
|
|
processed_count = 0
|
|||
|
|
for func_node in tqdm(nodes_with_raw_info, desc="生成函数摘要", unit="个"):
|
|||
|
|
func_node_id = func_node.id
|
|||
|
|
func_name = func_node.name
|
|||
|
|
|
|||
|
|
# 获取原始信息
|
|||
|
|
raw_info = function_raw_info_map[func_node_id]
|
|||
|
|
comment = raw_info['comment']
|
|||
|
|
code_snippet = raw_info['code_snippet']
|
|||
|
|
file_path = func_node.file_path
|
|||
|
|
func_node.raw_attributes.update({
|
|||
|
|
"comment": comment,
|
|||
|
|
"code_snippet": code_snippet,
|
|||
|
|
"includes": FILE_DEPENDENCIES.get(file_path, []),
|
|||
|
|
"qualified_name": _qualified_name(func_node),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 1. 生成流程逻辑
|
|||
|
|
try:
|
|||
|
|
logic = generate_code_logic(func_name, comment, code_snippet, file_path)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"为函数 {func_name} 生成逻辑失败: {e}")
|
|||
|
|
logic = "逻辑生成失败。"
|
|||
|
|
|
|||
|
|
# 2. 生成功能摘要
|
|||
|
|
try:
|
|||
|
|
summary = generate_code_summary(
|
|||
|
|
func_name, comment, logic, code_snippet, file_path
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"为函数 {func_name} 生成摘要失败: {e}")
|
|||
|
|
summary = "摘要生成失败。"
|
|||
|
|
|
|||
|
|
# 3. 更新节点信息
|
|||
|
|
func_node.summary = summary
|
|||
|
|
func_node.logic_flow = logic
|
|||
|
|
|
|||
|
|
# 获取调用关系
|
|||
|
|
called_functions = FUNCTION_CALL_GRAPH.get(func_name, [])
|
|||
|
|
caller_functions = CALLED_BY_GRAPH.get(func_name, [])
|
|||
|
|
func_node.raw_attributes.update({
|
|||
|
|
'calls': called_functions,
|
|||
|
|
'called_by': caller_functions,
|
|||
|
|
'includes': FILE_DEPENDENCIES.get(file_path, [])
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 4. 生成向量嵌入
|
|||
|
|
context_text_for_embedding = (
|
|||
|
|
f"函数名: {func_name}\n"
|
|||
|
|
f"功能摘要: {summary}\n"
|
|||
|
|
f"流程逻辑: {logic}\n"
|
|||
|
|
f"所在文件: {os.path.basename(file_path)}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
embedding_vector = get_qwen_embedding(context_text_for_embedding)
|
|||
|
|
embedding_values = (
|
|||
|
|
embedding_vector.tolist()
|
|||
|
|
if hasattr(embedding_vector, "tolist")
|
|||
|
|
else list(embedding_vector)
|
|||
|
|
)
|
|||
|
|
if not _has_real_embedding(embedding_values):
|
|||
|
|
raise RuntimeError("Embedding API returned an empty or zero vector.")
|
|||
|
|
func_node.embedding = embedding_values
|
|||
|
|
func_node.raw_attributes["embedding_available"] = True
|
|||
|
|
func_node.raw_attributes["embedding_error"] = ""
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"为函数 {func_name} 生成嵌入失败: {e}")
|
|||
|
|
func_node.embedding = [0.0] * 1024
|
|||
|
|
func_node.raw_attributes["embedding_available"] = False
|
|||
|
|
func_node.raw_attributes["embedding_error"] = str(e)[:500]
|
|||
|
|
|
|||
|
|
processed_count += 1
|
|||
|
|
# 每处理5个函数显示一次进度
|
|||
|
|
if processed_count % 5 == 0:
|
|||
|
|
print(f" 已处理 {processed_count}/{len(nodes_with_raw_info)} 个函数")
|
|||
|
|
|
|||
|
|
# 添加延迟避免API限流
|
|||
|
|
time.sleep(0.1)
|
|||
|
|
|
|||
|
|
print(f"✓ 完成 {len(nodes_with_raw_info)} 个函数的语义增强")
|
|||
|
|
else:
|
|||
|
|
logger.warning("没有找到任何有原始信息的函数节点，跳过摘要生成。")
|
|||
|
|
|
|||
|
|
# 3. 处理无原始信息的节点
|
|||
|
|
if nodes_without_raw_info:
|
|||
|
|
print(f"\n为 {len(nodes_without_raw_info)} 个无原始信息的节点生成基础信息...")
|
|||
|
|
for func_node in nodes_without_raw_info:
|
|||
|
|
func_node.summary = f"函数 {func_node.name}，位于 {func_node.file_path}"
|
|||
|
|
func_node.logic_flow = "无法生成逻辑流程（缺少原始代码信息）"
|
|||
|
|
func_node.embedding = [0.0] * 1024
|
|||
|
|
func_node.raw_attributes.update({
|
|||
|
|
"calls": FUNCTION_CALL_GRAPH.get(func_node.name, []),
|
|||
|
|
"called_by": CALLED_BY_GRAPH.get(func_node.name, []),
|
|||
|
|
"includes": FILE_DEPENDENCIES.get(func_node.file_path, []),
|
|||
|
|
"qualified_name": _qualified_name(func_node),
|
|||
|
|
"embedding_available": False,
|
|||
|
|
"embedding_error": "Missing source information.",
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_call_edges(all_nodes, all_edges, symbol_resolver):
|
|||
|
|
"""阶段二：构建函数调用边"""
|
|||
|
|
logger.info("阶段二：解析符号，构建函数调用图 (CALLS 边)...")
|
|||
|
|
|
|||
|
|
# 构建反向调用图
|
|||
|
|
CALLED_BY_GRAPH.clear()
|
|||
|
|
for caller_name, callee_names in FUNCTION_CALL_GRAPH.items():
|
|||
|
|
caller_node_id = f"Function:{caller_name}"
|
|||
|
|
# 确保调用者节点已存在
|
|||
|
|
if not any(n.id == caller_node_id for n in all_nodes):
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
for callee_name in callee_names:
|
|||
|
|
# 通过符号解析器将名称解析为节点ID
|
|||
|
|
callee_node_id = symbol_resolver.resolve_name(
|
|||
|
|
callee_name,
|
|||
|
|
{'current_file': next((n.file_path for n in all_nodes if n.id == caller_node_id), None)}
|
|||
|
|
)
|
|||
|
|
if callee_node_id:
|
|||
|
|
# 添加 CALLS 边
|
|||
|
|
all_edges.append(GraphEdge(
|
|||
|
|
source_id=caller_node_id,
|
|||
|
|
target_id=callee_node_id,
|
|||
|
|
type=EdgeType.CALLS
|
|||
|
|
))
|
|||
|
|
# 更新反向图
|
|||
|
|
if callee_name not in CALLED_BY_GRAPH:
|
|||
|
|
CALLED_BY_GRAPH[callee_name] = []
|
|||
|
|
CALLED_BY_GRAPH[callee_name].append(caller_name)
|
|||
|
|
|
|||
|
|
logger.info(f"阶段二完成。已构建边数: {len(all_edges)}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_results(all_nodes, all_edges, base_name, target_path, mode, output_dir=None):
|
|||
|
|
"""阶段四：持久化存储结果"""
|
|||
|
|
logger.info("阶段四：持久化图谱与向量索引...")
|
|||
|
|
|
|||
|
|
output_path = _normalize_output_dir(output_dir)
|
|||
|
|
graph_data_path = str(output_path / f"{base_name}_code_knowledge_graph.json")
|
|||
|
|
vector_db_path = str(output_path / f"{base_name}_rag.faiss")
|
|||
|
|
metadata_path = str(output_path / f"{base_name}_rag_metadata.json")
|
|||
|
|
|
|||
|
|
# 1. 保存为图谱文件
|
|||
|
|
nodes_dict = [node.dict() for node in all_nodes]
|
|||
|
|
edges_dict = [edge.dict() for edge in all_edges]
|
|||
|
|
|
|||
|
|
graph_data = {
|
|||
|
|
"metadata": {
|
|||
|
|
"project_root": target_path if mode == "full_project" else "single_file",
|
|||
|
|
"file_path": target_path if mode == "single_file" else None,
|
|||
|
|
"mode": mode,
|
|||
|
|
"generated_at": datetime.now().isoformat(),
|
|||
|
|
"total_nodes": len(all_nodes),
|
|||
|
|
"total_edges": len(all_edges)
|
|||
|
|
},
|
|||
|
|
"nodes": nodes_dict,
|
|||
|
|
"edges": edges_dict
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
with open(graph_data_path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(graph_data, f, indent=2, ensure_ascii=False, default=str)
|
|||
|
|
logger.info(f" 图谱结构已保存至: {graph_data_path}")
|
|||
|
|
print(f"✓ 图谱结构已保存: {graph_data_path}")
|
|||
|
|
|
|||
|
|
# 2. 构建并保存向量索引
|
|||
|
|
function_nodes = [node for node in all_nodes if node.type == NodeType.FUNCTION and node.embedding]
|
|||
|
|
if function_nodes:
|
|||
|
|
dimension = len(function_nodes[0].embedding)
|
|||
|
|
metadata_for_faiss = []
|
|||
|
|
|
|||
|
|
vectors = []
|
|||
|
|
for node in function_nodes:
|
|||
|
|
vectors.append(node.embedding)
|
|||
|
|
metadata_for_faiss.append(_function_metadata(node, dimension))
|
|||
|
|
|
|||
|
|
if vectors:
|
|||
|
|
if faiss is not None:
|
|||
|
|
if np is None:
|
|||
|
|
raise RuntimeError("numpy is required when writing a FAISS index.")
|
|||
|
|
vectors_np = np.array(vectors, dtype='float32')
|
|||
|
|
index = faiss.IndexFlatL2(dimension)
|
|||
|
|
index.add(vectors_np)
|
|||
|
|
faiss.write_index(index, vector_db_path)
|
|||
|
|
else:
|
|||
|
|
with open(vector_db_path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(
|
|||
|
|
{
|
|||
|
|
"format": "simple_l2_vector_index",
|
|||
|
|
"dimension": dimension,
|
|||
|
|
"vectors": [[float(value) for value in vector] for vector in vectors],
|
|||
|
|
},
|
|||
|
|
f,
|
|||
|
|
ensure_ascii=False,
|
|||
|
|
)
|
|||
|
|
with open(metadata_path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(metadata_for_faiss, f, indent=2, ensure_ascii=False)
|
|||
|
|
logger.info(f" 向量索引已保存。函数节点数: {len(function_nodes)}")
|
|||
|
|
logger.info(f" 向量数据库: {vector_db_path}")
|
|||
|
|
logger.info(f" 元数据: {metadata_path}")
|
|||
|
|
print(f"✓ 向量索引已保存: {vector_db_path}")
|
|||
|
|
else:
|
|||
|
|
logger.warning(" 没有可生成向量的函数节点，跳过FAISS索引构建。")
|
|||
|
|
print("⚠ 无向量数据，跳过FAISS索引构建")
|
|||
|
|
|
|||
|
|
return graph_data_path, vector_db_path, metadata_path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_code_knowledge_base(target_path, output_dir=None, semantic=True, base_name=None, embedding_dim=1024):
|
|||
|
|
"""Build a code knowledge base without interactive prompts."""
|
|||
|
|
target_path = str(Path(target_path).resolve())
|
|||
|
|
if not os.path.exists(target_path):
|
|||
|
|
raise FileNotFoundError(f"Target path does not exist: {target_path}")
|
|||
|
|
|
|||
|
|
all_nodes = []
|
|||
|
|
all_edges = []
|
|||
|
|
symbol_resolver = SymbolResolver()
|
|||
|
|
function_raw_info_map = {}
|
|||
|
|
|
|||
|
|
FUNCTION_CALL_GRAPH.clear()
|
|||
|
|
FILE_DEPENDENCIES.clear()
|
|||
|
|
CALLED_BY_GRAPH.clear()
|
|||
|
|
|
|||
|
|
if os.path.isdir(target_path):
|
|||
|
|
mode = "full_project"
|
|||
|
|
process_project_files(target_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)
|
|||
|
|
base_name = base_name or "project"
|
|||
|
|
else:
|
|||
|
|
mode = "single_file"
|
|||
|
|
process_single_file(target_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)
|
|||
|
|
base_name = base_name or f"file_{Path(target_path).stem}"
|
|||
|
|
|
|||
|
|
build_call_edges(all_nodes, all_edges, symbol_resolver)
|
|||
|
|
|
|||
|
|
if semantic:
|
|||
|
|
semantic_enhancement(all_nodes, all_edges, function_raw_info_map)
|
|||
|
|
else:
|
|||
|
|
for node in all_nodes:
|
|||
|
|
if node.type == NodeType.FUNCTION:
|
|||
|
|
raw_info = function_raw_info_map.get(node.id, {})
|
|||
|
|
node.summary = node.summary or f"Function {node.name}, located at {node.file_path}"
|
|||
|
|
node.logic_flow = node.logic_flow or "Semantic enhancement was skipped."
|
|||
|
|
node.embedding = node.embedding or [0.0] * embedding_dim
|
|||
|
|
node.raw_attributes.update({
|
|||
|
|
"comment": raw_info.get("comment", ""),
|
|||
|
|
"code_snippet": raw_info.get("code_snippet", ""),
|
|||
|
|
"calls": FUNCTION_CALL_GRAPH.get(node.name, []),
|
|||
|
|
"called_by": CALLED_BY_GRAPH.get(node.name, []),
|
|||
|
|
"includes": FILE_DEPENDENCIES.get(node.file_path, []),
|
|||
|
|
"qualified_name": _qualified_name(node),
|
|||
|
|
"embedding_available": False,
|
|||
|
|
"embedding_error": "Semantic enhancement was skipped.",
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return save_results(all_nodes, all_edges, base_name, target_path, mode, output_dir=output_dir)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_rag_database_interactive():
|
|||
|
|
"""交互式知识库构建主流程"""
|
|||
|
|
logger.info("=== 交互式知识图谱构建流程 ===")
|
|||
|
|
|
|||
|
|
# 显示模式选择菜单
|
|||
|
|
print("\n" + "=" * 50)
|
|||
|
|
print("知识库构建模式选择")
|
|||
|
|
print("=" * 50)
|
|||
|
|
print("1. 全项目模式 - 处理整个项目目录")
|
|||
|
|
print("2. 单文件模式 - 处理单个文件")
|
|||
|
|
print("3. 退出")
|
|||
|
|
print("=" * 50)
|
|||
|
|
|
|||
|
|
choice = input("请选择模式 (1-3): ").strip()
|
|||
|
|
|
|||
|
|
if choice == "3":
|
|||
|
|
logger.info("用户选择退出")
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
if choice == "1":
|
|||
|
|
# 全项目模式 - 直接引导用户输入
|
|||
|
|
print("\n" + "=" * 60)
|
|||
|
|
print("全项目知识库构建模式")
|
|||
|
|
print("=" * 60)
|
|||
|
|
|
|||
|
|
# 显示当前config.py中配置的默认路径
|
|||
|
|
current_default_path = str(PROJECT_ROOT) if PROJECT_ROOT else "未配置"
|
|||
|
|
if not current_default_path or current_default_path == ".":
|
|||
|
|
current_default_path = "当前工作目录"
|
|||
|
|
|
|||
|
|
# 直接要求用户输入
|
|||
|
|
while True:
|
|||
|
|
project_prompt = f"请直接输入要构建知识库的项目根目录完整路径\n(当前config.py中的默认路径: {current_default_path}): "
|
|||
|
|
user_input = input(project_prompt).strip()
|
|||
|
|
|
|||
|
|
if not user_input:
|
|||
|
|
# 用户没有输入，尝试使用config中的默认路径
|
|||
|
|
if PROJECT_ROOT and os.path.exists(PROJECT_ROOT) and os.path.isdir(PROJECT_ROOT):
|
|||
|
|
target_path = PROJECT_ROOT
|
|||
|
|
print(f"使用config.py中的默认路径: {target_path}")
|
|||
|
|
else:
|
|||
|
|
print("⚠ 您没有输入路径，且config.py中的默认路径无效。")
|
|||
|
|
retry = input("是否重新输入? (y/n): ").strip().lower()
|
|||
|
|
if retry != 'y':
|
|||
|
|
print("知识库构建取消。")
|
|||
|
|
return None, None
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
# 使用用户输入的路径
|
|||
|
|
target_path = Path(user_input)
|
|||
|
|
|
|||
|
|
# 验证路径
|
|||
|
|
if not os.path.exists(target_path):
|
|||
|
|
print(f"⚠ 路径不存在: {target_path}")
|
|||
|
|
retry = input("是否重新输入? (y/n): ").strip().lower()
|
|||
|
|
if retry != 'y':
|
|||
|
|
print("知识库构建取消。")
|
|||
|
|
return None, None
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
if not os.path.isdir(target_path):
|
|||
|
|
print(f"⚠ 路径不是目录: {target_path}")
|
|||
|
|
retry = input("是否重新输入? (y/n): ").strip().lower()
|
|||
|
|
if retry != 'y':
|
|||
|
|
print("知识库构建取消。")
|
|||
|
|
return None, None
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
break # 路径有效，退出循环
|
|||
|
|
|
|||
|
|
mode = "full_project"
|
|||
|
|
print(f"\n✓ 将处理整个项目: {target_path}")
|
|||
|
|
print("开始处理...")
|
|||
|
|
|
|||
|
|
elif choice == "2":
|
|||
|
|
# 单文件模式
|
|||
|
|
while True:
|
|||
|
|
file_path = input("\n请输入要处理的文件完整路径: ").strip()
|
|||
|
|
if not file_path:
|
|||
|
|
print("输入为空，退出单文件模式")
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
if not os.path.exists(file_path):
|
|||
|
|
print(f"错误: 文件不存在: {file_path}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
file_ext = Path(file_path).suffix.lower()
|
|||
|
|
if file_ext not in CPP_EXTENSIONS:
|
|||
|
|
print(f"警告: 文件类型 {file_ext} 可能不是C/C++文件，继续吗? (y/n): ", end="")
|
|||
|
|
confirm = input().strip().lower()
|
|||
|
|
if confirm != 'y':
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
target_path = file_path
|
|||
|
|
mode = "single_file"
|
|||
|
|
print(f"\n将处理单个文件: {target_path}")
|
|||
|
|
print("开始处理...")
|
|||
|
|
break
|
|||
|
|
else:
|
|||
|
|
print("无效的选择，退出")
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
# 确认开始
|
|||
|
|
print("\n即将开始构建知识库，确认继续? (y/n): ", end="")
|
|||
|
|
confirm = input().strip().lower()
|
|||
|
|
if confirm != 'y':
|
|||
|
|
print("操作已取消")
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
# --- 初始化数据结构 ---
|
|||
|
|
all_nodes = [] # 存储 GraphNode 对象
|
|||
|
|
all_edges = [] # 存储 GraphEdge 对象
|
|||
|
|
symbol_resolver = SymbolResolver() # 符号解析器
|
|||
|
|
function_raw_info_map = {} # 临时存储函数的原始信息
|
|||
|
|
|
|||
|
|
# 清空全局图
|
|||
|
|
FUNCTION_CALL_GRAPH.clear()
|
|||
|
|
FILE_DEPENDENCIES.clear()
|
|||
|
|
CALLED_BY_GRAPH.clear()
|
|||
|
|
|
|||
|
|
# ========================
|
|||
|
|
# 阶段一：静态分析，提取节点与硬性边
|
|||
|
|
# ========================
|
|||
|
|
logger.info("阶段一：静态分析，提取代码实体与结构关系...")
|
|||
|
|
|
|||
|
|
if choice == "1":
|
|||
|
|
# 全项目模式
|
|||
|
|
processed_files_count = process_project_files(target_path, all_nodes, all_edges,
|
|||
|
|
symbol_resolver, function_raw_info_map)
|
|||
|
|
logger.info(f"已处理 {processed_files_count} 个文件")
|
|||
|
|
else:
|
|||
|
|
# 单文件模式
|
|||
|
|
process_single_file(target_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)
|
|||
|
|
logger.info(f"已处理单个文件: {target_path}")
|
|||
|
|
|
|||
|
|
logger.info(f"阶段一完成。初步识别节点数: {len(all_nodes)}")
|
|||
|
|
|
|||
|
|
# ========================
|
|||
|
|
# 阶段二：解析符号，构建准确的调用边
|
|||
|
|
# ========================
|
|||
|
|
build_call_edges(all_nodes, all_edges, symbol_resolver)
|
|||
|
|
|
|||
|
|
# ========================
|
|||
|
|
# 阶段三：语义增强，生成摘要与向量
|
|||
|
|
# ========================
|
|||
|
|
print("\n" + "=" * 50)
|
|||
|
|
print("阶段三：语义增强")
|
|||
|
|
print("=" * 50)
|
|||
|
|
print("此阶段将调用LLM API为函数生成摘要和向量嵌入")
|
|||
|
|
print(f"需要处理的函数节点数: {len([n for n in all_nodes if n.type == NodeType.FUNCTION])}")
|
|||
|
|
print("这可能会消耗API调用额度并需要一些时间")
|
|||
|
|
print("-" * 50)
|
|||
|
|
|
|||
|
|
proceed = input("是否继续阶段三? (y/n): ").strip().lower()
|
|||
|
|
if proceed != 'y':
|
|||
|
|
logger.info("用户跳过阶段三（语义增强）")
|
|||
|
|
print("跳过阶段三，直接进入阶段四")
|
|||
|
|
|
|||
|
|
# 为未处理的函数节点设置默认值
|
|||
|
|
for node in all_nodes:
|
|||
|
|
if node.type == NodeType.FUNCTION and not node.summary:
|
|||
|
|
node.summary = f"函数 {node.name}，位于 {node.file_path}"
|
|||
|
|
node.logic_flow = "跳过语义增强阶段"
|
|||
|
|
node.embedding = [0.0] * 1024
|
|||
|
|
raw_info = function_raw_info_map.get(node.id, {})
|
|||
|
|
node.raw_attributes.update({
|
|||
|
|
"comment": raw_info.get("comment", ""),
|
|||
|
|
"code_snippet": raw_info.get("code_snippet", ""),
|
|||
|
|
"calls": FUNCTION_CALL_GRAPH.get(node.name, []),
|
|||
|
|
"called_by": CALLED_BY_GRAPH.get(node.name, []),
|
|||
|
|
"includes": FILE_DEPENDENCIES.get(node.file_path, []),
|
|||
|
|
"qualified_name": _qualified_name(node),
|
|||
|
|
"embedding_available": False,
|
|||
|
|
"embedding_error": "Semantic enhancement was skipped.",
|
|||
|
|
})
|
|||
|
|
else:
|
|||
|
|
semantic_enhancement(all_nodes, all_edges, function_raw_info_map)
|
|||
|
|
|
|||
|
|
# ========================
|
|||
|
|
# 阶段四：持久化存储
|
|||
|
|
# ========================
|
|||
|
|
print("\n" + "=" * 50)
|
|||
|
|
print("阶段四：持久化存储")
|
|||
|
|
print("=" * 50)
|
|||
|
|
|
|||
|
|
save_choice = input("是否保存结果? (y/n): ").strip().lower()
|
|||
|
|
if save_choice != 'y':
|
|||
|
|
logger.info("用户取消保存，退出构建流程")
|
|||
|
|
print("构建流程已取消")
|
|||
|
|
return all_nodes, all_edges
|
|||
|
|
|
|||
|
|
# 确定保存文件名
|
|||
|
|
if choice == "1":
|
|||
|
|
base_name = "project"
|
|||
|
|
else:
|
|||
|
|
file_name = Path(target_path).stem
|
|||
|
|
base_name = f"file_{file_name}"
|
|||
|
|
|
|||
|
|
graph_data_path, vector_db_path, metadata_path = save_results(
|
|||
|
|
all_nodes, all_edges, base_name, target_path, mode
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print("\n" + "=" * 50)
|
|||
|
|
print("构建完成！")
|
|||
|
|
print("=" * 50)
|
|||
|
|
print(f"总节点数: {len(all_nodes)}")
|
|||
|
|
print(f"总边数: {len(all_edges)}")
|
|||
|
|
print(f"图谱文件: {graph_data_path}")
|
|||
|
|
if os.path.exists(vector_db_path):
|
|||
|
|
print(f"向量文件: {vector_db_path}")
|
|||
|
|
|
|||
|
|
logger.info("=== 知识图谱构建流程全部完成 ===")
|
|||
|
|
return all_nodes, all_edges
|