RAG/sync/git_sync.py

628 lines
27 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
Git代码库同步子类
继承BaseSync实现代码拉取/增量同步/函数解析
"""
import os
from typing import List, Dict, Any, Set, Optional
from datetime import datetime
from loguru import logger
from config import BaseDataSourceConfig, GitDataSourceConfig, settings
from sync.base_sync import BaseSync
from sync.ast_parser import ASTParser
from utils.git_tool import GitTool
from utils.func_id_generator import generate_func_unique_id
class GitSync(BaseSync):
def __init__(self, config: GitDataSourceConfig, vector_store_manager=None):
"""
初始化Git同步器
Args:
config: Git数据源配置
vector_store_manager: 向量存储管理器
"""
super().__init__(config, vector_store_manager)
self.config = config
self.git_tools = [] # 存储多个GitTool实例
# 处理多仓库配置
if hasattr(config, 'git_mode') and config.git_mode == 'server' and hasattr(config, 'git_repositories') and config.git_repositories:
# Git服务器模式多仓库
host = getattr(config, 'git_server_host', 'localhost')
port = getattr(config, 'git_server_port', 9418)
username = getattr(config, 'git_server_username', '')
for repo_config in config.git_repositories:
# 获取仓库名称、路径和分支
if isinstance(repo_config, dict):
repo_name = repo_config.get('repository')
repo_path = repo_config.get('path', repo_name) # 优先使用path字段如果没有则使用repository
branch = repo_config.get('branch', '') # 如果没有指定分支,使用空字符串
else:
# 向后兼容:如果是字符串格式,使用默认分支
repo_name = repo_config
repo_path = repo_config
branch = getattr(config, 'git_branch', '') # 如果没有指定分支,使用空字符串
if not repo_name:
continue
# 构建git_url
if port == 9418:
# Git daemon协议
# 尝试两种格式:带.git后缀和不带.git后缀
git_url = f"git://{host}:{port}/{repo_path}"
# 同时支持带.git后缀的格式
if not repo_path.endswith('.git'):
git_url_with_suffix = f"git://{host}:{port}/{repo_path}.git"
else:
git_url_with_suffix = git_url
else:
# SSH协议
# 检查repo_path是否已经是完整路径或已经包含.git后缀
if repo_path.startswith('/') or repo_path.endswith('.git'):
# 如果是完整路径或已经包含.git后缀直接使用
if username:
git_url = f"ssh://{username}@{host}:{port}{repo_path}"
else:
git_url = f"ssh://{host}:{port}{repo_path}"
else:
# 否则,添加.git后缀
if username:
git_url = f"ssh://{username}@{host}:{port}/{repo_path}.git"
else:
git_url = f"ssh://{host}:{port}/{repo_path}.git"
# 初始化Git工具
git_tool = GitTool(
user_id="default", # 暂时使用默认用户ID
repo_id=f"{config.name}_{repo_name}",
git_config={
"git_url": git_url,
"branch": branch,
"ssh_key": config.ssh_key,
"https_token": config.https_token,
"password": getattr(config, 'git_server_password', None),
"local_repo_path": config.local_repo_path
}
)
# 存储仓库的分支信息,以便后续使用
git_tool.branch = branch
self.git_tools.append(git_tool)
else:
# 单个Git仓库
# 构建git_url
git_url = config.git_url
if hasattr(config, 'git_mode') and config.git_mode == 'server':
# 对于Git服务器模式构建git_url
host = getattr(config, 'git_server_host', 'localhost')
port = getattr(config, 'git_server_port', 9418)
repository = getattr(config, 'git_repository', '')
repo_path = repository
# 尝试从git_repositories中获取对应仓库的路径
if hasattr(config, 'git_repositories') and config.git_repositories:
for repo in config.git_repositories:
if isinstance(repo, dict) and repo.get('repository') == repository:
repo_path = repo.get('path', repository) # 优先使用path字段
break
if port == 9418:
# Git daemon协议
git_url = f"git://{host}:{port}/{repo_path}"
else:
# SSH协议
username = getattr(config, 'git_server_username', '')
# 检查repo_path是否已经是完整路径或已经包含.git后缀
if repo_path.startswith('/') or repo_path.endswith('.git'):
# 如果是完整路径或已经包含.git后缀直接使用
if username:
git_url = f"ssh://{username}@{host}:{port}{repo_path}"
else:
git_url = f"ssh://{host}:{port}{repo_path}"
else:
# 否则,添加.git后缀
if username:
git_url = f"ssh://{username}@{host}:{port}/{repo_path}.git"
else:
git_url = f"ssh://{host}:{port}/{repo_path}.git"
# 初始化Git工具
git_tool = GitTool(
user_id="default", # 暂时使用默认用户ID
repo_id=config.name,
git_config={
"git_url": git_url,
"branch": getattr(config, 'branch', ''), # 如果没有指定分支,使用空字符串
"ssh_key": config.ssh_key,
"https_token": config.https_token,
"password": getattr(config, 'git_server_password', None),
"local_repo_path": config.local_repo_path
}
)
self.git_tools.append(git_tool)
def fetch_all_documents(self) -> List[Dict[str, Any]]:
"""获取所有文档基于Git blob SHA进行文件级别去重"""
func_list = []
# 遍历所有Git工具实例支持多仓库
for git_tool in self.git_tools:
# 克隆/更新仓库
git_tool.clone_repo()
# 获取当前仓库所有文件的blob SHA
current_file_shas = git_tool.get_all_file_shas()
# 从ChromaDB获取已处理的文件SHA
processed_file_shas = self._get_processed_file_shas_from_chroma()
# 识别需要处理的新文件/修改文件
files_to_process = []
for file_path, current_sha in current_file_shas.items():
if file_path not in processed_file_shas or processed_file_shas[file_path] != current_sha:
files_to_process.append(file_path)
logger.info(f"仓库 {git_tool.repo_id} 文件去重结果: 总数{len(current_file_shas)}, 已处理{len(processed_file_shas)}, 待处理{len(files_to_process)}")
# 解析需要处理的文件
for file_path in files_to_process:
lang = ASTParser.detect_language(file_path)
if lang:
parser = ASTParser(file_path, lang)
try:
functions = parser.parse_functions()
for func in functions:
if func is None:
continue
func_id = generate_func_unique_id(
user_id="default",
repo_id=self.config.name,
branch=getattr(git_tool, 'branch', 'main'),
file_path=func["file_path"],
class_name=func.get("class_name"),
func_name=func["func_name"]
)
func['id'] = func_id
func['file_blob_sha'] = current_file_shas[file_path]
func_list.append(func)
except Exception as e:
logger.error(f"解析文件失败 {file_path}: {e}")
logger.info(f"获取到 {len(func_list)} 个函数")
return func_list
def _get_processed_file_shas_from_chroma(self, collection_key: str = 'default') -> Dict[str, str]:
"""从ChromaDB获取已处理的文件SHA映射"""
try:
if not self.vector_store_manager:
return {}
collection = self.vector_store_manager.collections.get(collection_key)
if not collection:
logger.warning(f"Collection {collection_key} not initialized")
return {}
results = collection.get(
where={
"$and": [
{"repo_id": {"$eq": self.config.name}},
{"branch": {"$eq": self.config.branch}}
]
}
)
processed_shas = {}
for metadata in results.get('metadatas', []):
if metadata and 'file_path' in metadata and 'file_blob_sha' in metadata:
file_path = metadata['file_path']
blob_sha = metadata['file_blob_sha']
processed_shas[file_path] = blob_sha
return processed_shas
except Exception as e:
logger.warning(f"从ChromaDB获取已处理文件SHA失败: {e}")
return {}
def doc_to_llamaindex_doc(self, doc: Dict) -> 'Document':
"""
转换函数信息为LlamaIndex Document
Args:
doc: 函数信息
Returns:
Document: LlamaIndex Document对象
"""
from llama_index.core import Document
# 生成函数唯一ID
func_id = doc.get('id')
# 生成函数描述
func_desc = self.generate_func_desc(doc)
# 获取函数体
func_body = doc.get("func_body", "")
# 创建Document对象
document = Document(
text=func_desc, # 使用描述作为文本
id_=func_id,
metadata={
"func_id": func_id,
"func_name": doc["func_name"],
"class_name": doc.get("class_name") if doc.get("class_name")!=None else "None",
"file_path": doc["file_path"],
"lang": doc["lang"],
"params": len(doc.get("params", [])), # 只存储参数数量,不存储完整参数列表
"return_type": doc.get("return_type") if doc.get("return_type")!=None else "None",
"docstring": (doc.get("docstring") or "")[:200], # 进一步限制文档字符串长度
"start_line": doc.get("start_line"),
"end_line": doc.get("end_line"),
"repo_id": self.config.name,
"branch": self.config.branch,
"func_body": doc["func_body"][:1000],
"file_blob_sha": doc.get("file_blob_sha", "")
}
)
return document
def generate_func_desc(self, func_info: Dict) -> str:
"""
生成函数描述
Args:
func_info: 函数信息
Returns:
str: 函数描述
"""
# 构建函数描述
parts = []
# 函数类型
if func_info.get("class_name"):
parts.append(f"{func_info['class_name']}类的{func_info['func_name']}方法")
else:
parts.append(f"{func_info['func_name']}函数")
# 参数信息
params = func_info.get("params", [])
if params:
param_str = []
for param in params:
if param.get("type"):
param_str.append(f"{param['name']}: {param['type']}")
else:
param_str.append(param['name'])
parts.append(f"接收参数: {', '.join(param_str)}")
# 返回值信息
return_type = func_info.get("return_type")
if return_type:
parts.append(f"返回类型: {return_type}")
# 文档字符串
docstring = func_info.get("docstring")
if docstring:
parts.append(f"功能描述: {docstring.strip()}")
else:
# NOTE如果没有文档字符串使用本地大模型根据函数体生成描述
# BUG: 同步服务终止后(删除对应配置 or 整个服务终止),仍在继续生成文档字符串
parts.append(f"功能描述: {self._generate_docstring_from_body(func_info.get('func_body', ''))}")
return ". ".join(parts)
def _generate_docstring_from_body(self, func_body: str) -> str:
"""
使用本地大模型根据函数体生成文档字符串
Args:
func_body: 函数体代码
Returns:
str: 生成的文档字符串
"""
if not func_body:
return "无文档字符串"
try:
from config import settings
from llama_index.llms.ollama import Ollama
# 初始化Ollama LLM
llm = Ollama(
model=settings.OLLAMA_MODEL,
base_url=settings.OLLAMA_BASE_URL,
temperature=0.3, # 降低温度,生成更确定的结果
request_timeout=300.0
)
# 构建提示词
prompt = f"""
请为以下函数生成简洁的文档字符串,描述其功能、参数和返回值:
{func_body}
要求:
1. 语言简洁明了不超过100字
2. 只返回文档字符串内容,不要包含其他内容
3. 重点描述函数的核心功能
"""
# 生成文档字符串
response = llm.complete(prompt)
generated_docstring = response.text.strip()
# 限制长度
if len(generated_docstring) > 200:
generated_docstring = generated_docstring[:200] + "..."
logger.debug(f"生成的文档字符串: {generated_docstring}")
return generated_docstring
except Exception as e:
logger.error(f"生成文档字符串失败: {e}")
# 降级方案:返回基于函数名的简单描述
return "执行相关操作的函数"
def fetch_new_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
"""
获取新文档(增量同步)
Args:
last_sync_time: 上次同步时间
Returns:
List[Dict[str, Any]]: 新函数信息列表
"""
func_list = []
# 遍历所有Git工具实例
for git_tool in self.git_tools:
# 检测远程更新
try:
has_update, local_commit, remote_commit = git_tool.detect_remote_update()
if not has_update:
logger.info(f"Git仓库 {git_tool.repo_id} 无更新")
continue
# 增量拉取
delta_files = git_tool.incremental_pull(local_commit, remote_commit)
#==== 处理重命名文件
rename_file_pairs = delta_files.get("RENAME", [])
if rename_file_pairs and self.vector_store_manager:
logger.info(f"处理重命名的文件: {len(rename_file_pairs)}")
for old_path, new_path in rename_file_pairs:
logger.info(f"重命名文件: {old_path} -> {new_path}")
# 从向量存储中更新元数据
try:
self.vector_store_manager.update_document_metadata(old_path, new_path)
except Exception as e:
logger.error(f"更新重命名文件元数据失败: {e}")
#==== 处理删除的文件
delete_files = delta_files.get("DELETE", [])
if delete_files and self.vector_store_manager:
logger.info(f"处理删除的文件: {len(delete_files)}")
for file_path in delete_files:
# 从向量存储中删除相关文档
try:
collection = self.vector_store_manager.collections.get('default')
if not collection:
logger.warning("Collection not initialized")
continue
# 获取所有文档的元数据
results = collection.get(include=['metadatas'])
metadatas = results.get('metadatas', [])
ids = results.get('ids', [])
# 找出需要删除的文档ID
to_delete_ids = []
for doc_id, metadata in zip(ids, metadatas):
if metadata and metadata.get('file_path') == file_path:
to_delete_ids.append(doc_id)
if to_delete_ids:
logger.info(f"删除文件 {file_path} 相关的 {len(to_delete_ids)} 个文档")
self.vector_store_manager.delete_documents(to_delete_ids)
except Exception as e:
logger.error(f"删除文件 {file_path} 相关文档失败: {e}")
#==== 解析新增/修改的文件
processed_files = set()
logger.info(f"增量更新处理文件: {len(delta_files.get('ADD', []))} 个新增, {len(delta_files.get('MODIFY', []))} 个修改")
for file_path in delta_files.get("ADD", []) + delta_files.get("MODIFY", []):
logger.debug(f"处理文件: {file_path}")
if file_path in processed_files:
logger.warning(f"文件 {file_path} 已被处理,跳过")
continue
processed_files.add(file_path)
lang = ASTParser.detect_language(file_path)
if lang:
parser = ASTParser(file_path, lang)
try:
functions = parser.parse_functions()
logger.debug(f"文件 {file_path} 解析出 {len(functions)} 个函数")
# 为每个函数生成doc_id并设置到字典中
for func in functions:
if func is None:
logger.warning(f"解析出空函数,跳过: {file_path}")
continue
# 生成唯一的文档ID
func_id = generate_func_unique_id(
user_id="default",
repo_id=self.config.name,
branch=getattr(git_tool, 'branch', 'main'),
file_path=func["file_path"],
class_name=func.get("class_name"),
func_name=func["func_name"]
)
logger.debug(f"生成函数ID: {func_id}")
func['id'] = func_id
func_list.extend(functions)
except Exception as e:
logger.error(f"解析文件失败 {file_path}: {e}")
except Exception as e:
logger.warning(f"处理Git仓库 {git_tool.repo_id} 增量同步失败: {e}")
return func_list
def get_synced_document_ids(self) -> Set[str]:
"""
获取已同步的文档ID
Returns:
Set[str]: 文档ID集合
"""
# 从向量存储中获取已同步的函数ID
if not self.vector_store_manager:
return set()
try:
# 获取所有已存在的文档ID
all_doc_ids = self.vector_store_manager.get_existing_doc_ids()
# 过滤出与当前Git仓库相关的文档ID
synced_ids = set()
# 获取所有文档的元数据,用于过滤
collection = self.vector_store_manager.collections.get('default')
if not collection:
logger.warning("Collection not initialized")
return set()
results = collection.get(include=['metadatas'])
metadatas = results.get('metadatas', [])
ids = results.get('ids', [])
for doc_id, metadata in zip(ids, metadatas):
if metadata and metadata.get('repo_id') == self.config.name:
synced_ids.add(doc_id)
logger.info(f"获取到 {len(synced_ids)} 个已同步的Git函数ID")
return synced_ids
except Exception as e:
logger.error(f"获取已同步文档ID失败: {e}")
return set()
def generate_doc_id(self, identifier: str) -> str:
"""
生成唯一的文档ID
Args:
identifier: 文档的唯一标识符(文件路径等)
Returns:
str: 唯一的文档ID
"""
from utils.func_id_generator import generate_func_unique_id
# 对于Git数据源使用函数唯一ID生成器
# 假设identifier是文件路径
return generate_func_unique_id(
user_id="default",
repo_id=self.config.name,
branch=self.config.branch,
file_path=identifier,
class_name="",
func_name=identifier.split('/')[-1].split('.')[0]
)
@staticmethod
def check_data_source_exists(config: BaseDataSourceConfig) -> bool:
"""
检查数据源是否存在
Args:
config: 数据源配置
Returns:
bool: 是否存在
"""
try:
# 构建git_url
git_url = config.git_url
if hasattr(config, 'git_mode') and config.git_mode == 'server':
# 对于Git服务器模式构建git_url
host = getattr(config, 'git_server_host', 'localhost')
port = getattr(config, 'git_server_port', 9418)
username = getattr(config, 'git_server_username', '')
repository = getattr(config, 'git_repository', '')
repo_path = repository
# 检查仓库名称是否存在
if not repository:
# 尝试使用git_repositories中的第一个仓库
if hasattr(config, 'git_repositories') and config.git_repositories:
first_repo = config.git_repositories[0]
if isinstance(first_repo, dict):
repository = first_repo.get('repository', '')
repo_path = first_repo.get('path', repository) # 优先使用path字段
else:
repository = first_repo
repo_path = repository
else:
# 尝试从git_repositories中获取对应仓库的路径
if hasattr(config, 'git_repositories') and config.git_repositories:
for repo in config.git_repositories:
if isinstance(repo, dict) and repo.get('repository') == repository:
repo_path = repo.get('path', repository) # 优先使用path字段
break
if port == 9418:
# Git daemon协议
git_url = f"git://{host}:{port}/{repo_path}"
else:
# SSH协议
# 检查repo_path是否已经是完整路径或已经包含.git后缀
if repo_path.startswith('/') or repo_path.endswith('.git'):
# 如果是完整路径或已经包含.git后缀直接使用
if username:
git_url = f"ssh://{username}@{host}{repo_path}"
else:
git_url = f"ssh://{host}{repo_path}"
else:
# 否则,添加.git后缀
if username:
git_url = f"ssh://{username}@{host}/{repo_path}.git"
else:
git_url = f"ssh://{host}/{repo_path}.git"
# 尝试克隆仓库
# 获取分支信息
branch = getattr(config, 'branch', '')
if not branch:
branch = getattr(config, 'git_branch', '')
# 尝试从git_repositories中获取第一个仓库的分支
if not branch and hasattr(config, 'git_repositories') and config.git_repositories:
first_repo = config.git_repositories[0]
if isinstance(first_repo, dict):
branch = first_repo.get('branch', '')
else:
branch = ''
git_tool = GitTool(
user_id="default",
repo_id=config.name,
git_config={
"git_url": git_url,
"branch": branch,
"ssh_key": config.ssh_key,
"https_token": config.https_token,
"password": getattr(config, 'git_server_password', None)
}
)
git_tool.clone_repo()
logger.info(f"Git数据源检查成功: {config.name}")
return True
except Exception as e:
logger.error(f"Git数据源检查失败: {e}")
return False