"""LightRAG 集成。负责管理实例、增量索引与问答。""" from __future__ import annotations import asyncio import logging import os from pathlib import Path import sys from contextvars import ContextVar from typing import Any, AsyncIterator, Dict, List, Optional from urllib.parse import urlparse import httpx import numpy as np from openai import AsyncOpenAI def _ensure_lightrag_available() -> None: """确保本地 LightRAG 代码可被 Python 找到.""" repo_root = Path(__file__).resolve().parents[3] local_pkg = repo_root / "LightRAG" if local_pkg.exists(): path_str = str(local_pkg) if path_str not in sys.path: sys.path.insert(0, path_str) logger = logging.getLogger(__name__) # 优先将本地 LightRAG 代码加入 sys.path，避免导入失败走占位实现 _ensure_lightrag_available() _EMBEDDING_MODEL_OVERRIDE: ContextVar[Optional[str]] = ContextVar( "lightrag_embedding_model_override", default=None, ) try: from lightrag import LightRAG, QueryParam from lightrag.kg.shared_storage import initialize_pipeline_status from lightrag.llm.ollama import ollama_model_complete, ollama_embed from lightrag.utils import logger as lightrag_logger _LIGHTRAG_AVAILABLE = True _IMPORT_ERROR: Optional[Exception] = None except Exception as exc: # pragma: no cover - 本地缺失或版本不兼容时使用占位实现 _LIGHTRAG_AVAILABLE = False _IMPORT_ERROR = exc lightrag_logger = logging.getLogger("lightrag_stub") # 占位的 ollama 方法，避免引用错误 async def ollama_model_complete(*_: Any, **__: Any) -> str: # type: ignore[override] return "" async def ollama_embed(*_: Any, **__: Any) -> list[list[float]]: # type: ignore[override] return [] class QueryParam: # type: ignore[override] def __init__(self, mode: str = "mix") -> None: self.mode = mode self.stream = True class LightRAG: # type: ignore[override] async def initialize_storages(self) -> None: return None async def ainsert(self, *_: Any, **__: Any) -> str: return "lightrag-skipped" async def aquery_llm(self, *_: Any, **__: Any) -> Dict[str, Any]: return { "llm_response": { "response_iterator": iter(()), "content": "LightRAG 已跳过", "is_streaming": False, }, "data": {"references": []}, "metadata": {"skipped": True}, } async def initialize_pipeline_status() -> None: # type: ignore[override] return None def gpt_4o_mini_complete(*_: Any, **__: Any) -> str: # type: ignore[override] return "" def openai_embed(*_: Any, **__: Any) -> List[float]: # type: ignore[override] return [] _ensure_lightrag_available() from app.config import settings from app.services.supabase_rest import supabase_rest class LightRAGService: """ 管理 LightRAG 单例、工作空间隔离与增量索引。 - 每个 workspace 映射为一个 LightRAG 实例（共享 Postgres） - 支持在同步/异步上下文中调用 """ def __init__(self) -> None: self.collection = settings.lightrag_collection self._instances: Dict[str, LightRAG] = {} self._locks: Dict[str, asyncio.Lock] = {} self._pipeline_ready = False self._configure_pg_env() if not _LIGHTRAG_AVAILABLE and _IMPORT_ERROR: logger.warning("LightRAG 不可用，使用占位实现：%s", _IMPORT_ERROR) self._working_dir = ( Path(__file__).resolve().parent.parent / "runtime" / "lightrag_cache" ) self._working_dir.mkdir(parents=True, exist_ok=True) self._availability_error: Optional[str] = None self._ollama_host = settings.ollama_base_url self._use_deepseek = bool(getattr(settings, "deepseek_api_key", "")) self._deepseek_client: Optional[AsyncOpenAI] = None if self._use_deepseek: self._deepseek_client = AsyncOpenAI( api_key=settings.deepseek_api_key, base_url=getattr(settings, "deepseek_base_url", None), ) self._llm_model_name = ( getattr(settings, "deepseek_model", "deepseek-chat") if self._use_deepseek else settings.lightrag_llm_model ) def _configure_pg_env(self) -> None: """根据配置将 pgvector 连接信息注入 LightRAG 需要的环境变量。""" parsed = urlparse(settings.lightrag_db_url) if parsed.scheme not in {"postgresql", "postgres"}: self._availability_error = "LIGHTRAG_DB_URL 必须是 Postgres 连接串" return if parsed.username: os.environ.setdefault("POSTGRES_USER", parsed.username) if parsed.password: os.environ.setdefault("POSTGRES_PASSWORD", parsed.password) if parsed.hostname: os.environ.setdefault("POSTGRES_HOST", parsed.hostname) if parsed.port: os.environ.setdefault("POSTGRES_PORT", str(parsed.port)) if parsed.path and len(parsed.path) > 1: os.environ.setdefault("POSTGRES_DATABASE", parsed.path.lstrip("/")) # Ollama 走本地 HTTP，无需 OpenAI key # 允许通过环境变量关闭 KG 抽取，避免大模型深拷贝异常 os.environ.setdefault("LIGHTRAG_DISABLE_ENTITY_RELATION", "true") # 为大维度向量配置 IVFFlat，避免 HNSW 2000 维限制 os.environ.setdefault("POSTGRES_VECTOR_INDEX_TYPE", "IVFFlat") os.environ.setdefault("EMBEDDING_DIM", str(settings.lightrag_embedding_dim if hasattr(settings, "lightrag_embedding_dim") else 4096)) def _namespace(self, workspace_id: Optional[str], user_id: str, *, prefer_deepseek: bool = False) -> str: """生成 LightRAG workspace 名称，优先 workspace，其次 user；根据模型标记区分实例。""" base = f"workspace_{workspace_id}" if workspace_id else f"user_{user_id}" if prefer_deepseek: return f"{base}_deepseek" return base def _is_available(self) -> tuple[bool, Optional[str]]: """ 判断 LightRAG 是否具备运行条件。 - import 失败或配置错误时返回 False 和原因 """ if not _LIGHTRAG_AVAILABLE: return False, f"LightRAG 导入失败: {self._availability_error or _IMPORT_ERROR}" if self._availability_error: # 缺少关键配置时也视为不可用 return False, self._availability_error return True, None async def _wait_with_timeout(self, coro: Any, *, timeout: float = 300.0) -> Any: """为外部调用包一层超时，避免卡住 worker / healthcheck。""" return await asyncio.wait_for(coro, timeout=timeout) async def _ollama_llm(self, *args: Any, **kwargs: Any) -> Any: """ 固定使用配置好的 Ollama LLM。 LightRAG 会传入 prompt/system_prompt/history_messages。 """ prompt = kwargs.pop("prompt", None) if prompt is None and args: prompt = args[0] if prompt is None: raise ValueError("缺少 prompt，无法调用 Ollama LLM") stream_flag = bool(kwargs.pop("stream", False)) timeout = kwargs.pop("timeout", None) return await ollama_model_complete( prompt=prompt, host=self._ollama_host, timeout=timeout, stream=stream_flag, **kwargs, ) async def _deepseek_llm(self, *args: Any, **kwargs: Any) -> Any: """ 使用 DeepSeek 在线模型，符合 LightRAG 的 llm_model_func 接口。 """ if not self._deepseek_client: raise RuntimeError("DeepSeek 客户端未初始化") prompt = kwargs.pop("prompt", None) if prompt is None and args: prompt = args[0] if prompt is None: raise ValueError("缺少 prompt，无法调用 DeepSeek LLM") system_prompt = kwargs.pop("system_prompt", None) history = kwargs.pop("history_messages", []) or [] stream_flag = bool(kwargs.pop("stream", False)) temperature = kwargs.pop("temperature", 0.2) max_tokens = kwargs.pop("max_tokens", 512) # 构建消息 messages = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) messages.extend(history) messages.append({"role": "user", "content": prompt}) client = self._deepseek_client if stream_flag: response = await client.chat.completions.create( model=self._llm_model_name, messages=messages, stream=True, temperature=temperature, max_tokens=max_tokens, ) async def _aiter(): async for chunk in response: delta = chunk.choices[0].delta.content or "" if delta: yield delta return _aiter() else: response = await client.chat.completions.create( model=self._llm_model_name, messages=messages, stream=False, temperature=temperature, max_tokens=max_tokens, ) return response.choices[0].message.content async def _ollama_embed(self, texts: list[str], **kwargs: Any) -> Any: """固定使用配置好的 Ollama embedding 模型。""" timeout = kwargs.pop("timeout", None) embed_model_override = kwargs.pop("embed_model", None) or _EMBEDDING_MODEL_OVERRIDE.get() embed_model = embed_model_override or settings.lightrag_embedding_model return await ollama_embed( texts, embed_model=embed_model, host=self._ollama_host, timeout=timeout, **kwargs, ) async def _embedding_rerank(self, query: str, documents: List[str], **_: Any) -> List[Dict[str, Any]]: """ 简易基于向量的 rerank：使用指定 embedding 模型计算 query/doc 向量并按余弦相似度排序。 若调用失败则返回原序。 """ if not documents: return [] try: query_vec = await self._ollama_embed([query]) doc_vecs = await self._ollama_embed(documents) q = np.array(query_vec[0], dtype=float) d = np.array(doc_vecs, dtype=float) q_norm = np.linalg.norm(q) + 1e-8 d_norm = np.linalg.norm(d, axis=1) + 1e-8 scores = (d @ q) / (d_norm * q_norm) order = np.argsort(-scores) return [ {"index": int(idx), "relevance_score": float(scores[idx])} for idx in order ] except Exception as exc: logger.warning("rerank 失败，回退原序：%s", exc) return [{"index": i, "relevance_score": 0.0} for i in range(len(documents))] async def _get_instance(self, workspace: str, *, prefer_deepseek: bool = False) -> LightRAG: if workspace in self._instances: return self._instances[workspace] lock = self._locks.setdefault(workspace, asyncio.Lock()) async with lock: if workspace in self._instances: return self._instances[workspace] llm_func = self._deepseek_llm if (self._use_deepseek and prefer_deepseek) else self._ollama_llm rag = LightRAG( working_dir=str(self._working_dir), workspace=workspace, kv_storage="PGKVStorage", vector_storage="PGVectorStorage", graph_storage="NetworkXStorage", doc_status_storage="PGDocStatusStorage", llm_model_func=llm_func, llm_model_name=self._llm_model_name if (self._use_deepseek and prefer_deepseek) else settings.lightrag_llm_model, llm_model_kwargs={ "options": { # 限制生成长度，避免本地大模型回答过慢 "num_predict": 512, "temperature": 0.2, } }, embedding_func=self._ollama_embed, rerank_model_func=self._embedding_rerank, ) await self._wait_with_timeout(rag.initialize_storages()) if not self._pipeline_ready: await self._wait_with_timeout(initialize_pipeline_status()) self._pipeline_ready = True self._instances[workspace] = rag lightrag_logger.info("LightRAG workspace %s ready", workspace) return rag async def index_document_async( self, *, document_id: str, user_id: str, workspace_id: Optional[str], text: str, title: Optional[str] = None, ) -> str: workspace = self._namespace(workspace_id, user_id) rag = await self._get_instance(workspace) clean_text = text.strip() if not clean_text: raise ValueError("空文本无法建立 LightRAG 索引") file_path = f"doc://{document_id}" if title: file_path = f"{file_path}?title={title}" track_id = await rag.ainsert( clean_text, ids=[document_id], file_paths=[file_path], ) return track_id def index_document( self, *, document_id: str, user_id: str, workspace_id: Optional[str], text: str, title: Optional[str] = None, ) -> str: """同步环境（如 Celery）调用的封装。""" return asyncio.run( self.index_document_async( document_id=document_id, user_id=user_id, workspace_id=workspace_id, text=text, title=title, ) ) async def query_async( self, *, query_text: str, user_id: str, workspace_id: Optional[str], stream: bool = True, mode: str = "naive", model_choice: Optional[str] = None, rag_settings: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: prefer_deepseek = model_choice == "deepseek" workspace = self._namespace(workspace_id, user_id, prefer_deepseek=prefer_deepseek) rag = await self._get_instance(workspace, prefer_deepseek=prefer_deepseek) param = QueryParam(mode=mode, top_k=6, chunk_top_k=6, max_total_tokens=4096) if rag_settings: try: if rag_settings.get("mode"): param.mode = str(rag_settings["mode"]) if isinstance(rag_settings.get("top_k"), (int, float)): param.top_k = int(rag_settings["top_k"]) if isinstance(rag_settings.get("chunk_top_k"), (int, float)): param.chunk_top_k = int(rag_settings["chunk_top_k"]) if isinstance(rag_settings.get("max_entity_tokens"), (int, float)): param.max_entity_tokens = int(rag_settings["max_entity_tokens"]) if isinstance(rag_settings.get("max_relation_tokens"), (int, float)): param.max_relation_tokens = int(rag_settings["max_relation_tokens"]) if isinstance(rag_settings.get("max_total_tokens"), (int, float)): param.max_total_tokens = int(rag_settings["max_total_tokens"]) if "enable_rerank" in rag_settings: param.enable_rerank = bool(rag_settings["enable_rerank"]) if rag_settings.get("user_prompt"): param.user_prompt = str(rag_settings["user_prompt"]) except Exception as exc: logger.warning("解析 RAG 参数失败，继续使用默认值: %s", exc) param.stream = stream embedding_token = None if rag_settings and isinstance(rag_settings.get("embedding_model"), str): embedding_token = _EMBEDDING_MODEL_OVERRIDE.set(str(rag_settings["embedding_model"])) try: result = await self._wait_with_timeout(rag.aquery_llm(query_text, param)) finally: if embedding_token is not None: _EMBEDDING_MODEL_OVERRIDE.reset(embedding_token) return result async def stream_answer( self, *, query_text: str, user_id: str, workspace_id: Optional[str], model_choice: Optional[str] = None, document_id: Optional[str] = None, ) -> Dict[str, Any]: """统一返回结构，包含 SSE 需要的 iterator 与引用。""" ok, reason = self._is_available() if not ok: return { "references": [], "iterator": iter(()), "content": f"LightRAG 未就绪：{reason}", "is_streaming": False, "metadata": {"skipped": True, "reason": reason}, } rag_settings = self._load_rag_settings(document_id) rag_mode = (rag_settings or {}).get("mode") result = await self.query_async( query_text=query_text, user_id=user_id, workspace_id=workspace_id, stream=True, model_choice=model_choice, mode=str(rag_mode) if isinstance(rag_mode, str) else "naive", rag_settings=rag_settings, ) llm_resp = result.get("llm_response", {}) references: List[Dict[str, Any]] = ( result.get("data", {}).get("references", []) or [] ) return { "references": references, "iterator": llm_resp.get("response_iterator"), "content": llm_resp.get("content"), "is_streaming": llm_resp.get("is_streaming", False), "metadata": result.get("metadata", {}), } def _load_rag_settings(self, document_id: Optional[str]) -> Optional[Dict[str, Any]]: if not document_id: return None try: doc = supabase_rest.select_one("documents", {"id": document_id}) except Exception as exc: logger.warning("读取文档 %s 的 RAG 配置失败：%s", document_id, exc) return None value = doc.get("rag_settings") if doc else None if isinstance(value, dict): return value return None async def answer_with_context( self, *, query_text: str, context: str, prefer_deepseek: bool = False, ) -> Dict[str, Any]: """基于外部上下文的简单问答，优先走 DeepSeek，无则回退 Ollama。""" prompt = ( "你是检索结果总结助手。根据以下搜索摘要与来源回答用户问题，" "答案需简洁且引用要点。保持中文输出，并在内容后附上引用编号。\n\n" f"【搜索摘要】\n{context}\n\n【用户问题】{query_text}" ) prefer_deepseek = prefer_deepseek and self._use_deepseek stream_flag = True iterator: Optional[AsyncIterator[str]] = None content: Optional[str] = None reason: Optional[str] = None try: if prefer_deepseek and self._deepseek_client: resp = await self._deepseek_llm(prompt, stream=stream_flag) iterator = resp if hasattr(resp, "__aiter__") else None content = None if iterator else str(resp) else: resp = await self._ollama_llm(prompt, stream=stream_flag) iterator = resp if hasattr(resp, "__aiter__") else None content = None if iterator else str(resp) except Exception as exc: # pragma: no cover - 运行时保护 iterator = None content = f"生成失败：{exc}" reason = str(exc) return { "references": [], "iterator": iterator, "content": content, "is_streaming": iterator is not None, "metadata": {"reason": reason} if reason else {}, } async def run_healthcheck(self) -> Dict[str, Any]: workspace = "__healthcheck__" ok, reason = self._is_available() if not ok: return {"ok": False, "workspace": workspace, "error": reason} try: rag = await self._get_instance(workspace) await self._wait_with_timeout(rag.doc_status.initialize(), timeout=5.0) return {"ok": True, "workspace": workspace} except Exception as exc: # pragma: no cover - 调试辅助 return {"ok": False, "workspace": workspace, "error": str(exc)} lightrag_service = LightRAGService()