343 lines
13 KiB
Python
343 lines
13 KiB
Python
"""
|
||
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
|
||
# 初始化Git工具
|
||
self.git_tool = GitTool(
|
||
user_id="default", # 暂时使用默认用户ID
|
||
repo_id=config.name,
|
||
git_config={
|
||
"git_url": config.git_url,
|
||
"branch": config.branch,
|
||
"ssh_key": config.ssh_key,
|
||
"https_token": config.https_token,
|
||
"local_repo_path": config.local_repo_path
|
||
}
|
||
)
|
||
|
||
def fetch_all_documents(self) -> List[Dict[str, Any]]:
|
||
"""获取所有文档，基于Git blob SHA进行文件级别去重"""
|
||
# 克隆/更新仓库
|
||
self.git_tool.clone_repo()
|
||
|
||
# 获取当前仓库所有文件的blob SHA
|
||
current_file_shas = self.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"文件去重结果: 总数{len(current_file_shas)}, 已处理{len(processed_file_shas)}, 待处理{len(files_to_process)}")
|
||
|
||
# 解析需要处理的文件
|
||
func_list = []
|
||
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
|
||
|
||
# 生成函数ID并设置文件SHA标识符
|
||
func_id = generate_func_unique_id(
|
||
user_id="default",
|
||
repo_id=self.config.name,
|
||
branch=self.config.branch,
|
||
file_path=func["file_path"],
|
||
class_name=func.get("class_name"),
|
||
func_name=func["func_name"]
|
||
)
|
||
func['id'] = func_id
|
||
# 添加文件SHA到函数元数据
|
||
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)} 个函数（来自 {len(files_to_process)} 个文件）")
|
||
return func_list
|
||
def _get_processed_file_shas_from_chroma(self) -> Dict[str, str]:
|
||
"""从ChromaDB获取已处理的文件SHA映射"""
|
||
try:
|
||
if not self.vector_store_manager or not self.vector_store_manager.collection:
|
||
return {}
|
||
|
||
# 使用正确的ChromaDB查询语法：每个字段需要使用操作符
|
||
# 或者使用 $and 操作符组合多个条件
|
||
results = self.vector_store_manager.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)
|
||
|
||
# 创建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", "")[: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], # 限制函数体长度，避免metadata过长
|
||
"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()}")
|
||
|
||
return ". ".join(parts)
|
||
|
||
def fetch_new_documents(self, last_sync_time=None) -> List[Dict[str, Any]]:
|
||
"""
|
||
获取新文档（增量同步）
|
||
|
||
Args:
|
||
last_sync_time: 上次同步时间
|
||
|
||
Returns:
|
||
List[Dict[str, Any]]: 新函数信息列表
|
||
"""
|
||
# 检测远程更新
|
||
has_update, local_commit, remote_commit = self.git_tool.detect_remote_update()
|
||
|
||
if not has_update:
|
||
logger.info("Git仓库无更新")
|
||
return []
|
||
|
||
# 增量拉取
|
||
delta_files = self.git_tool.incremental_pull(local_commit, remote_commit)
|
||
|
||
# 解析新增/修改的文件
|
||
processed_files = set()
|
||
func_list = []
|
||
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:
|
||
# 生成唯一的文档ID
|
||
func_id = generate_func_unique_id(
|
||
user_id="default",
|
||
repo_id=self.config.name,
|
||
branch=self.config.branch,
|
||
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}")
|
||
|
||
logger.info(f"增量同步获取到 {len(func_list)} 个函数")
|
||
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()
|
||
|
||
# 获取所有文档的元数据，用于过滤
|
||
results = self.vector_store_manager.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_tool = GitTool(
|
||
user_id="default",
|
||
repo_id=config.name,
|
||
git_config={
|
||
"git_url": config.git_url,
|
||
"branch": config.branch,
|
||
"ssh_key": config.ssh_key,
|
||
"https_token": config.https_token
|
||
}
|
||
)
|
||
git_tool.clone_repo()
|
||
logger.info(f"Git数据源检查成功: {config.name}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Git数据源检查失败: {e}")
|
||
return False |