"""
检索引擎模块
负责基于向量数据库的代码语义检索
"""
import os
import logging
from typing import List, Dict, Any, Optional, Tuple
import json
from .config import Config
from .indexer import CodeIndexer
logger = logging.getLogger(__name__)
class SearchEngine:
"""
检索引擎类
基于向量数据库的代码语义检索，提供高级搜索功能
"""
def __init__(self, config: Config, indexer: CodeIndexer):
"""
初始化检索引擎
Args:
config: 配置对象
indexer: 代码索引器对象
Returns:
无返回值
"""
self.config = config
self.indexer = indexer
logger.info("检索引擎初始化完成")
def search(self, query: str, project_ids: Optional[List[str]] = None,
filters: Optional[Dict[str, Any]] = None, limit: int = 10) -> List[Dict[str, Any]]:
"""
搜索代码
Args:
query: 查询字符串
project_ids: 项目ID列表，如果为None则搜索所有项目
filters: 过滤条件，如语言、文件类型等
limit: 返回结果数量限制
Returns:
代码块字典列表
"""
logger.info(f"搜索代码: {query}, 项目: {project_ids}, 过滤条件: {filters}")
# 如果未指定项目，获取所有已索引项目
if not project_ids:
indexed_projects = self.indexer.get_indexed_projects()
project_ids = [p["project_id"] for p in indexed_projects]
all_results = []
# 对每个项目执行搜索
for project_id in project_ids:
results = self.indexer.search(project_id, query, limit=limit)
# 应用过滤条件
if filters:
results = self._apply_filters(results, filters)
all_results.extend(results)
# 按相似度排序
all_results.sort(key=lambda x: x.get('similarity', 0), reverse=True)
# 限制结果数量
return all_results[:limit]
def _apply_filters(self, results: List[Dict[str, Any]],
filters: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
应用过滤条件
Args:
results: 搜索结果列表
filters: 过滤条件字典
Returns:
过滤后的结果列表
"""
filtered_results = []
for result in results:
match = True
# 语言过滤
if 'language' in filters and result.get('language') != filters['language']:
match = False
# 文件路径过滤
if 'file_path' in filters:
file_path = result.get('file_path', '')
if filters['file_path'] not in file_path:
match = False
# 代码类型过滤
if 'type' in filters and result.get('type') != filters['type']:
match = False
if match:
filtered_results.append(result)
return filtered_results
def search_by_file(self, file_path: str, project_id: Optional[str] = None,
limit: int = 10) -> List[Dict[str, Any]]:
"""
按文件路径搜索代码
Args:
file_path: 文件路径
project_id: 项目ID，如果为None则搜索所有项目
limit: 返回结果数量限制
Returns:
代码块字典列表
"""
# 构建过滤条件
filters = {'file_path': file_path}
# 如果未指定项目，获取所有已索引项目
project_ids = [project_id] if project_id else None
# 使用文件路径作为查询，并应用过滤条件
return self.search(os.path.basename(file_path), project_ids, filters, limit)
def search_by_function(self, function_name: str, project_ids: Optional[List[str]] = None,
limit: int = 10) -> List[Dict[str, Any]]:
"""
按函数名搜索代码
Args:
function_name: 函数名
project_ids: 项目ID列表，如果为None则搜索所有项目
limit: 返回结果数量限制
Returns:
代码块字典列表
"""
# 构建函数定义的查询
query = f"function {function_name}" if function_name else ""
# 执行搜索
return self.search(query, project_ids, None, limit)
def search_by_class(self, class_name: str, project_ids: Optional[List[str]] = None,
limit: int = 10) -> List[Dict[str, Any]]:
"""
按类名搜索代码
Args:
class_name: 类名
project_ids: 项目ID列表，如果为None则搜索所有项目
limit: 返回结果数量限制
Returns:
代码块字典列表
"""
# 构建类定义的查询
query = f"class {class_name}" if class_name else ""
# 执行搜索
return self.search(query, project_ids, None, limit)
def get_related_code(self, code_chunk: Dict[str, Any],
limit: int = 5) -> List[Dict[str, Any]]:
"""
获取与给定代码块相关的代码
Args:
code_chunk: 代码块字典
limit: 返回结果数量限制
Returns:
相关代码块字典列表
"""
# 使用代码块内容作为查询
query = code_chunk.get('content', '')
# 获取项目ID
file_path = code_chunk.get('file_path', '')
project_id = self._get_project_id_by_file_path(file_path)
if not project_id:
return []
# 执行搜索，排除自身
results = self.indexer.search(project_id, query, limit=limit+1)
# 过滤掉自身
filtered_results = []
for result in results:
if (result.get('file_path') != file_path or
result.get('start_line') != code_chunk.get('start_line')):
filtered_results.append(result)
return filtered_results[:limit]
def _get_project_id_by_file_path(self, file_path: str) -> Optional[str]:
"""
根据文件路径获取项目ID
Args:
file_path: 文件路径
Returns:
项目ID，如果未找到则返回None
"""
# 获取所有已索引项目
indexed_projects = self.indexer.get_indexed_projects()
for project in indexed_projects:
project_path = project.get("project_path", "")
if project_path and file_path.startswith(project_path):
return project.get("project_id")
return None
def get_code_context(self, file_path: str, line_number: int,
context_lines: int = 10) -> Dict[str, Any]:
"""
获取代码上下文
Args:
file_path: 文件路径
line_number: 行号
context_lines: 上下文行数
Returns:
代码上下文字典
"""
try:
# 读取文件内容
with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
lines = f.readlines()
# 计算上下文范围
start_line = max(1, line_number - context_lines)
end_line = min(len(lines), line_number + context_lines)
# 提取上下文
context_content = ''.join(lines[start_line-1:end_line])
return {
'file_path': file_path,
'start_line': start_line,
'end_line': end_line,
'target_line': line_number,
'content': context_content
}
except Exception as e:
logger.error(f"获取代码上下文失败 {file_path}:{line_number}: {str(e)}")
return {
'file_path': file_path,
'start_line': line_number,
'end_line': line_number,
'target_line': line_number,
'content': '',
'error': str(e)
}
def get_file_overview(self, file_path: str) -> Dict[str, Any]:
"""
获取文件概览
Args:
file_path: 文件路径
Returns:
文件概览字典
"""
try:
# 读取文件内容
with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
content = f.read()
# 获取文件大小
file_size = os.path.getsize(file_path)
# 获取行数
lines = content.split('\n')
line_count = len(lines)
# 获取文件扩展名
_, ext = os.path.splitext(file_path)
# 提取文件头部（最多100行）
header = '\n'.join(lines[:min(100, line_count)])
return {
'file_path': file_path,
'file_name': os.path.basename(file_path),
'file_size': file_size,
'line_count': line_count,
'extension': ext,
'header': header
}
except Exception as e:
logger.error(f"获取文件概览失败 {file_path}: {str(e)}")
return {
'file_path': file_path,
'file_name': os.path.basename(file_path),
'error': str(e)
}
def find_similar_code(self, code: str, language: str = None,
threshold: float = 0.7, limit: int = 5) -> List[Dict[str, Any]]:
"""
查找相似代码片段，使用改进的相似度算法
Args:
code: 代码片段
language: 编程语言（可选）
threshold: 相似度阈值（0-1之间）
limit: 返回结果数量限制
Returns:
相似代码片段列表，每个元素包含:
- content: 代码内容
- file_path: 文件路径
- start_line: 起始行
- end_line: 结束行
- similarity: 相似度分数
- matched_lines: 匹配的行
- similarity_details: 相似度详细信息
"""
try:
# 标准化代码，传入语言参数
normalized_code = self._normalize_code(code, language)
# 使用代码内容作为查询
results = []
indexed_projects = self.indexer.get_indexed_projects()
# 记录开始时间，用于性能监控
import time
start_time = time.time()
# 首先使用向量搜索找到候选结果
candidate_results = []
for project in indexed_projects:
project_id = project.get("project_id")
if not project_id:
continue
# 搜索相似代码，增加搜索范围以提高召回率
search_results = self.indexer.search(project_id, normalized_code, limit=limit*3)
candidate_results.extend(search_results)
# 对每个候选结果计算详细相似度
for result in candidate_results:
# 获取结果的语言
result_language = result.get('language', language)
# 标准化结果代码，传入语言参数
result_code = self._normalize_code(result.get('content', ''), result_language)
# 计算详细相似度，传入语言参数
similarity_info = self._calculate_detailed_similarity(
normalized_code,
result_code,
result_language
)
# 如果相似度超过阈值，添加到结果中
if similarity_info['overall_similarity'] >= threshold:
# 添加相似度信息
result.update({
'similarity': similarity_info['overall_similarity'],
'matched_lines': similarity_info['matched_lines'],
'similarity_details': {
'line_similarity': similarity_info['line_similarity'],
'structure_similarity': similarity_info['structure_similarity'],
'semantic_similarity': similarity_info['semantic_similarity'],
'control_flow_similarity': similarity_info['control_flow_similarity']
}
})
results.append(result)
# 按相似度排序
results.sort(key=lambda x: x.get('similarity', 0), reverse=True)
# 如果指定了语言，过滤结果
if language:
results = [r for r in results if r.get('language') == language]
# 记录搜索时间
search_time = time.time() - start_time
logger.info(f"相似代码搜索完成，耗时: {search_time:.2f}秒，找到 {len(results)} 个结果")
return results[:limit]
except Exception as e:
logger.error(f"查找相似代码失败: {str(e)}")
import traceback
logger.error(traceback.format_exc())
return []
def _normalize_code(self, code: str, language: str = None) -> str:
"""
标准化代码，移除空白字符和注释，并进行语言特定的标准化
Args:
code: 原始代码
language: 编程语言
Returns:
标准化后的代码
"""
if not code:
return ""
# 移除注释
lines = []
in_multiline_comment = False
in_python_docstring = False
# 语言特定的注释标记
single_line_comment_markers = {
'python': ['#'],
'javascript': ['//'],
'typescript': ['//'],
'java': ['//'],
'c': ['//'],
'cpp': ['//'],
'csharp': ['//'],
'php': ['//'],
'ruby': ['#'],
'go': ['//'],
'rust': ['//'],
'swift': ['//'],
'kotlin': ['//']
}
# 获取当前语言的单行注释标记
comment_markers = single_line_comment_markers.get(language, ['#', '//'])
for line in code.split('\n'):
processed_line = line.strip()
# 处理Python文档字符串
if language == 'python':
if '"""' in processed_line or "'''" in processed_line:
# 计算三引号的数量
triple_double_count = processed_line.count('"""')
triple_single_count = processed_line.count("'''")
# 如果数量为奇数，切换docstring状态
if triple_double_count % 2 == 1 or triple_single_count % 2 == 1:
in_python_docstring = not in_python_docstring
# 移除docstring部分
if '"""' in processed_line:
parts = processed_line.split('"""')
processed_line = parts[0] if not in_python_docstring else parts[-1]
if "'''" in processed_line:
parts = processed_line.split("'''")
processed_line = parts[0] if not in_python_docstring else parts[-1]
if in_python_docstring:
continue
# 处理多行注释
if '/*' in processed_line and not any(processed_line.startswith(m) for m in comment_markers):
in_multiline_comment = True
processed_line = processed_line[:processed_line.index('/*')].strip()
if '*/' in processed_line:
in_multiline_comment = False
if processed_line.index('*/') + 2 < len(processed_line):
processed_line = processed_line[processed_line.index('*/')+2:].strip()
else:
processed_line = ""
if in_multiline_comment:
continue
# 处理单行注释
for marker in comment_markers:
if marker in processed_line:
# 确保不是字符串中的注释标记
parts = []
in_string = False
string_char = None
i = 0
while i < len(processed_line):
if processed_line[i] in ['"', "'"]:
if not in_string:
in_string = True
string_char = processed_line[i]
elif processed_line[i] == string_char:
in_string = False
elif processed_line[i:i+len(marker)] == marker and not in_string:
processed_line = processed_line[:i].strip()
break
i += 1
# 语言特定的标准化
if language:
# 移除变量类型声明（适用于静态类型语言）
if language in ['java', 'c', 'cpp', 'csharp', 'typescript']:
# 简化类型声明，如 "int x = 5;" -> "x = 5;"
import re
processed_line = re.sub(r'\b(int|float|double|char|boolean|string|var|let|const|auto)\b\s+([a-zA-Z_][a-zA-Z0-9_]*)', r'\2', processed_line)
# 移除多余空白字符
processed_line = ' '.join(processed_line.split())
if processed_line:
lines.append(processed_line)
return '\n'.join(lines)
def _calculate_detailed_similarity(self, code1: str, code2: str, language: str = None) -> Dict[str, Any]:
"""
计算详细的代码相似度，使用多种算法综合评估
Args:
code1: 第一段代码
code2: 第二段代码
language: 编程语言
Returns:
相似度详细信息
"""
# 分割为行
lines1 = code1.split('\n')
lines2 = code2.split('\n')
# 1. 使用最长公共子序列算法找到匹配行
lcs_matrix = self._compute_lcs_matrix(lines1, lines2)
matched_pairs = self._extract_lcs_matches(lcs_matrix, lines1, lines2)
# 2. 计算行级别的匹配
matched_lines = []
total_line_similarity = 0
for i, j in matched_pairs:
# 计算行相似度
similarity = self._calculate_line_similarity(lines1[i], lines2[j])
if similarity >= 0.7: # 行匹配阈值
matched_lines.append({
'line1': i + 1,
'line2': j + 1,
'similarity': similarity
})
total_line_similarity += similarity
# 3. 计算结构相似度
structure_similarity = len(matched_pairs) / max(len(lines1), len(lines2)) if max(len(lines1), len(lines2)) > 0 else 0
# 4. 计算语义相似度 (使用词袋模型)
semantic_similarity = self._calculate_semantic_similarity(code1, code2)
# 5. 计算控制流相似度
control_flow_similarity = self._calculate_control_flow_similarity(code1, code2, language)
# 6. 综合计算整体相似度 (加权平均)
weights = {
'line': 0.4,
'structure': 0.3,
'semantic': 0.2,
'control_flow': 0.1
}
line_similarity_normalized = total_line_similarity / max(len(lines1), len(lines2)) if max(len(lines1), len(lines2)) > 0 else 0
overall_similarity = (
weights['line'] * line_similarity_normalized +
weights['structure'] * structure_similarity +
weights['semantic'] * semantic_similarity +
weights['control_flow'] * control_flow_similarity
)
return {
'overall_similarity': overall_similarity,
'matched_lines': matched_lines,
'line_similarity': line_similarity_normalized,
'structure_similarity': structure_similarity,
'semantic_similarity': semantic_similarity,
'control_flow_similarity': control_flow_similarity
}
def _compute_lcs_matrix(self, lines1: List[str], lines2: List[str]) -> List[List[int]]:
"""计算最长公共子序列矩阵"""
m, n = len(lines1), len(lines2)
lcs_matrix = [[0] * (n + 1) for _ in range(m + 1)]
for i in range(1, m + 1):
for j in range(1, n + 1):
if self._calculate_line_similarity(lines1[i-1], lines2[j-1]) >= 0.7:
lcs_matrix[i][j] = lcs_matrix[i-1][j-1] + 1
else:
lcs_matrix[i][j] = max(lcs_matrix[i-1][j], lcs_matrix[i][j-1])
return lcs_matrix
def _extract_lcs_matches(self, lcs_matrix: List[List[int]], lines1: List[str], lines2: List[str]) -> List[Tuple[int, int]]:
"""从LCS矩阵中提取匹配的行对"""
matches = []
i, j = len(lines1), len(lines2)
while i > 0 and j > 0:
if self._calculate_line_similarity(lines1[i-1], lines2[j-1]) >= 0.7 and lcs_matrix[i][j] == lcs_matrix[i-1][j-1] + 1:
matches.append((i-1, j-1))
i -= 1
j -= 1
elif lcs_matrix[i-1][j] >= lcs_matrix[i][j-1]:
i -= 1
else:
j -= 1
return matches[::-1] # 反转以获得正序
def _calculate_line_similarity(self, line1: str, line2: str) -> float:
"""
计算两行代码的相似度，使用改进的算法
Args:
line1: 第一行代码
line2: 第二行代码
Returns:
相似度分数 (0-1)
"""
if not line1 and not line2:
return 1.0
if not line1 or not line2:
return 0.0
# 1. 标记化
tokens1 = self._tokenize_code_line(line1)
tokens2 = self._tokenize_code_line(line2)
# 2. 计算Jaccard相似度
set1 = set(tokens1)
set2 = set(tokens2)
intersection = len(set1.intersection(set2))
union = len(set1.union(set2))
jaccard = intersection / union if union > 0 else 0
# 3. 计算序列相似度 (考虑顺序)
lcs_length = self._longest_common_subsequence(tokens1, tokens2)
sequence_sim = lcs_length / max(len(tokens1), len(tokens2)) if max(len(tokens1), len(tokens2)) > 0 else 0
# 4. 计算编辑距离相似度
edit_distance = self._levenshtein_distance(tokens1, tokens2)
max_length = max(len(tokens1), len(tokens2))
edit_sim = 1 - (edit_distance / max_length) if max_length > 0 else 0
# 5. 综合相似度 (加权平均)
return 0.3 * jaccard + 0.4 * sequence_sim + 0.3 * edit_sim
def _tokenize_code_line(self, line: str) -> List[str]:
"""将代码行分解为标记"""
import re
# 分割标识符、运算符、括号等
tokens = re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*|[0-9]+|[+\-*/=<>!&|^~%]+|[{}()\[\],.;:]', line)
return tokens
def _longest_common_subsequence(self, seq1: List[str], seq2: List[str]) -> int:
"""计算最长公共子序列长度"""
m, n = len(seq1), len(seq2)
dp = [[0] * (n + 1) for _ in range(m + 1)]
for i in range(1, m + 1):
for j in range(1, n + 1):
if seq1[i-1] == seq2[j-1]:
dp[i][j] = dp[i-1][j-1] + 1
else:
dp[i][j] = max(dp[i-1][j], dp[i][j-1])
return dp[m][n]
def _levenshtein_distance(self, seq1: List[str], seq2: List[str]) -> int:
"""计算编辑距离"""
m, n = len(seq1), len(seq2)
# 创建距离矩阵
dp = [[0] * (n + 1) for _ in range(m + 1)]
# 初始化第一行和第一列
for i in range(m + 1):
dp[i][0] = i
for j in range(n + 1):
dp[0][j] = j
# 填充矩阵
for i in range(1, m + 1):
for j in range(1, n + 1):
if seq1[i-1] == seq2[j-1]:
dp[i][j] = dp[i-1][j-1]
else:
dp[i][j] = min(
dp[i-1][j] + 1, # 删除
dp[i][j-1] + 1, # 插入
dp[i-1][j-1] + 1 # 替换
)
return dp[m][n]
def _calculate_semantic_similarity(self, code1: str, code2: str) -> float:
"""计算代码的语义相似度"""
# 提取所有标识符
import re
identifiers1 = set(re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', code1))
identifiers2 = set(re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', code2))
# 过滤掉关键字
keywords = {'if', 'else', 'for', 'while', 'return', 'function', 'class', 'var', 'let', 'const', 'import', 'export'}
identifiers1 = identifiers1 - keywords
identifiers2 = identifiers2 - keywords
# 计算标识符相似度
if not identifiers1 and not identifiers2:
return 1.0
if not identifiers1 or not identifiers2:
return 0.0
intersection = len(identifiers1.intersection(identifiers2))
union = len(identifiers1.union(identifiers2))
return intersection / union
def _calculate_control_flow_similarity(self, code1: str, code2: str, language: str = None) -> float:
"""计算控制流相似度"""
# 提取控制流结构
control_patterns = {
'if': r'\bif\b',
'else': r'\belse\b',
'for': r'\bfor\b',
'while': r'\bwhile\b',
'switch': r'\bswitch\b',
'case': r'\bcase\b',
'try': r'\btry\b',
'catch': r'\bcatch\b',
'return': r'\breturn\b'
}
# 计算每种控制结构的出现次数
import re
control_counts1 = {pattern: len(re.findall(regex, code1)) for pattern, regex in control_patterns.items()}
control_counts2 = {pattern: len(re.findall(regex, code2)) for pattern, regex in control_patterns.items()}
# 计算控制结构分布的相似度
total_structures1 = sum(control_counts1.values())
total_structures2 = sum(control_counts2.values())
if total_structures1 == 0 and total_structures2 == 0:
return 1.0
if total_structures1 == 0 or total_structures2 == 0:
return 0.0
# 计算每种控制结构的比例差异
similarity = 0.0
for pattern in control_patterns:
ratio1 = control_counts1[pattern] / total_structures1 if total_structures1 > 0 else 0
ratio2 = control_counts2[pattern] / total_structures2 if total_structures2 > 0 else 0
similarity += 1.0 - abs(ratio1 - ratio2)
return similarity / len(control_patterns)