mnote/wolai-backend/app/services/lightrag_service.py

"""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()