498 lines
21 KiB
Python
498 lines
21 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
卫星星务软件代码RAG知识库构建工具（严格保留Tree-sitter/Qwen调用方式）
|
||
- 仅优化关联关系处理逻辑
|
||
- 保留原始Tree-sitter解析和Qwen-V3调用流程
|
||
- 通过全局关联图确保关联关系正确性
|
||
- 适用于卫星软件项目（C/C++）
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import json
|
||
from collections import defaultdict
|
||
from pathlib import Path
|
||
import numpy as np
|
||
import faiss
|
||
import tree_sitter
|
||
import tree_sitter_c
|
||
import tree_sitter_cpp
|
||
from openai import OpenAI
|
||
from tree_sitter import Language, Parser
|
||
import requests
|
||
from tqdm import tqdm
|
||
import logging
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(levelname)s - %(message)s',
|
||
handlers=[
|
||
logging.FileHandler("rag_builder.log"),
|
||
logging.StreamHandler()
|
||
]
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# === 配置常量 ===
|
||
QWEN_API_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||
QWEN_API_KEY = ""
|
||
PROJECT_ROOT = "C:/Users/Administrator/Desktop\静态分析问题单过滤/gs-zk和扫描结果/PrjAttCtrlMng" # 替换为实际项目路径
|
||
VECTOR_DB_PATH = "satellite_rag.faiss"
|
||
METADATA_PATH = "satellite_rag_metadata.json"
|
||
MAX_CODE_LENGTH = 800 # 代码片段最大长度
|
||
|
||
|
||
# === 全局状态 ===
|
||
FUNCTION_CALL_GRAPH = defaultdict(list)
|
||
FILE_DEPENDENCIES = defaultdict(list)
|
||
CALLED_BY_GRAPH = defaultdict(list) # 反向图：callee -> [caller]
|
||
CPP_LANGUAGE = Language(tree_sitter_cpp.language())
|
||
parser = Parser()
|
||
parser.language=CPP_LANGUAGE
|
||
CPP_EXTENSIONS = {'.c', '.cpp', '.h', '.hpp', '.tcc'}
|
||
IGNORE_DIRS = {'build', 'cmake-build', '.git', 'vendor', 'lib', 'external','Debug'}
|
||
|
||
def generate_code_summary(func_name, comment, logic, code_snippet, file_path):
|
||
called_functions = FUNCTION_CALL_GRAPH.get(func_name, [])
|
||
caller_functions = CALLED_BY_GRAPH.get(func_name, []) # ← 直接获取，O(1)
|
||
included_headers = FILE_DEPENDENCIES.get(file_path, [])
|
||
|
||
prompt = f"""
|
||
你是一名资深的航天软件工程师，请总结以下C++函数的核心功能，**必须严格包含以下6点**：
|
||
|
||
1. 函数流程与逻辑:
|
||
{logic}
|
||
|
||
2. 函数目的:
|
||
→ 结合航天术语总结，必须根据上述“函数流程与逻辑”推导生成，不得依赖假设！
|
||
|
||
3. 输入参数（名称、类型、作用，用括号列出）
|
||
|
||
4. 返回值（类型和含义）
|
||
|
||
5. 与其他函数的关联关系：
|
||
- 被调用的函数: {', '.join(called_functions) or '无'}
|
||
- 调用此函数的函数: {', '.join(caller_functions) or '无'}
|
||
|
||
6. 与跨文件关联的函数（头文件）: {', '.join(included_headers) or '无'}
|
||
|
||
---
|
||
函数名: {func_name}
|
||
注释: {comment or '无'}
|
||
代码片段（截断至{MAX_CODE_LENGTH}字符）:
|
||
{code_snippet[:MAX_CODE_LENGTH]}
|
||
|
||
输出格式: "功能: [按上述6点组织的连贯段落，不要编号]"
|
||
"""
|
||
|
||
headers = {"Authorization": f"Bearer {QWEN_API_KEY}", "Content-Type": "application/json"}
|
||
payload = {
|
||
"model": "qwen-max",
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"temperature": 0.1
|
||
}
|
||
|
||
response = requests.post(
|
||
"https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
|
||
headers=headers,
|
||
json=payload
|
||
)
|
||
response.raise_for_status()
|
||
summary = response.json()['choices'][0]['message']['content'].strip()
|
||
return summary.replace("功能: ", "")
|
||
def generate_code_logic(func_name, comment, code_snippet, file_path):
|
||
prompt = f"""
|
||
你是一名资深的航天软件工程师，请生成以下C++函数的markdown格式的函数流程图，并总结函数的核心流程逻辑。
|
||
"""
|
||
# 你的提示词保持不变
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {QWEN_API_KEY}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
payload = {
|
||
"model": "qwen-max",
|
||
"messages": [ # ← 直接顶层字段，无 "input" 包裹
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
"temperature": 0.1 # 可选：提高摘要一致性
|
||
}
|
||
|
||
# 使用 OpenAI 兼容端点
|
||
response = requests.post(
|
||
"https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
|
||
headers=headers,
|
||
json=payload
|
||
)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
logic = result['choices'][0]['message']['content'].strip() # ← 注意路径变化！
|
||
return logic
|
||
def safe_decode(b: bytes) -> str:
|
||
for encoding in ['utf-8', 'gbk', 'gb2312', 'latin-1']:
|
||
try:
|
||
return b.decode(encoding)
|
||
except UnicodeDecodeError:
|
||
continue
|
||
return b.decode('utf-8', errors='replace')
|
||
|
||
# === 严格保留的代码提取辅助函数 ===
|
||
def extract_comment_before(node, source_lines):
|
||
"""严格保留注释提取逻辑（已验证正确）"""
|
||
start_line = node.start_point[0]
|
||
comment_lines = []
|
||
|
||
for i in range(start_line - 1, max(0, start_line - 10), -1):
|
||
line = safe_decode(source_lines[i]).strip()
|
||
if line.startswith('//') or line.startswith('/*'):
|
||
comment_lines.append(line)
|
||
|
||
return " ".join(comment_lines[::-1]) or "无注释"
|
||
|
||
|
||
def extract_code_snippet(tree, node, source_lines):
|
||
"""严格保留代码片段提取逻辑（已验证正确）"""
|
||
start_line = node.start_point[0]
|
||
end_line = node.end_point[0]
|
||
|
||
code_lines = []
|
||
for i in range(start_line, end_line + 1):
|
||
if i < len(source_lines):
|
||
code_lines.append(safe_decode(source_lines[i]))
|
||
|
||
return "\n".join(code_lines)
|
||
|
||
|
||
# === 严格保留的向量化嵌入 ===
|
||
def get_qwen_embedding(text: str) -> np.ndarray:
|
||
"""
|
||
调用 DashScope 的 text-embedding-v4 模型生成 1024 维文本嵌入向量。
|
||
使用 OpenAI 兼容 SDK 方式调用。
|
||
"""
|
||
try:
|
||
client = OpenAI(
|
||
api_key=QWEN_API_KEY,
|
||
base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||
)
|
||
response = client.embeddings.create(
|
||
model="text-embedding-v4",
|
||
input=text # OpenAI SDK 支持单个字符串或字符串列表
|
||
)
|
||
embedding = response.data[0].embedding
|
||
return np.array(embedding, dtype='float32')
|
||
except Exception as e:
|
||
# 假设 logger 已定义；若未定义，可临时用 print 替代
|
||
try:
|
||
logger.error(f"获取嵌入失败: {e}")
|
||
except NameError:
|
||
print(f"获取嵌入失败: {e}")
|
||
return np.zeros(1024, dtype='float32')
|
||
def extract_function_name(func_def_node):
|
||
"""
|
||
从 function_definition 节点中提取完整的函数名（含命名空间/类作用域）
|
||
"""
|
||
declarator = func_def_node.child_by_field_name("declarator")
|
||
if not declarator:
|
||
return "<unknown>"
|
||
|
||
# 处理 A::B::func 这种形式：可能是 field_expression 嵌套
|
||
def get_qualified_name(node):
|
||
if node.type == "identifier":
|
||
return node.text.decode('utf-8')
|
||
elif node.type == "field_expression":
|
||
# left.field -> left 可能是 identifier 或 field_expression
|
||
left = node.child_by_field_name("argument") or node.children[0]
|
||
right = node.child_by_field_name("field") or node.children[-1]
|
||
left_name = get_qualified_name(left) if left else ""
|
||
right_name = get_qualified_name(right) if right else ""
|
||
return f"{left_name}::{right_name}" if left_name and right_name else (left_name or right_name)
|
||
elif node.type == "function_declarator":
|
||
# 函数声明器：取其 declarator 部分
|
||
func_declarator = node.child_by_field_name("declarator")
|
||
if func_declarator:
|
||
return get_qualified_name(func_declarator)
|
||
# 兜底：直接取文本
|
||
return node.text.decode('utf-8').split('(')[0].strip()
|
||
|
||
return get_qualified_name(declarator)
|
||
def extract_qualified_name_from_field(field_node):
|
||
"""递归解析 field_expression，如 A::B::foo"""
|
||
def _extract(n):
|
||
if n.type == "identifier":
|
||
return n.text.decode('utf-8')
|
||
elif n.type == "field_expression":
|
||
# left.field -> argument.field
|
||
left = n.child_by_field_name("argument") or (n.children[0] if n.children else None)
|
||
right = n.child_by_field_name("field") or (n.children[-1] if n.children else None)
|
||
left_str = _extract(left) if left else ""
|
||
right_str = _extract(right) if right else ""
|
||
if left_str and right_str:
|
||
return f"{left_str}::{right_str}"
|
||
return left_str or right_str
|
||
else:
|
||
# 兜底：取文本，截断参数列表
|
||
text = n.text.decode('utf-8')
|
||
return text.split('(')[0].strip()
|
||
return _extract(field_node)
|
||
|
||
|
||
def extract_full_function_name(func_def_node):
|
||
"""
|
||
从 function_definition 节点中提取干净的函数名（支持 A::foo）
|
||
"""
|
||
declarator = func_def_node.child_by_field_name("declarator")
|
||
if not declarator:
|
||
return "<unknown>"
|
||
|
||
# 情况1: declarator 是 function_declarator（最常见）
|
||
if declarator.type == "function_declarator":
|
||
inner_declarator = declarator.child_by_field_name("declarator")
|
||
if not inner_declarator:
|
||
return "<unknown>"
|
||
if inner_declarator.type == "identifier":
|
||
return inner_declarator.text.decode('utf-8')
|
||
elif inner_declarator.type == "field_expression":
|
||
return extract_qualified_name_from_field(inner_declarator)
|
||
else:
|
||
# 兜底处理（如数组、指针等，但函数名通常不会这样）
|
||
raw = inner_declarator.text.decode('utf-8')
|
||
return raw.split('(')[0].strip()
|
||
|
||
# 情况2: declarator 直接是 identifier（旧式C或简单情况）
|
||
elif declarator.type == "identifier":
|
||
return declarator.text.decode('utf-8')
|
||
|
||
# 情况3: 其他类型（如 pointer_declarator），尝试兜底
|
||
else:
|
||
raw = declarator.text.decode('utf-8')
|
||
return raw.split('(')[0].strip()
|
||
def traverse_type(node, target_type):
|
||
"""递归遍历 AST 节点，查找指定类型"""
|
||
if node.type == target_type:
|
||
yield node
|
||
for child in node.children:
|
||
yield from traverse_type(child, target_type)
|
||
# === 主处理流程（仅优化关联关系处理）===
|
||
def build_rag_database():
|
||
"""严格保留Tree-sitter/Qwen调用流程，仅优化关联关系处理"""
|
||
global FUNCTION_CALL_GRAPH, FILE_DEPENDENCIES, CALLED_BY_GRAPH
|
||
|
||
# 清空旧图（避免多次调用累积）
|
||
FUNCTION_CALL_GRAPH.clear()
|
||
FILE_DEPENDENCIES.clear()
|
||
CALLED_BY_GRAPH.clear()
|
||
|
||
dimension = 1024
|
||
index = faiss.IndexFlatL2(dimension)
|
||
metadata_list = []
|
||
|
||
import re
|
||
|
||
# 全局图（确保在阶段1前清空）
|
||
FUNCTION_CALL_GRAPH = {}
|
||
CALLED_BY_GRAPH = {}
|
||
FILE_DEPENDENCIES = {}
|
||
|
||
def extract_callee_name(func_node):
|
||
"""从 call_expression 的 function 节点中提取函数名标识符"""
|
||
if func_node.type == "identifier":
|
||
return func_node.text.decode('utf-8')
|
||
elif func_node.type == "field_expression": # 如 A::func 或 a.b
|
||
field = func_node.child_by_field_name("field")
|
||
if field and field.type == "identifier":
|
||
return field.text.decode('utf-8')
|
||
# 递归处理复杂命名空间
|
||
text = func_node.text.decode('utf-8')
|
||
parts = re.split(r'[.:]+', text)
|
||
return parts[-1].split('(')[0].strip() if parts else "<unknown>"
|
||
elif func_node.type == "member_expression": # C++ 风格: obj->method 或 obj.method
|
||
member = func_node.child_by_field_name("member")
|
||
if member and member.type == "identifier":
|
||
return member.text.decode('utf-8')
|
||
# 兜底
|
||
text = func_node.text.decode('utf-8')
|
||
for sep in ['->', '.']:
|
||
if sep in text:
|
||
return text.split(sep)[-1].split('(')[0].strip()
|
||
return text.split('(')[0].strip()
|
||
else:
|
||
# 兜底：从文本中提取最可能的函数名
|
||
text = func_node.text.decode('utf-8')
|
||
# 移除模板 <...>
|
||
text = re.sub(r'<[^>]*>', '', text)
|
||
# 取括号前部分
|
||
base = text.split('(')[0].strip()
|
||
# 处理 a::b::c -> c
|
||
for sep in ['::', '->', '.']:
|
||
if sep in base:
|
||
base = base.split(sep)[-1]
|
||
# 简单合法性检查
|
||
if base and base[0].isalpha() and all(c.isalnum() or c in '_$' for c in base):
|
||
return base
|
||
return "<unknown>"
|
||
|
||
def collect_call_expressions(node):
|
||
"""递归遍历 AST 节点，收集所有 call_expression 中的 callee 名称（去重）"""
|
||
calls = set()
|
||
|
||
def traverse(n):
|
||
if n.type == "call_expression":
|
||
func_node = n.child_by_field_name("function")
|
||
if func_node:
|
||
callee = extract_callee_name(func_node)
|
||
if callee != "<unknown>":
|
||
calls.add(callee)
|
||
# 递归子节点
|
||
for child in n.children:
|
||
traverse(child)
|
||
|
||
traverse(node)
|
||
return list(calls)
|
||
|
||
# ========================
|
||
# 阶段1：构建调用图与依赖图
|
||
# ========================
|
||
logger.info("阶段1: 构建函数调用图和文件依赖图")
|
||
|
||
total_funcs = 0
|
||
|
||
for root, _, files in os.walk(PROJECT_ROOT):
|
||
if any(ignore_dir in Path(root).parts for ignore_dir in IGNORE_DIRS):
|
||
continue
|
||
for file in files:
|
||
if Path(file).suffix.lower() in CPP_EXTENSIONS:
|
||
file_path = os.path.join(root, file)
|
||
logger.debug(f"解析文件构建关联图: {file_path}")
|
||
try:
|
||
with open(file_path, 'rb') as f:
|
||
source_bytes = f.read()
|
||
tree = parser.parse(source_bytes)
|
||
|
||
# --- 1. 收集头文件依赖 ---
|
||
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
|
||
|
||
# --- 2. 遍历顶层节点，找函数定义 ---
|
||
for node in tree.root_node.named_children:
|
||
if node.type == "function_definition":
|
||
func_name = extract_full_function_name(node)
|
||
if func_name == "<unknown>":
|
||
continue
|
||
|
||
total_funcs += 1
|
||
FUNCTION_CALL_GRAPH[func_name] = [] # 初始化
|
||
|
||
# 获取函数体（body）
|
||
body = node.child_by_field_name("body")
|
||
if body:
|
||
callees = collect_call_expressions(body)
|
||
FUNCTION_CALL_GRAPH[func_name] = callees
|
||
|
||
except Exception as e:
|
||
logger.error(f"解析文件失败 {file_path}: {e}", exc_info=True)
|
||
|
||
# --- 3. 构建反向调用图（谁调用了我）---
|
||
CALLED_BY_GRAPH.clear()
|
||
for caller, callees in FUNCTION_CALL_GRAPH.items():
|
||
for callee in callees:
|
||
if callee not in CALLED_BY_GRAPH:
|
||
CALLED_BY_GRAPH[callee] = []
|
||
CALLED_BY_GRAPH[callee].append(caller)
|
||
|
||
logger.info(f"阶段1完成。共识别 {total_funcs} 个函数，{len(FILE_DEPENDENCIES)} 个文件依赖。")
|
||
# === 阶段2: 生成摘要并构建RAG ===
|
||
logger.info("阶段2: 生成函数摘要并构建RAG")
|
||
for root, _, files in os.walk(PROJECT_ROOT):
|
||
if any(ignore_dir in Path(root).parts for ignore_dir in IGNORE_DIRS):
|
||
continue
|
||
for file in files:
|
||
if Path(file).suffix.lower() in CPP_EXTENSIONS:
|
||
file_path = os.path.join(root, file)
|
||
logger.debug(f"处理文件: {file_path}")
|
||
try:
|
||
with open(file_path, 'rb') as f:
|
||
source = f.read()
|
||
tree = parser.parse(source)
|
||
source_lines = source.split(b'\n')
|
||
|
||
for node in tree.root_node.named_children:
|
||
if node.type == "function_definition":
|
||
# === 关键修正：使用统一的函数名提取逻辑 ===
|
||
func_name = extract_full_function_name(node)
|
||
if func_name == "<unknown>":
|
||
logger.warning(f"无法解析函数名，跳过节点: {node.text[:100]}...")
|
||
continue
|
||
|
||
# 提取前置注释（保持不变）
|
||
comment = extract_comment_before(node, source_lines)
|
||
|
||
# 提取代码片段（保持不变）
|
||
func_code = extract_code_snippet(tree, node, source_lines)
|
||
print(func_code)
|
||
# 生成摘要（保持不变）
|
||
logic = generate_code_logic(func_name, comment, func_code, file_path)
|
||
summary = generate_code_summary(func_name, comment,logic,func_code, file_path)
|
||
print(summary)
|
||
# 从全局图中获取关系（确保 func_name 与阶段1一致！）
|
||
called_functions = FUNCTION_CALL_GRAPH.get(func_name, [])
|
||
caller_functions = CALLED_BY_GRAPH.get(func_name, [])
|
||
included_headers = FILE_DEPENDENCIES.get(file_path, [])
|
||
|
||
# 构建上下文文本
|
||
context_text = (
|
||
f"【实体类型】函数\n"
|
||
f"【函数名】{func_name}\n"
|
||
f"【所在文件】{os.path.basename(file_path)}\n"
|
||
f"【代码注释】{comment.strip() if comment else '无'}\n"
|
||
f"【功能摘要】{summary}\n"
|
||
f"【流程逻辑】{logic}\n"
|
||
f"【调用的函数】{', '.join(called_functions) if called_functions else '无'}\n"
|
||
f"【被以下函数调用】{', '.join(caller_functions) if caller_functions else '无'}\n"
|
||
f"【包含的头文件】{', '.join(included_headers) if included_headers else '无'}"
|
||
)
|
||
|
||
# 生成嵌入向量
|
||
vector = get_qwen_embedding(context_text)
|
||
|
||
# 构建元数据
|
||
metadata = {
|
||
"file": file_path,
|
||
"line": node.start_point[0] + 1,
|
||
"type": "function",
|
||
"name": func_name,
|
||
"summary": summary,
|
||
"logic": logic,
|
||
"calls": called_functions,
|
||
"called_by": caller_functions,
|
||
"includes": included_headers,
|
||
"comment": comment or ""
|
||
}
|
||
|
||
# 添加到向量库
|
||
index.add(np.array([vector], dtype='float32'))
|
||
metadata_list.append(metadata)
|
||
|
||
except Exception as e:
|
||
logger.error(f"处理文件 {file_path} 失败: {str(e)}", exc_info=True)
|
||
# 保存结果
|
||
faiss.write_index(index, VECTOR_DB_PATH)
|
||
with open(METADATA_PATH, 'w', encoding='utf-8') as f:
|
||
json.dump(metadata_list, f, indent=2, ensure_ascii=False)
|
||
|
||
logger.info(f"知识库构建完成! 总实体: {len(metadata_list)}")
|
||
logger.info(f"向量数据库保存至: {VECTOR_DB_PATH}")
|
||
logger.info(f"元数据保存至: {METADATA_PATH}")
|
||
|
||
return index, metadata_list # 如果你确实需要返回，可以放这里
|
||
|
||
|
||
if __name__ == "__main__":
|
||
build_rag_database() |