import os import json import time import traceback import shutil import asyncio from pathlib import Path from typing import Dict, List, Any, Optional from datetime import datetime from lightrag import LightRAG, QueryParam from lightrag.llm.openai import openai_complete_if_cache, openai_embed from lightrag.utils import EmbeddingFunc, setup_logger from lightrag.kg.shared_storage import initialize_pipeline_status from src import config from src.utils import logger, hashstr, get_docker_safe_url from src.plugins import ocr work_dir = os.path.join(config.save_dir, "lightrag_data") log_dir = os.path.join(work_dir, "logs", "lightrag") setup_logger("lightrag", log_file_path=os.path.join(log_dir, f"lightrag_{datetime.now().strftime('%Y-%m-%d')}.log")) class LightRagBasedKB: """基于 LightRAG 的知识库管理类""" def __init__(self) -> None: # 存储 LightRAG 实例映射 {db_id: LightRAG} self.instances: Dict[str, LightRAG] = {} # 数据库元信息存储 {db_id: metadata} self.databases_meta: Dict[str, dict] = {} # 文件信息存储 {file_id: file_info} self.files_meta: Dict[str, dict] = {} # 工作目录 self.work_dir = os.path.join(config.save_dir, "lightrag_data") os.makedirs(self.work_dir, exist_ok=True) # 加载已有的元数据 self._load_metadata() logger.info("LightRagBasedKB initialized") def _load_metadata(self): """加载元数据""" meta_file = os.path.join(self.work_dir, "metadata.json") if os.path.exists(meta_file): try: with open(meta_file, 'r', encoding='utf-8') as f: data = json.load(f) self.databases_meta = data.get("databases", {}) self.files_meta = data.get("files", {}) logger.info(f"Loaded metadata for {len(self.databases_meta)} databases") except Exception as e: logger.error(f"Failed to load metadata: {e}") def _save_metadata(self): """保存元数据""" # 确保工作目录存在 os.makedirs(self.work_dir, exist_ok=True) meta_file = os.path.join(self.work_dir, "metadata.json") try: data = { "databases": self.databases_meta, "files": self.files_meta } with open(meta_file, 'w', encoding='utf-8') as f: json.dump(data, f, ensure_ascii=False, indent=2) except Exception as e: logger.error(f"Failed to save metadata: {e}") async def _get_lightrag_instance(self, db_id: str) -> Optional[LightRAG]: """获取或创建 LightRAG 实例""" logger.info(f"Getting or creating LightRAG instance for {db_id}") if db_id in self.instances: return self.instances[db_id] if db_id not in self.databases_meta: return None llm_info = self.databases_meta[db_id].get("llm_info", {}) embed_info = self.databases_meta[db_id].get("embed_info", {}) logger.info(f"LLM info: {llm_info}") logger.info(f"Embed info: {embed_info}") # 创建 LightRAG 实例 working_dir = os.path.join(self.work_dir, db_id) os.makedirs(working_dir, exist_ok=True) try: # 使用配置的 LLM 和 embedding 函数 rag = LightRAG( working_dir=working_dir, llm_model_func=self._get_llm_func(llm_info), embedding_func=self._get_embedding_func(embed_info), vector_storage="MilvusVectorDBStorage", kv_storage="JsonKVStorage", graph_storage="PGGraphStorage", doc_status_storage="JsonDocStatusStorage", log_file_path=os.path.join(self.work_dir, db_id, "lightrag.log"), ) # 异步初始化存储 await self._initialize_rag_storages(rag) self.instances[db_id] = rag return rag except Exception as e: logger.error(f"Failed to create LightRAG instance for {db_id}: {e}") logger.error(f"Traceback: {traceback.format_exc()}") return None async def _initialize_rag_storages(self, rag: LightRAG): """异步初始化 LightRAG 存储""" logger.info(f"Initializing LightRAG storages for {rag.working_dir}") await rag.initialize_storages() await initialize_pipeline_status() def _get_llm_func(self, llm_info: dict): """获取 LLM 函数""" # llm_info = llm_info | { # "model_name": "qwen3-1.7b", # "provider": "custom" # } # provider_info = config.model_names[llm_info.get("provider")] # api_key = os.getenv(provider_info.get("env")[0] or "OPENAI_API_KEY") or "no_api_key" # base_url = get_docker_safe_url(provider_info.get("base_url", "http://localhost:8081/v1")) from src.models import get_custom_model llm_info = get_custom_model("qwen3:32b-RFnC") async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs): return await openai_complete_if_cache( llm_info.get("name", "qwen3-1.7b"), prompt, system_prompt=system_prompt, history_messages=history_messages, api_key=llm_info.get("api_key"), base_url=get_docker_safe_url(llm_info.get("api_base")), extra_body={"enable_thinking": False}, **kwargs, ) return llm_model_func def _get_embedding_func(self, embed_info: dict): """获取 embedding 函数""" api_key = os.getenv(embed_info.get("api_key", "OPENAI_API_KEY")) or "no_api_key" base_url = embed_info.get("base_url", "http://localhost:8081/v1").replace("/embeddings", "") return EmbeddingFunc( embedding_dim=embed_info.get("dimension") or 1024, max_token_size=4096, func=lambda texts: openai_embed( texts=texts, model=embed_info.get("model_name") or "Qwen3-Embedding-0.6B", api_key=api_key, base_url=get_docker_safe_url(base_url) ), ) async def _process_file_to_markdown(self, file_path: str, params: dict | None = None) -> str: """将不同类型的文件转换为 markdown 格式""" file_path_obj = Path(file_path) file_ext = file_path_obj.suffix.lower() if file_ext == '.pdf': # 使用 OCR 处理 PDF from src.core.indexing import parse_pdf_async text = await parse_pdf_async(str(file_path_obj), params=params) return f"Using OCR to process {file_path_obj.name}\n\n{text}" elif file_ext in ['.txt', '.md']: # 直接读取文本文件 with open(file_path_obj, 'r', encoding='utf-8') as f: content = f.read() return f"# {file_path_obj.name}\n\n{content}" elif file_ext in ['.doc', '.docx']: # 处理 Word 文档 from docx import Document # type: ignore doc = Document(file_path_obj) text = '\n'.join([para.text for para in doc.paragraphs]) return f"# {file_path_obj.name}\n\n{text}" elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp']: # 使用 OCR 处理图片 text = ocr.process_image(str(file_path_obj)) return f"# {file_path_obj.name}\n\n{text}" else: # 尝试作为文本文件读取 import textract # type: ignore text = textract.process(file_path_obj) return f"# {file_path_obj.name}\n\n{text}" async def _process_url_to_markdown(self, url: str, params: dict | None = None) -> str: """将 URL 转换为 markdown 格式""" import requests from bs4 import BeautifulSoup response = requests.get(url, timeout=30) soup = BeautifulSoup(response.content, 'html.parser') text_content = soup.get_text() return f"# {url}\n\n{text_content}" # ============================================================================= # data_router.py 中使用的核心方法 # ============================================================================= def get_databases(self): """获取所有数据库信息 - data_router.py 使用""" databases = [] for db_id, meta in self.databases_meta.items(): db_dict = meta.copy() db_dict["db_id"] = db_id # 获取文件信息 db_files = {} for file_id, file_info in self.files_meta.items(): if file_info.get("database_id") == db_id: db_files[file_id] = { "file_id": file_id, "filename": file_info.get("filename", ""), "path": file_info.get("path", ""), "type": file_info.get("file_type", ""), "status": file_info.get("status", "done"), "created_at": file_info.get("created_at", time.time()) } db_dict["files"] = db_files db_dict["row_count"] = len(db_files) db_dict["status"] = "已连接" databases.append(db_dict) return {"databases": databases} def create_database(self, database_name, description, embed_info: dict | None = None, **kwargs): """创建数据库 - data_router.py 使用""" db_id = f"kb_{hashstr(database_name, with_salt=True)}" # 创建数据库记录 self.databases_meta[db_id] = { "name": database_name, "description": description, "embed_info": embed_info, "metadata": kwargs, "created_at": datetime.now().isoformat() } self._save_metadata() # 创建工作目录 working_dir = os.path.join(self.work_dir, db_id) os.makedirs(working_dir, exist_ok=True) # 返回数据库信息 db_dict = self.databases_meta[db_id].copy() db_dict["db_id"] = db_id db_dict["files"] = {} return db_dict def delete_database(self, db_id): """删除数据库 - data_router.py 使用""" # TODO 删除数据库时，需要删除文件记录，并删除 LightRAG 中的文件 if db_id in self.databases_meta: # 删除相关文件记录 files_to_delete = [fid for fid, finfo in self.files_meta.items() if finfo.get("database_id") == db_id] for file_id in files_to_delete: del self.files_meta[file_id] # 删除数据库记录 del self.databases_meta[db_id] # 删除 LightRAG 实例 if db_id in self.instances: del self.instances[db_id] self._save_metadata() # 删除工作目录 working_dir = os.path.join(self.work_dir, db_id) if os.path.exists(working_dir): try: shutil.rmtree(working_dir) except Exception as e: logger.error(f"Error deleting working directory {working_dir}: {e}") return {"message": "删除成功"} async def add_content(self, db_id, items, params: dict | None = None): """通用的内容添加方法 - 支持文件和URL""" if db_id not in self.databases_meta: raise ValueError(f"Database {db_id} not found") rag = await self._get_lightrag_instance(db_id) if not rag: raise ValueError(f"Failed to get LightRAG instance for {db_id}") content_type = params.get('content_type', 'file') if params else 'file' processed_items_info = [] for item in items: # 根据内容类型生成不同的ID和文件名 if content_type == "file": file_path = Path(item) file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}" file_type = file_path.suffix.lower().replace(".", "") filename = file_path.name item_path = str(file_path) else: # URL file_id = f"url_{hashstr(item + str(time.time()), 6)}" file_type = "url" filename = f"webpage_{hashstr(item, 6)}.md" item_path = item # 添加文件记录 file_record = { "database_id": db_id, "filename": filename, "path": item_path, "file_type": file_type, "status": "processing", "created_at": time.time() } self.files_meta[file_id] = file_record self._save_metadata() # 添加 file_id 到返回数据 file_record = file_record.copy() file_record["file_id"] = file_id try: # 根据内容类型处理内容 if content_type == "file": markdown_content = await self._process_file_to_markdown(item, params=params) logger.info(f"Markdown content: {markdown_content[:100].replace('\n', ' ')}...") else: # URL markdown_content = await self._process_url_to_markdown(item, params=params) # 使用 LightRAG 插入内容 await rag.ainsert( input=markdown_content, ids=file_id, file_paths=item_path ) logger.info(f"Inserted {content_type} {item} into LightRAG. Done.") # 更新状态为完成 self.files_meta[file_id]["status"] = "done" self._save_metadata() file_record['status'] = "done" except Exception as e: logger.error(f"处理{content_type} {item} 失败: {e}, {traceback.format_exc()}") self.files_meta[file_id]["status"] = "failed" self._save_metadata() file_record['status'] = "failed" processed_items_info.append(file_record) return processed_items_info def get_database_info(self, db_id): """获取数据库详细信息 - data_router.py 使用""" if db_id not in self.databases_meta: return None meta = self.databases_meta[db_id].copy() meta["db_id"] = db_id # 获取文件信息 db_files = {} for file_id, file_info in self.files_meta.items(): if file_info.get("database_id") == db_id: db_files[file_id] = { "file_id": file_id, "filename": file_info.get("filename", ""), "path": file_info.get("path", ""), "type": file_info.get("file_type", ""), "status": file_info.get("status", "done"), "created_at": file_info.get("created_at", time.time()) } meta["files"] = db_files meta["row_count"] = len(db_files) meta["status"] = "已连接" return meta async def delete_file(self, db_id, file_id): """删除文件 - data_router.py 使用""" # TODO 删除文件时，需要删除文件记录，并删除 LightRAG 中的文件 rag = await self._get_lightrag_instance(db_id) if rag: try: # 使用 LightRAG 删除文档 await rag.adelete_by_doc_id(file_id) except Exception as e: logger.error(f"Error deleting file {file_id} from LightRAG: {e}") # 删除文件记录 if file_id in self.files_meta: del self.files_meta[file_id] self._save_metadata() async def get_file_info(self, db_id, file_id): """获取文件信息和其 chunks - data_router.py 使用""" if file_id not in self.files_meta: raise Exception(f"File not found: {file_id}") # 使用 LightRAG 获取 chunks rag = await self._get_lightrag_instance(db_id) if rag: try: # 获取文档的所有 chunks assert hasattr(rag.text_chunks, 'get_all'), "text_chunks does not have get_all method" all_chunks = await rag.text_chunks.get_all() # 筛选属于该文档的 chunks doc_chunks = [] for chunk_id, chunk_data in all_chunks.items(): if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == file_id: chunk_data["id"] = chunk_id chunk_data["content_vector"] = [] doc_chunks.append(chunk_data) # 按 chunk_order_index 排序 doc_chunks.sort(key=lambda x: x.get("chunk_order_index", 0)) # logger.debug(f"All chunks: {doc_chunks}") return {"lines": doc_chunks} except Exception as e: logger.error(f"Error getting chunks for file {file_id}: {e}") return {"lines": []} def get_db_upload_path(self, db_id=None): """获取数据库上传路径 - data_router.py 使用""" if db_id: uploads_folder = os.path.join(self.work_dir, db_id, "uploads") os.makedirs(uploads_folder, exist_ok=True) return uploads_folder general_uploads = os.path.join(self.work_dir, "uploads") os.makedirs(general_uploads, exist_ok=True) return general_uploads def update_database(self, db_id, name, description): """更新数据库 - data_router.py 使用""" if db_id not in self.databases_meta: raise ValueError(f"数据库 {db_id} 不存在") self.databases_meta[db_id]["name"] = name self.databases_meta[db_id]["description"] = description self._save_metadata() # 返回更新后的数据库信息 return self.get_database_info(db_id) # ============================================================================= # 为了系统兼容性需要的其他方法 # ============================================================================= def query(self, query_text, db_id, **kwargs): logger.warning("query is deprecated, use aquery instead") return asyncio.run(self.aquery(query_text, db_id, **kwargs)) async def aquery(self, query_text, db_id, **kwargs): """查询知识库 - 用于检索器""" rag = await self._get_lightrag_instance(db_id) if not rag: raise ValueError(f"Database {db_id} not found") try: # 设置查询参数 params_dict = { "mode": "mix", "only_need_context": True, "top_k": 10, } | kwargs param = QueryParam(**params_dict) # 执行查询 response = await rag.aquery(query_text, param) logger.debug(f"Query response: {response}") return response except Exception as e: logger.error(f"Query error: {e}, {traceback.format_exc()}") return "" def get_retrievers(self): """获取所有检索器 - 用于工具系统""" retrievers = {} for db_id, meta in self.databases_meta.items(): def make_retriever(db_id): async def retriever(query_text): return await self.aquery(query_text, db_id) return retriever retrievers[db_id] = { "name": meta["name"], "description": meta["description"], "retriever": make_retriever(db_id), } return retrievers